// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        corev1 "k8s.io/api/core/v1"
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/runtime"
)

// Condition types for EmbeddingServer (reuses common conditions from MCPServer)
// ConditionPodTemplateValid is shared with MCPServer

const (
        // ConditionModelReady indicates whether the embedding model is downloaded and ready
        ConditionModelReady = "ModelReady"

        // ConditionVolumeReady indicates whether the PVC for model caching is ready
        ConditionVolumeReady = "VolumeReady"
)

// Condition reasons for EmbeddingServer
// Image validation and PodTemplate reasons are shared with MCPServer

const (
        // ConditionReasonModelDownloading indicates the model is being downloaded
        ConditionReasonModelDownloading = "ModelDownloading"
        // ConditionReasonModelReady indicates the model is downloaded and ready
        ConditionReasonModelReady = "ModelReady"
        // ConditionReasonModelFailed indicates the model download or initialization failed
        ConditionReasonModelFailed = "ModelFailed"

        // ConditionReasonVolumeCreating indicates the PVC is being created
        ConditionReasonVolumeCreating = "VolumeCreating"
        // ConditionReasonVolumeReady indicates the PVC is ready
        ConditionReasonVolumeReady = "VolumeReady"
        // ConditionReasonVolumeFailed indicates the PVC creation failed
        ConditionReasonVolumeFailed = "VolumeFailed"
)

// EmbeddingServerSpec defines the desired state of EmbeddingServer
type EmbeddingServerSpec struct {
        // Model is the HuggingFace embedding model to use (e.g., "sentence-transformers/all-MiniLM-L6-v2")
        // +kubebuilder:default="BAAI/bge-small-en-v1.5"
        // +optional
        Model string `json:"model,omitempty"`

        // HFTokenSecretRef is a reference to a Kubernetes Secret containing the huggingface token.
        // If provided, the secret value will be provided to the embedding server for authentication with huggingface.
        // +optional
        HFTokenSecretRef *SecretKeyRef `json:"hfTokenSecretRef,omitempty"`

        // Image is the container image for the embedding inference server.
        // Images must be from HuggingFace Text Embeddings Inference (https://github.com/huggingface/text-embeddings-inference).
        // +kubebuilder:default="ghcr.io/huggingface/text-embeddings-inference:cpu-latest"
        // +optional
        Image string `json:"image,omitempty"`

        // ImagePullPolicy defines the pull policy for the container image
        // +kubebuilder:validation:Enum=Always;Never;IfNotPresent
        // +kubebuilder:default="IfNotPresent"
        // +optional
        ImagePullPolicy corev1.PullPolicy `json:"imagePullPolicy,omitempty"`

        // Port is the port to expose the embedding service on
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:validation:Maximum=65535
        // +kubebuilder:default=8080
        Port int32 `json:"port,omitempty"`

        // Args are additional arguments to pass to the embedding inference server
        // +listType=atomic
        // +optional
        Args []string `json:"args,omitempty"`

        // Env are environment variables to set in the container
        // +listType=map
        // +listMapKey=name
        // +optional
        Env []EnvVar `json:"env,omitempty"`

        // Resources defines compute resources for the embedding server
        // +optional
        Resources ResourceRequirements `json:"resources,omitempty"`

        // ModelCache configures persistent storage for downloaded models
        // When enabled, models are cached in a PVC and reused across pod restarts
        // +optional
        ModelCache *ModelCacheConfig `json:"modelCache,omitempty"`

        // PodTemplateSpec allows customizing the pod (node selection, tolerations, etc.)
        // This field accepts a PodTemplateSpec object as JSON/YAML.
        // Note that to modify the specific container the embedding server runs in, you must specify
        // the 'embedding' container name in the PodTemplateSpec.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        PodTemplateSpec *runtime.RawExtension `json:"podTemplateSpec,omitempty"`

        // ResourceOverrides allows overriding annotations and labels for resources created by the operator
        // +optional
        ResourceOverrides *EmbeddingResourceOverrides `json:"resourceOverrides,omitempty"`

        // Replicas is the number of embedding server replicas to run
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:default=1
        // +optional
        Replicas *int32 `json:"replicas,omitempty"`
}

// ModelCacheConfig configures persistent storage for model caching
type ModelCacheConfig struct {
        // Enabled controls whether model caching is enabled
        // +kubebuilder:default=true
        // +optional
        Enabled bool `json:"enabled,omitempty"`

        // StorageClassName is the storage class to use for the PVC
        // If not specified, uses the cluster's default storage class
        // +optional
        StorageClassName *string `json:"storageClassName,omitempty"`

        // Size is the size of the PVC for model caching (e.g., "10Gi")
        // +kubebuilder:default="10Gi"
        // +optional
        Size string `json:"size,omitempty"`

        // AccessMode is the access mode for the PVC
        // +kubebuilder:default="ReadWriteOnce"
        // +kubebuilder:validation:Enum=ReadWriteOnce;ReadWriteMany;ReadOnlyMany
        // +optional
        AccessMode string `json:"accessMode,omitempty"`
}

// EmbeddingResourceOverrides defines overrides for annotations and labels on created resources
type EmbeddingResourceOverrides struct {
        // StatefulSet defines overrides for the StatefulSet resource
        // +optional
        StatefulSet *EmbeddingStatefulSetOverrides `json:"statefulSet,omitempty"`

        // Service defines overrides for the Service resource
        // +optional
        Service *ResourceMetadataOverrides `json:"service,omitempty"`

        // PersistentVolumeClaim defines overrides for the PVC resource
        // +optional
        PersistentVolumeClaim *ResourceMetadataOverrides `json:"persistentVolumeClaim,omitempty"`
}

// EmbeddingStatefulSetOverrides defines overrides specific to the embedding statefulset
type EmbeddingStatefulSetOverrides struct {
        // ResourceMetadataOverrides is embedded to inherit annotations and labels fields
        ResourceMetadataOverrides `json:",inline"` // nolint:revive

        // PodTemplateMetadataOverrides defines metadata overrides for the pod template
        // +optional
        PodTemplateMetadataOverrides *ResourceMetadataOverrides `json:"podTemplateMetadataOverrides,omitempty"`
}

// EmbeddingServerStatus defines the observed state of EmbeddingServer
type EmbeddingServerStatus struct {
        // Conditions represent the latest available observations of the EmbeddingServer's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // Phase is the current phase of the EmbeddingServer
        // +optional
        Phase EmbeddingServerPhase `json:"phase,omitempty"`

        // Message provides additional information about the current phase
        // +optional
        Message string `json:"message,omitempty"`

        // URL is the URL where the embedding service can be accessed
        // +optional
        URL string `json:"url,omitempty"`

        // ReadyReplicas is the number of ready replicas
        // +optional
        ReadyReplicas int32 `json:"readyReplicas,omitempty"`

        // ObservedGeneration reflects the generation most recently observed by the controller
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`
}

// EmbeddingServerPhase is the phase of the EmbeddingServer
// +kubebuilder:validation:Enum=Pending;Downloading;Ready;Failed;Terminating
type EmbeddingServerPhase string

const (
        // EmbeddingServerPhasePending means the EmbeddingServer is being created
        EmbeddingServerPhasePending EmbeddingServerPhase = "Pending"

        // EmbeddingServerPhaseDownloading means the model is being downloaded
        EmbeddingServerPhaseDownloading EmbeddingServerPhase = "Downloading"

        // EmbeddingServerPhaseReady means the EmbeddingServer is ready
        EmbeddingServerPhaseReady EmbeddingServerPhase = "Ready"

        // EmbeddingServerPhaseFailed means the EmbeddingServer failed to start
        EmbeddingServerPhaseFailed EmbeddingServerPhase = "Failed"

        // EmbeddingServerPhaseTerminating means the EmbeddingServer is being deleted
        EmbeddingServerPhaseTerminating EmbeddingServerPhase = "Terminating"
)

//+kubebuilder:object:root=true
//+kubebuilder:storageversion
//+kubebuilder:subresource:status
//+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
//+kubebuilder:resource:shortName=emb;embedding,categories=toolhive
//+kubebuilder:printcolumn:name="Status",type="string",JSONPath=".status.phase"
//+kubebuilder:printcolumn:name="Model",type="string",JSONPath=".spec.model"
//+kubebuilder:printcolumn:name="Ready",type="integer",JSONPath=".status.readyReplicas"
//+kubebuilder:printcolumn:name="URL",type="string",JSONPath=".status.url"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"

// EmbeddingServer is the Schema for the embeddingservers API
type EmbeddingServer struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   EmbeddingServerSpec   `json:"spec,omitempty"`
        Status EmbeddingServerStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// EmbeddingServerList contains a list of EmbeddingServer
type EmbeddingServerList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []EmbeddingServer `json:"items"`
}

// GetName returns the name of the EmbeddingServer
func (e *EmbeddingServer) GetName() string {
        return e.Name
}

// GetNamespace returns the namespace of the EmbeddingServer
func (e *EmbeddingServer) GetNamespace() string {
        return e.Namespace
}

// GetPort returns the port of the EmbeddingServer
func (e *EmbeddingServer) GetPort() int32 {
        if e.Spec.Port > 0 {
                return e.Spec.Port
        }
        return 8080
}

// GetReplicas returns the number of replicas for the EmbeddingServer
func (e *EmbeddingServer) GetReplicas() int32 {
        if e.Spec.Replicas != nil {
                return *e.Spec.Replicas
        }
        return 1
}

// IsModelCacheEnabled returns whether model caching is enabled
func (e *EmbeddingServer) IsModelCacheEnabled() bool {
        if e.Spec.ModelCache == nil {
                return false
        }
        return e.Spec.ModelCache.Enabled
}

// GetImagePullPolicy returns the image pull policy for the EmbeddingServer
func (e *EmbeddingServer) GetImagePullPolicy() corev1.PullPolicy {
        if e.Spec.ImagePullPolicy != "" {
                return e.Spec.ImagePullPolicy
        }
        return corev1.PullIfNotPresent
}

func init() {
        SchemeBuilder.Register(&EmbeddingServer{}, &EmbeddingServerList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package v1beta1 contains API Schema definitions for the toolhive v1beta1 API group
// +kubebuilder:object:generate=true
// +groupName=toolhive.stacklok.dev
package v1beta1

import (
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/runtime"
        "k8s.io/apimachinery/pkg/runtime/schema"
)

var (
        // GroupVersion is group version used to register these objects
        GroupVersion = schema.GroupVersion{Group: "toolhive.stacklok.dev", Version: "v1beta1"}

        // SchemeBuilder is used to add go types to the GroupVersionKind scheme
        SchemeBuilder = newSchemeBuilder(GroupVersion)

        // AddToScheme adds the types in this group-version to the given scheme.
        AddToScheme = SchemeBuilder.AddToScheme
)

type schemeBuilder struct {
        runtime.SchemeBuilder
        GroupVersion schema.GroupVersion
}

func newSchemeBuilder(groupVersion schema.GroupVersion) *schemeBuilder {
        builder := &schemeBuilder{GroupVersion: groupVersion}
        builder.SchemeBuilder.Register(func(scheme *runtime.Scheme) error {
                metav1.AddToGroupVersion(scheme, groupVersion)
                return nil
        })
        return builder
}

func (builder *schemeBuilder) Register(objects ...runtime.Object) *schemeBuilder {
        builder.SchemeBuilder.Register(func(scheme *runtime.Scheme) error {
                scheme.AddKnownTypes(builder.GroupVersion, objects...)
                return nil
        })
        return builder
}

func (builder *schemeBuilder) Build() (*runtime.Scheme, error) {
        scheme := runtime.NewScheme()
        return scheme, builder.AddToScheme(scheme)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        "fmt"

        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        runtime "k8s.io/apimachinery/pkg/runtime"

        "github.com/stacklok/toolhive/pkg/authz/authorizers"
)

// Condition type and reasons for MCPAuthzConfig status (RFC-0023)
const (
        // ConditionTypeAuthzConfigValid indicates whether the MCPAuthzConfig configuration is valid
        ConditionTypeAuthzConfigValid = ConditionTypeValid

        // ConditionReasonAuthzConfigValid indicates spec validation passed
        ConditionReasonAuthzConfigValid = "ConfigValid"

        // ConditionReasonAuthzConfigInvalid indicates spec validation failed
        ConditionReasonAuthzConfigInvalid = "ConfigInvalid"
)

// MCPAuthzConfigSpec defines the desired state of MCPAuthzConfig.
// MCPAuthzConfig resources are namespace-scoped and can only be referenced by
// MCPServer, MCPRemoteProxy, or VirtualMCPServer resources in the same namespace.
type MCPAuthzConfigSpec struct {
        // Type identifies the authorizer backend (e.g., "cedarv1", "httpv1").
        // Must match a registered authorizer type in the factory registry.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Type string `json:"type"`

        // Config contains the backend-specific authorization configuration.
        // The structure depends on the Type field:
        //   - cedarv1: policies ([]string), entities_json (string), primary_upstream_provider (string), group_claim_name (string)
        //   - httpv1: http ({url, timeout, insecure_skip_verify}), context ({include_args, include_operation}), claim_mapping (string)
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        Config runtime.RawExtension `json:"config"`
}

// MCPAuthzConfigStatus defines the observed state of MCPAuthzConfig
type MCPAuthzConfigStatus struct {
        // Conditions represent the latest available observations of the MCPAuthzConfig's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration is the most recent generation observed for this MCPAuthzConfig.
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // ConfigHash is a hash of the current configuration for change detection
        // +optional
        ConfigHash string `json:"configHash,omitempty"`
}

// +kubebuilder:object:root=true
// +kubebuilder:storageversion
// +kubebuilder:subresource:status
// +kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
// +kubebuilder:resource:shortName=authzcfg,categories=toolhive
// +kubebuilder:printcolumn:name="Type",type=string,JSONPath=`.spec.type`
// +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status`
// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp`

// MCPAuthzConfig is the Schema for the mcpauthzconfigs API.
// MCPAuthzConfig resources are namespace-scoped and can only be referenced by
// MCPServer, MCPRemoteProxy, or VirtualMCPServer resources within the same namespace.
// Cross-namespace references are not supported for security and isolation reasons.
type MCPAuthzConfig struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPAuthzConfigSpec   `json:"spec,omitempty"`
        Status MCPAuthzConfigStatus `json:"status,omitempty"`
}

// +kubebuilder:object:root=true

// MCPAuthzConfigList contains a list of MCPAuthzConfig
type MCPAuthzConfigList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPAuthzConfig `json:"items"`
}

// MCPAuthzConfigReference references a shared MCPAuthzConfig resource.
// The referenced MCPAuthzConfig must be in the same namespace as the referencing workload.
type MCPAuthzConfigReference struct {
        // Name is the name of the MCPAuthzConfig resource in the same namespace.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Name string `json:"name"`
}

// Validate performs structural validation on the MCPAuthzConfig spec.
// This method is called by the controller during reconciliation.
//
// Validate provides defense-in-depth alongside CEL validation rules: CEL catches
// issues at API admission time, but this method also validates stored objects to
// catch any that bypassed CEL or were stored before CEL rules were added.
//
// Validate consults the authorizer factory registry (pkg/authz/authorizers) to
// reject spec.type values that no backend has registered. Backend-specific
// schema validation (interpreting spec.config) still lives in the controller,
// where the registered factory's ValidateConfig is invoked.
func (r *MCPAuthzConfig) Validate() error {
        if r.Spec.Type == "" {
                return fmt.Errorf("type must not be empty")
        }
        if !authorizers.IsRegistered(r.Spec.Type) {
                return fmt.Errorf(
                        "type %q is not a registered authorizer backend (registered: %v)",
                        r.Spec.Type, authorizers.RegisteredTypes(),
                )
        }
        if len(r.Spec.Config.Raw) == 0 {
                return fmt.Errorf("config must not be empty")
        }
        return nil
}

func init() {
        SchemeBuilder.Register(&MCPAuthzConfig{}, &MCPAuthzConfigList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        "fmt"
        "slices"
        "sort"

        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

        "github.com/stacklok/toolhive/pkg/authserver"
        "github.com/stacklok/toolhive/pkg/authserver/oauthparams"
        "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange"
)

// External auth configuration types
const (
        // ExternalAuthTypeTokenExchange is the type for RFC-8693 token exchange
        ExternalAuthTypeTokenExchange ExternalAuthType = "tokenExchange"

        // ExternalAuthTypeHeaderInjection is the type for custom header injection
        ExternalAuthTypeHeaderInjection ExternalAuthType = "headerInjection"

        // ExternalAuthTypeBearerToken is the type for bearer token authentication
        // This allows authenticating to remote MCP servers using bearer tokens stored in Kubernetes Secrets
        ExternalAuthTypeBearerToken ExternalAuthType = "bearerToken"

        // ExternalAuthTypeUnauthenticated is the type for no authentication
        // This should only be used for backends on trusted networks (e.g., localhost, VPC)
        // or when authentication is handled by network-level security
        ExternalAuthTypeUnauthenticated ExternalAuthType = "unauthenticated"

        // ExternalAuthTypeEmbeddedAuthServer is the type for embedded OAuth2/OIDC authorization server
        // This enables running an embedded auth server that delegates to upstream IDPs
        ExternalAuthTypeEmbeddedAuthServer ExternalAuthType = "embeddedAuthServer"

        // ExternalAuthTypeAWSSts is the type for AWS STS authentication
        ExternalAuthTypeAWSSts ExternalAuthType = "awsSts"

        // ExternalAuthTypeUpstreamInject is the type for upstream token injection
        // This injects an upstream IdP access token as the Authorization: Bearer header
        ExternalAuthTypeUpstreamInject ExternalAuthType = "upstreamInject"

        // ExternalAuthTypeOBO is the type for on-behalf-of (OBO) flows.
        // This type requires a build with an OBO handler registered via
        // controllerutil.RegisterOBOHandler; an upstream-only build surfaces
        // status.conditions[Valid] = False with Reason: EnterpriseRequired
        // when an obo-typed MCPExternalAuthConfig is applied.
        ExternalAuthTypeOBO ExternalAuthType = "obo"

        // ExternalAuthTypeXAA is the type for XAA (Cross-Application Access) auth.
        // XAA performs a two-step token exchange to obtain access tokens for target services:
        //   - IdP exchange (RFC 8693): Exchange the user's ID token at their IdP for an ID-JAG JWT
        //   - Target grant (RFC 7523): Exchange the ID-JAG at the target app's AS for an access token
        ExternalAuthTypeXAA ExternalAuthType = "xaa"
)

// ExternalAuthType represents the type of external authentication
type ExternalAuthType string

// MCPExternalAuthConfigSpec defines the desired state of MCPExternalAuthConfig.
// MCPExternalAuthConfig resources are namespace-scoped and can only be referenced by
// MCPServer resources in the same namespace.
//
// +kubebuilder:validation:XValidation:rule="self.type == 'tokenExchange' ? has(self.tokenExchange) : !has(self.tokenExchange)",message="tokenExchange configuration must be set if and only if type is 'tokenExchange'"
// +kubebuilder:validation:XValidation:rule="self.type == 'headerInjection' ? has(self.headerInjection) : !has(self.headerInjection)",message="headerInjection configuration must be set if and only if type is 'headerInjection'"
// +kubebuilder:validation:XValidation:rule="self.type == 'bearerToken' ? has(self.bearerToken) : !has(self.bearerToken)",message="bearerToken configuration must be set if and only if type is 'bearerToken'"
// +kubebuilder:validation:XValidation:rule="self.type == 'embeddedAuthServer' ? has(self.embeddedAuthServer) : !has(self.embeddedAuthServer)",message="embeddedAuthServer configuration must be set if and only if type is 'embeddedAuthServer'"
// +kubebuilder:validation:XValidation:rule="self.type == 'awsSts' ? has(self.awsSts) : !has(self.awsSts)",message="awsSts configuration must be set if and only if type is 'awsSts'"
// +kubebuilder:validation:XValidation:rule="self.type == 'upstreamInject' ? has(self.upstreamInject) : !has(self.upstreamInject)",message="upstreamInject configuration must be set if and only if type is 'upstreamInject'"
// +kubebuilder:validation:XValidation:rule="self.type == 'obo' ? has(self.obo) : !has(self.obo)",message="obo configuration must be set if and only if type is 'obo'"
// +kubebuilder:validation:XValidation:rule="self.type == 'xaa' ? has(self.xaa) : !has(self.xaa)",message="xaa configuration must be set if and only if type is 'xaa'"
// +kubebuilder:validation:XValidation:rule="self.type == 'unauthenticated' ? (!has(self.tokenExchange) && !has(self.headerInjection) && !has(self.bearerToken) && !has(self.embeddedAuthServer) && !has(self.awsSts) && !has(self.upstreamInject) && !has(self.obo) && !has(self.xaa)) : true",message="no configuration must be set when type is 'unauthenticated'"
//
//nolint:lll // CEL validation rules exceed line length limit
type MCPExternalAuthConfigSpec struct {
        // Type is the type of external authentication to configure.
        // When set to "obo", the cluster must run a build that has registered an
        // OBO handler via controllerutil.RegisterOBOHandler; upstream-only builds
        // surface status.conditions[Valid] = False with Reason: EnterpriseRequired
        // for obo-typed configs.
        // +kubebuilder:validation:Enum=tokenExchange;headerInjection;bearerToken;unauthenticated;embeddedAuthServer;awsSts;upstreamInject;obo;xaa
        // +kubebuilder:validation:Required
        Type ExternalAuthType `json:"type"`

        // TokenExchange configures RFC-8693 OAuth 2.0 Token Exchange
        // Only used when Type is "tokenExchange"
        // +optional
        TokenExchange *TokenExchangeConfig `json:"tokenExchange,omitempty"`

        // HeaderInjection configures custom HTTP header injection
        // Only used when Type is "headerInjection"
        // +optional
        HeaderInjection *HeaderInjectionConfig `json:"headerInjection,omitempty"`

        // BearerToken configures bearer token authentication
        // Only used when Type is "bearerToken"
        // +optional
        BearerToken *BearerTokenConfig `json:"bearerToken,omitempty"`

        // EmbeddedAuthServer configures an embedded OAuth2/OIDC authorization server
        // Only used when Type is "embeddedAuthServer"
        // +optional
        EmbeddedAuthServer *EmbeddedAuthServerConfig `json:"embeddedAuthServer,omitempty"`

        // AWSSts configures AWS STS authentication with SigV4 request signing
        // Only used when Type is "awsSts"
        // +optional
        AWSSts *AWSStsConfig `json:"awsSts,omitempty"`

        // UpstreamInject configures upstream token injection for backend requests.
        // Only used when Type is "upstreamInject".
        // +optional
        UpstreamInject *UpstreamInjectSpec `json:"upstreamInject,omitempty"`

        // OBO configures On-Behalf-Of (OBO) authentication.
        // Only used when Type is "obo". Setting this field on an upstream-only build
        // causes the MCPExternalAuthConfig to transition to
        // status.conditions[Valid] = False with Reason: EnterpriseRequired, because
        // no OBO handler is registered. See OBOConfig for the field-to-runtime
        // contract mapping.
        // +optional
        OBO *OBOConfig `json:"obo,omitempty"`

        // XAA configures XAA (Cross-Application Access) auth for backend requests.
        // Only used when Type is "xaa".
        // +optional
        XAA *XAASpec `json:"xaa,omitempty"`
}

// OBOConfig holds configuration for the On-Behalf-Of (OBO) external auth type.
// Only used when Type is "obo".
//
// This is the user-facing CRD surface for the Microsoft Entra OBO flow. It is
// structurally valid in upstream (OSS) builds but inert: an upstream-only build
// returns obo.ErrEnterpriseRequired at reconcile (Valid=False, Reason:
// EnterpriseRequired) because no OBO handler is registered via
// controllerutil.RegisterOBOHandler. A build with the enterprise OBO handler
// translates these fields into the runtime wire contract obo.MiddlewareParameters,
// so the field names and semantics here track that contract rather than the
// upstream TokenExchangeConfig (which uses different names, e.g.
// subjectProviderName / externalTokenHeaderName). In particular there is no
// externalTokenHeaderName: the OBO subject is sourced from the authenticated
// Identity, never from an inbound request header.
//
// Field-to-contract mapping performed by the operator's OBO handler:
//   - tenantId (+ optional authority) → tokenUrl
//     (https://login.microsoftonline.com/<tenantId>/oauth2/v2.0/token, or the
//     configured authority base joined with the tenant for sovereign clouds)
//   - clientSecretRef → resolved into a pod env var; only the env var name
//     travels in the contract, as clientSecretEnvVar
//   - audience / scopes → collapsed to a single exchange target by
//     obo.MiddlewareParameters.ExchangeTarget() (space-joined scopes win,
//     otherwise audience)
//   - cacheSkew → the contract's integer-seconds cacheSkewSeconds
//
// Every field is optional at the CRD level, and the schema deliberately carries
// no required field and no cross-field CEL rule. spec.obo shipped as an empty
// placeholder ({}) in earlier releases, so adding a required field or an
// admission rule that rejects {} would be a backward-incompatible narrowing of
// an already-stored, round-trippable object. Presence and combination
// requirements — a tenant, a client-auth credential, and at least one of
// audience/scopes — are therefore enforced by the registered OBO handler at
// reconcile, which reports a violation as Valid=False / Reason=InvalidConfig.
// The per-field patterns below still apply, but only to a value that is present.
//
//nolint:lll // the tenantId GUID/domain pattern exceeds the line length limit
type OBOConfig struct {
        // TenantID is the Microsoft Entra (Azure AD) directory (tenant) identifier.
        // Optional at the CRD level (see the type doc); the operator enforces its
        // presence, since an OBO confidential-client exchange must target a specific
        // tenant. When set, it must be one of the two forms the Entra v2.0 token
        // endpoint addresses: a directory GUID, or a verified domain name (e.g.
        // contoso.onmicrosoft.com). Well-known aliases such as "common",
        // "organizations", and "consumers" are NOT accepted. The operator
        // interpolates it into the token endpoint
        // (<authority>/<tenantId>/oauth2/v2.0/token), so the value is constrained to
        // the GUID/domain shape (no path metacharacters); the pattern and 253-char
        // cap mirror the enterprise exchanger's validateTenant, so any tenantId
        // admitted here is one the runtime can consume.
        // +kubebuilder:validation:MaxLength=253
        // +kubebuilder:validation:Pattern=`^([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}|([a-zA-Z0-9]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,})$`
        // +optional
        TenantID string `json:"tenantId,omitempty"`

        // Authority overrides the default Entra login host
        // (https://login.microsoftonline.com) for sovereign or national clouds, e.g.
        // https://login.microsoftonline.us (US Gov) or
        // https://login.partner.microsoftonline.cn (China). When set, the operator
        // builds the token endpoint by joining <authority>, <tenantId>, and the
        // v2.0 token path. Must be an HTTPS URL with no userinfo, query, fragment,
        // or trailing slash; a path IS permitted and is prefixed before the tenant
        // segment, as some sovereign / B2C / CIAM endpoints require. The OBO exchange
        // POSTs the client secret and the end-user assertion to this host, so it is a
        // credential trust boundary: HTTPS is required and userinfo (user@host) is
        // rejected to prevent host confusion (per RFC 3986 the real host follows the
        // "@", so https://login.microsoftonline.com@attacker.example targets
        // attacker.example). This is intentionally stricter than the downstream
        // exchanger's validateHTTPSURL, which also accepts http for loopback hosts
        // and tolerates a trailing slash — rejecting those at admission is the safe
        // direction.
        // +kubebuilder:validation:Pattern=`^https://[^\s?#@]+[^/\s?#@]$`
        // +optional
        Authority string `json:"authority,omitempty"`

        // ClientID is the confidential client's application (client) ID registered
        // in Entra. Emitted verbatim as the runtime contract's clientId.
        // Optional at the CRD level so future client-authentication methods (e.g.
        // certificate or workload-identity credentials, planned fast-follows) can be
        // added without a breaking schema change. The operator enforces that clientId
        // and clientSecretRef are both present for the v1 shared-secret flow.
        // +optional
        ClientID string `json:"clientId,omitempty"`

        // ClientSecretRef references a Kubernetes Secret containing the confidential
        // client's secret. v1 supports a shared client secret only. The operator
        // injects the resolved value into the proxyrunner pod as an environment
        // variable and emits only that variable's name in the runtime contract, as
        // clientSecretEnvVar — the secret value never travels in the contract.
        // Optional at the CRD level for the same forward-compatibility reason as
        // clientId (a certificate/workload-identity flow needs no client secret);
        // the operator enforces presence for the v1 shared-secret flow.
        // +optional
        ClientSecretRef *SecretKeyRef `json:"clientSecretRef,omitempty"`

        // Audience is the backend target identifier requested in the exchanged
        // token. Used as the exchange target when Scopes is empty. At least one of
        // audience or scopes must be set; the operator enforces that at reconcile
        // (it is not an admission-time rule — see the type doc).
        // +optional
        Audience string `json:"audience,omitempty"`

        // Scopes are the delegated scopes to request for the exchanged token, e.g.
        // ["api://<backend>/.default"]. When non-empty they take precedence over
        // Audience. At least one of audience or scopes must be set; the operator
        // enforces that at reconcile. The MaxItems and per-item length caps are
        // defensive bounds on an otherwise unbounded list.
        // +kubebuilder:validation:MaxItems=20
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=256
        // +listType=atomic
        // +optional
        Scopes []string `json:"scopes,omitempty"`

        // SubjectTokenProviderName selects the source of the OBO subject (assertion)
        // token from the request's authenticated Identity:
        //   - Omitted: use the inbound end-user token the client presented
        //     (Identity.Token) — the deployment with no embedded auth server, where
        //     the client holds an Entra token directly.
        //   - Set: use the named upstream provider's token
        //     (Identity.UpstreamTokens[<name>]) — the embedded-auth-server
        //     deployment, where the inbound token is the proxy's own session token.
        //     The value must match a configured upstream provider name.
        // The subject is always sourced from the authenticated Identity, never from
        // an inbound request header, so the upstream auth middleware must run first.
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=63
        // +kubebuilder:validation:Pattern=`^[a-z0-9]([a-z0-9-]*[a-z0-9])?$`
        // +optional
        SubjectTokenProviderName string `json:"subjectTokenProviderName,omitempty"`

        // CacheSkew overrides the OBO token cache's default expiry skew (the margin
        // by which a cached token is treated as expired before its real expiry),
        // e.g. "30s". The operator converts it to the runtime contract's
        // integer-seconds cacheSkewSeconds. Should not be negative, but the schema
        // does not enforce that — metav1.Duration carries no numeric minimum — and
        // upstream builds do not reject it. A negative value is rejected only by an
        // enterprise build's OBO handler once that handler validates the converted
        // parameters; it is not enforced at admission or in upstream-only builds.
        // When omitted, the cache default applies.
        // +optional
        CacheSkew *metav1.Duration `json:"cacheSkew,omitempty"`
}

// TokenExchangeConfig holds configuration for RFC-8693 OAuth 2.0 Token Exchange.
// This configuration is used to exchange incoming authentication tokens for tokens
// that can be used with external services.
// The structure matches the tokenexchange.Config from pkg/oauthproto/tokenexchange/middleware.go
type TokenExchangeConfig struct {
        // TokenURL is the OAuth 2.0 token endpoint URL for token exchange
        // +kubebuilder:validation:Required
        TokenURL string `json:"tokenUrl"`

        // ClientID is the OAuth 2.0 client identifier
        // Optional for some token exchange flows (e.g., Google Cloud Workforce Identity)
        // +optional
        ClientID string `json:"clientId,omitempty"`

        // ClientSecretRef is a reference to a secret containing the OAuth 2.0 client secret
        // Optional for some token exchange flows (e.g., Google Cloud Workforce Identity)
        // +optional
        ClientSecretRef *SecretKeyRef `json:"clientSecretRef,omitempty"`

        // Audience is the target audience for the exchanged token
        // +kubebuilder:validation:Required
        Audience string `json:"audience"`

        // Scopes is a list of OAuth 2.0 scopes to request for the exchanged token
        // +listType=atomic
        // +optional
        Scopes []string `json:"scopes,omitempty"`

        // SubjectTokenType is the type of the incoming subject token.
        // Accepts short forms: "access_token" (default), "id_token", "jwt"
        // Or full URNs: "urn:ietf:params:oauth:token-type:access_token",
        //               "urn:ietf:params:oauth:token-type:id_token",
        //               "urn:ietf:params:oauth:token-type:jwt"
        // For Google Workload Identity Federation with OIDC providers (like Okta), use "id_token"
        // +kubebuilder:validation:Pattern=`^(access_token|id_token|jwt|urn:ietf:params:oauth:token-type:(access_token|id_token|jwt))?$`
        // +optional
        SubjectTokenType string `json:"subjectTokenType,omitempty"`

        // ExternalTokenHeaderName is the name of the custom header to use for the exchanged token.
        // If set, the exchanged token will be added to this custom header (e.g., "X-Upstream-Token").
        // If empty or not set, the exchanged token will replace the Authorization header (default behavior).
        // +optional
        ExternalTokenHeaderName string `json:"externalTokenHeaderName,omitempty"`

        // SubjectProviderName is the name of the upstream provider whose token is used as the
        // RFC 8693 subject token instead of identity.Token when performing token exchange.
        // When left empty and an embedded authorization server is configured on the VirtualMCPServer,
        // the controller automatically populates this field with the first configured upstream
        // provider name. Set it explicitly to override that default or to select a specific
        // provider when multiple upstreams are configured.
        // +optional
        SubjectProviderName string `json:"subjectProviderName,omitempty"`
}

// HeaderInjectionConfig holds configuration for custom HTTP header injection authentication.
// This allows injecting a secret-based header value into requests to backend MCP servers.
// For security reasons, only secret references are supported (no plaintext values).
type HeaderInjectionConfig struct {
        // HeaderName is the name of the HTTP header to inject
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        HeaderName string `json:"headerName"`

        // ValueSecretRef references a Kubernetes Secret containing the header value
        // +kubebuilder:validation:Required
        ValueSecretRef *SecretKeyRef `json:"valueSecretRef"`
}

// BearerTokenConfig holds configuration for bearer token authentication.
// This allows authenticating to remote MCP servers using bearer tokens stored in Kubernetes Secrets.
// For security reasons, only secret references are supported (no plaintext values).
type BearerTokenConfig struct {
        // TokenSecretRef references a Kubernetes Secret containing the bearer token
        // +kubebuilder:validation:Required
        TokenSecretRef *SecretKeyRef `json:"tokenSecretRef"`
}

// DelegateClientConfig configures a pre-provisioned confidential OAuth client
// for RFC 8693 token exchange. Its secret is referenced from a Kubernetes
// Secret and is never represented inline.
//
// +kubebuilder:validation:XValidation:rule="has(self.clientSecretRef) && size(self.clientSecretRef.name) > 0 && size(self.clientSecretRef.key) > 0",message="clientSecretRef.name and clientSecretRef.key are required and must be non-empty"
//
//nolint:lll // CEL validation rule exceeds line length limit
type DelegateClientConfig struct {
        // ClientID is the OAuth client_id presented at the token endpoint.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=256
        ClientID string `json:"clientId"`

        // ClientSecretRef references the Kubernetes Secret key containing the client secret.
        // +kubebuilder:validation:Required
        ClientSecretRef *SecretKeyRef `json:"clientSecretRef"`

        // Scopes is the narrowed set of OAuth scopes this client may request.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=10
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=256
        // +listType=atomic
        Scopes []string `json:"scopes"`

        // Audiences is the narrowed set of RFC 8707 resources this client may request.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=10
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=2048
        // +listType=atomic
        Audiences []string `json:"audiences"`
}

// TrustedIssuerConfig configures an external OIDC issuer whose tokens are
// accepted as RFC 8693 subject tokens or RFC 7523 JWT-bearer assertions. Trust
// fields remain top-level; canonical grant policy references this declaration by
// Name. The embedded grant-policy fields are retained for released CRD compatibility.
//
// +kubebuilder:validation:XValidation:rule="!has(self.allowedDelegateClients) || !('*' in self.allowedDelegateClients) || size(self.allowedDelegateClients) == 1",message="allowedDelegateClients must not combine the wildcard \"*\" with specific client IDs"
// +kubebuilder:validation:XValidation:rule="!(has(self.allowMayAct) && self.allowMayAct && '*' in self.allowedDelegateClients)",message="allowMayAct must not be enabled when allowedDelegateClients contains the wildcard \"*\""
// +kubebuilder:validation:XValidation:rule="!has(self.actorClaim) || !(self.actorClaim in ['sub', 'iss', 'aud', 'exp', 'iat', 'nbf', 'jti', 'name', 'email', 'scope', 'scp', 'may_act'])",message="actorClaim must name a readable claim; use client_id or a non-reserved claim such as azp, appid, or cid"
// +kubebuilder:validation:XValidation:rule="!(has(self.allowPrivateIPs) && self.allowPrivateIPs) || (has(self.jwksUrl) && self.jwksUrl != \"\")",message="allowPrivateIPs requires jwksUrl to be set explicitly"
// +kubebuilder:validation:XValidation:rule="!((has(self.actorClaim) && size(self.actorClaim) > 0) || (has(self.actorMatcher) && size(self.actorMatcher) > 0) || (has(self.allowMayAct) && self.allowMayAct) || (has(self.allowedActors) && size(self.allowedActors) > 0) || (has(self.allowedDelegateClients) && size(self.allowedDelegateClients) > 0)) || (has(self.expectedAudience) && size(self.expectedAudience) > 0)",message="expectedAudience is required when legacy RFC 8693 policy is configured"
// +kubebuilder:validation:XValidation:rule="!((has(self.expectedAudience) && size(self.expectedAudience) > 0) || (has(self.actorClaim) && size(self.actorClaim) > 0) || (has(self.actorMatcher) && size(self.actorMatcher) > 0) || (has(self.allowMayAct) && self.allowMayAct) || (has(self.allowedActors) && size(self.allowedActors) > 0)) || (has(self.allowedDelegateClients) && size(self.allowedDelegateClients) > 0)",message="allowedDelegateClients is required when legacy RFC 8693 policy is configured"
//
//nolint:lll // CEL validation rules exceed line length limit
type TrustedIssuerConfig struct {
        // Name optionally identifies this trust declaration for canonical issuerRef references.
        // Names must be unique within trustedIssuers when set.
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        // +optional
        Name string `json:"name,omitempty"`

        // IssuerURL is the expected "iss" claim value (exact match).
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=2048
        IssuerURL string `json:"issuerUrl"`

        // ExpectedAudience is the expected "aud" claim value that must appear in
        // an RFC 8693 subject token's audience list. It is not used by an RFC 7523
        // JWT-bearer assertion, whose audience is the token endpoint.
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=2048
        // +optional
        // This legacy field is deprecated; configure RFC 8693 policy under
        // inboundGrants.tokenExchange.issuerPolicies.
        ExpectedAudience string `json:"expectedAudience,omitempty"`

        // JWKSURL is the URL to fetch the issuer's JSON Web Key Set from. If
        // empty, it is resolved via OIDC discovery at
        // {issuerUrl}/.well-known/openid-configuration.
        // +optional
        // +kubebuilder:validation:MaxLength=2048
        JWKSURL string `json:"jwksUrl,omitempty"`

        // InsecureAllowHTTP permits plain-HTTP OIDC discovery and JWKS fetches
        // for THIS issuer only. Development and testing only — never set in
        // production.
        // +optional
        InsecureAllowHTTP bool `json:"insecureAllowHTTP,omitempty"`

        // AllowPrivateIPs permits OIDC discovery and JWKS fetches for THIS issuer
        // to resolve to a private or loopback address. Use only when the issuer
        // is hosted inside the same cluster and has no public endpoint. Requires
        // jwksUrl to be set explicitly (enforced at admission and by shared
        // validation), since otherwise OIDC discovery — fetched from the external
        // issuer itself — would choose the private dial target.
        // +optional
        AllowPrivateIPs bool `json:"allowPrivateIPs,omitempty"`

        // CABundleRef references a ConfigMap containing PEM CA certificates used when
        // fetching this issuer's OIDC discovery document and JWKS. The bundle is added
        // to the system roots for this issuer's client only; public roots still apply
        // and other issuers are unaffected. Write access to the referenced ConfigMap is
        // equivalent to controlling this issuer's trust anchor for subject-token
        // validation — restrict it with the same care as a signing-key Secret.
        // +optional
        CABundleRef *CABundleSource `json:"caBundleRef,omitempty"`

        // ActorClaim names the claim identifying the client that requested the
        // subject token from this external issuer (used by allowedActors below).
        // Defaults to "azp" when empty; use "appid" for Microsoft Entra v1, "cid"
        // for Okta. The special value "client_id" reads the subject token's
        // client_id claim instead.
        // +optional
        // +kubebuilder:validation:MaxLength=64
        // This legacy field is deprecated; configure RFC 8693 policy under
        // inboundGrants.tokenExchange.issuerPolicies.
        ActorClaim string `json:"actorClaim,omitempty"`

        // AllowedActors is the allowlist of actorClaim values authorized to
        // exchange a subject token from this issuer when it carries no
        // "may_act" claim, in addition to (not instead of) actorMatcher below —
        // either signal is sufficient. Empty denies every token unless
        // actorMatcher is set, or allowMayAct is true and the token carries a
        // permitted may_act claim.
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=256
        // +listType=atomic
        // +optional
        // This legacy field is deprecated; configure RFC 8693 policy under
        // inboundGrants.tokenExchange.issuerPolicies.
        AllowedActors []string `json:"allowedActors,omitempty"`

        // ActorMatcher is an admin-authored CEL expression evaluated against the
        // subject token's complete signature-verified claims map (bound as
        // "claims") to authorize a class of external actors, in addition to (not
        // instead of) allowedActors — either signal is sufficient. Must evaluate
        // to a boolean; a non-boolean result denies the token at evaluation time,
        // not at reconcile time. A syntactically invalid expression fails
        // reconciliation (surfaced via the AuthServerConfigValidated condition),
        // not admission — there is no validating webhook for this field.
        // +optional
        // +kubebuilder:validation:MaxLength=4096
        // This legacy field is deprecated; configure RFC 8693 policy under
        // inboundGrants.tokenExchange.issuerPolicies.
        ActorMatcher string `json:"actorMatcher,omitempty"`

        // AllowedDelegateClients restricts which ToolHive client IDs may exchange
        // an RFC 8693 subject token from this issuer. Required unless only
        // jwtBearerGrant is configured; set it to ["*"] to permit any confidential
        // client holding the token-exchange grant, or list specific client IDs to
        // bind delegation to them.
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=256
        // +listType=atomic
        // +optional
        // This legacy field is deprecated; configure RFC 8693 policy under
        // inboundGrants.tokenExchange.issuerPolicies.
        AllowedDelegateClients []string `json:"allowedDelegateClients,omitempty"`

        // AllowMayAct permits this external issuer's may_act claim to authorize
        // delegation. Defaults to false; external issuers must be opted in
        // explicitly because may_act bypasses allowedActors and actorMatcher.
        // Does not affect self-issued subject tokens. The wildcard is never
        // permitted alongside specific allowedDelegateClients, regardless of
        // this setting.
        // +kubebuilder:default=false
        // +optional
        // This legacy field is deprecated; configure RFC 8693 policy under
        // inboundGrants.tokenExchange.issuerPolicies.
        AllowMayAct bool `json:"allowMayAct,omitempty"`

        // JWTBearerGrant enables the plain RFC 7523 JWT-bearer grant for this
        // issuer. It is independent of RFC 8693 delegation policy.
        // +optional
        // This legacy field is deprecated; configure RFC 7523 policy under
        // inboundGrants.jwtBearer.issuerPolicies.
        JWTBearerGrant *JWTBearerGrantConfig `json:"jwtBearerGrant,omitempty"`
}

// JWTBearerGrantConfig limits RFC 7523 JWT-bearer assertions for one trusted
// issuer. Each assertion subject must have an exact binding and request exactly
// one of that binding's allowed resources.
//
// +kubebuilder:validation:XValidation:rule="duration(self.maxAssertionAge) > duration('0s')",message="maxAssertionAge must be greater than zero"
// +kubebuilder:validation:XValidation:rule="self.subjectBindings.all(binding, self.subjectBindings.filter(other, other.subject == binding.subject).size() == 1)",message="subjectBindings must not contain duplicate subjects"
//
//nolint:lll // CEL validation rules exceed line length limit
type JWTBearerGrantConfig struct {
        // MaxAssertionAge caps the exp-iat interval independently of exp.
        // +kubebuilder:validation:Required
        MaxAssertionAge *metav1.Duration `json:"maxAssertionAge"`

        // SubjectBindings maps an exact external subject to allowed RFC 8707
        // resources.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=50
        // +listType=atomic
        SubjectBindings []JWTBearerSubjectBinding `json:"subjectBindings"`

        // AcceptedAudiences identifies this authorization server's accepted
        // assertion audiences. When omitted, runtime validation defaults to the
        // token endpoint.
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=2048
        // +kubebuilder:validation:items:Pattern=`^https?://[^[:space:]]+$`
        // +listType=atomic
        // +optional
        AcceptedAudiences []string `json:"acceptedAudiences,omitempty"`
}

// JWTBearerSubjectBinding configures the exact subject and allowed resources
// for one RFC 7523 JWT-bearer assertion identity.
type JWTBearerSubjectBinding struct {
        // Subject is an exact assertion sub value. Wildcards are not supported.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=256
        // +kubebuilder:validation:Pattern=`^[^*]+$`
        Subject string `json:"subject"`

        // AllowedResources is the exact set of RFC 8707 resources this subject may
        // request.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=2048
        // +kubebuilder:validation:items:Pattern=`^https?://[^[:space:]]+$`
        // +listType=atomic
        AllowedResources []string `json:"allowedResources"`
}

// SPIFFEAuthenticationMethod identifies the credential type permitted for a
// SPIFFE workload. Mirrors authserver.SPIFFEAuthenticationMethod field-for-field;
// runtime parsing remains authoritative.
type SPIFFEAuthenticationMethod string

// SPIFFEBundleSourceType identifies the selected trust-bundle source. Mirrors
// authserver.SPIFFEBundleSourceType.
type SPIFFEBundleSourceType string

// SPIFFEBundleEndpointProfile identifies how a SPIFFE Bundle Endpoint's TLS
// connection is authenticated. Mirrors authserver.SPIFFEBundleEndpointProfile.
type SPIFFEBundleEndpointProfile string

const (
        // SPIFFEAuthenticationMethodX509 authenticates a workload with an X.509-SVID.
        SPIFFEAuthenticationMethodX509 SPIFFEAuthenticationMethod = "spiffe_x509"
        // SPIFFEAuthenticationMethodJWT authenticates a workload with a JWT-SVID.
        SPIFFEAuthenticationMethodJWT SPIFFEAuthenticationMethod = "spiffe_jwt"

        // SPIFFEBundleSourceTypeEndpoint selects a HTTPS SPIFFE Bundle Endpoint.
        SPIFFEBundleSourceTypeEndpoint SPIFFEBundleSourceType = "bundle_endpoint"
        // SPIFFEBundleSourceTypeFile selects a ConfigMap-mounted SPIFFE trust bundle.
        SPIFFEBundleSourceTypeFile SPIFFEBundleSourceType = "file"
        // SPIFFEBundleSourceTypeWorkloadAPI selects the local SPIFFE Workload API.
        SPIFFEBundleSourceTypeWorkloadAPI SPIFFEBundleSourceType = "workload_api"

        // SPIFFEBundleEndpointProfileHTTPSWeb authenticates the bundle endpoint's
        // TLS connection with a Web PKI certificate.
        SPIFFEBundleEndpointProfileHTTPSWeb SPIFFEBundleEndpointProfile = "https_web"
        // SPIFFEBundleEndpointProfileHTTPSSPIFFE authenticates the bundle
        // endpoint's TLS connection with a separately distributed X.509-SVID root.
        SPIFFEBundleEndpointProfileHTTPSSPIFFE SPIFFEBundleEndpointProfile = "https_spiffe"
)

// SPIFFEBundleSourceConfig is a discriminated bundle-source declaration. Type
// determines which, and only which, source payload may be set. The file-bundle
// loader is implemented in isolation, but SPIFFE authentication is not wired
// to any bundle source; endpoint and Workload API loading are not implemented.
//
// +kubebuilder:validation:XValidation:rule="self.type == 'bundle_endpoint' ? has(self.endpoint) : !has(self.endpoint)",message="endpoint configuration must be set if and only if type is 'bundle_endpoint'"
// +kubebuilder:validation:XValidation:rule="self.type == 'file' ? has(self.file) : !has(self.file)",message="file configuration must be set if and only if type is 'file'"
// +kubebuilder:validation:XValidation:rule="self.type == 'workload_api' ? has(self.workloadAPI) : !has(self.workloadAPI)",message="workloadAPI configuration must be set if and only if type is 'workload_api'"
//
//nolint:lll // CEL validation rules exceed line length limit
type SPIFFEBundleSourceConfig struct {
        // Type selects the trust-bundle source.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Enum=bundle_endpoint;file;workload_api
        Type SPIFFEBundleSourceType `json:"type"`

        // Endpoint declares a HTTPS SPIFFE Bundle Endpoint. Required when Type is
        // "bundle_endpoint".
        // +optional
        Endpoint *SPIFFEBundleEndpointSourceConfig `json:"endpoint,omitempty"`

        // File declares a SPIFFE trust bundle projected from a ConfigMap.
        // Required when Type is "file".
        // +optional
        File *SPIFFEFileBundleSourceConfig `json:"file,omitempty"`

        // WorkloadAPI selects the local SPIFFE Workload API. Required when Type
        // is "workload_api".
        // +optional
        WorkloadAPI *SPIFFEWorkloadAPIBundleSourceConfig `json:"workloadAPI,omitempty"`
}

// SPIFFEFileBundleSourceConfig selects one ConfigMap key containing a SPIFFE
// trust-bundle JWKS document, typically projected as a mounted file. The
// ConfigMap name and key follow Kubernetes ConfigMap naming constraints.
type SPIFFEFileBundleSourceConfig struct {
        // ConfigMapName is the name of the ConfigMap containing the trust bundle.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        // +kubebuilder:validation:Pattern=`^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$`
        ConfigMapName string `json:"configMapName"`

        // ConfigMapKey is the key within the ConfigMap holding the SPIFFE JWKS
        // trust-bundle document. It must not be "." or start with "..".
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        // +kubebuilder:validation:Pattern=`^[A-Za-z0-9._-]+$`
        // +kubebuilder:validation:XValidation:rule="self != '.' && !self.startsWith('..')"
        ConfigMapKey string `json:"configMapKey"`
}

// SPIFFEBundleEndpointSourceConfig declares a HTTPS SPIFFE Bundle Endpoint.
type SPIFFEBundleEndpointSourceConfig struct {
        // URL is the HTTPS SPIFFE Bundle Endpoint URL.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=2048
        URL string `json:"url"`

        // Profile selects how the endpoint's TLS connection is authenticated:
        // SPIFFEBundleEndpointProfileHTTPSWeb (Web PKI) or
        // SPIFFEBundleEndpointProfileHTTPSSPIFFE (a separately distributed
        // X.509-SVID root).
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Enum=https_web;https_spiffe
        Profile SPIFFEBundleEndpointProfile `json:"profile"`
}

// SPIFFEWorkloadAPIBundleSourceConfig selects the local SPIFFE Workload API.
// It deliberately has no payload; loading and deployment details are
// deferred to the bundle-loading implementation.
type SPIFFEWorkloadAPIBundleSourceConfig struct{}

// SPIFFETrustDomainConfig declares one SPIFFE trust domain accepted by the
// embedded authorization server. Configuration is not authentication: no
// live X.509-SVID or JWT-SVID validation exists yet, so a declared trust
// domain does not by itself let any workload authenticate — RunConfig.Validate
// (pkg/authserver/config.go) currently hard-rejects any non-empty
// spiffeTrustDomains at authserver startup via validateSPIFFENotYetEnforced,
// a deliberate placeholder until real SVID verification lands.
type SPIFFETrustDomainConfig struct {
        // Name uniquely identifies this declaration and is referenced by
        // inboundGrants.spiffeClientAuth[].trustDomainRef.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        Name string `json:"name"`

        // TrustDomain is the SPIFFE trust domain accepted by this declaration.
        // This pattern is a best-effort CRD-level approximation of the SPIFFE
        // trust-domain grammar; runtime parsing via
        // spiffeid.TrustDomainFromString remains authoritative.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=255
        // +kubebuilder:validation:Pattern=`^([a-z0-9_]|[a-z0-9_]([a-z0-9_-]|\.[a-z0-9_-])*[a-z0-9_])$`
        TrustDomain string `json:"trustDomain"`

        // Methods explicitly enables the supported credential types for this
        // trust domain.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=2
        // +kubebuilder:validation:items:Enum=spiffe_x509;spiffe_jwt
        // +listType=set
        Methods []SPIFFEAuthenticationMethod `json:"methods"`

        // BundleSource declares exactly one trust-bundle source and is validated for
        // shape here. The standalone file loader exists, but is not wired into SPIFFE
        // authentication.
        // +kubebuilder:validation:Required
        BundleSource SPIFFEBundleSourceConfig `json:"bundleSource"`
}

// SPIFFEClientConfig associates one SPIFFE principal pattern from a declared
// trust domain with an explicit OAuth client identity and permissions.
// Configuration is not authentication: configured SPIFFE clients remain
// non-public OAuth clients without a secret until live SPIFFE credential
// validation is implemented (see SPIFFETrustDomainConfig's doc comment).
//
// GrantTypes is deliberately not exposed here: the runtime only accepts
// exactly the RFC 8693 token-exchange grant for a SPIFFE client
// (validateSPIFFEGrants in pkg/authserver/spiffe_trust.go), so the converter
// always supplies it instead of letting it be configured.
//
// +kubebuilder:validation:XValidation:rule="self.principalPattern.split('/').all(segment, segment != '.' && segment != '..')",message="principalPattern path must not contain . or .. segments"
//
//nolint:lll // CEL validation rule exceeds line length limit
type SPIFFEClientConfig struct {
        // TrustDomainRef references spiffeTrustDomains[].name.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        TrustDomainRef string `json:"trustDomainRef"`

        // PrincipalPattern is a concrete SPIFFE ID or a terminal /* pattern within
        // the declared trust domain. This pattern is a best-effort CRD-level
        // approximation; runtime parsing via spiffeid.FromString remains
        // authoritative.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MaxLength=2048
        // +kubebuilder:validation:Pattern=`^spiffe://[a-z0-9._-]+((/[a-zA-Z0-9._-]+)+(/\*)?|/\*)$`
        PrincipalPattern string `json:"principalPattern"`

        // ClientID is the explicit OAuth client_id. It is never derived from a
        // SPIFFE ID.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        ClientID string `json:"clientId"`

        // Methods are the credential types this association may authenticate
        // with. Must be a subset of the referenced trust domain's methods.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=2
        // +kubebuilder:validation:items:Enum=spiffe_x509;spiffe_jwt
        // +listType=set
        Methods []SPIFFEAuthenticationMethod `json:"methods"`

        // Resources are RFC 8707 resource indicators this association may
        // request. Must be a subset of the server's allowed_audiences allowlist,
        // which is derived at reconcile time and not available on this CRD, so
        // allowlist membership is validated at reconcile time, not admission.
        // Shape (a well-formed absolute HTTP(S) URI) is independent of that
        // derived allowlist and is validated here. Distinct from Audiences: a
        // resource permission does not imply the same value is also a permitted
        // token audience, or vice versa.
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=2048
        // +kubebuilder:validation:items:Pattern=`^https?://[^@#[:space:]]+$`
        // +listType=set
        // +optional
        Resources []string `json:"resources,omitempty"`

        // Audiences are RFC 8693 token audiences this association may request.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=2048
        // +listType=set
        Audiences []string `json:"audiences"`

        // Scopes are OAuth scopes granted to this association. Must be a subset
        // of the server's effective supported scopes.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=256
        // +listType=set
        Scopes []string `json:"scopes"`
}

// InboundGrantsConfig groups canonical inbound OAuth grant-family configuration.
type InboundGrantsConfig struct {
        // SPIFFEClientAuth associates SPIFFE principal patterns with explicit
        // OAuth client identities and permissions. A sibling of TokenExchange and
        // JWTBearer below, not nested under either: client authentication does
        // not by itself confer a grant. See SPIFFEClientConfig.
        // +kubebuilder:validation:MaxItems=100
        // +listType=atomic
        // +optional
        SPIFFEClientAuth []SPIFFEClientConfig `json:"spiffeClientAuth,omitempty"`

        // TokenExchange configures RFC 8693 clients and issuer policies.
        // +optional
        TokenExchange *TokenExchangeInboundGrantConfig `json:"tokenExchange,omitempty"`

        // JWTBearer configures RFC 7523 issuer policies.
        // +optional
        JWTBearer *JWTBearerInboundGrantConfig `json:"jwtBearer,omitempty"`
}

// TokenExchangeInboundGrantConfig configures canonical RFC 8693 inbound grants.
type TokenExchangeInboundGrantConfig struct {
        // DelegateClients configures pre-provisioned confidential clients.
        // +kubebuilder:validation:MaxItems=10
        // +listType=atomic
        // +optional
        DelegateClients []DelegateClientConfig `json:"delegateClients,omitempty"`

        // IssuerPolicies binds RFC 8693 policy to named trusted issuers.
        // +kubebuilder:validation:MaxItems=20
        // +listType=atomic
        // +optional
        IssuerPolicies []TokenExchangeIssuerPolicyConfig `json:"issuerPolicies,omitempty"`
}

// TokenExchangeIssuerPolicyConfig binds RFC 8693 policy to a named trusted issuer.
//
// +kubebuilder:validation:XValidation:rule="!('*' in self.allowedDelegateClients) || size(self.allowedDelegateClients) == 1",message="allowedDelegateClients must not combine the wildcard \"*\" with specific client IDs"
// +kubebuilder:validation:XValidation:rule="!(has(self.allowMayAct) && self.allowMayAct && '*' in self.allowedDelegateClients)",message="allowMayAct must not be enabled when allowedDelegateClients contains the wildcard \"*\""
// +kubebuilder:validation:XValidation:rule="!has(self.actorClaim) || !(self.actorClaim in ['sub', 'iss', 'aud', 'exp', 'iat', 'nbf', 'jti', 'name', 'email', 'scope', 'scp', 'may_act'])",message="actorClaim must name a readable claim; use client_id or a non-reserved claim such as azp, appid, or cid"
//
//nolint:lll // CEL validation rules exceed line length limits.
type TokenExchangeIssuerPolicyConfig struct {
        // IssuerRef references trustedIssuers[].name.
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        IssuerRef string `json:"issuerRef"`

        // ExpectedAudience is the required RFC 8693 subject-token audience.
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=2048
        ExpectedAudience string `json:"expectedAudience"`

        // ActorClaim names the claim containing the external actor identity.
        // +kubebuilder:validation:MaxLength=64
        // +optional
        ActorClaim string `json:"actorClaim,omitempty"`

        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=256
        // +listType=atomic
        // +optional
        AllowedActors []string `json:"allowedActors,omitempty"`

        // +kubebuilder:validation:MaxLength=4096
        // +optional
        ActorMatcher string `json:"actorMatcher,omitempty"`

        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=50
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:MaxLength=256
        // +listType=atomic
        AllowedDelegateClients []string `json:"allowedDelegateClients"`

        // +optional
        AllowMayAct bool `json:"allowMayAct,omitempty"`
}

// JWTBearerInboundGrantConfig configures canonical RFC 7523 inbound grants.
type JWTBearerInboundGrantConfig struct {
        // IssuerPolicies binds RFC 7523 policy to named trusted issuers.
        // +kubebuilder:validation:MaxItems=20
        // +listType=atomic
        // +optional
        IssuerPolicies []JWTBearerIssuerPolicyConfig `json:"issuerPolicies,omitempty"`
}

// JWTBearerIssuerPolicyConfig binds RFC 7523 policy to a named trusted issuer.
type JWTBearerIssuerPolicyConfig struct {
        // IssuerRef references trustedIssuers[].name.
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        IssuerRef string `json:"issuerRef"`

        JWTBearerGrantConfig `json:",inline"` // nolint:revive
}

// EmbeddedAuthServerConfig holds configuration for the embedded OAuth2/OIDC authorization server.
// This enables running an authorization server that delegates authentication to upstream IDPs
// or accepts token exchange without an interactive authorization flow.
// This type is shared by MCPExternalAuthConfig.Spec.EmbeddedAuthServer and
// VirtualMCPServer.Spec.AuthServerConfig, so the XValidation rules below are
// enforced at admission for both CRDs. CEL only requires the explicit opt-in for
// delegate clients using an HTTP issuer; the shared Go validator performs the
// precise loopback-host security check.
//
// +kubebuilder:validation:XValidation:rule="(has(self.upstreamProviders) && size(self.upstreamProviders) > 0) || (has(self.delegateClients) && size(self.delegateClients) > 0) || (has(self.trustedIssuers) && self.trustedIssuers.exists(issuer, has(issuer.jwtBearerGrant))) || (has(self.inboundGrants) && (has(self.inboundGrants.tokenExchange) || has(self.inboundGrants.jwtBearer) || (has(self.inboundGrants.spiffeClientAuth) && size(self.inboundGrants.spiffeClientAuth) > 0)))",message="at least one upstream provider or inbound grant family is required"
//
// +kubebuilder:validation:XValidation:rule="((has(self.spiffeTrustDomains) && size(self.spiffeTrustDomains) > 0) == (has(self.inboundGrants) && has(self.inboundGrants.spiffeClientAuth) && size(self.inboundGrants.spiffeClientAuth) > 0))",message="spiffeTrustDomains and inboundGrants.spiffeClientAuth must be configured together"
// +kubebuilder:validation:XValidation:rule="!has(self.spiffeTrustDomains) || self.spiffeTrustDomains.all(domain, self.spiffeTrustDomains.filter(other, other.name == domain.name).size() == 1)",message="spiffeTrustDomains must not contain duplicate names"
// +kubebuilder:validation:XValidation:rule="!has(self.spiffeTrustDomains) || self.spiffeTrustDomains.all(domain, self.spiffeTrustDomains.filter(other, other.trustDomain == domain.trustDomain).size() == 1)",message="spiffeTrustDomains must not contain duplicate trust domains"
// +kubebuilder:validation:XValidation:rule="!has(self.spiffeTrustDomains) || self.spiffeTrustDomains.all(domain, has(self.inboundGrants) && has(self.inboundGrants.spiffeClientAuth) && self.inboundGrants.spiffeClientAuth.exists(client, client.trustDomainRef == domain.name))",message="every SPIFFE trust domain must be referenced by a SPIFFE client"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.spiffeClientAuth) || self.inboundGrants.spiffeClientAuth.all(client, has(self.spiffeTrustDomains) && self.spiffeTrustDomains.exists(domain, domain.name == client.trustDomainRef))",message="every SPIFFE client trustDomainRef must reference a declared trust domain"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.spiffeClientAuth) || self.inboundGrants.spiffeClientAuth.all(client, !has(self.spiffeTrustDomains) || !self.spiffeTrustDomains.exists(domain, domain.name == client.trustDomainRef) || self.spiffeTrustDomains.exists(domain, domain.name == client.trustDomainRef && client.methods.all(method, method in domain.methods)))",message="spiffeClientAuth methods must be enabled by the referenced trust domain"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.spiffeClientAuth) || self.inboundGrants.spiffeClientAuth.all(client, !has(self.spiffeTrustDomains) || !self.spiffeTrustDomains.exists(domain, domain.name == client.trustDomainRef) || self.spiffeTrustDomains.exists(domain, domain.name == client.trustDomainRef && client.principalPattern.startsWith('spiffe://' + domain.trustDomain + '/')))",message="spiffeClientAuth principalPattern trust domain must match the referenced trust domain"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.spiffeClientAuth) || self.inboundGrants.spiffeClientAuth.all(client, self.inboundGrants.spiffeClientAuth.filter(other, other.clientId == client.clientId).size() == 1)",message="spiffeClientAuth must not contain duplicate client IDs"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.spiffeClientAuth) || self.inboundGrants.spiffeClientAuth.all(client, self.inboundGrants.spiffeClientAuth.filter(other, other.principalPattern == client.principalPattern).size() == 1)",message="spiffeClientAuth must not contain duplicate principal patterns"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.spiffeClientAuth) || self.inboundGrants.spiffeClientAuth.all(client, !client.clientId.startsWith('synthetic:'))",message="spiffeClientAuth clientId must not use the reserved synthetic: prefix"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.spiffeClientAuth) || self.inboundGrants.spiffeClientAuth.all(client, !client.clientId.matches('^[A-Za-z][A-Za-z0-9+.-]*://.+'))",message="spiffeClientAuth clientId must not be an absolute URL"
//
// +kubebuilder:validation:XValidation:rule="!(has(self.insecureAllowHTTP) && self.insecureAllowHTTP && ((has(self.allowConfidentialClientRegistration) && self.allowConfidentialClientRegistration) || (has(self.delegateClients) && size(self.delegateClients) > 0) || (has(self.inboundGrants) && has(self.inboundGrants.tokenExchange) && has(self.inboundGrants.tokenExchange.delegateClients) && size(self.inboundGrants.tokenExchange.delegateClients) > 0)))",message="insecureAllowHTTP cannot be combined with confidential client registration or delegateClients; client secrets would be issued or used in cleartext over an unauthenticated endpoint"
// +kubebuilder:validation:XValidation:rule="(!has(self.forceConfidentialRedirectUris) || size(self.forceConfidentialRedirectUris) == 0) || (has(self.allowConfidentialClientRegistration) && self.allowConfidentialClientRegistration)",message="forceConfidentialRedirectUris requires allowConfidentialClientRegistration to be true"
// +kubebuilder:validation:XValidation:rule="!self.issuer.startsWith('http://') || (has(self.insecureAllowConfidentialOverLoopbackHTTP) && self.insecureAllowConfidentialOverLoopbackHTTP) || ((!has(self.allowConfidentialClientRegistration) || !self.allowConfidentialClientRegistration) && (!has(self.delegateClients) || size(self.delegateClients) == 0) && (!has(self.inboundGrants) || !has(self.inboundGrants.tokenExchange) || !has(self.inboundGrants.tokenExchange.delegateClients) || size(self.inboundGrants.tokenExchange.delegateClients) == 0))",message="confidential client registration or delegateClients with an HTTP issuer require insecureAllowConfidentialOverLoopbackHTTP to be explicitly enabled; the issuer must still be loopback"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.tokenExchange) || ((!has(self.delegateClients) || size(self.delegateClients) == 0) && (!has(self.trustedIssuers) || !self.trustedIssuers.exists(issuer, has(issuer.expectedAudience) || has(issuer.actorClaim) || has(issuer.allowedActors) || has(issuer.actorMatcher) || has(issuer.allowedDelegateClients) || (has(issuer.allowMayAct) && issuer.allowMayAct))))",message="canonical tokenExchange conflicts with legacy delegateClients or RFC 8693 trusted issuer policy"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.jwtBearer) || !has(self.trustedIssuers) || !self.trustedIssuers.exists(issuer, has(issuer.jwtBearerGrant))",message="canonical jwtBearer conflicts with legacy jwtBearerGrant"
// +kubebuilder:validation:XValidation:rule="!has(self.trustedIssuers) || self.trustedIssuers.all(issuer, !has(issuer.name) || self.trustedIssuers.filter(other, has(other.name) && other.name == issuer.name).size() == 1)",message="trustedIssuers must not contain duplicate names"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.tokenExchange) || !has(self.inboundGrants.tokenExchange.issuerPolicies) || self.inboundGrants.tokenExchange.issuerPolicies.all(policy, has(self.trustedIssuers) && self.trustedIssuers.exists(issuer, has(issuer.name) && issuer.name == policy.issuerRef))",message="every tokenExchange issuerRef must reference a named trusted issuer"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.tokenExchange) || !has(self.inboundGrants.tokenExchange.issuerPolicies) || self.inboundGrants.tokenExchange.issuerPolicies.all(policy, self.inboundGrants.tokenExchange.issuerPolicies.filter(other, other.issuerRef == policy.issuerRef).size() == 1)",message="tokenExchange issuerPolicies must not contain duplicate issuerRef values"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.jwtBearer) || !has(self.inboundGrants.jwtBearer.issuerPolicies) || self.inboundGrants.jwtBearer.issuerPolicies.all(policy, has(self.trustedIssuers) && self.trustedIssuers.exists(issuer, has(issuer.name) && issuer.name == policy.issuerRef))",message="every jwtBearer issuerRef must reference a named trusted issuer"
// +kubebuilder:validation:XValidation:rule="!has(self.inboundGrants) || !has(self.inboundGrants.jwtBearer) || !has(self.inboundGrants.jwtBearer.issuerPolicies) || self.inboundGrants.jwtBearer.issuerPolicies.all(policy, self.inboundGrants.jwtBearer.issuerPolicies.filter(other, other.issuerRef == policy.issuerRef).size() == 1)",message="jwtBearer issuerPolicies must not contain duplicate issuerRef values"
//
// The shared Go-level ValidateConfidentialClientTransport validator remains the
// source of truth for confidential-client transport and loopback policy,
// including delegate clients. Trusted issuer endpoint shape is validated by
// ValidateInboundGrants; audience and outbound DNS/private-IP checks remain
// runtime-only.
//
//nolint:lll // CEL validation rules exceed line length limits.
type EmbeddedAuthServerConfig struct {
        // Issuer is the issuer identifier for this authorization server.
        // This will be included in the "iss" claim of issued tokens.
        // Must be a valid HTTPS URL (or HTTP for localhost, or HTTP for trusted in-cluster hosts when
        // insecureAllowHTTP is true) without query, fragment, or trailing slash (per RFC 8414).
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https?://[^\s?#]+[^/\s?#]$`
        Issuer string `json:"issuer"`

        // AuthorizationEndpointBaseURL overrides the base URL used for the authorization_endpoint
        // in the OAuth discovery document. When set, the discovery document will advertise
        // `{authorizationEndpointBaseUrl}/oauth/authorize` instead of `{issuer}/oauth/authorize`.
        // All other endpoints (token, registration, JWKS) remain derived from the issuer.
        // This is useful when the browser-facing authorization endpoint needs to be on a
        // different host than the issuer used for backend-to-backend calls.
        // The upstream callback (`redirectUri` on each upstream provider, defaulting to
        // `{resourceUrl}/oauth/callback`) must share this hostname: the callback only
        // completes a login for the browser that started it, bound by a host-only cookie
        // set on this endpoint, so a different callback host rejects every browser login.
        // Must be a valid HTTPS URL (or HTTP for localhost, or HTTP for trusted in-cluster hosts
        // when insecureAllowHTTP is true) without query, fragment, or trailing slash.
        // +kubebuilder:validation:Pattern=`^https?://[^\s?#]+[^/\s?#]$`
        // +optional
        AuthorizationEndpointBaseURL string `json:"authorizationEndpointBaseUrl,omitempty"`

        // SigningKeySecretRefs references Kubernetes Secrets containing signing keys for JWT operations.
        // Supports key rotation by allowing multiple keys (oldest keys are used for verification only).
        // If not specified, an ephemeral signing key will be auto-generated (development only -
        // JWTs will be invalid after restart).
        // +kubebuilder:validation:MaxItems=5
        // +listType=atomic
        // +optional
        SigningKeySecretRefs []SecretKeyRef `json:"signingKeySecretRefs,omitempty"`

        // HMACSecretRefs references Kubernetes Secrets containing symmetric secrets for signing
        // authorization codes and refresh tokens (opaque tokens).
        // Current secret must be at least 32 bytes and cryptographically random.
        // Supports secret rotation via multiple entries (first is current, rest are for verification).
        // If not specified, an ephemeral secret will be auto-generated (development only -
        // auth codes and refresh tokens will be invalid after restart).
        // +listType=atomic
        // +optional
        HMACSecretRefs []SecretKeyRef `json:"hmacSecretRefs,omitempty"`

        // TokenLifespans configures the duration that various tokens are valid.
        // If not specified, defaults are applied (access: 1h, refresh: 7d, authCode: 10m).
        // +optional
        TokenLifespans *TokenLifespanConfig `json:"tokenLifespans,omitempty"`

        // SPIFFETrustDomains declares SPIFFE trust domains for
        // inboundGrants.spiffeClientAuth associations. See SPIFFETrustDomainConfig's
        // doc comment for why declaring a domain does not by itself enable
        // authentication in this build.
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=50
        // +listType=atomic
        // +optional
        SPIFFETrustDomains []SPIFFETrustDomainConfig `json:"spiffeTrustDomains,omitempty"`

        // InboundGrants configures canonical inbound OAuth grant families.
        // +optional
        InboundGrants *InboundGrantsConfig `json:"inboundGrants,omitempty"`

        // UpstreamProviders configures connections to upstream Identity Providers.
        // When configured, the embedded auth server delegates interactive authentication
        // to these providers. It may be omitted only when delegateClients or a trusted
        // issuer with jwtBearerGrant enables token-only operation.
        // MCPServer and MCPRemoteProxy support a single upstream; VirtualMCPServer supports multiple.
        // +listType=map
        // +listMapKey=name
        // +optional
        UpstreamProviders []UpstreamProviderConfig `json:"upstreamProviders,omitempty"`

        // PrimaryUpstreamProvider names the upstream IDP whose access token Cedar
        // should read claims from when authorising a request. Must match the name
        // of one of the entries in UpstreamProviders. When empty, the controller
        // auto-selects the first entry of UpstreamProviders.
        //
        // Only meaningful on VirtualMCPServer, where multiple upstream providers
        // can be configured and Cedar needs to pick which token's claims to
        // evaluate. The VirtualMCPServer controller validates this field against
        // UpstreamProviders at admission and rejects unresolvable values.
        //
        // On MCPServer and MCPRemoteProxy this field is structurally present (the
        // EmbeddedAuthServerConfig struct is shared) but has no runtime effect:
        // those CRDs are restricted to a single upstream so there is no choice to
        // make. Setting it on those CRDs is silently ignored.
        // +optional
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=63
        // +kubebuilder:validation:Pattern=`^[a-z0-9]([a-z0-9-]*[a-z0-9])?$`
        PrimaryUpstreamProvider string `json:"primaryUpstreamProvider,omitempty"`

        // Storage configures the storage backend for the embedded auth server.
        // If not specified, defaults to in-memory storage.
        // +optional
        Storage *AuthServerStorageConfig `json:"storage,omitempty"`

        // DisableUpstreamTokenInjection prevents the embedded auth server from injecting
        // upstream IdP tokens into requests forwarded to the backend MCP server.
        // When true, the embedded auth server still handles OAuth flows for clients,
        // but instead of swapping ToolHive JWTs for upstream tokens the proxy STRIPS
        // the client's credential headers (Authorization, Cookie, Proxy-Authorization)
        // after validating the JWT — the backend receives an unauthenticated request.
        // Use headerForward to attach static credentials (e.g. an API key) if the
        // backend needs them. Cannot be combined with token exchange, AWS STS, or OBO
        // middleware, which would re-add credentials after the strip.
        // This is useful when the backend MCP server does not require authentication
        // (e.g., public documentation servers) but you still want client authentication.
        // +kubebuilder:default=false
        // +optional
        DisableUpstreamTokenInjection bool `json:"disableUpstreamTokenInjection,omitempty"`

        // InsecureAllowHTTP permits an http:// issuer URL for non-localhost hosts.
        // Only set this for in-cluster Kubernetes deployments where traffic between
        // pods traverses a trusted network (e.g. the in-cluster service mesh).
        // Production deployments reachable outside the cluster MUST use https://.
        //
        // On VirtualMCPServer: when false (the default), http:// issuers for non-localhost
        // hosts are rejected at reconcile time with an AuthServerConfigValidated=False condition.
        //
        // On MCPServer and MCPRemoteProxy (via MCPExternalAuthConfig): this field is
        // structurally present but enforcement is deferred to pod startup via Config.Validate();
        // a misconfigured issuer will cause the pod to crash at startup rather than surface
        // as an operator condition.
        //
        // One combination is rejected at admission on all three CRDs regardless of the
        // above: setting this field alongside confidential client registration or
        // delegate clients, which would issue or use client secrets in cleartext over
        // an unauthenticated endpoint (see the XValidation rule on
        // EmbeddedAuthServerConfig).
        // +kubebuilder:default=false
        // +optional
        InsecureAllowHTTP bool `json:"insecureAllowHTTP,omitempty"`

        // BaselineClientScopes is a baseline set of OAuth 2.0 scopes guaranteed to be
        // included in every client registration. The embedded auth server unions these
        // scopes into the registered set returned by RFC 7591 Dynamic Client
        // Registration, so a client that narrows the `scope` field at /oauth/register
        // can still request the baseline scopes at /oauth/authorize. All values must
        // be present in the upstream-derived scopesSupported set; the auth server
        // fails to start if any value is missing.
        //
        // Security: every client registered via /oauth/register will gain the
        // ability to request these scopes at /oauth/authorize, regardless of what
        // the client itself requested. Keep the baseline narrow (typically
        // "openid" and "offline_access"). Adding a privileged scope here — e.g.
        // "admin:read" — would grant it to every DCR-registered client, including
        // public clients like Claude Code, Cursor, and VS Code.
        // When cimd.enabled is true, every dynamically resolved CIMD client will
        // also gain the ability to request these scopes, including third-party
        // clients resolved from arbitrary HTTPS URLs.
        // +kubebuilder:validation:MaxItems=10
        // +kubebuilder:validation:items:MinLength=1
        // +kubebuilder:validation:items:Pattern=`^[\x21\x23-\x5B\x5D-\x7E]+$`
        // +listType=atomic
        // +optional
        BaselineClientScopes []string `json:"baselineClientScopes,omitempty"`

        // AllowConfidentialClientRegistration permits RFC 7591 Dynamic Client
        // Registration of confidential clients: when true, /oauth/register
        // accepts token_endpoint_auth_method values client_secret_basic and
        // client_secret_post in addition to "none" (still the default on
        // omission) and mints a client_secret returned exactly once.
        // Confidential registrations are restricted to https non-loopback
        // redirect URIs, and on the Redis storage backend all DCR-issued
        // registrations are evicted after 30 days of inactivity and must
        // re-register. This gates registration only: disabling it does not
        // revoke or reject already-minted secrets at the token endpoint.
        //
        // Security: registration is unauthenticated, so enabling this lets any
        // caller who can reach the endpoint obtain a client credential.
        // Combining it with insecureAllowHTTP is rejected at validation.
        // +kubebuilder:default=false
        // +optional
        AllowConfidentialClientRegistration bool `json:"allowConfidentialClientRegistration,omitempty"`

        // AllowPrivateKeyJWTRegistration permits Dynamic Client Registration of
        // clients using private_key_jwt authentication. Registration behavior is
        // intentionally configured separately from confidential-client registration.
        //
        // Security: registration is unauthenticated, so enabling this lets any
        // caller who can reach the endpoint register a private_key_jwt client.
        // Unlike allowConfidentialClientRegistration, this is NOT rejected when
        // combined with insecureAllowHTTP: registration never returns a secret
        // for a private_key_jwt client, so there is nothing for cleartext HTTP
        // to expose.
        // +kubebuilder:default=false
        // +optional
        AllowPrivateKeyJWTRegistration bool `json:"allowPrivateKeyJWTRegistration,omitempty"`

        // InsecureAllowConfidentialOverLoopbackHTTP opts in to confidential
        // Dynamic Client Registration (DCR) and delegate clients when issuer is a
        // plain-HTTP loopback URL (e.g. "http://localhost:8080"). Without this
        // flag, that combination is rejected at reconcile time: a loopback http://
        // issuer is normally fine for local development since the traffic never
        // leaves the machine, but confidential clients send secrets over cleartext.
        // Forcing TLS onto every loopback deployment instead would just push
        // operators toward insecureAllowHTTP, which is worse: that also disables
        // the non-loopback host check. Has no effect when there are no confidential
        // clients or issuer is https.
        //
        // private_key_jwt registration has no equivalent flag or transport
        // restriction: unlike confidential registration, it never returns a
        // client_secret (or any other secret) in the DCR response, so there is
        // nothing here for cleartext HTTP to expose.
        // +kubebuilder:default=false
        // +optional
        InsecureAllowConfidentialOverLoopbackHTTP bool `json:"insecureAllowConfidentialOverLoopbackHTTP,omitempty"`

        // DelegateClients configures pre-provisioned confidential clients for RFC 8693
        // token exchange. Each secret is referenced from a Kubernetes Secret; no
        // plaintext secret, redirect URI, or grant selection is accepted here. The
        // operator always supplies the token-exchange grant when it converts this
        // configuration to the runtime contract.
        //
        // This is independent of allowConfidentialClientRegistration: it neither
        // enables nor requires unauthenticated confidential dynamic client
        // registration.
        // +kubebuilder:validation:MaxItems=10
        // +listType=atomic
        // +optional
        // This legacy field is deprecated; use inboundGrants.tokenExchange.delegateClients.
        DelegateClients []DelegateClientConfig `json:"delegateClients,omitempty"`

        // TrustedIssuers configures external OIDC issuers whose tokens are
        // accepted as RFC 8693 subject tokens during token exchange, in addition
        // to self-issued subject tokens. Empty (the default) means only
        // self-issued subject tokens are accepted. See
        // docs/arch/17-token-exchange-delegation.md for the trust model.
        // +kubebuilder:validation:MaxItems=20
        // +listType=atomic
        // +optional
        TrustedIssuers []TrustedIssuerConfig `json:"trustedIssuers,omitempty"`

        // ForceConfidentialRedirectURIs lists redirect URIs that must be
        // registered as confidential clients regardless of the
        // token_endpoint_auth_method the DCR request declares. A registration
        // whose redirectUris contains an EXACT match for one of these entries is
        // issued a real client_secret and reported back as
        // token_endpoint_auth_method "client_secret_post", even if the request
        // said "none" or omitted the field.
        //
        // Intended for MCP clients that declare themselves public
        // (token_endpoint_auth_method: "none") per RFC 7591 but then refuse to
        // proceed because the response carries no client_secret — a
        // self-contradictory request. RFC 7591 §3.2.1 permits the server to
        // substitute client metadata, so this takes such a client at its word
        // that it wants a secret. Remove an entry once the client is fixed to
        // handle "none" registrations correctly.
        //
        // Exact matching is deliberate: an attacker who registers with someone
        // else's callback URI is issued a secret for a client whose
        // authorization codes are delivered to that someone else's redirect
        // endpoint, not to the attacker, so this is not a way to obtain a usable
        // credential for another client.
        //
        // Requires allowConfidentialClientRegistration to be true. Every entry
        // must be an https non-loopback URI — a loopback client is a public
        // client by construction (OAuth 2.1 §2.1) and must not be issued a
        // secret; this is enforced at reconcile time since CEL cannot express
        // the loopback-hostname check.
        // +kubebuilder:validation:MaxItems=10
        // +kubebuilder:validation:items:Pattern=`^https://[^\s?#]+$`
        // +listType=atomic
        // +optional
        ForceConfidentialRedirectURIs []string `json:"forceConfidentialRedirectUris,omitempty"`

        // CIMD configures Client ID Metadata Document support. When omitted, CIMD is disabled.
        // +optional
        CIMD *EmbeddedAuthServerCIMDConfig `json:"cimd,omitempty"`

        // RegistrationRateLimit bounds the unauthenticated /oauth/register
        // endpoint with a token bucket: up to maxTokens registrations at once,
        // refilled at maxTokens per refillPeriod. For a sustained rate of R per
        // second with a burst of B, set maxTokens to B and refillPeriod to B/R;
        // for example, 10 per minute with a burst of 3 is maxTokens 3,
        // refillPeriod 18s. If omitted, defaults to 1 registration per second
        // with a burst of 5 (maxTokens 5, refillPeriod 5s). The limit is per
        // replica, so the effective rate is the configured rate multiplied by the
        // replica count. Every successful registration stores a client, so a
        // higher rate fills storage.maxDCRClients sooner, and with maxDCRClients
        // 0 this limit is the only bound on stored clients.
        // +kubebuilder:validation:XValidation:rule="duration(self.refillPeriod) > duration('0s')",message="registrationRateLimit.refillPeriod must be positive"
        // +optional
        RegistrationRateLimit *RateLimitBucket `json:"registrationRateLimit,omitempty"`
}

// effectiveInboundGrants is the folded validation view of legacy and canonical
// inbound grant configuration. Its slice fields are cloned from the source
// EmbeddedAuthServerConfig when built, so mutating them cannot affect it.
type effectiveInboundGrants struct {
        tokenExchange   bool
        jwtBearer       bool
        delegateClients []DelegateClientConfig
        trustedIssuers  []tokenexchange.TrustedIssuer
}

// EffectiveGrants folds legacy and canonical inbound grants using the same
// conflict and overlay semantics as authserver.NormalizeInboundGrants.
func (c *EmbeddedAuthServerConfig) EffectiveGrants() (effectiveInboundGrants, error) {
        if err := c.validateJWTBearerMaxAssertionAges(); err != nil {
                return effectiveInboundGrants{}, err
        }
        result := effectiveInboundGrants{
                tokenExchange:   c.InboundGrants == nil,
                delegateClients: cloneDelegateClientConfigs(c.DelegateClients),
                trustedIssuers:  buildTrustedIssuerConfigs(c.TrustedIssuers),
        }
        issuerByName, legacyTokenExchange, legacyJWTBearer, err := c.indexEffectiveTrustedIssuers(result.trustedIssuers)
        if err != nil {
                return effectiveInboundGrants{}, err
        }
        result.jwtBearer = legacyJWTBearer
        if c.InboundGrants == nil {
                return result, nil
        }
        if err := applyEffectiveTokenExchange(&result, c.InboundGrants.TokenExchange, issuerByName, legacyTokenExchange); err != nil {
                return effectiveInboundGrants{}, err
        }
        if err := applyEffectiveJWTBearer(&result, c.InboundGrants.JWTBearer, issuerByName, legacyJWTBearer); err != nil {
                return effectiveInboundGrants{}, err
        }
        return result, nil
}

func (c *EmbeddedAuthServerConfig) validateJWTBearerMaxAssertionAges() error {
        for i, issuer := range c.TrustedIssuers {
                if issuer.JWTBearerGrant != nil && issuer.JWTBearerGrant.MaxAssertionAge == nil {
                        return fmt.Errorf("trustedIssuers[%d].jwtBearerGrant.maxAssertionAge is required", i)
                }
        }
        if c.InboundGrants == nil || c.InboundGrants.JWTBearer == nil {
                return nil
        }
        for i, policy := range c.InboundGrants.JWTBearer.IssuerPolicies {
                if policy.MaxAssertionAge == nil {
                        return fmt.Errorf("inboundGrants.jwtBearer.issuerPolicies[%d].maxAssertionAge is required", i)
                }
        }
        return nil
}

func (c *EmbeddedAuthServerConfig) indexEffectiveTrustedIssuers(
        issuers []tokenexchange.TrustedIssuer,
) (map[string]int, bool, bool, error) {
        byName := make(map[string]int, len(issuers))
        byURL := make(map[string]int, len(issuers))
        legacyTokenExchange := len(c.DelegateClients) > 0
        legacyJWTBearer := false
        for i, issuer := range issuers {
                if previous, ok := byURL[issuer.IssuerURL]; ok {
                        return nil, false, false, fmt.Errorf(
                                "trustedIssuers[%d].issuerUrl duplicates trustedIssuers[%d].issuerUrl", i, previous)
                }
                byURL[issuer.IssuerURL] = i
                if issuer.Name != "" {
                        if previous, ok := byName[issuer.Name]; ok {
                                return nil, false, false, fmt.Errorf(
                                        "trustedIssuers[%d].name duplicates trustedIssuers[%d].name %q", i, previous, issuer.Name)
                        }
                        byName[issuer.Name] = i
                }
                legacyTokenExchange = legacyTokenExchange || hasLegacyTokenExchangePolicy(c.TrustedIssuers[i])
                legacyJWTBearer = legacyJWTBearer || issuer.JWTBearerGrant != nil
        }
        return byName, legacyTokenExchange, legacyJWTBearer, nil
}

func applyEffectiveTokenExchange(
        result *effectiveInboundGrants,
        config *TokenExchangeInboundGrantConfig,
        issuerByName map[string]int,
        legacy bool,
) error {
        if config == nil {
                result.tokenExchange = legacy
                return nil
        }
        if legacy {
                return fmt.Errorf("inboundGrants.tokenExchange conflicts with legacy delegateClients or RFC 8693 policy in trustedIssuers")
        }
        result.tokenExchange = true
        result.delegateClients = cloneDelegateClientConfigs(config.DelegateClients)
        seen := make(map[string]int, len(config.IssuerPolicies))
        for i, policy := range config.IssuerPolicies {
                issuerIndex, err := resolveEffectiveIssuerRef(issuerByName, seen, policy.IssuerRef,
                        fmt.Sprintf("inboundGrants.tokenExchange.issuerPolicies[%d]", i))
                if err != nil {
                        return err
                }
                seen[policy.IssuerRef] = i
                issuer := &result.trustedIssuers[issuerIndex]
                issuer.ExpectedAudience = policy.ExpectedAudience
                issuer.ActorClaim = policy.ActorClaim
                issuer.AllowedActors = slices.Clone(policy.AllowedActors)
                issuer.ActorMatcher = policy.ActorMatcher
                issuer.AllowedDelegateClients = slices.Clone(policy.AllowedDelegateClients)
                issuer.AllowMayAct = policy.AllowMayAct
        }
        return nil
}

func applyEffectiveJWTBearer(
        result *effectiveInboundGrants,
        config *JWTBearerInboundGrantConfig,
        issuerByName map[string]int,
        legacy bool,
) error {
        if config == nil {
                return nil
        }
        if legacy {
                return fmt.Errorf("inboundGrants.jwtBearer conflicts with legacy trustedIssuers[*].jwtBearerGrant")
        }
        result.jwtBearer = true
        seen := make(map[string]int, len(config.IssuerPolicies))
        for i, policy := range config.IssuerPolicies {
                issuerIndex, err := resolveEffectiveIssuerRef(issuerByName, seen, policy.IssuerRef,
                        fmt.Sprintf("inboundGrants.jwtBearer.issuerPolicies[%d]", i))
                if err != nil {
                        return err
                }
                seen[policy.IssuerRef] = i
                result.trustedIssuers[issuerIndex].JWTBearerGrant = buildJWTBearerGrantPolicy(&policy.JWTBearerGrantConfig)
        }
        return nil
}

// ValidateInboundGrants validates the effective grant projection shared by
// MCPExternalAuthConfig and inline VirtualMCPServer auth configuration.
func (c *EmbeddedAuthServerConfig) ValidateInboundGrants() error {
        grants, err := c.EffectiveGrants()
        if err != nil {
                return err
        }
        if err := authserver.ValidateConfidentialClientTransport(
                c.AllowConfidentialClientRegistration || len(grants.delegateClients) > 0,
                c.InsecureAllowHTTP,
                c.Issuer,
                c.InsecureAllowConfidentialOverLoopbackHTTP,
        ); err != nil {
                return err
        }
        if err := tokenexchange.ValidateTrustedIssuers(grants.trustedIssuers, c.Issuer, nil); err != nil {
                return fmt.Errorf("trustedIssuers: %w", err)
        }
        return nil
}

func hasLegacyTokenExchangePolicy(issuer TrustedIssuerConfig) bool {
        return issuer.ExpectedAudience != "" || issuer.ActorClaim != "" || len(issuer.AllowedActors) > 0 ||
                issuer.ActorMatcher != "" || len(issuer.AllowedDelegateClients) > 0 || issuer.AllowMayAct
}

func resolveEffectiveIssuerRef(byName, seen map[string]int, ref, path string) (int, error) {
        if ref == "" {
                return 0, fmt.Errorf("%s.issuerRef is required", path)
        }
        if previous, ok := seen[ref]; ok {
                return 0, fmt.Errorf("%s.issuerRef duplicates issuer policy [%d] for %q", path, previous, ref)
        }
        index, ok := byName[ref]
        if !ok {
                return 0, fmt.Errorf("%s.issuerRef references unknown or unnamed trusted issuer %q", path, ref)
        }
        return index, nil
}

func cloneDelegateClientConfigs(clients []DelegateClientConfig) []DelegateClientConfig {
        cloned := slices.Clone(clients)
        for i := range cloned {
                cloned[i].Scopes = slices.Clone(cloned[i].Scopes)
                cloned[i].Audiences = slices.Clone(cloned[i].Audiences)
                if cloned[i].ClientSecretRef != nil {
                        secretRef := *cloned[i].ClientSecretRef
                        cloned[i].ClientSecretRef = &secretRef
                }
        }
        return cloned
}

// TokenLifespanConfig holds configuration for token lifetimes.
type TokenLifespanConfig struct {
        // AccessTokenLifespan is the duration that access tokens are valid.
        // Format: Go duration string (e.g., "1h", "30m", "24h").
        // If empty, defaults to 1 hour.
        // +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
        // +optional
        AccessTokenLifespan string `json:"accessTokenLifespan,omitempty"`

        // RefreshTokenLifespan is the duration that refresh tokens are valid.
        // Format: Go duration string (e.g., "168h", "7d" as "168h").
        // If empty, defaults to 7 days (168h).
        // +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
        // +optional
        RefreshTokenLifespan string `json:"refreshTokenLifespan,omitempty"`

        // AuthCodeLifespan is the duration that authorization codes are valid.
        // Format: Go duration string (e.g., "10m", "5m").
        // If empty, defaults to 10 minutes.
        // +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
        // +optional
        AuthCodeLifespan string `json:"authCodeLifespan,omitempty"`
}

// EmbeddedAuthServerCIMDConfig configures Client ID Metadata Document (CIMD) support
// on the embedded authorization server. When enabled, the AS accepts HTTPS URLs as
// client_id values and resolves them via the CIMD protocol, allowing clients such as
// VS Code to authenticate without prior Dynamic Client Registration.
type EmbeddedAuthServerCIMDConfig struct {
        // Enabled activates CIMD client lookup. When false (the default), the AS only
        // accepts client_id values that were registered via DCR.
        // +kubebuilder:default=false
        Enabled bool `json:"enabled"`

        // CacheMaxSize is the maximum number of CIMD documents held in the LRU cache.
        // Defaults to 256 when Enabled is true and this field is omitted.
        // +kubebuilder:validation:Minimum=1
        // +optional
        CacheMaxSize int `json:"cacheMaxSize,omitempty"`

        // CacheFallbackTTL is the fixed TTL applied to every cached CIMD document.
        // Cache-Control header parsing is not yet implemented; all entries use this value.
        // Format: Go duration string (e.g. "5m", "10m", "1h").
        // Defaults to 5 minutes when Enabled is true and this field is omitted.
        // +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
        // +optional
        CacheFallbackTTL string `json:"cacheFallbackTtl,omitempty"`
}

// UpstreamProviderType identifies the type of upstream Identity Provider.
type UpstreamProviderType string

const (
        // UpstreamProviderTypeOIDC is for OIDC providers with discovery support
        UpstreamProviderTypeOIDC UpstreamProviderType = "oidc"

        // UpstreamProviderTypeOAuth2 is for pure OAuth 2.0 providers with explicit endpoints
        UpstreamProviderTypeOAuth2 UpstreamProviderType = "oauth2"
)

// UpstreamProviderConfig defines configuration for an upstream Identity Provider.
//
// Exactly one of OIDCConfig or OAuth2Config must be set and must match the
// declared Type: oidc-typed providers set OIDCConfig, oauth2-typed providers
// set OAuth2Config. The CEL rule below enforces the pairing at admission; the
// matching Go-level check in validateUpstreamProvider provides defense-in-depth
// for stored objects.
//
// The rule is structured as a chain of equality checks ending in an explicit
// `false`, so adding a new UpstreamProviderType value without extending this
// rule fails admission instead of silently demanding the OAuth2 shape. When
// adding a new type, extend both this rule and validateUpstreamProvider.
//
// +kubebuilder:validation:XValidation:rule="self.type == 'oidc' ? (has(self.oidcConfig) && !has(self.oauth2Config)) : self.type == 'oauth2' ? (has(self.oauth2Config) && !has(self.oidcConfig)) : false",message="type must be 'oidc' or 'oauth2'; oidcConfig must be set when type is 'oidc' and oauth2Config must be set when type is 'oauth2' (and the other must not be set)"
//
//nolint:lll // CEL validation rules exceed line length limit
type UpstreamProviderConfig struct {
        // Name uniquely identifies this upstream provider.
        // Used for routing decisions and session binding in multi-upstream scenarios.
        // Must be lowercase alphanumeric with hyphens (DNS-label-like).
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=63
        // +kubebuilder:validation:Pattern=`^[a-z0-9]([a-z0-9-]*[a-z0-9])?$`
        Name string `json:"name"`

        // Type specifies the provider type: "oidc" or "oauth2"
        // +kubebuilder:validation:Enum=oidc;oauth2
        // +kubebuilder:validation:Required
        Type UpstreamProviderType `json:"type"`

        // OIDCConfig contains OIDC-specific configuration.
        // Required when Type is "oidc", must be nil when Type is "oauth2".
        // +optional
        OIDCConfig *OIDCUpstreamConfig `json:"oidcConfig,omitempty"`

        // OAuth2Config contains OAuth 2.0-specific configuration.
        // Required when Type is "oauth2", must be nil when Type is "oidc".
        // +optional
        OAuth2Config *OAuth2UpstreamConfig `json:"oauth2Config,omitempty"`
}

// OIDCUpstreamConfig contains configuration for OIDC providers.
// +kubebuilder:validation:XValidation:rule="(has(self.clientId) && size(self.clientId) > 0) ? !has(self.dcrConfig) : has(self.dcrConfig)",message="exactly one of clientId or dcrConfig must be set"
// +kubebuilder:validation:XValidation:rule="!(has(self.dcrConfig) && has(self.clientSecretRef))",message="clientSecretRef must not be set when dcrConfig is set; the client_secret is obtained at runtime via Dynamic Client Registration"
// +kubebuilder:validation:XValidation:rule="!(has(self.dcrConfig) && has(self.dcrConfig.discoveryUrl) && has(self.dcrConfig.registrationEndpoint))",message="at most one of discoveryUrl or registrationEndpoint may be set"
//
//nolint:lll // CEL validation rules exceed line length limit
type OIDCUpstreamConfig struct {
        // IssuerURL is the OIDC issuer URL for automatic endpoint discovery.
        // Must be a valid HTTPS URL.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https://.*$`
        IssuerURL string `json:"issuerUrl"`

        // ClientID is the OAuth 2.0 client identifier registered with the upstream IdP.
        // +optional
        ClientID string `json:"clientId,omitempty"`

        // DCRConfig enables RFC 7591 Dynamic Client Registration. When set, ClientID
        // and ClientSecretRef must be omitted. If neither discoveryUrl nor
        // registrationEndpoint is set, discovery is derived from issuerUrl.
        // +optional
        DCRConfig *DCRUpstreamConfig `json:"dcrConfig,omitempty"`

        // ClientSecretRef references a Kubernetes Secret containing the OAuth 2.0 client secret.
        // Optional for public clients using PKCE instead of client secret.
        // +optional
        ClientSecretRef *SecretKeyRef `json:"clientSecretRef,omitempty"`

        // RedirectURI is the callback URL where the upstream IdP will redirect after authentication.
        // When not specified, defaults to `{resourceUrl}/oauth/callback` where `resourceUrl` is the
        // URL associated with the resource (e.g., MCPServer or vMCP) using this config.
        // Its hostname must match the browser-facing authorization endpoint (`issuer`, or
        // `authorizationEndpointBaseUrl` when set): the callback is bound to the browser that
        // started the login by a host-only cookie, so a different host rejects every browser login.
        // +optional
        RedirectURI string `json:"redirectUri,omitempty"`

        // Scopes are the OAuth scopes to request from the upstream IdP.
        // If not specified, defaults to ["openid", "offline_access"].
        // When using additionalAuthorizationParams with provider-specific refresh token
        // mechanisms (e.g., Google's access_type=offline), set explicit scopes to avoid
        // sending both offline_access and the provider-specific parameter.
        // +listType=atomic
        // +optional
        Scopes []string `json:"scopes,omitempty"`

        // UserInfoOverride allows customizing UserInfo fetching behavior for OIDC providers.
        // By default, the UserInfo endpoint is discovered automatically via OIDC discovery.
        // Use this to override the endpoint URL, HTTP method, or field mappings for providers
        // that return non-standard claim names in their UserInfo response.
        // +optional
        UserInfoOverride *UserInfoConfig `json:"userInfoOverride,omitempty"`

        // AdditionalAuthorizationParams are extra query parameters to include in
        // authorization requests sent to the upstream provider.
        // This is useful for providers that require custom parameters, such as
        // Google's access_type=offline for obtaining refresh tokens.
        // Note: when using access_type=offline, also set explicit scopes to avoid
        // the default offline_access scope being sent alongside it.
        // Framework-managed parameters (response_type, client_id, redirect_uri,
        // scope, state, code_challenge, code_challenge_method, nonce) are not allowed.
        // +kubebuilder:validation:MaxProperties=16
        // +optional
        AdditionalAuthorizationParams map[string]string `json:"additionalAuthorizationParams,omitempty"`

        // AdditionalTokenParams are extra form-body parameters to include in
        // token requests (authorization code exchange and refresh) sent to the
        // upstream provider's token endpoint.
        // This is useful for providers that enforce RFC 8707 resource indicators
        // on token requests, where the resource parameter must accompany the code
        // exchange and refresh, not only the authorization request.
        // Framework-managed parameters (grant_type, code, redirect_uri, client_id,
        // client_secret, code_verifier, refresh_token, scope) are not allowed.
        // +kubebuilder:validation:MaxProperties=16
        // +optional
        AdditionalTokenParams map[string]string `json:"additionalTokenParams,omitempty"`

        // OmitScopeOnRefresh stops the scope parameter from being sent on
        // refresh_token grant requests to the upstream provider's token endpoint.
        // Scopes are still sent on the authorization request.
        // By default (false), the configured scopes are sent on refresh as well,
        // which some providers require to avoid narrowing the refreshed token.
        // Set it for providers that reject scope on refresh, such as Salesforce
        // ("invalid_request - scope parameter not supported").
        // +optional
        OmitScopeOnRefresh bool `json:"omitScopeOnRefresh,omitempty"`

        // SubjectClaim names the validated ID-token claim to use as the upstream
        // subject. Defaults to "sub" when empty. Set it for IdPs where "sub" isn't
        // stable per user — e.g. Entra/Azure AD, whose "sub" rotates per application
        // and whose stable identifier is "oid".
        //
        // The value is looked up verbatim as a top-level claim name, so it is
        // constrained to a claim-name shape: it must start with a letter or
        // underscore and contain only letters, digits, and underscores. This rejects
        // dotted, colon-namespaced, or whitespace-containing values at admission
        // rather than letting a typo silently miss the claim at login, and keeps the
        // field aligned with the directory service's per-issuer bindingClaim.
        //
        // Changing this on a live deployment re-keys existing users (the value
        // resolves to the internal user ID), so treat it as immutable once users
        // exist.
        //
        // Per-IdP notes:
        //   - Entra/Azure AD: use "oid"; it is only emitted when the upstream scopes
        //     include "profile". "oid" is unique within a single tenant — multi-tenant
        //     apps need oid+tid, which this single-claim field cannot express.
        //   - Okta: the org auth server already puts the stable id in "sub" (default
        //     works). A custom auth server's "sub" is the mutable login/email and the
        //     stable "uid" lives only in the access token, not the ID token — map a
        //     custom ID-token claim and point subjectClaim at it.
        // The pattern matches the claim-name shape and allows empty (defaults to
        // "sub"). Using Pattern rather than a CEL XValidation rule keeps this off the
        // CRD's CEL cost budget — a single-field format check via CEL is rejected by
        // the apiserver as too expensive once multiplied across the upstreams list.
        // +optional
        // +kubebuilder:validation:MaxLength=128
        // +kubebuilder:validation:Pattern=`^([a-zA-Z_][a-zA-Z0-9_]*)?$`
        SubjectClaim string `json:"subjectClaim,omitempty"`

        // CABundleRef references a ConfigMap containing a CA bundle added to the
        // system roots when connecting to this upstream; it does not restrict trust
        // to this bundle or disable public-root trust. The selected key is projected
        // as ca.crt.
        // +optional
        CABundleRef *CABundleSource `json:"caBundleRef,omitempty"`

        // AllowPrivateIPs permits the upstream provider's HTTP client to connect to
        // private IP ranges (RFC-1918, link-local). Use only when the upstream is
        // hosted inside the same cluster and has no public endpoint.
        // +optional
        AllowPrivateIPs bool `json:"allowPrivateIPs,omitempty"`
}

// OAuth2UpstreamConfig contains configuration for pure OAuth 2.0 providers.
// OAuth 2.0 providers require explicit endpoint configuration.
//
// Exactly one of ClientID or DCRConfig must be set: ClientID is used when the
// client is pre-provisioned out of band, DCRConfig enables RFC 7591 Dynamic
// Client Registration at runtime.
//
// ClientSecretRef is mutually exclusive with DCRConfig: when DCRConfig is set,
// the client_secret is obtained from the registration response (RFC 7591
// §3.2.1) and any static ClientSecretRef would be either dead config or a
// competing source of truth. The XValidation rule below rejects the
// combination at admission; ValidateOAuth2DCRConfig is the matching
// reconcile-time backstop.
//
// Layered XOR behavior: the ClientID/DCRConfig rule treats `clientId: ""` as
// absent (size>0) but treats `dcrConfig: {}` as present (has() is true for an
// empty object). For input `{ clientId: "", dcrConfig: {} }` the outer rule
// passes and the inner DCRUpstreamConfig XOR fires with "exactly one of
// discoveryUrl or registrationEndpoint must be set". This is intentional —
// adding a non-empty subfield check (e.g.,
// `has(self.dcrConfig.discoveryUrl) || has(self.dcrConfig.registrationEndpoint)`)
// would inflate CEL cost on an already-budget-bound rule, and the inner
// message is still actionable.
//
// +kubebuilder:validation:XValidation:rule="(has(self.clientId) && size(self.clientId) > 0) ? !has(self.dcrConfig) : has(self.dcrConfig)",message="exactly one of clientId or dcrConfig must be set"
// +kubebuilder:validation:XValidation:rule="!has(self.dcrConfig) || (has(self.dcrConfig.discoveryUrl) != has(self.dcrConfig.registrationEndpoint))",message="exactly one of discoveryUrl or registrationEndpoint must be set when dcrConfig is set"
// +kubebuilder:validation:XValidation:rule="!(has(self.dcrConfig) && has(self.clientSecretRef))",message="clientSecretRef must not be set when dcrConfig is set; the client_secret is obtained at runtime via Dynamic Client Registration"
//
//nolint:lll // CEL validation rules exceed line length limit
type OAuth2UpstreamConfig struct {
        // AuthorizationEndpoint is the URL for the OAuth authorization endpoint.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https?://.*$`
        AuthorizationEndpoint string `json:"authorizationEndpoint"`

        // TokenEndpoint is the URL for the OAuth token endpoint.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https?://.*$`
        TokenEndpoint string `json:"tokenEndpoint"`

        // UserInfo contains configuration for fetching user information from the upstream provider.
        // When omitted and IdentityFromToken is also unset, the embedded auth server runs in
        // synthesis mode for this upstream: a non-PII subject derived from the access token, no
        // Name/Email. Use this shape for upstreams with no userinfo surface and no identity in
        // the token response (e.g., MCP authorization servers per the MCP spec). When
        // IdentityFromToken is set instead, identity is resolved from the token response body
        // (e.g., Snowflake's "username" field, Slack's "authed_user.id"); the userinfo HTTP call
        // is skipped entirely.
        // +optional
        UserInfo *UserInfoConfig `json:"userInfo,omitempty"`

        // ClientID is the OAuth 2.0 client identifier registered with the upstream IDP.
        // Mutually exclusive with DCRConfig: when DCRConfig is set, ClientID is obtained
        // at runtime via RFC 7591 Dynamic Client Registration and must be left empty.
        // +optional
        ClientID string `json:"clientId,omitempty"`

        // ClientSecretRef references a Kubernetes Secret containing the OAuth 2.0 client secret.
        // Optional for public clients using PKCE instead of client secret.
        // +optional
        ClientSecretRef *SecretKeyRef `json:"clientSecretRef,omitempty"`

        // TokenEndpointAuthMethod selects how the client authenticates at the OAuth token
        // endpoint. When empty, credentials are sent in the request body (the historical
        // client_secret_post-shaped default). Set this to client_secret_basic explicitly
        // for providers that require HTTP Basic auth. Public clients without a secret use
        // the "none" method.
        // +kubebuilder:validation:Enum=none;client_secret_basic;client_secret_post
        // +optional
        TokenEndpointAuthMethod string `json:"tokenEndpointAuthMethod,omitempty"`

        // RedirectURI is the callback URL where the upstream IdP will redirect after authentication.
        // When not specified, defaults to `{resourceUrl}/oauth/callback` where `resourceUrl` is the
        // URL associated with the resource (e.g., MCPServer or vMCP) using this config.
        // Its hostname must match the browser-facing authorization endpoint (`issuer`, or
        // `authorizationEndpointBaseUrl` when set): the callback is bound to the browser that
        // started the login by a host-only cookie, so a different host rejects every browser login.
        // +optional
        RedirectURI string `json:"redirectUri,omitempty"`

        // Scopes are the OAuth scopes to request from the upstream IdP.
        // +listType=atomic
        // +optional
        Scopes []string `json:"scopes,omitempty"`

        // TokenResponseMapping configures custom field extraction from non-standard token responses.
        // Some OAuth providers (e.g., GovSlack) nest token fields under non-standard paths
        // instead of returning them at the top level. When set, ToolHive performs the token
        // exchange HTTP call directly and extracts fields using the configured dot-notation paths.
        // If nil, standard OAuth 2.0 token response parsing is used.
        // For extracting user identity from the token response, see IdentityFromToken.
        // +optional
        TokenResponseMapping *TokenResponseMapping `json:"tokenResponseMapping,omitempty"`

        // IdentityFromToken extracts user identity (subject, name, email) directly
        // from the OAuth2 token-endpoint response body using gjson dot-notation paths.
        // When set, the embedded auth server skips the userinfo HTTP call entirely
        // and resolves identity from the token response. See IdentityFromTokenConfig
        // for trust-model and uniqueness considerations.
        // +optional
        IdentityFromToken *IdentityFromTokenConfig `json:"identityFromToken,omitempty"`

        // AdditionalAuthorizationParams are extra query parameters to include in
        // authorization requests sent to the upstream provider.
        // This is useful for providers that require custom parameters, such as
        // Google's access_type=offline for obtaining refresh tokens.
        // Framework-managed parameters (response_type, client_id, redirect_uri,
        // scope, state, code_challenge, code_challenge_method, nonce) are not allowed.
        // +kubebuilder:validation:MaxProperties=16
        // +optional
        AdditionalAuthorizationParams map[string]string `json:"additionalAuthorizationParams,omitempty"`

        // AdditionalTokenParams are extra form-body parameters to include in
        // token requests (authorization code exchange and refresh) sent to the
        // upstream provider's token endpoint.
        // This is useful for providers that enforce RFC 8707 resource indicators
        // on token requests, where the resource parameter must accompany the code
        // exchange and refresh, not only the authorization request.
        // Framework-managed parameters (grant_type, code, redirect_uri, client_id,
        // client_secret, code_verifier, refresh_token, scope) are not allowed.
        // +kubebuilder:validation:MaxProperties=16
        // +optional
        AdditionalTokenParams map[string]string `json:"additionalTokenParams,omitempty"`

        // OmitScopeOnRefresh stops the scope parameter from being sent on
        // refresh_token grant requests to the upstream provider's token endpoint.
        // Scopes are still sent on the authorization request.
        // By default (false), the configured scopes are sent on refresh as well,
        // which some providers require to avoid narrowing the refreshed token.
        // Set it for providers that reject scope on refresh, such as Salesforce
        // ("invalid_request - scope parameter not supported").
        // +optional
        OmitScopeOnRefresh bool `json:"omitScopeOnRefresh,omitempty"`

        // CABundleRef references a ConfigMap containing a CA bundle added to the
        // system roots when connecting to this upstream; it does not restrict trust
        // to this bundle or disable public-root trust. The selected key is projected
        // as ca.crt.
        // +optional
        CABundleRef *CABundleSource `json:"caBundleRef,omitempty"`

        // InsecureAllowHTTP permits plain-HTTP authorization and token endpoint URLs
        // for this upstream. Only for in-cluster development environments (e.g. an
        // OAuth2 provider served over HTTP in a kind cluster) where TLS is not
        // available. Never set this in production.
        // +optional
        InsecureAllowHTTP bool `json:"insecureAllowHTTP,omitempty"`

        // AllowPrivateIPs permits the upstream provider's HTTP client to connect to
        // private IP ranges (RFC-1918, link-local). Use only when the upstream is
        // hosted inside the same cluster and has no public endpoint. HTTP-scheme
        // restrictions are unchanged — HTTPS is still required for non-localhost
        // hosts unless InsecureAllowHTTP is set. Defaults to false.
        // +optional
        AllowPrivateIPs bool `json:"allowPrivateIPs,omitempty"`

        // DCRConfig enables RFC 7591 Dynamic Client Registration against the upstream
        // authorization server. When set, the client credentials are obtained at
        // runtime rather than being pre-provisioned, and ClientID must be left empty.
        // Mutually exclusive with ClientID.
        // +optional
        DCRConfig *DCRUpstreamConfig `json:"dcrConfig,omitempty"`
}

// DCRUpstreamConfig configures RFC 7591 Dynamic Client Registration for an
// OAuth2 or OIDC upstream. OAuth2 requires one of DiscoveryURL or
// RegistrationEndpoint; OIDC may omit both and derive discovery from IssuerURL.
//
// Exactly one of DiscoveryURL or RegistrationEndpoint must be set. DiscoveryURL
// points at an RFC 8414 / OIDC Discovery document from which the registration
// endpoint is resolved; RegistrationEndpoint is used directly when the upstream
// does not publish discovery metadata.
//
// The XOR rule uses has() alone (not has() + size() > 0) to keep the rule's
// estimated CEL cost under the apiserver's per-rule static budget. With
// `omitempty` on both fields, an unset field is absent on the wire and has()
// returns false; the explicit-empty-string edge case is rejected at reconcile
// time by ValidateOAuth2DCRConfig.
//
// +kubebuilder:validation:XValidation:rule="!(has(self.discoveryUrl) && has(self.registrationEndpoint))",message="at most one of discoveryUrl or registrationEndpoint may be set"
//
//nolint:lll // CEL validation rules exceed line length limit
type DCRUpstreamConfig struct {
        // DiscoveryURL is the RFC 8414 / OIDC Discovery document URL. The resolver
        // issues a single GET against this URL (no well-known-path fallback) and
        // reads registration_endpoint, authorization_endpoint, token_endpoint,
        // token_endpoint_auth_methods_supported, and scopes_supported from the
        // response.
        // Mutually exclusive with RegistrationEndpoint.
        // HTTPS is required because the registration endpoint resolved from this
        // document carries the initial access token and the issued client_secret
        // (RFC 7591 §3, RFC 8414 §3). MaxLength is a defensive size cap (etcd
        // object budget, regex evaluation cost) and matches the conventional URL
        // length cap.
        // +optional
        // +kubebuilder:validation:Pattern=`^https://[^\s?#]+[^/\s?#]$`
        // +kubebuilder:validation:MaxLength=2048
        DiscoveryURL string `json:"discoveryUrl,omitempty"`

        // RegistrationEndpoint is the RFC 7591 registration endpoint URL used
        // directly, bypassing discovery. When using this field, the caller is
        // expected to also supply AuthorizationEndpoint, TokenEndpoint, and an
        // explicit Scopes list on the parent OAuth2UpstreamConfig.
        // Mutually exclusive with DiscoveryURL.
        // HTTPS is required because the registration endpoint carries the initial
        // access token and the issued client_secret (RFC 7591 §3, RFC 8414 §3).
        // MaxLength is a defensive size cap (etcd object budget, regex evaluation
        // cost) and matches the conventional URL length cap.
        // +optional
        // +kubebuilder:validation:Pattern=`^https://[^\s?#]+[^/\s?#]$`
        // +kubebuilder:validation:MaxLength=2048
        RegistrationEndpoint string `json:"registrationEndpoint,omitempty"`

        // RotationKey is an optional operator-chosen label that is part of the
        // cached registration's identity but is never sent to the upstream.
        // Change it to make the proxies register a fresh client at their next
        // start instead of reusing the cached one: for example after the
        // upstream has dropped the registration, or when a rotation is required.
        // At start ToolHive checks the cached client at the upstream (see
        // RegistrationProbe) and registers a new one by itself when the upstream
        // rejects it. Where the configured probe cannot tell, the stale client_id
        // may only be rejected in the user's browser, and changing this key is
        // the remedy. The previous registration's row stays in storage, and
        // upstream refresh tokens bound to the previous client_id stop
        // refreshing, so affected users sign in again.
        // +optional
        // +kubebuilder:validation:Pattern=`^[A-Za-z0-9._-]+$`
        // +kubebuilder:validation:MaxLength=64
        RotationKey string `json:"rotationKey,omitempty"`

        // RegistrationProbe selects how ToolHive checks the cached upstream client
        // at start. tokenEndpoint (the default) redeems a random authorization
        // code with the cached client and reads invalid_client as "registration
        // gone"; it decides only for upstreams that authenticate the client
        // before validating the grant. authorizationEndpoint sends an
        // authorization request for the cached client without following
        // redirects and reads a direct 400, 401 or 403 that carries an RFC 6749
        // JSON error object (error invalid_client, unauthorized_client or
        // invalid_request) as "gone" (RFC 6749 §4.1.2.1), for upstreams whose
        // token endpoint answers invalid_grant for any client_id; an HTML page or
        // a bare gateway status is not a verdict. A false rejection re-registers
        // at every start and, when registration also fails, fails startup. none
        // skips the check; the registration is then retired only through
        // RotationKey.
        // +optional
        // +kubebuilder:validation:Enum=tokenEndpoint;authorizationEndpoint;none
        RegistrationProbe string `json:"registrationProbe,omitempty"`

        // InitialAccessTokenRef is an optional reference to a Kubernetes Secret
        // carrying an RFC 7591 §3 initial access token. When set, the resolver
        // presents the token value as a Bearer credential on the registration
        // request. Mirrors the ClientSecretRef pattern.
        // +optional
        InitialAccessTokenRef *SecretKeyRef `json:"initialAccessTokenRef,omitempty"`

        // SoftwareID is the RFC 7591 "software_id" registration metadata value,
        // identifying the client software independent of any particular
        // registration instance. Typically a UUID or short identifier.
        // +optional
        // +kubebuilder:validation:MaxLength=255
        SoftwareID string `json:"softwareId,omitempty"`

        // SoftwareStatement is the RFC 7591 "software_statement" JWT asserting
        // metadata about the client software, signed by a party the authorization
        // server trusts.
        //
        // Stored inline on the CR. The JWT is signed but not encrypted, so its
        // contents are visible to anyone with get/list/watch on this resource and
        // appear in etcd backups in plaintext. Treat the value as non-confidential
        // (signed attestation, not a secret). Operators that rotate software
        // statements like bearer credentials should keep them at the authorization
        // server side and rely on the registration endpoint's initial access
        // token (see InitialAccessTokenRef) instead of placing them on the CR.
        //
        // Bounded to 16384 characters as a defensive size cap (etcd object
        // budget, regex evaluation cost). Real-world signed statements with
        // embedded x5c certificate chains, JWKS keys, or OIDC-Federation
        // trust-framework metadata routinely exceed 4 KB.
        // +optional
        // +kubebuilder:validation:MaxLength=16384
        SoftwareStatement string `json:"softwareStatement,omitempty"`
}

// IdentityFromTokenConfig extracts user identity (subject, name, email) directly from the
// OAuth2 token-endpoint response body using gjson dot-notation paths. When configured on an
// OAuth2UpstreamConfig, the embedded auth server skips the userinfo HTTP call entirely and
// resolves identity from the token response.
//
// Paths use gjson dot-notation, where each segment is a JSON object key. For example,
// "username" extracts a top-level field, and "authed_user.id" extracts a nested field.
//
// Trust-model warning: Identity claims extracted via this block are not cryptographically
// verified — they are trusted only via the TLS connection to the token endpoint. Prefer
// OIDC + ID token validation when verifiable claims are required.
//
// Subject uniqueness is scoped to the upstream provider entry. To keep identity namespaces
// separate across multiple instances of the same provider (e.g., two Snowflake accounts),
// use distinct upstream provider entries.
type IdentityFromTokenConfig struct {
        // SubjectPath is the dot-notation path to the subject (user ID) field in the token response.
        // Warning: claims read from the token response are trusted only via TLS, not
        // cryptographically verified; prefer OIDC ID tokens when verifiable claims are required.
        // Example: "authed_user.id" for Slack (top-level token-response field). For providers
        // whose token response embeds the access token as a JWT (e.g. Snowflake), use the
        // "@upstreamjwt" modifier to decode the payload, e.g. "access_token|@upstreamjwt|sub".
        // The "@upstreamjwt" modifier performs no signature verification either.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=256
        SubjectPath string `json:"subjectPath"`

        // NamePath is the dot-notation path to the display name field in the token response.
        // If not specified or if the path does not resolve to a string, the display name is omitted.
        // Omit the field entirely rather than setting it to an empty string.
        // +optional
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=256
        NamePath string `json:"namePath,omitempty"`

        // EmailPath is the dot-notation path to the email address field in the token response.
        // If not specified or if the path does not resolve to a string, the email is omitted.
        // Omit the field entirely rather than setting it to an empty string.
        // +optional
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=256
        EmailPath string `json:"emailPath,omitempty"`
}

// TokenResponseMapping maps non-standard token response fields to standard OAuth 2.0 fields
// using dot-notation JSON paths. This supports upstream providers like GovSlack that nest
// the access token under paths like "authed_user.access_token".
//
// For extracting user identity from the token response, see IdentityFromToken.
type TokenResponseMapping struct {
        // AccessTokenPath is the dot-notation path to the access token in the response.
        // Example: "authed_user.access_token"
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        AccessTokenPath string `json:"accessTokenPath"`

        // ScopePath is the dot-notation path to the scope string in the response.
        // If not specified, defaults to "scope".
        // +optional
        ScopePath string `json:"scopePath,omitempty"`

        // RefreshTokenPath is the dot-notation path to the refresh token in the response.
        // If not specified, defaults to "refresh_token".
        // +optional
        RefreshTokenPath string `json:"refreshTokenPath,omitempty"`

        // ExpiresInPath is the dot-notation path to the expires_in value (in seconds).
        // If not specified, defaults to "expires_in".
        // +optional
        ExpiresInPath string `json:"expiresInPath,omitempty"`
}

// UserInfoConfig contains configuration for fetching user information from an upstream provider.
// This supports both standard OIDC UserInfo endpoints and custom provider-specific endpoints
// like GitHub's /user API. For providers that do not expose a usable userinfo endpoint but
// include identity in the OAuth2 token response, use IdentityFromToken on OAuth2UpstreamConfig
// instead.
type UserInfoConfig struct {
        // EndpointURL is the URL of the userinfo endpoint.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https?://.*$`
        EndpointURL string `json:"endpointUrl"`

        // HTTPMethod is the HTTP method to use for the userinfo request.
        // If not specified, defaults to GET.
        // +kubebuilder:validation:Enum=GET;POST
        // +optional
        HTTPMethod string `json:"httpMethod,omitempty"`

        // AdditionalHeaders contains extra headers to include in the userinfo request.
        // Useful for providers that require specific headers (e.g., GitHub's Accept header).
        // +optional
        AdditionalHeaders map[string]string `json:"additionalHeaders,omitempty"`

        // FieldMapping contains custom field mapping configuration for non-standard providers.
        // If nil, standard OIDC field names are used ("sub", "name", "email").
        // +optional
        FieldMapping *UserInfoFieldMapping `json:"fieldMapping,omitempty"`
}

// UserInfoFieldMapping maps provider-specific field names to standard UserInfo fields.
// This allows adapting non-standard provider responses to the canonical UserInfo structure.
// Each field supports an ordered list of claim names to try. The first non-empty value
// found will be used.
//
// Example for GitHub:
//
//        fieldMapping:
//          subjectFields: ["id", "login"]
//          nameFields: ["name", "login"]
//          emailFields: ["email"]
type UserInfoFieldMapping struct {
        // SubjectFields is an ordered list of field names to try for the user ID.
        // The first non-empty value found will be used.
        // Default: ["sub"]
        // +listType=atomic
        // +optional
        SubjectFields []string `json:"subjectFields,omitempty"`

        // NameFields is an ordered list of field names to try for the display name.
        // The first non-empty value found will be used.
        // Default: ["name"]
        // +listType=atomic
        // +optional
        NameFields []string `json:"nameFields,omitempty"`

        // EmailFields is an ordered list of field names to try for the email address.
        // The first non-empty value found will be used.
        // Default: ["email"]
        // +listType=atomic
        // +optional
        EmailFields []string `json:"emailFields,omitempty"`
}

// Auth server storage types
const (
        // AuthServerStorageTypeMemory is the in-memory storage backend (default)
        AuthServerStorageTypeMemory AuthServerStorageType = "memory"

        // AuthServerStorageTypeRedis is the Redis storage backend
        AuthServerStorageTypeRedis AuthServerStorageType = "redis"
)

// AuthServerStorageType represents the type of storage backend for the embedded auth server
type AuthServerStorageType string

// AuthServerStorageConfig configures the storage backend for the embedded auth server.
type AuthServerStorageConfig struct {
        // Type specifies the storage backend type.
        // Valid values: "memory" (default), "redis".
        // +kubebuilder:validation:Enum=memory;redis
        // +kubebuilder:default=memory
        Type AuthServerStorageType `json:"type,omitempty"`

        // Redis configures the Redis storage backend.
        // Required when type is "redis".
        // +optional
        Redis *RedisStorageConfig `json:"redis,omitempty"`

        // MaxDCRClients caps the number of live clients created by Dynamic Client
        // Registration (/oauth/register). At the cap, registration fails with HTTP
        // 503 until a client expires; no existing client is evicted. If omitted,
        // defaults to 10000. Operator-configured and CIMD clients are not counted.
        // On Redis the count is shared by every replica; with memory storage each
        // replica keeps its own count.
        //
        // A client never used for a token exchange lives 24 hours, so a caller
        // registering at the full registrationRateLimit keeps up to
        // rate x replicas x 86400 unused clients alive; at the default rate one
        // replica fills the default cap in about 2.8 hours.
        //
        // 0 means unlimited. The endpoint is unauthenticated, so 0 leaves stored
        // clients bounded only by the per-replica registrationRateLimit, about
        // 86400 per replica per day at the default rate. Use 0 only temporarily,
        // for example while an existing client population above the cap drains;
        // a large finite value, such as two to three times the observed live
        // client count, keeps some protection.
        // +kubebuilder:validation:Minimum=0
        // +optional
        MaxDCRClients *int32 `json:"maxDCRClients,omitempty"`
}

// RedisStorageConfig configures Redis connection for auth server storage.
// Exactly one of addr or sentinelConfig must be set. Set clusterMode to true when
// addr points to a Redis Cluster discovery endpoint (GCP Memorystore Cluster,
// AWS ElastiCache cluster mode enabled).
//
// +kubebuilder:validation:XValidation:rule="(has(self.addr) && self.addr.size() > 0) != has(self.sentinelConfig)",message="exactly one of addr or sentinelConfig must be set"
// +kubebuilder:validation:XValidation:rule="!(has(self.clusterMode) && self.clusterMode) || (has(self.addr) && self.addr.size() > 0)",message="clusterMode requires addr to be set"
//
//nolint:lll // CEL validation rules exceed line length limit
type RedisStorageConfig struct {
        // Addr is the Redis server address (host:port). Required for standalone and cluster modes.
        // Use for managed Redis services that expose a single endpoint (GCP Memorystore basic tier,
        // AWS ElastiCache without cluster mode, or cluster-mode services when clusterMode is true).
        // Mutually exclusive with sentinelConfig.
        // +optional
        Addr string `json:"addr,omitempty"`

        // ClusterMode enables the Redis Cluster protocol. Set to true when addr points to a
        // Redis Cluster discovery endpoint (e.g., GCP Memorystore Cluster, AWS ElastiCache
        // cluster mode enabled). Requires addr to be set.
        // +optional
        ClusterMode bool `json:"clusterMode,omitempty"`

        // SentinelConfig holds Redis Sentinel configuration.
        // Use for self-managed Redis with Sentinel-based HA. Mutually exclusive with addr.
        // +optional
        SentinelConfig *RedisSentinelConfig `json:"sentinelConfig,omitempty"`

        // ACLUserConfig configures Redis ACL user authentication.
        // +kubebuilder:validation:Required
        ACLUserConfig *RedisACLUserConfig `json:"aclUserConfig"`

        // DialTimeout is the timeout for establishing connections.
        // Format: Go duration string (e.g., "5s", "1m").
        // +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
        // +kubebuilder:default="5s"
        // +optional
        DialTimeout string `json:"dialTimeout,omitempty"`

        // ReadTimeout is the timeout for socket reads.
        // Format: Go duration string (e.g., "3s", "1m").
        // +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
        // +kubebuilder:default="3s"
        // +optional
        ReadTimeout string `json:"readTimeout,omitempty"`

        // WriteTimeout is the timeout for socket writes.
        // Format: Go duration string (e.g., "3s", "1m").
        // +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
        // +kubebuilder:default="3s"
        // +optional
        WriteTimeout string `json:"writeTimeout,omitempty"`

        // TLS configures TLS for connections to the Redis/Valkey master or cluster nodes.
        // Presence of this field enables TLS. Omit to use plaintext.
        // +optional
        TLS *RedisTLSConfig `json:"tls,omitempty"`

        // SentinelTLS configures TLS for connections to Sentinel instances.
        // Only applies when sentinelConfig is set. Presence of this field enables TLS.
        // +optional
        SentinelTLS *RedisTLSConfig `json:"sentinelTls,omitempty"`
}

// RedisSentinelConfig configures Redis Sentinel connection.
type RedisSentinelConfig struct {
        // MasterName is the name of the Redis master monitored by Sentinel.
        // +kubebuilder:validation:Required
        MasterName string `json:"masterName"`

        // SentinelAddrs is a list of Sentinel host:port addresses.
        // Mutually exclusive with SentinelService.
        // +listType=atomic
        // +optional
        SentinelAddrs []string `json:"sentinelAddrs,omitempty"`

        // SentinelService enables automatic discovery from a Kubernetes Service.
        // Mutually exclusive with SentinelAddrs.
        // +optional
        SentinelService *SentinelServiceRef `json:"sentinelService,omitempty"`

        // DB is the Redis database number.
        // +kubebuilder:default=0
        // +optional
        DB int32 `json:"db,omitempty"`
}

// SentinelServiceRef references a Kubernetes Service for Sentinel discovery.
type SentinelServiceRef struct {
        // Name of the Sentinel Service.
        // +kubebuilder:validation:Required
        Name string `json:"name"`

        // Namespace of the Sentinel Service (defaults to same namespace).
        // +optional
        Namespace string `json:"namespace,omitempty"`

        // Port of the Sentinel service.
        // +kubebuilder:default=26379
        // +optional
        Port int32 `json:"port,omitempty"`
}

// RedisTLSConfig configures TLS for Redis connections.
// Presence of this struct on a connection type enables TLS for that connection.
type RedisTLSConfig struct {
        // InsecureSkipVerify disables server certificate and hostname verification.
        // The connection is encrypted but not authenticated, so anyone able to
        // impersonate the server can read the credentials and data sent over it.
        // Intended for testing only; for self-signed or private CAs set
        // caCertSecretRef instead.
        // +optional
        InsecureSkipVerify bool `json:"insecureSkipVerify,omitempty"`

        // CACertSecretRef references a Secret containing a PEM-encoded CA certificate
        // for verifying the server. When not specified, system root CAs are used.
        // If the Secret or key does not exist, the pod cannot start (it stays in
        // ContainerCreating) rather than connecting without verification.
        // +optional
        CACertSecretRef *SecretKeyRef `json:"caCertSecretRef,omitempty"`
}

// RedisACLUserConfig configures Redis ACL user authentication.
type RedisACLUserConfig struct {
        // UsernameSecretRef references a Secret containing the Redis ACL username.
        // When omitted, connections use legacy password-only AUTH. Omit for managed
        // Redis tiers that do not support ACL users (e.g. GCP Memorystore Basic/Standard
        // HA, Azure Cache for Redis). Set for services that support ACL users (e.g. AWS
        // ElastiCache non-cluster with Redis 6+ RBAC).
        // +optional
        UsernameSecretRef *SecretKeyRef `json:"usernameSecretRef,omitempty"`

        // PasswordSecretRef references a Secret containing the Redis ACL password.
        // +kubebuilder:validation:Required
        PasswordSecretRef *SecretKeyRef `json:"passwordSecretRef"`
}

// SecretKeyRef is a reference to a key within a Secret
type SecretKeyRef struct {
        // Name is the name of the secret
        // +kubebuilder:validation:Required
        Name string `json:"name"`

        // Key is the key within the secret
        // +kubebuilder:validation:Required
        Key string `json:"key"`
}

// AWSStsConfig holds configuration for AWS STS authentication with SigV4 request signing.
// This configuration exchanges incoming authentication tokens (typically OIDC JWT) for AWS STS
// temporary credentials, then signs requests to AWS services using SigV4.
type AWSStsConfig struct {
        // Region is the AWS region for the STS endpoint and service (e.g., "us-east-1", "eu-west-1")
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:Pattern=`^[a-z]{2}(-[a-z]+)+-\d+$`
        Region string `json:"region"`

        // Service is the AWS service name for SigV4 signing
        // Defaults to "aws-mcp" for AWS MCP Server endpoints
        // +kubebuilder:default="aws-mcp"
        // +optional
        Service string `json:"service,omitempty"`

        // FallbackRoleArn is the IAM role ARN to assume when no role mappings match
        // Used as the default role when RoleMappings is empty or no mapping matches
        // At least one of FallbackRoleArn or RoleMappings must be configured (enforced by webhook)
        // +kubebuilder:validation:Pattern=`^arn:(aws|aws-cn|aws-us-gov):iam::\d{12}:role/[\w+=,.@\-_/]+$`
        // +optional
        FallbackRoleArn string `json:"fallbackRoleArn,omitempty"`

        // RoleMappings defines claim-based role selection rules
        // Allows mapping JWT claims (e.g., groups, roles) to specific IAM roles
        // Lower priority values are evaluated first (higher priority)
        // +listType=atomic
        // +optional
        RoleMappings []RoleMapping `json:"roleMappings,omitempty"`

        // RoleClaim is the JWT claim to use for role mapping evaluation
        // Defaults to "groups" to match common OIDC group claims
        // +kubebuilder:default="groups"
        // +optional
        RoleClaim string `json:"roleClaim,omitempty"`

        // SessionDuration is the duration in seconds for the STS session
        // Must be between 900 (15 minutes) and 43200 (12 hours)
        // Defaults to 3600 (1 hour) if not specified
        // +kubebuilder:validation:Minimum=900
        // +kubebuilder:validation:Maximum=43200
        // +kubebuilder:default=3600
        // +optional
        SessionDuration *int32 `json:"sessionDuration,omitempty"`

        // SessionNameClaim is the JWT claim to use for role session name
        // Defaults to "sub" to use the subject claim
        // +kubebuilder:default="sub"
        // +optional
        SessionNameClaim string `json:"sessionNameClaim,omitempty"`

        // SubjectProviderName is the name of the upstream provider whose access token
        // is used as the web identity token for STS AssumeRoleWithWebIdentity.
        // This field is used exclusively by VirtualMCPServer, where there is no
        // upstream swap middleware to replace the bearer token before the strategy runs.
        // When left empty and an embedded authorization server is configured on the
        // VirtualMCPServer, the controller automatically populates this field with
        // the first configured upstream provider name. Set it explicitly to override
        // that default or to select a specific provider when multiple upstreams are
        // configured.
        // When no embedded auth server is present, the bearer token from the incoming
        // request's Authorization header is used instead.
        // +optional
        SubjectProviderName string `json:"subjectProviderName,omitempty"`
}

// RoleMapping defines a rule for mapping JWT claims to IAM roles.
// Mappings are evaluated in priority order (lower number = higher priority), and the first
// matching rule determines which IAM role to assume.
// Exactly one of Claim or Matcher must be specified.
type RoleMapping struct {
        // Claim is a simple claim value to match against
        // The claim type is specified by AWSStsConfig.RoleClaim
        // For example, if RoleClaim is "groups", this would be a group name
        // Internally compiled to a CEL expression: "<claim_value>" in claims["<role_claim>"]
        // Mutually exclusive with Matcher
        // +kubebuilder:validation:MinLength=1
        // +optional
        Claim string `json:"claim,omitempty"`

        // Matcher is a CEL expression for complex matching against JWT claims
        // The expression has access to a "claims" variable containing all JWT claims as map[string]any
        // Examples:
        //   - "admins" in claims["groups"]
        //   - claims["sub"] == "user123" && !("act" in claims)
        // Mutually exclusive with Claim
        // +kubebuilder:validation:MinLength=1
        // +optional
        Matcher string `json:"matcher,omitempty"`

        // RoleArn is the IAM role ARN to assume when this mapping matches
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^arn:(aws|aws-cn|aws-us-gov):iam::\d{12}:role/[\w+=,.@\-_/]+$`
        RoleArn string `json:"roleArn"`

        // Priority determines evaluation order (lower values = higher priority)
        // Allows fine-grained control over role selection precedence
        // When omitted, this mapping has the lowest possible priority and
        // configuration order acts as tie-breaker via stable sort
        // +kubebuilder:validation:Minimum=0
        // +optional
        Priority *int32 `json:"priority,omitempty"`
}

// UpstreamInjectSpec holds configuration for upstream token injection.
// This strategy injects an upstream IdP access token obtained by the embedded
// authorization server into backend requests as the Authorization: Bearer header.
type UpstreamInjectSpec struct {
        // ProviderName is the name of the upstream IdP provider whose access token
        // should be injected as the Authorization: Bearer header.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        ProviderName string `json:"providerName"`
}

// Condition types specific to MCPExternalAuthConfig and the inline embedded
// auth server config it shares with VirtualMCPServer.
const (
        // ConditionTypeIdentitySynthesized is an advisory set to True when at
        // least one OAuth2 upstream has no real-identity source configured
        // (neither userInfo nor identityFromToken). The embedded auth server
        // then synthesizes its subject from the access token, with no
        // Name/Email claims. Surfaces on resources that own the upstream
        // declaration so a missing identity source is visible in
        // `kubectl describe` instead of only in proxyrunner logs.
        ConditionTypeIdentitySynthesized = "IdentitySynthesized"

        // ConditionTypeDeprecatedInboundGrantConfiguration reports whether released
        // legacy inbound grant fields remain populated.
        ConditionTypeDeprecatedInboundGrantConfiguration = "DeprecatedInboundGrantConfiguration"
)

// Condition reasons for the deprecated inbound grant advisory.
const (
        ConditionReasonLegacyInboundGrantFields           = "LegacyInboundGrantFields"
        ConditionReasonCanonicalInboundGrantConfiguration = "CanonicalInboundGrantConfiguration"
)

// Condition reasons for ConditionTypeIdentitySynthesized.
const (
        // ConditionReasonIdentitySynthesizedActive: one or more OAuth2 upstreams
        // have neither userInfo nor identityFromToken configured. The condition
        // message names the affected upstream(s).
        ConditionReasonIdentitySynthesizedActive = "OAuth2UpstreamWithoutUserInfo"

        // ConditionReasonIdentitySynthesizedInactive: every OAuth2 upstream has
        // a real-identity source (userInfo or identityFromToken); real identity
        // is being resolved.
        ConditionReasonIdentitySynthesizedInactive = "AllUpstreamsHaveUserInfo"
)

// Condition reasons for ConditionTypeValid on OBO-typed configs. These
// literals are part of the user-facing contract — external consumers and
// downstream tooling pattern-match on them.
const (
        // ConditionReasonEnterpriseRequired: an obo-typed MCPExternalAuthConfig
        // requires an enterprise build that has registered an OBO handler via
        // controllerutil.RegisterOBOHandler. Upstream-only builds surface this
        // reason for every obo-typed config.
        ConditionReasonEnterpriseRequired = "EnterpriseRequired"

        // ConditionReasonInvalidConfig: an obo-typed MCPExternalAuthConfig is
        // well-formed at the CRD level but fails the registered OBO handler's
        // Validate() with an error other than the enterprise-required sentinel.
        // Used by out-of-tree handlers; unreachable in upstream-only builds.
        ConditionReasonInvalidConfig = "InvalidConfig"

        // ConditionReasonInvalidCABundle: a referenced CA bundle ConfigMap is
        // missing its key or holds content that is not a PEM certificate.
        //
        // Distinct from ConditionReasonInvalidConfig because the failing input is
        // ConfigMap *content*, which is covered by neither metadata.generation nor
        // the referenced config's spec hash. The guards that hold a terminal
        // ConditionReasonInvalidConfig steady across reconciles key off those two
        // values, so reusing that reason here would pin the failure in place even
        // after the ConfigMap is repaired.
        ConditionReasonInvalidCABundle = "InvalidCABundle"
)

// XAASpec holds configuration for the XAA (Cross-Application Access) auth strategy.
// XAA implements draft-ietf-oauth-identity-assertion-authz-grant (ID-JAG) — a
// two-step token exchange to obtain access tokens for target services:
//   - IdP exchange (RFC 8693): Exchange the user's ID token at their IdP for an ID-JAG JWT
//   - Target grant (RFC 7523): Exchange the ID-JAG at the target app's AS for an access token
type XAASpec struct {
        // IDPTokenURL is the IdP token endpoint for IdP exchange (RFC 8693).
        // Must be a valid HTTPS URL.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https://.*$`
        IDPTokenURL string `json:"idpTokenUrl"`

        // IDPClientID is the OAuth client ID at the IdP for IdP exchange.
        // +optional
        IDPClientID string `json:"idpClientId,omitempty"`

        // IDPClientSecretRef references a Kubernetes Secret containing the IdP client secret.
        // +optional
        IDPClientSecretRef *SecretKeyRef `json:"idpClientSecretRef,omitempty"`

        // TargetTokenURL is the target AS token endpoint for target grant (RFC 7523).
        // +kubebuilder:validation:Required
        TargetTokenURL string `json:"targetTokenUrl"`

        // InsecureTargetTokenURL allows plain HTTP for TargetTokenURL.
        // WARNING: this is insecure and must only be set for in-cluster or
        // development/testing endpoints — never in production.
        // +optional
        InsecureTargetTokenURL bool `json:"insecureTargetTokenUrl,omitempty"`

        // TargetClientID is the OAuth client ID at the target AS for target grant.
        // ID-JAG draft §9.1 RECOMMENDS confidential clients for target grant; most
        // conformant target authorization servers will reject an unauthenticated
        // JWT-bearer grant per the §4.4.1 client_id continuity requirement.
        // +optional
        TargetClientID string `json:"targetClientId,omitempty"`

        // TargetClientSecretRef references a Kubernetes Secret for the target AS client secret.
        // +optional
        TargetClientSecretRef *SecretKeyRef `json:"targetClientSecretRef,omitempty"`

        // TargetAudience is the resource AS URL for the ID-JAG audience claim.
        // +kubebuilder:validation:Required
        TargetAudience string `json:"targetAudience"`

        // TargetResource is the RFC 8707 resource indicator sent as the `resource`
        // parameter in IdP exchange (RFC 8693, draft §4.3, OPTIONAL). It
        // identifies the target resource server — not the access-token audience, which
        // is governed by TargetAudience. For MCP backends, set to the MCP server URL.
        // Some authorization servers (e.g. Okta's early ID-JAG implementation) require
        // this parameter in practice despite the draft marking it optional — set it
        // when your IdP needs it.
        // +optional
        TargetResource string `json:"targetResource,omitempty"`

        // Scopes are the requested scopes for the XAA exchange (IdP exchange and target grant).
        // +listType=atomic
        // +optional
        Scopes []string `json:"scopes,omitempty"`

        // SubjectProviderName selects which upstream provider's ID token to use.
        // When left empty and an embedded authorization server is configured,
        // the controller automatically populates this field with the first configured
        // upstream provider name.
        // +optional
        SubjectProviderName string `json:"subjectProviderName,omitempty"`

        // SubjectTokenType is the token-type URN of the upstream subject token
        // used in IdP exchange. Defaults to "urn:ietf:params:oauth:token-type:id_token"
        // when empty.
        // +kubebuilder:validation:Enum="urn:ietf:params:oauth:token-type:id_token"
        // +optional
        SubjectTokenType string `json:"subjectTokenType,omitempty"`
}

// MCPExternalAuthConfigStatus defines the observed state of MCPExternalAuthConfig
type MCPExternalAuthConfigStatus struct {
        // Conditions represent the latest available observations of the MCPExternalAuthConfig's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration is the most recent generation observed for this MCPExternalAuthConfig.
        // It corresponds to the MCPExternalAuthConfig's generation, which is updated on mutation by the API Server.
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // ConfigHash is a hash of the current configuration for change detection
        // +optional
        ConfigHash string `json:"configHash,omitempty"`
}

// +kubebuilder:object:root=true
// +kubebuilder:storageversion
// +kubebuilder:subresource:status
// +kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
// +kubebuilder:resource:shortName=extauth;mcpextauth,categories=toolhive
// +kubebuilder:printcolumn:name="Type",type=string,JSONPath=`.spec.type`
// +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status`
// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp`

// MCPExternalAuthConfig is the Schema for the mcpexternalauthconfigs API.
// MCPExternalAuthConfig resources are namespace-scoped and can only be referenced by
// MCPServer resources within the same namespace. Cross-namespace references
// are not supported for security and isolation reasons.
type MCPExternalAuthConfig struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPExternalAuthConfigSpec   `json:"spec,omitempty"`
        Status MCPExternalAuthConfigStatus `json:"status,omitempty"`
}

// +kubebuilder:object:root=true

// MCPExternalAuthConfigList contains a list of MCPExternalAuthConfig
type MCPExternalAuthConfigList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPExternalAuthConfig `json:"items"`
}

// Validate performs validation on the MCPExternalAuthConfig spec.
// This method is called by the controller during reconciliation.
//
// Note: These validations provide defense-in-depth alongside the
// +kubebuilder:validation:XValidation markers on MCPExternalAuthConfigSpec.
// CEL catches issues at API admission time, but this method also validates stored objects
// to catch any that bypassed CEL or were stored before CEL rules were added.
func (r *MCPExternalAuthConfig) Validate() error {
        // First, validate type/config consistency (defense-in-depth with CEL)
        if err := r.validateTypeConfigConsistency(); err != nil {
                return err
        }

        // Then perform type-specific complex validation
        switch r.Spec.Type {
        case ExternalAuthTypeEmbeddedAuthServer:
                return r.validateEmbeddedAuthServer()
        case ExternalAuthTypeAWSSts:
                return r.validateAWSSts()
        case ExternalAuthTypeUpstreamInject:
                if r.Spec.UpstreamInject == nil || r.Spec.UpstreamInject.ProviderName == "" {
                        return fmt.Errorf("upstreamInject requires a non-empty providerName")
                }
                return nil
        case ExternalAuthTypeXAA:
                if r.Spec.XAA == nil {
                        return fmt.Errorf("xaa requires configuration")
                }
                return nil
        case ExternalAuthTypeTokenExchange,
                ExternalAuthTypeHeaderInjection,
                ExternalAuthTypeBearerToken,
                ExternalAuthTypeUnauthenticated:
                // No complex validation needed for these types
                return nil
        case ExternalAuthTypeOBO:
                // Structural validation (the OBO field is set iff Type is "obo")
                // has already run via r.validateTypeConfigConsistency() at the top
                // of this method, so this arm is reached only when the structural
                // invariant holds — and the matching CEL rule on the spec catches
                // it at admission time. OBOConfig carries no required field and no
                // cross-field CEL rule: spec.obo shipped as an empty placeholder in an
                // earlier release, so a required field or a rule rejecting {} would be a
                // backward-incompatible narrowing. The kubebuilder markers check only
                // per-field shape (tenantId GUID/domain, authority URL,
                // subjectTokenProviderName, scopes/cacheSkew bounds) at admission, and
                // only for values that are present. All presence and combination
                // requirements (a tenant, a client-auth credential, at least one of
                // audience/scopes) and protocol-level checks are owned by the registered
                // handler, not the upstream type. That handler runs at reconcile
                // time via the controllerutil.OBOValidate function-pointer hook:
                // upstream-only builds return obo.ErrEnterpriseRequired, which the
                // reconciler maps to status.conditions[Valid] = False / Reason:
                // EnterpriseRequired. Out-of-tree builds that register a handler via
                // controllerutil.RegisterOBOHandler short-circuit the sentinel and
                // run their own protocol-level checks. Splitting the tiers this way
                // keeps the upstream CRD schema stable across builds.
                return nil
        default:
                // Unknown type - should be caught by enum validation, but handle defensively
                return fmt.Errorf("unsupported auth type: %s", r.Spec.Type)
        }
}

// typeConfigEntry maps an ExternalAuthType to its corresponding config field presence check.
type typeConfigEntry struct {
        authType  ExternalAuthType
        fieldName string
        isSet     bool
}

// validateTypeConfigConsistency validates that the correct config is set for the selected type.
// This mirrors the CEL validation rules but provides defense-in-depth for stored objects.
//
// Each ExternalAuthType with a config sub-field is one row in the typeConfigEntry
// table, checked by a uniform biconditional in the loop below ("config set iff
// type matches"). OBO is checked separately because its handler is registered
// out-of-tree; the unauthenticated guard then asserts no config is set for the
// unauthenticated type. Each row still maps one-to-one onto the corresponding
// CEL XValidation rule on MCPExternalAuthConfigSpec, so a reviewer can audit the
// structural-validation contract by comparing the table against the markers.
// See issue #5329 for the broader discussion.
func (r *MCPExternalAuthConfig) validateTypeConfigConsistency() error {
        entries := []typeConfigEntry{
                {ExternalAuthTypeTokenExchange, "tokenExchange", r.Spec.TokenExchange != nil},
                {ExternalAuthTypeHeaderInjection, "headerInjection", r.Spec.HeaderInjection != nil},
                {ExternalAuthTypeBearerToken, "bearerToken", r.Spec.BearerToken != nil},
                {ExternalAuthTypeEmbeddedAuthServer, "embeddedAuthServer", r.Spec.EmbeddedAuthServer != nil},
                {ExternalAuthTypeAWSSts, "awsSts", r.Spec.AWSSts != nil},
                {ExternalAuthTypeUpstreamInject, "upstreamInject", r.Spec.UpstreamInject != nil},
                {ExternalAuthTypeXAA, "xaa", r.Spec.XAA != nil},
        }
        if (r.Spec.OBO == nil) == (r.Spec.Type == ExternalAuthTypeOBO) {
                return fmt.Errorf("obo configuration must be set if and only if type is 'obo'")
        }

        for _, e := range entries {
                wantSet := r.Spec.Type == e.authType
                if e.isSet != wantSet {
                        return fmt.Errorf("%s configuration must be set if and only if type is '%s'", e.fieldName, e.authType)
                }
        }

        // Redundant with the per-type biconditionals above — each fires first for
        // Type=Unauthenticated with any non-nil field — but retained as a single
        // readable invariant so a contributor adding a new ExternalAuthType extends
        // the "no configuration must be set" check here too.
        if r.Spec.Type == ExternalAuthTypeUnauthenticated {
                if r.Spec.OBO != nil {
                        return fmt.Errorf("no configuration must be set when type is 'unauthenticated'")
                }
                for _, e := range entries {
                        if e.isSet {
                                return fmt.Errorf("no configuration must be set when type is 'unauthenticated'")
                        }
                }
        }

        return nil
}

// validateEmbeddedAuthServer validates embeddedAuthServer type configuration.
// This performs complex business logic validation that CEL cannot express.
func (r *MCPExternalAuthConfig) validateEmbeddedAuthServer() error {
        // Validate upstream providers
        cfg := r.Spec.EmbeddedAuthServer
        if cfg == nil {
                return nil
        }

        if len(cfg.UpstreamProviders) == 0 && len(cfg.DelegateClients) == 0 && cfg.InboundGrants == nil &&
                !hasJWTBearerTrustedIssuer(cfg.TrustedIssuers) {
                return fmt.Errorf("at least one upstream provider or inbound grant family is required")
        }
        // Note: multi-upstream is accepted at the CRD level. Consumer controllers
        // (MCPServer, MCPRemoteProxy) enforce single-upstream restrictions;
        // VirtualMCPServer allows multiple upstreams.

        // Defense-in-depth with the shared runtime policy. This folds legacy and
        // canonical grant declarations before checking confidential transport and
        // trusted issuer policy.
        if err := cfg.ValidateInboundGrants(); err != nil {
                return err
        }

        // The "requires allowConfidentialClientRegistration" half is also
        // enforced by the type-level XValidation rule above (defense-in-depth,
        // same reasoning as ValidateConfidentialClientTransport). The
        // https-non-loopback-per-entry check has no CEL equivalent here since it
        // needs the loopback-hostname helper, so it lives only in Go.
        //
        // The permission-shaped half of SPIFFE trust/client-auth config (grant
        // types, scopes, resources, audiences) has no equivalent Go-level
        // pre-check here, matching the DelegateClients precedent below — a
        // meaningful check needs AllowedAudiences/ScopesSupported, which are only
        // known once derived at reconcile time. controllerutil.
        // validateDelegateClientsAndTrustedIssuers revalidates the full SPIFFE
        // trust config (via RunConfig.Validate) once those derived values exist.
        //
        // The two checks below are admission-time-safe: they depend only on this
        // object's own spec, so they run here instead of waiting for reconcile.
        if err := validateSPIFFEBundleEndpoints(cfg.SPIFFETrustDomains); err != nil {
                return err
        }
        if err := validateSPIFFEPrincipalPatternOverlap(cfg.InboundGrants); err != nil {
                return err
        }
        if err := authserver.ValidateForceConfidentialRedirectURIs(
                cfg.ForceConfidentialRedirectURIs, cfg.AllowConfidentialClientRegistration,
        ); err != nil {
                return err
        }

        seen := make(map[string]bool, len(cfg.UpstreamProviders))
        for i, provider := range cfg.UpstreamProviders {
                if seen[provider.Name] {
                        return fmt.Errorf("upstreamProviders[%d]: duplicate name %q", i, provider.Name)
                }
                seen[provider.Name] = true

                if err := r.validateUpstreamProvider(i, &provider); err != nil {
                        return err
                }
        }

        for i := range cfg.TrustedIssuers {
                if err := validateUpstreamCABundleRef(cfg.TrustedIssuers[i].CABundleRef); err != nil {
                        return fmt.Errorf("trustedIssuers[%d] (%q) caBundleRef: %w", i, cfg.TrustedIssuers[i].IssuerURL, err)
                }
        }
        return nil
}

// validateSPIFFEBundleEndpoints rejects a structurally invalid SPIFFE Bundle
// Endpoint URL (wrong scheme, an IP-literal or loopback host, credentials,
// a query, or a fragment) at admission time. Not CEL-expressible: it needs
// real IP-literal and loopback-range detection, not a best-effort regex.
func validateSPIFFEBundleEndpoints(domains []SPIFFETrustDomainConfig) error {
        for i, domain := range domains {
                if domain.BundleSource.Type != SPIFFEBundleSourceTypeEndpoint || domain.BundleSource.Endpoint == nil {
                        continue
                }
                endpoint := domain.BundleSource.Endpoint
                if err := authserver.ValidateSPIFFEBundleEndpoint(authserver.SPIFFEBundleEndpointSourceRunConfig{
                        URL:     endpoint.URL,
                        Profile: authserver.SPIFFEBundleEndpointProfile(endpoint.Profile),
                }); err != nil {
                        return fmt.Errorf("spiffeTrustDomains[%d].bundleSource: %w", i, err)
                }
        }
        return nil
}

// validateSPIFFEPrincipalPatternOverlap validates each SPIFFE client-auth
// entry's own principalPattern and resources, then rejects two entries whose
// principalPattern values overlap (e.g. a "/agent/*" wildcard and a concrete
// "/agent/one"), all at admission time. Not CEL-expressible: real overlap
// semantics need the go-spiffe parser for normalization, the same
// runtime-parser dependency this epic has deliberately kept in Go elsewhere;
// resource shape is folded in here (rather than its own top-level call) to
// keep validateEmbeddedAuthServer's own branching within the complexity
// budget, since both are per-entry, admission-time-safe checks over the same
// slice.
func validateSPIFFEPrincipalPatternOverlap(inboundGrants *InboundGrantsConfig) error {
        if inboundGrants == nil {
                return nil
        }
        entries := inboundGrants.SPIFFEClientAuth
        // Validate each entry's own pattern and resources first, in its own
        // index, so a malformed value is always attributed to the entry that
        // owns it rather than to whichever earlier index the pairwise loop below
        // happened to be comparing it against.
        for i := range entries {
                if err := authserver.ValidateSPIFFEPrincipalPattern(entries[i].PrincipalPattern); err != nil {
                        return fmt.Errorf("inboundGrants.spiffeClientAuth[%d].principalPattern: %w", i, err)
                }
                if err := authserver.ValidateResourceIndicators(
                        entries[i].Resources, fmt.Sprintf("inboundGrants.spiffeClientAuth[%d].resources", i),
                ); err != nil {
                        return err
                }
        }
        for i := range entries {
                for j := i + 1; j < len(entries); j++ {
                        // Patterns were already validated above, so the error return
                        // here is unreachable in practice.
                        overlaps, err := authserver.SPIFFEPatternsOverlap(entries[i].PrincipalPattern, entries[j].PrincipalPattern)
                        if err != nil {
                                return fmt.Errorf("inboundGrants.spiffeClientAuth[%d].principalPattern: %w", i, err)
                        }
                        if overlaps {
                                return fmt.Errorf(
                                        "inboundGrants.spiffeClientAuth[%d].principalPattern %q overlaps spiffeClientAuth[%d].principalPattern %q",
                                        i, entries[i].PrincipalPattern, j, entries[j].PrincipalPattern,
                                )
                        }
                }
        }
        return nil
}

func hasJWTBearerTrustedIssuer(issuers []TrustedIssuerConfig) bool {
        return slices.ContainsFunc(issuers, func(issuer TrustedIssuerConfig) bool {
                return issuer.JWTBearerGrant != nil
        })
}

// buildTrustedIssuerConfigs converts CRD entries to the authoritative runtime
// type without sharing caller-owned slices.
func buildTrustedIssuerConfigs(issuers []TrustedIssuerConfig) []tokenexchange.TrustedIssuer {
        configs := make([]tokenexchange.TrustedIssuer, len(issuers))
        for i, issuer := range issuers {
                configs[i] = tokenexchange.TrustedIssuer{
                        Name:                   issuer.Name,
                        IssuerURL:              issuer.IssuerURL,
                        ExpectedAudience:       issuer.ExpectedAudience,
                        JWKSURL:                issuer.JWKSURL,
                        InsecureAllowHTTP:      issuer.InsecureAllowHTTP,
                        AllowPrivateIPs:        issuer.AllowPrivateIPs,
                        ActorClaim:             issuer.ActorClaim,
                        AllowedActors:          slices.Clone(issuer.AllowedActors),
                        ActorMatcher:           issuer.ActorMatcher,
                        AllowedDelegateClients: slices.Clone(issuer.AllowedDelegateClients),
                        AllowMayAct:            issuer.AllowMayAct,
                        JWTBearerGrant:         buildJWTBearerGrantPolicy(issuer.JWTBearerGrant),
                }
        }
        return configs
}

// buildJWTBearerGrantPolicy converts the CRD's JWT-bearer grant config to the
// runtime policy type without sharing caller-owned slices.
func buildJWTBearerGrantPolicy(grant *JWTBearerGrantConfig) *tokenexchange.JWTBearerGrantPolicy {
        if grant == nil {
                return nil
        }
        bindings := make([]tokenexchange.JWTBearerSubjectBinding, len(grant.SubjectBindings))
        for i, binding := range grant.SubjectBindings {
                bindings[i] = tokenexchange.JWTBearerSubjectBinding{
                        Subject:          binding.Subject,
                        AllowedResources: slices.Clone(binding.AllowedResources),
                }
        }
        return &tokenexchange.JWTBearerGrantPolicy{
                MaxAssertionAge:   grant.MaxAssertionAge.Duration.String(),
                SubjectBindings:   bindings,
                AcceptedAudiences: slices.Clone(grant.AcceptedAudiences),
        }
}

// validateUpstreamProvider validates a single upstream provider configuration.
// The discriminator check mirrors the combined CEL XValidation rule on
// UpstreamProviderConfig: a single boolean test produces a single message that
// matches what admission emits, so reconcile-time and admission-time errors
// stay aligned.
func (*MCPExternalAuthConfig) validateUpstreamProvider(index int, provider *UpstreamProviderConfig) error {
        prefix := fmt.Sprintf("upstreamProviders[%d]", index)

        typeOK := provider.Type == UpstreamProviderTypeOIDC || provider.Type == UpstreamProviderTypeOAuth2
        configOK := (provider.Type == UpstreamProviderTypeOIDC && provider.OIDCConfig != nil && provider.OAuth2Config == nil) ||
                (provider.Type == UpstreamProviderTypeOAuth2 && provider.OAuth2Config != nil && provider.OIDCConfig == nil)
        if !typeOK || !configOK {
                return fmt.Errorf("%s: type must be 'oidc' or 'oauth2'; oidcConfig must be set when type is 'oidc' "+
                        "and oauth2Config must be set when type is 'oauth2' (and the other must not be set)", prefix)
        }

        if err := validateUpstreamProviderCABundle(prefix, provider); err != nil {
                return err
        }

        if err := validateUpstreamProviderTypeConfig(prefix, provider); err != nil {
                return err
        }

        // Validate additionalAuthorizationParams does not contain reserved keys
        if err := ValidateAdditionalAuthorizationParams(prefix, provider.AdditionalAuthorizationParams()); err != nil {
                return err
        }

        // Validate additionalTokenParams does not contain reserved keys
        return ValidateAdditionalTokenParams(prefix, provider.AdditionalTokenParams())
}

// validateUpstreamProviderTypeConfig runs the checks that only apply to the
// provider's declared type. Split out of validateUpstreamProvider to keep that
// function under the cyclomatic limit, in the same shape as
// validateUpstreamProviderCABundle below.
func validateUpstreamProviderTypeConfig(prefix string, provider *UpstreamProviderConfig) error {
        // The discriminator checked by the caller guarantees the type-matched
        // config is non-nil here.
        switch provider.Type {
        case UpstreamProviderTypeOIDC:
                if err := ValidateOIDCDCRConfig(provider.OIDCConfig); err != nil {
                        return fmt.Errorf("%s: %w", prefix, err)
                }
        case UpstreamProviderTypeOAuth2:
                if err := ValidateOAuth2DCRConfig(provider.OAuth2Config); err != nil {
                        return fmt.Errorf("%s: %w", prefix, err)
                }
                if err := validateOAuth2UpstreamConfig(provider.OAuth2Config); err != nil {
                        return fmt.Errorf("%s: %w", prefix, err)
                }
        }
        return nil
}

func validateUpstreamProviderCABundle(prefix string, provider *UpstreamProviderConfig) error {
        field := "oidcConfig.caBundleRef"
        if provider.Type == UpstreamProviderTypeOAuth2 {
                field = "oauth2Config.caBundleRef"
        }
        if err := validateUpstreamCABundleRef(provider.CABundleRef()); err != nil {
                return fmt.Errorf("%s: %s: %w", prefix, field, err)
        }
        return nil
}

// Length caps for DCR-related string fields. Mirror the
// +kubebuilder:validation:MaxLength markers on DCRUpstreamConfig so that
// ValidateOAuth2DCRConfig is a true reconcile-time backstop for length
// constraints, not just for the XOR rules.
const (
        // MaxDCRURLLength matches the MaxLength marker on
        // DCRUpstreamConfig.DiscoveryURL and DCRUpstreamConfig.RegistrationEndpoint.
        MaxDCRURLLength = 2048

        // MaxSoftwareStatementLength matches the MaxLength marker on
        // DCRUpstreamConfig.SoftwareStatement.
        MaxSoftwareStatementLength = 16384
)

// validateUpstreamCABundleRef validates the source shape shared by OIDC and
// OAuth2 upstream CA bundles. ConfigMap existence and key contents are
// resolved by Kubernetes when the generated Pod is scheduled.
func validateUpstreamCABundleRef(ref *CABundleSource) error {
        if ref == nil {
                return nil
        }
        if ref.ConfigMapRef == nil {
                return fmt.Errorf("configMapRef must be specified")
        }
        if ref.ConfigMapRef.Name == "" {
                return fmt.Errorf("configMapRef.name must not be empty")
        }
        return nil
}

// ValidateOIDCDCRConfig enforces OIDC client/DCR exclusivity and DCR field
// limits. OIDC permits an empty endpoint selector because the issuer-derived
// well-known discovery URL is used at runtime.
func ValidateOIDCDCRConfig(cfg *OIDCUpstreamConfig) error {
        hasClientID := cfg.ClientID != ""
        hasDCR := cfg.DCRConfig != nil
        if hasClientID == hasDCR {
                return fmt.Errorf("oidcConfig: exactly one of clientId or dcrConfig must be set")
        }
        if !hasDCR {
                return nil
        }
        if cfg.ClientSecretRef != nil {
                return fmt.Errorf(
                        "oidcConfig: clientSecretRef must not be set when dcrConfig is set; " +
                                "the client_secret is obtained at runtime via Dynamic Client Registration")
        }
        if cfg.DCRConfig.DiscoveryURL != "" && cfg.DCRConfig.RegistrationEndpoint != "" {
                return fmt.Errorf("oidcConfig.dcrConfig: at most one of discoveryUrl or registrationEndpoint may be set")
        }
        if l := len(cfg.DCRConfig.DiscoveryURL); l > MaxDCRURLLength {
                return fmt.Errorf("oidcConfig.dcrConfig.discoveryUrl: length %d exceeds maximum %d", l, MaxDCRURLLength)
        }
        if l := len(cfg.DCRConfig.RegistrationEndpoint); l > MaxDCRURLLength {
                return fmt.Errorf("oidcConfig.dcrConfig.registrationEndpoint: length %d exceeds maximum %d", l, MaxDCRURLLength)
        }
        if l := len(cfg.DCRConfig.SoftwareStatement); l > MaxSoftwareStatementLength {
                return fmt.Errorf("oidcConfig.dcrConfig.softwareStatement: length %d exceeds maximum %d", l, MaxSoftwareStatementLength)
        }
        return nil
}

// ValidateOAuth2DCRConfig enforces the mutual exclusivity between ClientID and
// DCRConfig, between ClientSecretRef and DCRConfig, and (when DCRConfig is
// present) between DiscoveryURL and RegistrationEndpoint. It also enforces the
// MaxLength caps declared on DCRUpstreamConfig so reconcile-time matches
// admission-time. These rules mirror the CEL validation on OAuth2UpstreamConfig
// and DCRUpstreamConfig and provide defense-in-depth for stored objects (e.g.,
// objects stored before CEL rules were added or validated through code paths
// that bypass admission).
//
// Errors are scoped to "oauth2Config[.dcrConfig[.field]]" so callers can wrap
// with their own outer scope (e.g. "upstreamProviders[i]: %w" or
// "upstream %q: %w") without duplicating the field name.
//
// Exported so the controllerutil conversion layer can reuse the same
// invariants when building runtime configs, rejecting malformed objects at
// reconcile time rather than waiting until the authserver process parses them.
func ValidateOAuth2DCRConfig(cfg *OAuth2UpstreamConfig) error {
        hasClientID := cfg.ClientID != ""
        hasDCR := cfg.DCRConfig != nil

        if hasClientID == hasDCR {
                return fmt.Errorf("oauth2Config: exactly one of clientId or dcrConfig must be set")
        }

        if !hasDCR {
                return nil
        }

        if cfg.ClientSecretRef != nil {
                return fmt.Errorf(
                        "oauth2Config: clientSecretRef must not be set when dcrConfig is set; " +
                                "the client_secret is obtained at runtime via Dynamic Client Registration")
        }

        hasDiscoveryURL := cfg.DCRConfig.DiscoveryURL != ""
        hasRegistrationEndpoint := cfg.DCRConfig.RegistrationEndpoint != ""
        if hasDiscoveryURL == hasRegistrationEndpoint {
                return fmt.Errorf("oauth2Config.dcrConfig: exactly one of discoveryUrl or registrationEndpoint must be set")
        }

        if l := len(cfg.DCRConfig.DiscoveryURL); l > MaxDCRURLLength {
                return fmt.Errorf("oauth2Config.dcrConfig.discoveryUrl: length %d exceeds maximum %d", l, MaxDCRURLLength)
        }
        if l := len(cfg.DCRConfig.RegistrationEndpoint); l > MaxDCRURLLength {
                return fmt.Errorf("oauth2Config.dcrConfig.registrationEndpoint: length %d exceeds maximum %d", l, MaxDCRURLLength)
        }
        if l := len(cfg.DCRConfig.SoftwareStatement); l > MaxSoftwareStatementLength {
                return fmt.Errorf(
                        "oauth2Config.dcrConfig.softwareStatement: length %d exceeds maximum %d",
                        l, MaxSoftwareStatementLength)
        }
        return nil
}

// validateOAuth2UpstreamConfig validates OAuth2-specific upstream provider configuration
// beyond the DCR / ClientID constraints handled by ValidateOAuth2DCRConfig. Errors are
// scoped to "oauth2Config[.field]" so callers can wrap with their own outer scope (e.g.
// "upstreamProviders[i]: %w") without duplicating the field name, matching the contract
// of ValidateOAuth2DCRConfig.
func validateOAuth2UpstreamConfig(cfg *OAuth2UpstreamConfig) error {
        if cfg.IdentityFromToken != nil && cfg.IdentityFromToken.SubjectPath == "" {
                return fmt.Errorf("oauth2Config.identityFromToken.subjectPath must not be empty when identityFromToken is set")
        }
        return nil
}

// AdditionalAuthorizationParams returns the additional authorization parameters
// from whichever upstream config is set, or nil if none.
func (p *UpstreamProviderConfig) AdditionalAuthorizationParams() map[string]string {
        if p.OIDCConfig != nil {
                return p.OIDCConfig.AdditionalAuthorizationParams
        }
        if p.OAuth2Config != nil {
                return p.OAuth2Config.AdditionalAuthorizationParams
        }
        return nil
}

// AdditionalTokenParams returns the additional token-request parameters
// from whichever upstream config is set, or nil if none.
func (p *UpstreamProviderConfig) AdditionalTokenParams() map[string]string {
        if p.OIDCConfig != nil {
                return p.OIDCConfig.AdditionalTokenParams
        }
        if p.OAuth2Config != nil {
                return p.OAuth2Config.AdditionalTokenParams
        }
        return nil
}

// CABundleRef returns the CA bundle reference for the provider's configured
// type, or nil when the type-matched config or the reference is absent.
func (p *UpstreamProviderConfig) CABundleRef() *CABundleSource {
        switch p.Type {
        case UpstreamProviderTypeOIDC:
                if p.OIDCConfig != nil {
                        return p.OIDCConfig.CABundleRef
                }
        case UpstreamProviderTypeOAuth2:
                if p.OAuth2Config != nil {
                        return p.OAuth2Config.CABundleRef
                }
        }
        return nil
}

// SyntheticIdentityUpstreams returns the names of OAuth2 upstreams running
// in synthesis mode (neither userInfo nor identityFromToken configured),
// sorted lexically for deterministic condition messages. OIDC upstreams are
// skipped — they always have an ID-token-derived subject. Upstreams with
// identityFromToken are also skipped — the subject is extracted from the
// token response, not synthesized. Source of truth for the
// ConditionTypeIdentitySynthesized advisory.
func (c *EmbeddedAuthServerConfig) SyntheticIdentityUpstreams() []string {
        if c == nil {
                return nil
        }
        var names []string
        for i := range c.UpstreamProviders {
                p := &c.UpstreamProviders[i]
                if p.Type != UpstreamProviderTypeOAuth2 || p.OAuth2Config == nil {
                        continue
                }
                if p.OAuth2Config.UserInfo == nil && p.OAuth2Config.IdentityFromToken == nil {
                        names = append(names, p.Name)
                }
        }
        sort.Strings(names)
        return names
}

// ValidateAdditionalAuthorizationParams checks that no reserved OAuth2 parameters
// are present in the additional authorization params map.
func ValidateAdditionalAuthorizationParams(prefix string, params map[string]string) error {
        if err := oauthparams.Validate(params); err != nil {
                return fmt.Errorf("%s.additionalAuthorizationParams: %w", prefix, err)
        }
        return nil
}

// ValidateAdditionalTokenParams checks that no reserved OAuth2 token-request
// parameters are present in the additional token params map.
func ValidateAdditionalTokenParams(prefix string, params map[string]string) error {
        if err := oauthparams.ValidateTokenParams(params); err != nil {
                return fmt.Errorf("%s.additionalTokenParams: %w", prefix, err)
        }
        return nil
}

// validateAWSSts validates awsSts type configuration.
// This performs complex business logic validation that CEL cannot express.
func (r *MCPExternalAuthConfig) validateAWSSts() error {
        cfg := r.Spec.AWSSts
        if cfg == nil {
                return nil
        }

        // Region is required
        if cfg.Region == "" {
                return fmt.Errorf("awsSts.region is required")
        }

        // At least one of fallbackRoleArn or roleMappings must be configured
        // Both can be set: fallbackRoleArn is used when no mapping matches
        hasRoleArn := cfg.FallbackRoleArn != ""
        hasRoleMappings := len(cfg.RoleMappings) > 0

        if !hasRoleArn && !hasRoleMappings {
                return fmt.Errorf("awsSts: at least one of fallbackRoleArn or roleMappings must be configured")
        }

        // Validate role mappings if present
        for i, mapping := range cfg.RoleMappings {
                if mapping.RoleArn == "" {
                        return fmt.Errorf("awsSts.roleMappings[%d].roleArn is required", i)
                }
                // Exactly one of claim or matcher must be set
                if mapping.Claim == "" && mapping.Matcher == "" {
                        return fmt.Errorf("awsSts.roleMappings[%d]: exactly one of claim or matcher must be set", i)
                }
                if mapping.Claim != "" && mapping.Matcher != "" {
                        return fmt.Errorf("awsSts.roleMappings[%d]: claim and matcher are mutually exclusive", i)
                }
        }

        // Validate session duration if set
        // Bounds match AWS STS limits: 900s (15 min) to 43200s (12 hours)
        if cfg.SessionDuration != nil {
                duration := *cfg.SessionDuration
                const (
                        minSessionDuration int32 = 900   // 15 minutes
                        maxSessionDuration int32 = 43200 // 12 hours
                )
                if duration < minSessionDuration || duration > maxSessionDuration {
                        return fmt.Errorf("awsSts.sessionDuration must be between %d and %d seconds",
                                minSessionDuration, maxSessionDuration)
                }
        }

        return nil
}

func init() {
        SchemeBuilder.Register(&MCPExternalAuthConfig{}, &MCPExternalAuthConfigList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// MCPGroupSpec defines the desired state of MCPGroup
type MCPGroupSpec struct {
        // Description provides human-readable context
        // +optional
        Description string `json:"description,omitempty"`
}

// MCPGroupStatus defines observed state
type MCPGroupStatus struct {
        // ObservedGeneration reflects the generation most recently observed by the controller
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // Phase indicates current state
        // +optional
        // +kubebuilder:default=Pending
        Phase MCPGroupPhase `json:"phase,omitempty"`

        // Servers lists MCPServer names in this group
        // +listType=set
        // +optional
        Servers []string `json:"servers,omitempty"`

        // ServerCount is the number of MCPServers
        // +optional
        ServerCount int32 `json:"serverCount,omitempty"`

        // RemoteProxies lists MCPRemoteProxy names in this group
        // +listType=set
        // +optional
        RemoteProxies []string `json:"remoteProxies,omitempty"`

        // RemoteProxyCount is the number of MCPRemoteProxies
        // +optional
        RemoteProxyCount int32 `json:"remoteProxyCount,omitempty"`

        // Entries lists MCPServerEntry names in this group
        // +listType=set
        // +optional
        Entries []string `json:"entries,omitempty"`

        // EntryCount is the number of MCPServerEntries
        // +optional
        EntryCount int32 `json:"entryCount,omitempty"`

        // Conditions represent observations
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`
}

// MCPGroupPhase represents the lifecycle phase of an MCPGroup
// +kubebuilder:validation:Enum=Ready;Pending;Failed
type MCPGroupPhase string

const (
        // MCPGroupPhaseReady indicates the MCPGroup is ready
        MCPGroupPhaseReady MCPGroupPhase = "Ready"

        // MCPGroupPhasePending indicates the MCPGroup is pending
        MCPGroupPhasePending MCPGroupPhase = "Pending"

        // MCPGroupPhaseFailed indicates the MCPGroup has failed
        MCPGroupPhaseFailed MCPGroupPhase = "Failed"
)

// Condition types for MCPGroup
const (
        ConditionTypeMCPServersChecked = "MCPServersChecked"
)

// MCPGroupConditionReason represents the reason for a condition's last transition
const (
        ConditionReasonListMCPServersFailed    = "ListMCPServersCheckFailed"
        ConditionReasonListMCPServersSucceeded = "ListMCPServersCheckSucceeded"
)

//+kubebuilder:object:root=true
//+kubebuilder:storageversion
//+kubebuilder:subresource:status
//+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
//+kubebuilder:resource:shortName=mcpg,categories=toolhive
//+kubebuilder:printcolumn:name="Servers",type="integer",JSONPath=".status.serverCount"
//+kubebuilder:printcolumn:name="Proxies",type="integer",JSONPath=".status.remoteProxyCount"
//+kubebuilder:printcolumn:name="Phase",type="string",JSONPath=".status.phase"
//+kubebuilder:printcolumn:name="Checked",type="string",JSONPath=".status.conditions[?(@.type=='MCPServersChecked')].status"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"

// MCPGroup is the Schema for the mcpgroups API
type MCPGroup struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPGroupSpec   `json:"spec,omitempty"`
        Status MCPGroupStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// MCPGroupList contains a list of MCPGroup
type MCPGroupList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPGroup `json:"items"`
}

func init() {
        SchemeBuilder.Register(&MCPGroup{}, &MCPGroupList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        "fmt"

        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// OIDC configuration source types for MCPOIDCConfig
const (
        // MCPOIDCConfigTypeKubernetesServiceAccount is the type for Kubernetes service account token validation
        MCPOIDCConfigTypeKubernetesServiceAccount MCPOIDCConfigSourceType = "kubernetesServiceAccount"

        // MCPOIDCConfigTypeInline is the type for inline OIDC configuration
        MCPOIDCConfigTypeInline MCPOIDCConfigSourceType = "inline"
)

// Condition type and reasons for MCPOIDCConfig status (RFC-0023)
const (
        // ConditionTypeOIDCConfigValid indicates whether the MCPOIDCConfig configuration is valid
        ConditionTypeOIDCConfigValid = ConditionTypeValid

        // ConditionReasonOIDCConfigValid indicates spec validation passed
        ConditionReasonOIDCConfigValid = "ConfigValid"

        // ConditionReasonOIDCConfigInvalid indicates spec validation failed
        ConditionReasonOIDCConfigInvalid = "ConfigInvalid"
)

// MCPOIDCConfigSourceType represents the type of OIDC configuration source for MCPOIDCConfig
type MCPOIDCConfigSourceType string

// MCPOIDCConfigSpec defines the desired state of MCPOIDCConfig.
// MCPOIDCConfig resources are namespace-scoped and can only be referenced by
// MCPServer resources in the same namespace.
//
// +kubebuilder:validation:XValidation:rule="self.type == 'kubernetesServiceAccount' ? has(self.kubernetesServiceAccount) : !has(self.kubernetesServiceAccount)",message="kubernetesServiceAccount must be set when type is 'kubernetesServiceAccount', and must not be set otherwise"
// +kubebuilder:validation:XValidation:rule="self.type == 'inline' ? has(self.inline) : !has(self.inline)",message="inline must be set when type is 'inline', and must not be set otherwise"
//
//nolint:lll // CEL validation rules exceed line length limit
type MCPOIDCConfigSpec struct {
        // Type is the type of OIDC configuration source
        // +kubebuilder:validation:Enum=kubernetesServiceAccount;inline
        // +kubebuilder:validation:Required
        Type MCPOIDCConfigSourceType `json:"type"`

        // KubernetesServiceAccount configures OIDC for Kubernetes service account token validation.
        // Only used when Type is "kubernetesServiceAccount".
        // +optional
        KubernetesServiceAccount *KubernetesServiceAccountOIDCConfig `json:"kubernetesServiceAccount,omitempty"`

        // Inline contains direct OIDC configuration.
        // Only used when Type is "inline".
        // +optional
        Inline *InlineOIDCSharedConfig `json:"inline,omitempty"`
}

// KubernetesServiceAccountOIDCConfig configures OIDC for Kubernetes service account token validation.
// This contains shared fields without audience, which is specified per-server via MCPOIDCConfigReference.
type KubernetesServiceAccountOIDCConfig struct {
        // ServiceAccount is the name of the service account to validate tokens for.
        // If empty, uses the pod's service account.
        // +optional
        ServiceAccount string `json:"serviceAccount,omitempty"`

        // Namespace is the namespace of the service account.
        // If empty, uses the MCPServer's namespace.
        // +optional
        Namespace string `json:"namespace,omitempty"`

        // Issuer is the OIDC issuer URL.
        // +kubebuilder:default="https://kubernetes.default.svc"
        // +optional
        Issuer string `json:"issuer,omitempty"`

        // JWKSURL is the URL to fetch the JWKS from.
        // If empty, OIDC discovery will be used to automatically determine the JWKS URL.
        // +optional
        JWKSURL string `json:"jwksUrl,omitempty"`

        // IntrospectionURL is the URL for token introspection endpoint.
        // If empty, OIDC discovery will be used to automatically determine the introspection URL.
        // +optional
        IntrospectionURL string `json:"introspectionUrl,omitempty"`

        // UseClusterAuth enables using the Kubernetes cluster's CA bundle and service account token.
        // When true, uses /var/run/secrets/kubernetes.io/serviceaccount/ca.crt for TLS verification
        // and /var/run/secrets/kubernetes.io/serviceaccount/token for bearer token authentication.
        // Defaults to true if not specified.
        // +optional
        UseClusterAuth *bool `json:"useClusterAuth"`
}

// InlineOIDCSharedConfig contains direct OIDC configuration.
// This contains shared fields without audience and scopes, which are specified per-server
// via MCPOIDCConfigReference.
type InlineOIDCSharedConfig struct {
        // Issuer is the OIDC issuer URL
        // +kubebuilder:validation:Required
        Issuer string `json:"issuer"`

        // JWKSURL is the URL to fetch the JWKS from
        // +optional
        JWKSURL string `json:"jwksUrl,omitempty"`

        // IntrospectionURL is the URL for token introspection endpoint
        // +optional
        IntrospectionURL string `json:"introspectionUrl,omitempty"`

        // ClientID is the OIDC client ID
        // +optional
        ClientID string `json:"clientId,omitempty"`

        // ClientSecretRef is a reference to a Kubernetes Secret containing the client secret
        // +optional
        ClientSecretRef *SecretKeyRef `json:"clientSecretRef,omitempty"`

        // CABundleRef references a ConfigMap containing the CA certificate bundle.
        // When specified, ToolHive auto-mounts the ConfigMap and auto-computes ThvCABundlePath.
        // +optional
        CABundleRef *CABundleSource `json:"caBundleRef,omitempty"`

        // JWKSAuthTokenPath is the path to file containing bearer token for JWKS/OIDC requests
        // +optional
        JWKSAuthTokenPath string `json:"jwksAuthTokenPath,omitempty"`

        // JWKSAllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses.
        // Note: at runtime, if either JWKSAllowPrivateIP or ProtectedResourceAllowPrivateIP
        // is true, private IPs are allowed for all OIDC HTTP requests (JWKS, discovery, introspection).
        // +kubebuilder:default=false
        // +optional
        JWKSAllowPrivateIP bool `json:"jwksAllowPrivateIP"`

        // ProtectedResourceAllowPrivateIP allows protected resource endpoint on private IP addresses.
        // Note: at runtime, if either ProtectedResourceAllowPrivateIP or JWKSAllowPrivateIP
        // is true, private IPs are allowed for all OIDC HTTP requests (JWKS, discovery, introspection).
        // +kubebuilder:default=false
        // +optional
        ProtectedResourceAllowPrivateIP bool `json:"protectedResourceAllowPrivateIP"`

        // InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuer and JWKS URLs for development/testing.
        // WARNING: This is insecure and should NEVER be used in production.
        // +kubebuilder:default=false
        // +optional
        InsecureAllowHTTP bool `json:"insecureAllowHTTP"`
}

// Well-known WorkloadReference Kind values.
const (
        WorkloadKindMCPServer        = "MCPServer"
        WorkloadKindVirtualMCPServer = "VirtualMCPServer"
        WorkloadKindMCPRemoteProxy   = "MCPRemoteProxy"
)

// WorkloadReference identifies a workload that references a shared configuration resource.
// Namespace is implicit — cross-namespace references are not supported.
type WorkloadReference struct {
        // Kind is the type of workload resource
        // +kubebuilder:validation:Enum=MCPServer;VirtualMCPServer;MCPRemoteProxy
        // +kubebuilder:validation:Required
        Kind string `json:"kind"`

        // Name is the name of the workload resource
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Name string `json:"name"`
}

// MCPOIDCConfigStatus defines the observed state of MCPOIDCConfig
type MCPOIDCConfigStatus struct {
        // Conditions represent the latest available observations of the MCPOIDCConfig's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration is the most recent generation observed for this MCPOIDCConfig.
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // ConfigHash is a hash of the current configuration for change detection
        // +optional
        ConfigHash string `json:"configHash,omitempty"`
}

// +kubebuilder:object:root=true
// +kubebuilder:storageversion
// +kubebuilder:subresource:status
// +kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
// +kubebuilder:resource:shortName=mcpoidc,categories=toolhive
// +kubebuilder:printcolumn:name="Source",type=string,JSONPath=`.spec.type`
// +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status`
// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp`

// MCPOIDCConfig is the Schema for the mcpoidcconfigs API.
// MCPOIDCConfig resources are namespace-scoped and can only be referenced by
// MCPServer resources within the same namespace. Cross-namespace references
// are not supported for security and isolation reasons.
type MCPOIDCConfig struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPOIDCConfigSpec   `json:"spec,omitempty"`
        Status MCPOIDCConfigStatus `json:"status,omitempty"`
}

// +kubebuilder:object:root=true

// MCPOIDCConfigList contains a list of MCPOIDCConfig
type MCPOIDCConfigList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPOIDCConfig `json:"items"`
}

// MCPOIDCConfigReference is a reference to an MCPOIDCConfig resource with per-server overrides.
// The referenced MCPOIDCConfig must be in the same namespace as the MCPServer.
type MCPOIDCConfigReference struct {
        // Name is the name of the MCPOIDCConfig resource
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Name string `json:"name"`

        // Audience is the expected audience for token validation.
        // This MUST be unique per server to prevent token replay attacks.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Audience string `json:"audience"`

        // Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728).
        // If empty, defaults to ["openid"].
        // +listType=atomic
        // +optional
        Scopes []string `json:"scopes,omitempty"`

        // ResourceURL is the public URL for OAuth protected resource metadata (RFC 9728).
        // When the server is exposed via Ingress or gateway, set this to the external
        // URL that MCP clients connect to. If not specified, defaults to the internal
        // Kubernetes service URL.
        // +optional
        ResourceURL string `json:"resourceUrl,omitempty"`
}

// Validate performs validation on the MCPOIDCConfig spec.
// This method is called by the controller during reconciliation.
//
// Note: These validations provide defense-in-depth alongside CEL validation rules.
// CEL catches issues at API admission time, but this method also validates stored objects
// to catch any that bypassed CEL or were stored before CEL rules were added.
func (r *MCPOIDCConfig) Validate() error {
        return r.validateTypeConfigConsistency()
}

// validateTypeConfigConsistency validates that the correct config is set for the selected type.
// This mirrors the CEL validation rules but provides defense-in-depth for stored objects.
func (r *MCPOIDCConfig) validateTypeConfigConsistency() error {
        if (r.Spec.KubernetesServiceAccount == nil) == (r.Spec.Type == MCPOIDCConfigTypeKubernetesServiceAccount) {
                return fmt.Errorf("kubernetesServiceAccount configuration must be set if and only if type is 'kubernetesServiceAccount'")
        }
        if (r.Spec.Inline == nil) == (r.Spec.Type == MCPOIDCConfigTypeInline) {
                return fmt.Errorf("inline configuration must be set if and only if type is 'inline'")
        }
        return nil
}

func init() {
        SchemeBuilder.Register(&MCPOIDCConfig{}, &MCPOIDCConfigList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        "encoding/json"
        "fmt"

        corev1 "k8s.io/api/core/v1"
        apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/runtime"
)

// MCPRegistrySpec defines the desired state of MCPRegistry
type MCPRegistrySpec struct {
        // ConfigYAML is the complete registry server config.yaml content.
        // The operator creates a ConfigMap from this string and mounts it
        // at /config/config.yaml in the registry-api container.
        // The operator does NOT parse, validate, or transform this content —
        // configuration validation is the registry server's responsibility.
        //
        // Security note: this content is stored in a ConfigMap, not a Secret.
        // Do not inline credentials (passwords, tokens, client secrets) in this
        // field. Instead, reference credentials via file paths and mount the
        // actual secrets using the Volumes and VolumeMounts fields. For database
        // passwords, use PGPassSecretRef.
        //
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        ConfigYAML string `json:"configYAML"`

        // Volumes defines additional volumes to add to the registry API pod.
        // Each entry is a standard Kubernetes Volume object (JSON/YAML).
        // The operator appends them to the pod spec alongside its own config volume.
        //
        // Use these to mount:
        //   - Secrets (git auth tokens, OAuth client secrets, CA certs)
        //   - ConfigMaps (registry data files)
        //   - PersistentVolumeClaims (registry data on persistent storage)
        //   - Any other volume type the registry server needs
        //
        // +optional
        // +listType=atomic
        // +kubebuilder:pruning:PreserveUnknownFields
        Volumes []apiextensionsv1.JSON `json:"volumes,omitempty"`

        // VolumeMounts defines additional volume mounts for the registry-api container.
        // Each entry is a standard Kubernetes VolumeMount object (JSON/YAML).
        // The operator appends them to the container's volume mounts alongside the config mount.
        //
        // Mount paths must match the file paths referenced in configYAML.
        // For example, if configYAML references passwordFile: /secrets/git-creds/token,
        // a corresponding volume mount must exist with mountPath: /secrets/git-creds.
        //
        // +optional
        // +listType=atomic
        // +kubebuilder:pruning:PreserveUnknownFields
        VolumeMounts []apiextensionsv1.JSON `json:"volumeMounts,omitempty"`

        // PGPassSecretRef references a Secret containing a pre-created pgpass file.
        //
        // Why this is a dedicated field instead of a regular volume/volumeMount:
        // PostgreSQL's libpq rejects pgpass files that aren't mode 0600. Kubernetes
        // secret volumes mount files as root-owned, and the registry-api container
        // runs as non-root (UID 65532). A root-owned 0600 file is unreadable by
        // UID 65532, and using fsGroup changes permissions to 0640 which libpq also
        // rejects. The only solution is an init container that copies the file to an
        // emptyDir as the app user and runs chmod 0600. This cannot be expressed
        // through volumes/volumeMounts alone -- it requires an init container, two
        // extra volumes (secret + emptyDir), a subPath mount, and an environment
        // variable, all wired together correctly.
        //
        // When specified, the operator generates all of that plumbing invisibly.
        // The user creates the Secret with pgpass-formatted content; the operator
        // handles only the Kubernetes permission mechanics.
        //
        // Example Secret:
        //
        //        apiVersion: v1
        //        kind: Secret
        //        metadata:
        //          name: my-pgpass
        //        stringData:
        //          .pgpass: |
        //            postgres:5432:registry:db_app:mypassword
        //            postgres:5432:registry:db_migrator:otherpassword
        //
        // Then reference it:
        //
        //        pgpassSecretRef:
        //          name: my-pgpass
        //          key: .pgpass
        //
        // +optional
        PGPassSecretRef *corev1.SecretKeySelector `json:"pgpassSecretRef,omitempty"`

        // DisplayName is a human-readable name for the registry.
        // +optional
        DisplayName string `json:"displayName,omitempty"`

        // PodTemplateSpec defines the pod template to use for the registry API server.
        // This allows for customizing the pod configuration beyond what is provided by the other fields.
        // Note that to modify the specific container the registry API server runs in, you must specify
        // the `registry-api` container name in the PodTemplateSpec.
        // This field accepts a PodTemplateSpec object as JSON/YAML.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        PodTemplateSpec *runtime.RawExtension `json:"podTemplateSpec,omitempty"`

        // ImagePullSecrets allows specifying image pull secrets for the registry API workload.
        // These are applied to both the registry-api Deployment's PodSpec.ImagePullSecrets
        // and to the operator-managed ServiceAccount the registry API runs as, so private
        // images are pullable through either path.
        //
        // Use this field for new manifests.
        //
        // Important: this is the ONLY way to attach image-pull credentials to the
        // operator-managed ServiceAccount. The legacy
        // spec.podTemplateSpec.spec.imagePullSecrets path populates the Deployment's pod
        // spec ONLY — it does NOT touch the ServiceAccount. On managed Kubernetes
        // platforms that rely on ServiceAccount-level credential injection (for example
        // GKE Workload Identity, OpenShift's per-SA dockercfg secrets, EKS IRSA), using
        // only the legacy PodTemplateSpec path can fail to pull private images even when
        // the secret exists in the namespace. Always set spec.imagePullSecrets when
        // SA-level credentials matter.
        //
        // Precedence with PodTemplateSpec:
        //   - This field is applied first as the controller-generated default.
        //   - Values set under spec.podTemplateSpec.spec.imagePullSecrets are user overrides
        //     and win on overlap. If the user supplies imagePullSecrets via PodTemplateSpec,
        //     those replace the default list on the Deployment (the list is treated atomically).
        //   - The ServiceAccount is always populated from this field — PodTemplateSpec does not
        //     affect the ServiceAccount.
        //
        // An omitted field and an explicitly empty list are equivalent: both leave the
        // ServiceAccount's existing ImagePullSecrets unchanged. This preserves
        // platform-managed pull secrets (for example OpenShift's per-SA dockercfg
        // entries) when overlays or patches emit an empty list. Truly clearing the
        // ServiceAccount's pull secrets requires recreating the resource.
        //
        // +listType=atomic
        // +optional
        ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,omitempty"`
}

// MCPRegistryStatus defines the observed state of MCPRegistry
type MCPRegistryStatus struct {
        // Conditions represent the latest available observations of the MCPRegistry's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration reflects the generation most recently observed by the controller
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // Phase represents the current overall phase of the MCPRegistry
        // +optional
        Phase MCPRegistryPhase `json:"phase,omitempty"`

        // Message provides additional information about the current phase
        // +optional
        Message string `json:"message,omitempty"`

        // URL is the URL where the registry API can be accessed
        // +optional
        URL string `json:"url,omitempty"`

        // ReadyReplicas is the number of ready registry API replicas
        // +optional
        ReadyReplicas int32 `json:"readyReplicas,omitempty"`
}

// MCPRegistryPhase represents the phase of the MCPRegistry
// +kubebuilder:validation:Enum=Pending;Ready;Failed;Terminating
type MCPRegistryPhase string

const (
        // MCPRegistryPhasePending means the MCPRegistry is being initialized
        MCPRegistryPhasePending MCPRegistryPhase = "Pending"

        // MCPRegistryPhaseReady means the MCPRegistry is ready and operational
        MCPRegistryPhaseReady MCPRegistryPhase = "Ready"

        // MCPRegistryPhaseFailed means the MCPRegistry has failed
        MCPRegistryPhaseFailed MCPRegistryPhase = "Failed"

        // MCPRegistryPhaseTerminating means the MCPRegistry is being deleted
        MCPRegistryPhaseTerminating MCPRegistryPhase = "Terminating"
)

// Condition reasons for MCPRegistry
const (
        // ConditionReasonRegistryReady indicates the MCPRegistry is ready
        ConditionReasonRegistryReady = "Ready"

        // ConditionReasonRegistryNotReady indicates the MCPRegistry is not ready
        ConditionReasonRegistryNotReady = "NotReady"
)

// Developer note: the MCPRegistry deprecation is expressed as prose in the type
// doc comment below, NOT via the Go "Deprecated:" convention. The operator still
// reconciles this type, so the staticcheck SA1019 analyzer must not flag the
// operator's own internal uses. The kubectl-visible deprecation comes from the
// +kubebuilder:deprecatedversion:warning marker below.

// MCPRegistry is the Schema for the mcpregistries API.
//
// The MCPRegistry CRD is deprecated and will be removed in a future release.
// Install the ToolHive registry server via the toolhive-registry-server Helm chart
// instead: https://github.com/stacklok/toolhive-registry-server
//
// +kubebuilder:object:root=true
// +kubebuilder:storageversion
// +kubebuilder:deprecatedversion:warning="MCPRegistry is deprecated and will be removed in a future release; install the ToolHive registry server via the toolhive-registry-server Helm chart (https://github.com/stacklok/toolhive-registry-server) instead"
// +kubebuilder:subresource:status
// +kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
// +kubebuilder:printcolumn:name="Status",type="string",JSONPath=".status.phase"
// +kubebuilder:printcolumn:name="Ready",type="string",JSONPath=".status.conditions[?(@.type=='Ready')].status"
// +kubebuilder:printcolumn:name="Replicas",type="integer",JSONPath=".status.readyReplicas"
// +kubebuilder:printcolumn:name="URL",type="string",JSONPath=".status.url"
// +kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"
// +kubebuilder:resource:shortName=mcpreg;registry,scope=Namespaced,categories=toolhive
//
//nolint:lll // kubebuilder deprecatedversion marker cannot be line-wrapped
type MCPRegistry struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPRegistrySpec   `json:"spec,omitempty"`
        Status MCPRegistryStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// MCPRegistryList contains a list of MCPRegistry
type MCPRegistryList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPRegistry `json:"items"`
}

// GetAPIResourceName returns the base name for registry API resources (deployment, service)
func (r *MCPRegistry) GetAPIResourceName() string {
        return fmt.Sprintf("%s-api", r.Name)
}

func init() {
        SchemeBuilder.Register(&MCPRegistry{}, &MCPRegistryList{})
}

// HasPodTemplateSpec returns true if the MCPRegistry has a PodTemplateSpec
func (r *MCPRegistry) HasPodTemplateSpec() bool {
        return r.Spec.PodTemplateSpec != nil
}

// GetPodTemplateSpecRaw returns the raw PodTemplateSpec
func (r *MCPRegistry) GetPodTemplateSpecRaw() *runtime.RawExtension {
        return r.Spec.PodTemplateSpec
}

// ParseVolumes deserializes the raw JSON Volumes into typed corev1.Volume objects.
// Returns an empty slice if Volumes is nil or empty.
func (s *MCPRegistrySpec) ParseVolumes() ([]corev1.Volume, error) {
        volumes := make([]corev1.Volume, 0, len(s.Volumes))
        for i, raw := range s.Volumes {
                var vol corev1.Volume
                if err := json.Unmarshal(raw.Raw, &vol); err != nil {
                        return nil, fmt.Errorf("failed to unmarshal volumes[%d]: %w", i, err)
                }
                volumes = append(volumes, vol)
        }
        return volumes, nil
}

// ParseVolumeMounts deserializes the raw JSON VolumeMounts into typed corev1.VolumeMount objects.
// Returns an empty slice if VolumeMounts is nil or empty.
func (s *MCPRegistrySpec) ParseVolumeMounts() ([]corev1.VolumeMount, error) {
        mounts := make([]corev1.VolumeMount, 0, len(s.VolumeMounts))
        for i, raw := range s.VolumeMounts {
                var mount corev1.VolumeMount
                if err := json.Unmarshal(raw.Raw, &mount); err != nil {
                        return nil, fmt.Errorf("failed to unmarshal volumeMounts[%d]: %w", i, err)
                }
                mounts = append(mounts, mount)
        }
        return mounts, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/runtime"
)

// HeaderForwardConfig defines header forward configuration for remote servers.
type HeaderForwardConfig struct {
        // AddPlaintextHeaders is a map of header names to literal values to inject into requests.
        // WARNING: Values are stored in plaintext and visible via kubectl commands.
        // Use addHeadersFromSecret for sensitive data like API keys or tokens.
        // +optional
        AddPlaintextHeaders map[string]string `json:"addPlaintextHeaders,omitempty"`

        // AddHeadersFromSecret references Kubernetes Secrets for sensitive header values.
        // +listType=map
        // +listMapKey=headerName
        // +optional
        AddHeadersFromSecret []HeaderFromSecret `json:"addHeadersFromSecret,omitempty"`
}

// HeaderFromSecret defines a header whose value comes from a Kubernetes Secret.
type HeaderFromSecret struct {
        // HeaderName is the HTTP header name (e.g., "X-API-Key")
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=255
        HeaderName string `json:"headerName"`

        // ValueSecretRef references the Secret and key containing the header value
        // +kubebuilder:validation:Required
        ValueSecretRef *SecretKeyRef `json:"valueSecretRef"`
}

// MCPRemoteProxySpec defines the desired state of MCPRemoteProxy
//
// +kubebuilder:validation:XValidation:rule="!(has(self.authzConfig) && has(self.authzConfigRef))",message="authzConfig and authzConfigRef are mutually exclusive; use authzConfigRef to reference a shared MCPAuthzConfig"
//
//nolint:lll // CEL validation rules exceed line length limit
type MCPRemoteProxySpec struct {
        // RemoteURL is the URL of the remote MCP server to proxy
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https?://`
        RemoteURL string `json:"remoteUrl"`

        // AllowPrivateEndpoint permits RemoteURL to point at a private or
        // in-cluster endpoint: RFC 1918 and IPv6 unique-local addresses, and
        // hostnames ending in "cluster.local". (Bare ".svc" hostnames are never
        // blocked, since no DNS resolution is performed.) Enable this to reach a
        // co-located in-cluster backend in-mesh so the backend's workload-identity
        // authorization policy still applies. Loopback, link-local, cloud-metadata,
        // and kubernetes.default endpoints remain blocked regardless of this
        // setting.
        // +kubebuilder:default=false
        // +optional
        AllowPrivateEndpoint bool `json:"allowPrivateEndpoint,omitempty"`

        // ProxyPort is the port to expose the MCP proxy on
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:validation:Maximum=65535
        // +kubebuilder:default=8080
        ProxyPort int32 `json:"proxyPort,omitempty"`

        // Transport is the transport method for the remote proxy (sse or streamable-http)
        // +kubebuilder:validation:Enum=sse;streamable-http
        // +kubebuilder:default=streamable-http
        Transport string `json:"transport,omitempty"`

        // OIDCConfigRef references a shared MCPOIDCConfig resource for OIDC authentication.
        // The referenced MCPOIDCConfig must exist in the same namespace as this MCPRemoteProxy.
        // Per-server overrides (audience, scopes) are specified here; shared provider config
        // lives in the MCPOIDCConfig resource.
        //
        // SECURITY: if this field is omitted and no other authentication source is configured,
        // the proxy runs UNAUTHENTICATED. It accepts every request that can reach its port and
        // forwards it to the remote MCP server under a synthetic local-user identity, with no
        // token or credential check. Set this field to enforce identity-based access control
        // per request.
        // +optional
        OIDCConfigRef *MCPOIDCConfigReference `json:"oidcConfigRef,omitempty"`

        // ExternalAuthConfigRef references a MCPExternalAuthConfig resource for token exchange.
        // When specified, the proxy will exchange validated incoming tokens for remote service tokens.
        // The referenced MCPExternalAuthConfig must exist in the same namespace as this MCPRemoteProxy.
        // +optional
        ExternalAuthConfigRef *ExternalAuthConfigRef `json:"externalAuthConfigRef,omitempty"`

        // AuthServerRef optionally references a resource that configures an embedded
        // OAuth 2.0/OIDC authorization server to authenticate MCP clients.
        // Currently the only supported kind is MCPExternalAuthConfig (type: embeddedAuthServer).
        // +optional
        AuthServerRef *AuthServerRef `json:"authServerRef,omitempty"`

        // HeaderForward configures headers to inject into requests to the remote MCP server.
        // Use this to add custom headers like X-Tenant-ID or correlation IDs.
        // +optional
        HeaderForward *HeaderForwardConfig `json:"headerForward,omitempty"`

        // AuthzConfig defines authorization policy configuration for the proxy.
        // AuthzConfig and AuthzConfigRef are mutually exclusive.
        // +optional
        AuthzConfig *AuthzConfigRef `json:"authzConfig,omitempty"`

        // AuthzConfigRef references a shared MCPAuthzConfig resource for authorization.
        // The referenced MCPAuthzConfig must exist in the same namespace as this MCPRemoteProxy.
        // Mutually exclusive with authzConfig.
        // +optional
        AuthzConfigRef *MCPAuthzConfigReference `json:"authzConfigRef,omitempty"`

        // Audit defines audit logging configuration for the proxy
        // +optional
        Audit *AuditConfig `json:"audit,omitempty"`

        // ToolConfigRef references a MCPToolConfig resource for tool filtering and renaming.
        // The referenced MCPToolConfig must exist in the same namespace as this MCPRemoteProxy.
        // Cross-namespace references are not supported for security and isolation reasons.
        // If specified, this allows filtering and overriding tools from the remote MCP server.
        // +optional
        ToolConfigRef *ToolConfigRef `json:"toolConfigRef,omitempty"`

        // TelemetryConfigRef references an MCPTelemetryConfig resource for shared telemetry configuration.
        // The referenced MCPTelemetryConfig must exist in the same namespace as this MCPRemoteProxy.
        // Cross-namespace references are not supported for security and isolation reasons.
        // +optional
        TelemetryConfigRef *MCPTelemetryConfigReference `json:"telemetryConfigRef,omitempty"`

        // Resources defines the resource requirements for the proxy container
        // +optional
        Resources ResourceRequirements `json:"resources,omitempty"`

        // ServiceAccount is the name of an already existing service account to use by the proxy.
        // If not specified, a ServiceAccount will be created automatically and used by the proxy.
        // +optional
        ServiceAccount *string `json:"serviceAccount,omitempty"`

        // PodTemplateSpec defines the pod template to use for the MCPRemoteProxy
        // This allows for customizing the pod configuration beyond what is provided by the other fields.
        // Note that to modify the specific container the remote proxy runs in, you must specify
        // the `toolhive` container name in the PodTemplateSpec.
        // This field accepts a PodTemplateSpec object as JSON/YAML.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        PodTemplateSpec *runtime.RawExtension `json:"podTemplateSpec,omitempty"`

        // TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies
        // When enabled, the proxy will use X-Forwarded-Proto, X-Forwarded-Host, X-Forwarded-Port,
        // and X-Forwarded-Prefix headers to construct endpoint URLs
        // +kubebuilder:default=false
        // +optional
        TrustProxyHeaders bool `json:"trustProxyHeaders,omitempty"`

        // EndpointPrefix is the path prefix to prepend to SSE endpoint URLs.
        // This is used to handle path-based ingress routing scenarios where the ingress
        // strips a path prefix before forwarding to the backend.
        // +optional
        EndpointPrefix string `json:"endpointPrefix,omitempty"`

        // ResourceOverrides allows overriding annotations and labels for resources created by the operator
        // +optional
        ResourceOverrides *ResourceOverrides `json:"resourceOverrides,omitempty"`

        // GroupRef references the MCPGroup this proxy belongs to.
        // The referenced MCPGroup must be in the same namespace.
        // +optional
        GroupRef *MCPGroupRef `json:"groupRef,omitempty"`

        // SessionAffinity controls whether the Service routes repeated client connections to the same pod.
        // MCP protocols (SSE, streamable-http) are stateful, so ClientIP is the default.
        // Set to "None" for stateless servers or when using an external load balancer with its own affinity.
        //
        // Interaction with sessionStorage: when running multiple replicas with
        // sessionStorage.provider "redis", set this to "None" so requests are
        // distributed across replicas and sessions resolve via the shared store.
        // Conversely, "None" without Redis-backed sessionStorage breaks session
        // continuity — any request landing on a different pod fails with
        // "Session not found".
        // +kubebuilder:validation:Enum=ClientIP;None
        // +kubebuilder:default=ClientIP
        // +optional
        SessionAffinity string `json:"sessionAffinity,omitempty"`

        // Replicas is the desired number of proxy pod replicas.
        // MCPRemoteProxy creates a single Deployment for the proxy process, so there
        // is only one replicas field (mirrors VirtualMCPServer.spec.replicas).
        // When nil, the operator does not set Deployment.Spec.Replicas, leaving replica
        // management to an HPA or other external controller.
        // When set above 1, also configure sessionStorage with the redis provider and
        // sessionAffinity: "None" so sessions resolve across replicas; otherwise a
        // SessionStorageWarning condition is surfaced on the resource status.
        // +kubebuilder:validation:Minimum=0
        // +optional
        Replicas *int32 `json:"replicas,omitempty"`

        // SessionStorage configures session storage for stateful horizontal scaling.
        // When nil, no session storage is configured and the proxy falls back to
        // pod-local in-memory session state — incompatible with multi-replica
        // deployments behind load balancers that don't preserve client-IP affinity
        // (e.g. AWS ALB across multiple AZs).
        //
        // The transparent proxy validates `Mcp-Session-Id` against this store on
        // every non-initialize request (see pkg/transport/proxy/transparent/
        // transparent_proxy.go) and rewrites client-facing session IDs to backend
        // session IDs using session metadata. Both lookups require shared state
        // across replicas.
        //
        // When using the Redis provider, also set sessionAffinity to "None" so the
        // Service routes requests round-robin and all replicas rely on the shared
        // session store rather than pod-local state.
        //
        // Mirrors MCPServer.spec.sessionStorage and VirtualMCPServer.spec.sessionStorage.
        // +optional
        SessionStorage *SessionStorageConfig `json:"sessionStorage,omitempty"`
}

// MCPRemoteProxyStatus defines the observed state of MCPRemoteProxy
type MCPRemoteProxyStatus struct {
        // Phase is the current phase of the MCPRemoteProxy
        // +optional
        Phase MCPRemoteProxyPhase `json:"phase,omitempty"`

        // URL is the internal cluster URL where the proxy can be accessed
        // +optional
        URL string `json:"url,omitempty"`

        // ExternalURL is the external URL where the proxy can be accessed (if exposed externally)
        // +optional
        ExternalURL string `json:"externalUrl,omitempty"`

        // ObservedGeneration reflects the generation of the most recently observed MCPRemoteProxy
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // Conditions represent the latest available observations of the MCPRemoteProxy's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ToolConfigHash stores the hash of the referenced ToolConfig for change detection
        // +optional
        ToolConfigHash string `json:"toolConfigHash,omitempty"`

        // TelemetryConfigHash stores the hash of the referenced MCPTelemetryConfig for change detection
        // +optional
        TelemetryConfigHash string `json:"telemetryConfigHash,omitempty"`

        // ExternalAuthConfigHash is the hash of the referenced MCPExternalAuthConfig spec
        // +optional
        ExternalAuthConfigHash string `json:"externalAuthConfigHash,omitempty"`

        // AuthServerConfigHash is the hash of the referenced authServerRef spec,
        // used to detect configuration changes and trigger reconciliation.
        // +optional
        AuthServerConfigHash string `json:"authServerConfigHash,omitempty"`

        // AuthzConfigHash is the hash of the referenced MCPAuthzConfig spec for change detection
        // +optional
        AuthzConfigHash string `json:"authzConfigHash,omitempty"`

        // OIDCConfigHash is the hash of the referenced MCPOIDCConfig spec for change detection
        // +optional
        OIDCConfigHash string `json:"oidcConfigHash,omitempty"`

        // Message provides additional information about the current phase
        // +optional
        Message string `json:"message,omitempty"`
}

// MCPRemoteProxyPhase is a label for the condition of a MCPRemoteProxy at the current time
// +kubebuilder:validation:Enum=Pending;Ready;Failed;Terminating
type MCPRemoteProxyPhase string

const (
        // MCPRemoteProxyPhasePending means the proxy is being created
        MCPRemoteProxyPhasePending MCPRemoteProxyPhase = "Pending"

        // MCPRemoteProxyPhaseReady means the proxy is ready and operational
        MCPRemoteProxyPhaseReady MCPRemoteProxyPhase = "Ready"

        // MCPRemoteProxyPhaseFailed means the proxy failed to start or encountered an error
        MCPRemoteProxyPhaseFailed MCPRemoteProxyPhase = "Failed"

        // MCPRemoteProxyPhaseTerminating means the proxy is being deleted
        MCPRemoteProxyPhaseTerminating MCPRemoteProxyPhase = "Terminating"
)

// Condition types for MCPRemoteProxy
const (
        // ConditionTypeReady indicates overall readiness of the proxy
        ConditionTypeReady = "Ready"

        // ConditionTypeRemoteAvailable indicates whether the remote MCP server is reachable
        ConditionTypeRemoteAvailable = "RemoteAvailable"

        // ConditionTypeAuthConfigured indicates whether authentication is properly configured
        ConditionTypeAuthConfigured = "AuthConfigured"

        // ConditionTypeMCPRemoteProxyGroupRefValidated indicates whether the GroupRef is valid
        ConditionTypeMCPRemoteProxyGroupRefValidated = "GroupRefValidated"

        // ConditionTypeMCPRemoteProxyToolConfigValidated indicates whether the ToolConfigRef is valid
        ConditionTypeMCPRemoteProxyToolConfigValidated = "ToolConfigValidated"

        // ConditionTypeMCPRemoteProxyTelemetryConfigRefValidated indicates whether the TelemetryConfigRef is valid
        ConditionTypeMCPRemoteProxyTelemetryConfigRefValidated = "TelemetryConfigRefValidated"

        // ConditionTypeMCPRemoteProxyExternalAuthConfigValidated indicates whether the ExternalAuthConfigRef is valid
        ConditionTypeMCPRemoteProxyExternalAuthConfigValidated = "ExternalAuthConfigValidated"

        // ConditionTypeMCPRemoteProxyAuthServerRefValidated indicates whether the AuthServerRef is valid
        ConditionTypeMCPRemoteProxyAuthServerRefValidated = "AuthServerRefValidated"

        // ConditionTypeMCPRemoteProxyPodTemplateValid indicates whether the PodTemplateSpec is valid
        ConditionTypeMCPRemoteProxyPodTemplateValid = "PodTemplateValid"

        // ConditionTypeMCPRemoteProxyCABundleRefValidated indicates whether the OIDC CA bundle reference is valid
        ConditionTypeMCPRemoteProxyCABundleRefValidated = "CABundleRefValidated"

        // ConditionTypeConfigurationValid indicates whether the proxy spec has passed all pre-deployment validation checks
        ConditionTypeConfigurationValid = "ConfigurationValid"
)

// Condition reasons for MCPRemoteProxy
const (
        // ConditionReasonDeploymentReady indicates the deployment is ready
        ConditionReasonDeploymentReady = "DeploymentReady"

        // ConditionReasonDeploymentNotReady indicates the deployment is not ready
        ConditionReasonDeploymentNotReady = "DeploymentNotReady"

        // ConditionReasonRemoteURLReachable indicates the remote URL is reachable
        ConditionReasonRemoteURLReachable = "RemoteURLReachable"

        // ConditionReasonRemoteURLUnreachable indicates the remote URL is unreachable
        ConditionReasonRemoteURLUnreachable = "RemoteURLUnreachable"

        // ConditionReasonAuthValid indicates authentication configuration is valid
        ConditionReasonAuthValid = "AuthValid"

        // ConditionReasonAuthInvalid indicates authentication configuration is invalid
        ConditionReasonAuthInvalid = "AuthInvalid"

        // ConditionReasonMissingOIDCConfig indicates OIDCConfig is not specified
        ConditionReasonMissingOIDCConfig = "MissingOIDCConfig"

        // ConditionReasonMCPRemoteProxyGroupRefValidated indicates the GroupRef is valid
        ConditionReasonMCPRemoteProxyGroupRefValidated = "GroupRefIsValid"

        // ConditionReasonMCPRemoteProxyGroupRefNotFound indicates the GroupRef is invalid
        ConditionReasonMCPRemoteProxyGroupRefNotFound = "GroupRefNotFound"

        // ConditionReasonMCPRemoteProxyGroupRefNotReady indicates the referenced MCPGroup is not in the Ready state
        ConditionReasonMCPRemoteProxyGroupRefNotReady = "GroupRefNotReady"

        // ConditionReasonMCPRemoteProxyToolConfigValid indicates the ToolConfigRef is valid
        ConditionReasonMCPRemoteProxyToolConfigValid = "ToolConfigValid"

        // ConditionReasonMCPRemoteProxyToolConfigNotFound indicates the referenced MCPToolConfig was not found
        ConditionReasonMCPRemoteProxyToolConfigNotFound = "ToolConfigNotFound"

        // ConditionReasonMCPRemoteProxyToolConfigFetchError indicates an error occurred fetching the MCPToolConfig
        ConditionReasonMCPRemoteProxyToolConfigFetchError = "ToolConfigFetchError"

        // ConditionReasonMCPRemoteProxyTelemetryConfigRefValid indicates the TelemetryConfigRef is valid
        ConditionReasonMCPRemoteProxyTelemetryConfigRefValid = "TelemetryConfigRefValid"

        // ConditionReasonMCPRemoteProxyTelemetryConfigRefNotFound indicates the referenced MCPTelemetryConfig was not found
        ConditionReasonMCPRemoteProxyTelemetryConfigRefNotFound = "TelemetryConfigRefNotFound"

        // ConditionReasonMCPRemoteProxyTelemetryConfigRefInvalid indicates the referenced MCPTelemetryConfig is invalid
        ConditionReasonMCPRemoteProxyTelemetryConfigRefInvalid = "TelemetryConfigRefInvalid"

        // ConditionReasonMCPRemoteProxyTelemetryConfigRefFetchError indicates an error occurred fetching the MCPTelemetryConfig
        ConditionReasonMCPRemoteProxyTelemetryConfigRefFetchError = "TelemetryConfigRefFetchError"

        // ConditionReasonMCPRemoteProxyExternalAuthConfigValid indicates the ExternalAuthConfigRef is valid
        ConditionReasonMCPRemoteProxyExternalAuthConfigValid = "ExternalAuthConfigValid"

        // ConditionReasonMCPRemoteProxyExternalAuthConfigNotFound indicates the referenced MCPExternalAuthConfig was not found
        ConditionReasonMCPRemoteProxyExternalAuthConfigNotFound = "ExternalAuthConfigNotFound"

        // ConditionReasonMCPRemoteProxyExternalAuthConfigFetchError indicates an error occurred fetching the MCPExternalAuthConfig
        ConditionReasonMCPRemoteProxyExternalAuthConfigFetchError = "ExternalAuthConfigFetchError"

        // ConditionReasonMCPRemoteProxyExternalAuthConfigMultiUpstream indicates multi-upstream is not supported
        // for MCPRemoteProxy (use VirtualMCPServer for multi-upstream).
        ConditionReasonMCPRemoteProxyExternalAuthConfigMultiUpstream = "MultiUpstreamNotSupported"

        // ConditionReasonMCPRemoteProxyAuthServerRefValid indicates the AuthServerRef is valid
        ConditionReasonMCPRemoteProxyAuthServerRefValid = "AuthServerRefValid"

        // ConditionReasonMCPRemoteProxyAuthServerRefNotFound indicates the referenced auth server config was not found
        ConditionReasonMCPRemoteProxyAuthServerRefNotFound = "AuthServerRefNotFound"

        // ConditionReasonMCPRemoteProxyAuthServerRefFetchError indicates an error occurred fetching the auth server config
        ConditionReasonMCPRemoteProxyAuthServerRefFetchError = "AuthServerRefFetchError"

        // ConditionReasonMCPRemoteProxyAuthServerRefInvalidKind indicates the authServerRef kind is not supported
        ConditionReasonMCPRemoteProxyAuthServerRefInvalidKind = "AuthServerRefInvalidKind"

        // ConditionReasonMCPRemoteProxyAuthServerRefInvalidType indicates the referenced config is not an embeddedAuthServer
        ConditionReasonMCPRemoteProxyAuthServerRefInvalidType = "AuthServerRefInvalidType"

        // ConditionReasonMCPRemoteProxyAuthServerRefInvalidConfig indicates the
        // referenced embedded auth server configuration is invalid.
        ConditionReasonMCPRemoteProxyAuthServerRefInvalidConfig = "AuthServerRefInvalidConfig"

        // ConditionReasonMCPRemoteProxyAuthServerRefMultiUpstream indicates multi-upstream is not supported
        ConditionReasonMCPRemoteProxyAuthServerRefMultiUpstream = "MultiUpstreamNotSupported"

        // ConditionReasonMCPRemoteProxyPodTemplateValid indicates PodTemplateSpec validation succeeded
        ConditionReasonMCPRemoteProxyPodTemplateValid = "ValidPodTemplateSpec"

        // ConditionReasonMCPRemoteProxyPodTemplateInvalid indicates PodTemplateSpec validation failed
        ConditionReasonMCPRemoteProxyPodTemplateInvalid = "InvalidPodTemplateSpec"

        // ConditionReasonMCPRemoteProxyCABundleRefValid indicates the CA bundle ref is valid and the ConfigMap exists
        ConditionReasonMCPRemoteProxyCABundleRefValid = "CABundleRefValid"

        // ConditionReasonMCPRemoteProxyCABundleRefNotFound indicates the referenced CA bundle ConfigMap was not found
        ConditionReasonMCPRemoteProxyCABundleRefNotFound = "CABundleRefNotFound"

        // ConditionReasonMCPRemoteProxyCABundleRefInvalid indicates the CA bundle ref configuration is invalid
        ConditionReasonMCPRemoteProxyCABundleRefInvalid = "CABundleRefInvalid"

        // ConditionReasonConfigurationValid indicates all configuration validations passed
        ConditionReasonConfigurationValid = "ConfigurationValid"

        // ConditionReasonOIDCIssuerInsecure indicates the OIDC issuer URL uses HTTP instead of HTTPS
        ConditionReasonOIDCIssuerInsecure = "OIDCIssuerInsecure"

        // ConditionReasonOIDCIssuerInvalid indicates the OIDC issuer URL is malformed
        ConditionReasonOIDCIssuerInvalid = "OIDCIssuerInvalid"

        // ConditionReasonAuthzPolicySyntaxInvalid indicates an inline Cedar policy has a syntax error
        ConditionReasonAuthzPolicySyntaxInvalid = "AuthzPolicySyntaxInvalid"

        // ConditionReasonAuthzConfigMapNotFound indicates the referenced authz ConfigMap was not found.
        // Shared with VirtualMCPServer (see virtualmcpserver_types.go); both reconcilers use this
        // reason when the ConfigMap itself is absent.
        ConditionReasonAuthzConfigMapNotFound = "AuthzConfigMapNotFound"

        // ConditionReasonAuthzConfigMapInvalid indicates the referenced authz ConfigMap was found
        // but its payload is missing/empty/malformed, fails validation, or does not contain a
        // Cedar-flavoured config. Shared with VirtualMCPServer; both reconcilers use this reason
        // to distinguish a malformed payload from a missing ConfigMap.
        ConditionReasonAuthzConfigMapInvalid = "AuthzConfigMapInvalid"

        // ConditionReasonHeaderSecretNotFound indicates a referenced header Secret was not found
        ConditionReasonHeaderSecretNotFound = "HeaderSecretNotFound"

        // ConditionReasonRemoteURLInvalid indicates the remoteUrl is malformed or has an invalid scheme
        ConditionReasonRemoteURLInvalid = "RemoteURLInvalid"

        // ConditionReasonJWKSURLInvalid indicates the JWKS URL is malformed or has an invalid scheme
        ConditionReasonJWKSURLInvalid = "JWKSURLInvalid"
)

//+kubebuilder:object:root=true
//+kubebuilder:storageversion
//+kubebuilder:subresource:status
//+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
//+kubebuilder:resource:shortName=rp;mcprp,categories=toolhive
//+kubebuilder:printcolumn:name="Phase",type="string",JSONPath=".status.phase"
//+kubebuilder:printcolumn:name="Remote URL",type="string",JSONPath=".spec.remoteUrl"
//+kubebuilder:printcolumn:name="URL",type="string",JSONPath=".status.url"
//+kubebuilder:printcolumn:name="Ready",type="string",JSONPath=".status.conditions[?(@.type=='Ready')].status"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"

// MCPRemoteProxy is the Schema for the mcpremoteproxies API
// It enables proxying remote MCP servers with authentication, authorization, audit logging, and tool filtering
type MCPRemoteProxy struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPRemoteProxySpec   `json:"spec,omitempty"`
        Status MCPRemoteProxyStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// MCPRemoteProxyList contains a list of MCPRemoteProxy
type MCPRemoteProxyList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPRemoteProxy `json:"items"`
}

func init() {
        SchemeBuilder.Register(&MCPRemoteProxy{}, &MCPRemoteProxyList{})
}

// GetName returns the name of the MCPRemoteProxy
func (m *MCPRemoteProxy) GetName() string {
        return m.Name
}

// GetNamespace returns the namespace of the MCPRemoteProxy
func (m *MCPRemoteProxy) GetNamespace() string {
        return m.Namespace
}

// GetProxyPort returns the proxy port of the MCPRemoteProxy
func (m *MCPRemoteProxy) GetProxyPort() int32 {
        if m.Spec.ProxyPort > 0 {
                return m.Spec.ProxyPort
        }
        return 8080
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        corev1 "k8s.io/api/core/v1"
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/runtime"

        ratelimittypes "github.com/stacklok/toolhive/pkg/ratelimit/types"
)

// Condition types for MCPServer
// Note: ConditionTypeReady is shared across multiple resources and defined in mcpremoteproxy_types.go
const (
        // ConditionGroupRefValidated indicates whether the GroupRef is valid
        ConditionGroupRefValidated = "GroupRefValidated"

        // ConditionPodTemplateValid indicates whether the PodTemplateSpec is valid
        ConditionPodTemplateValid = "PodTemplateValid"
)

const (
        // ConditionReasonReady indicates the MCPServer is ready
        ConditionReasonReady = "Ready"

        // ConditionReasonNotReady indicates the MCPServer is not ready
        ConditionReasonNotReady = "NotReady"
)

const (
        // ConditionReasonGroupRefValidated indicates the GroupRef is valid
        ConditionReasonGroupRefValidated = "GroupRefIsValid"

        // ConditionReasonGroupRefNotFound indicates the GroupRef is invalid
        ConditionReasonGroupRefNotFound = "GroupRefNotFound"

        // ConditionReasonGroupRefNotReady indicates the referenced MCPGroup is not in the Ready state
        ConditionReasonGroupRefNotReady = "GroupRefNotReady"
)

const (
        // ConditionReasonPodTemplateValid indicates PodTemplateSpec validation succeeded
        ConditionReasonPodTemplateValid = "ValidPodTemplateSpec"

        // ConditionReasonPodTemplateInvalid indicates PodTemplateSpec validation failed
        ConditionReasonPodTemplateInvalid = "InvalidPodTemplateSpec"
)

// Condition type for CA bundle validation
const (
        // ConditionCABundleRefValidated indicates whether the CABundleRef is valid
        ConditionCABundleRefValidated = "CABundleRefValidated"
)

// Condition type for MCPOIDCConfig reference validation
const (
        // ConditionOIDCConfigRefValidated indicates whether the OIDCConfigRef is valid
        ConditionOIDCConfigRefValidated = "OIDCConfigRefValidated"

        // ConditionAuthzConfigRefValidated indicates whether the AuthzConfigRef is valid
        ConditionAuthzConfigRefValidated = "AuthzConfigRefValidated"
)

const (
        // ConditionReasonOIDCConfigRefValid indicates the referenced MCPOIDCConfig is valid and ready
        ConditionReasonOIDCConfigRefValid = "OIDCConfigRefValid"

        // ConditionReasonOIDCConfigRefNotFound indicates the referenced MCPOIDCConfig was not found
        ConditionReasonOIDCConfigRefNotFound = "OIDCConfigRefNotFound"

        // ConditionReasonOIDCConfigRefNotValid indicates the referenced MCPOIDCConfig is not valid
        ConditionReasonOIDCConfigRefNotValid = "OIDCConfigRefNotValid"

        // ConditionReasonOIDCConfigRefError indicates an error occurred validating the OIDCConfigRef
        ConditionReasonOIDCConfigRefError = "OIDCConfigRefError"
)

const (
        // ConditionReasonAuthzConfigRefValid indicates the referenced MCPAuthzConfig is valid and ready
        ConditionReasonAuthzConfigRefValid = "AuthzConfigRefValid"

        // ConditionReasonAuthzConfigRefNotFound indicates the referenced MCPAuthzConfig was not found
        ConditionReasonAuthzConfigRefNotFound = "AuthzConfigRefNotFound"

        // ConditionReasonAuthzConfigRefNotValid indicates the referenced MCPAuthzConfig is not valid
        ConditionReasonAuthzConfigRefNotValid = "AuthzConfigRefNotValid"

        // ConditionReasonAuthzConfigRefError indicates an error occurred validating the AuthzConfigRef
        ConditionReasonAuthzConfigRefError = "AuthzConfigRefError"
)

const (
        // ConditionReasonCABundleRefValid indicates the CABundleRef is valid and the ConfigMap exists
        ConditionReasonCABundleRefValid = "CABundleRefValid"

        // ConditionReasonCABundleRefNotFound indicates the referenced ConfigMap was not found
        ConditionReasonCABundleRefNotFound = "CABundleRefNotFound"

        // ConditionReasonCABundleRefInvalid indicates the CABundleRef configuration is invalid
        ConditionReasonCABundleRefInvalid = "CABundleRefInvalid"
)

const (
        // ConditionTypeExternalAuthConfigValidated indicates whether the ExternalAuthConfig is valid
        ConditionTypeExternalAuthConfigValidated = "ExternalAuthConfigValidated"

        // ConditionTypeWebhookConfigValidated indicates whether the WebhookConfig is valid
        ConditionTypeWebhookConfigValidated = "WebhookConfigValidated"

        // ConditionTypeAuthzPrimaryUpstreamProviderIgnored is an advisory condition set
        // when spec.authzConfig.inline.primaryUpstreamProvider is non-empty on a CR type
        // that has no embedded auth server (MCPServer / MCPRemoteProxy). The field has
        // no effect on those resources and is documented as VirtualMCPServer-only.
        //
        // Tied to the deprecated InlineAuthzConfig.PrimaryUpstreamProvider field
        // (see mcpserver_types.go). When that field is removed at end of the
        // deprecation cycle, this condition and ConditionReasonAuthzPrimaryUpstreamProviderIgnored
        // below should be removed in the same change: there is no other path that
        // fires this advisory.
        ConditionTypeAuthzPrimaryUpstreamProviderIgnored = "AuthzPrimaryUpstreamProviderIgnored"
)

const (
        // ConditionReasonExternalAuthConfigMultiUpstream indicates the ExternalAuthConfig has multiple upstreams,
        // which is not supported for MCPServer (use VirtualMCPServer for multi-upstream).
        ConditionReasonExternalAuthConfigMultiUpstream = "MultiUpstreamNotSupported"

        // ConditionReasonExternalAuthConfigValid indicates the referenced
        // MCPExternalAuthConfig passed every check the MCPServer controller applies
        // to it. Mirrors ConditionReasonMCPRemoteProxyExternalAuthConfigValid.
        ConditionReasonExternalAuthConfigValid = "ExternalAuthConfigValid"

        // ConditionReasonWebhookConfigInvalid indicates the referenced webhook config is invalid or missing
        ConditionReasonWebhookConfigInvalid = "WebhookConfigInvalid"

        // ConditionReasonAuthzPrimaryUpstreamProviderIgnored indicates that
        // primaryUpstreamProvider is set on a CR type without an embedded auth server,
        // where the field has no runtime effect.
        ConditionReasonAuthzPrimaryUpstreamProviderIgnored = "PrimaryUpstreamProviderIgnored"
)

const (
        // ConditionTypeAuthServerRefValidated indicates whether the AuthServerRef is valid
        ConditionTypeAuthServerRefValidated = "AuthServerRefValidated"
)

const (
        // ConditionReasonAuthServerRefValid indicates the referenced auth server config is valid
        ConditionReasonAuthServerRefValid = "AuthServerRefValid"

        // ConditionReasonAuthServerRefNotFound indicates the referenced auth server config was not found
        ConditionReasonAuthServerRefNotFound = "AuthServerRefNotFound"

        // ConditionReasonAuthServerRefFetchError indicates an error occurred fetching the auth server config
        ConditionReasonAuthServerRefFetchError = "AuthServerRefFetchError"

        // ConditionReasonAuthServerRefInvalidKind indicates the authServerRef kind is not supported
        ConditionReasonAuthServerRefInvalidKind = "AuthServerRefInvalidKind"

        // ConditionReasonAuthServerRefInvalidType indicates the referenced config is not an embeddedAuthServer
        ConditionReasonAuthServerRefInvalidType = "AuthServerRefInvalidType"

        // ConditionReasonAuthServerRefInvalidConfig indicates the referenced embedded auth server configuration is invalid.
        ConditionReasonAuthServerRefInvalidConfig = "AuthServerRefInvalidConfig"

        // ConditionReasonAuthServerRefMultiUpstream indicates multi-upstream is not supported
        ConditionReasonAuthServerRefMultiUpstream = "MultiUpstreamNotSupported"
)

// ConditionTelemetryConfigRefValidated indicates whether the TelemetryConfigRef is valid
const ConditionTelemetryConfigRefValidated = "TelemetryConfigRefValidated"

const (
        // ConditionReasonTelemetryConfigRefValid indicates the referenced MCPTelemetryConfig is valid
        ConditionReasonTelemetryConfigRefValid = "TelemetryConfigRefValid"

        // ConditionReasonTelemetryConfigRefNotFound indicates the referenced MCPTelemetryConfig was not found
        ConditionReasonTelemetryConfigRefNotFound = "TelemetryConfigRefNotFound"

        // ConditionReasonTelemetryConfigRefInvalid indicates the referenced MCPTelemetryConfig is not valid
        ConditionReasonTelemetryConfigRefInvalid = "TelemetryConfigRefInvalid"

        // ConditionReasonTelemetryConfigRefError indicates a transient error occurred fetching the config
        ConditionReasonTelemetryConfigRefError = "TelemetryConfigRefError"
)

// ConditionStdioReplicaCapped indicates spec.replicas was capped at 1 for stdio transport.
const ConditionStdioReplicaCapped = "StdioReplicaCapped"

const (
        // ConditionReasonStdioReplicaCapped is set when spec.replicas > 1 for a stdio transport.
        ConditionReasonStdioReplicaCapped = "StdioTransportCapAt1"
        // ConditionReasonStdioReplicaCapNotActive is set when the stdio replica cap does not apply.
        ConditionReasonStdioReplicaCapNotActive = "StdioReplicaCapNotActive"
)

// ConditionSessionStorageWarning indicates replicas > 1 but no Redis session storage is configured.
const ConditionSessionStorageWarning = "SessionStorageWarning"

const (
        // ConditionReasonSessionStorageMissing is set when replicas > 1 and no Redis session storage is configured.
        ConditionReasonSessionStorageMissing = "SessionStorageMissingForReplicas"
        // ConditionReasonSessionStorageConfigured is set when replicas > 1 and Redis session storage is configured.
        ConditionReasonSessionStorageConfigured = "SessionStorageConfigured"
        // ConditionReasonSessionStorageNotApplicable is set when replicas is nil or <= 1 and the warning is not active.
        ConditionReasonSessionStorageNotApplicable = "SessionStorageWarningNotApplicable"
)

// ConditionRateLimitConfigValid indicates whether the rate limit configuration is valid.
const ConditionRateLimitConfigValid = "RateLimitConfigValid"

const (
        // ConditionReasonRateLimitConfigValid indicates the rate limit configuration is valid.
        ConditionReasonRateLimitConfigValid = "RateLimitConfigValid"
        // ConditionReasonRateLimitPerUserRequiresAuth indicates perUser rate limiting requires authentication.
        ConditionReasonRateLimitPerUserRequiresAuth = "PerUserRequiresAuth"
        // ConditionReasonRateLimitNotApplicable indicates rate limiting is not configured.
        ConditionReasonRateLimitNotApplicable = "RateLimitNotApplicable"
)

// SessionStorageProviderRedis is the provider name for Redis-backed session storage.
const SessionStorageProviderRedis = "redis"

// MCPServerSpec defines the desired state of MCPServer
//
// +kubebuilder:validation:XValidation:rule="!(has(self.authzConfig) && has(self.authzConfigRef))",message="authzConfig and authzConfigRef are mutually exclusive; use authzConfigRef to reference a shared MCPAuthzConfig"
// +kubebuilder:validation:XValidation:rule="!has(self.rateLimiting) || (has(self.sessionStorage) && self.sessionStorage.provider == 'redis')",message="rateLimiting requires sessionStorage with provider 'redis'"
// +kubebuilder:validation:XValidation:rule="!(has(self.rateLimiting) && has(self.rateLimiting.perUser)) || has(self.oidcConfigRef) || has(self.externalAuthConfigRef)",message="rateLimiting.perUser requires authentication (oidcConfigRef or externalAuthConfigRef)"
// +kubebuilder:validation:XValidation:rule="!has(self.rateLimiting) || !has(self.rateLimiting.tools) || self.rateLimiting.tools.all(t, !has(t.perUser)) || has(self.oidcConfigRef) || has(self.externalAuthConfigRef)",message="per-tool perUser rate limiting requires authentication (oidcConfigRef or externalAuthConfigRef)"
//
//nolint:lll // CEL validation rules exceed line length limit
type MCPServerSpec struct {
        // Image is the container image for the MCP server
        // +kubebuilder:validation:Required
        Image string `json:"image"`

        // Transport is the transport method for the MCP server (stdio, streamable-http or sse)
        // +kubebuilder:validation:Enum=stdio;streamable-http;sse
        // +kubebuilder:default=stdio
        Transport string `json:"transport,omitempty"`

        // ProxyMode is the proxy mode for stdio transport (sse or streamable-http)
        // This setting is ONLY applicable when Transport is "stdio".
        // For direct transports (sse, streamable-http), this field is ignored.
        // The default value is applied by Kubernetes but will be ignored for non-stdio transports.
        // +kubebuilder:validation:Enum=sse;streamable-http
        // +kubebuilder:default=streamable-http
        // +optional
        ProxyMode string `json:"proxyMode,omitempty"`

        // ProxyPort is the port to expose the proxy runner on
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:validation:Maximum=65535
        // +kubebuilder:default=8080
        ProxyPort int32 `json:"proxyPort,omitempty"`

        // MCPPort is the port that MCP server listens to
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:validation:Maximum=65535
        // +optional
        MCPPort int32 `json:"mcpPort,omitempty"`

        // Args are additional arguments to pass to the MCP server
        // +listType=atomic
        // +optional
        Args []string `json:"args,omitempty"`

        // Env are environment variables to set in the MCP server container
        // +listType=map
        // +listMapKey=name
        // +optional
        Env []EnvVar `json:"env,omitempty"`

        // Volumes are volumes to mount in the MCP server container
        // +listType=map
        // +listMapKey=name
        // +optional
        Volumes []Volume `json:"volumes,omitempty"`

        // Resources defines the resource requirements for the MCP server container
        // +optional
        Resources ResourceRequirements `json:"resources,omitempty"`

        // Secrets are references to secrets to mount in the MCP server container
        // +listType=map
        // +listMapKey=name
        // +listMapKey=key
        // +optional
        Secrets []SecretRef `json:"secrets,omitempty"`

        // ServiceAccount is the name of an already existing service account to use by the MCP server.
        // If not specified, a ServiceAccount will be created automatically and used by the MCP server.
        // +optional
        ServiceAccount *string `json:"serviceAccount,omitempty"`

        // PermissionProfile is deprecated and retained for API compatibility. It is
        // not supported by the Kubernetes operator and has no effect.
        // +optional
        PermissionProfile *PermissionProfileRef `json:"permissionProfile,omitempty"`

        // PodTemplateSpec defines the pod template to use for the MCP server
        // This allows for customizing the pod configuration beyond what is provided by the other fields.
        // Note that to modify the specific container the MCP server runs in, you must specify
        // the `mcp` container name in the PodTemplateSpec.
        // This field accepts a PodTemplateSpec object as JSON/YAML.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        PodTemplateSpec *runtime.RawExtension `json:"podTemplateSpec,omitempty"`

        // ResourceOverrides allows overriding annotations and labels for resources created by the operator
        // +optional
        ResourceOverrides *ResourceOverrides `json:"resourceOverrides,omitempty"`

        // OIDCConfigRef references a shared MCPOIDCConfig resource for OIDC authentication.
        // The referenced MCPOIDCConfig must exist in the same namespace as this MCPServer.
        // Per-server overrides (audience, scopes) are specified here; shared provider config
        // lives in the MCPOIDCConfig resource.
        //
        // SECURITY: if this field is omitted and no other authentication source is configured,
        // the proxy runs UNAUTHENTICATED. It accepts every request that can reach its port and
        // forwards it to the MCP server under a synthetic local-user identity, with no token or
        // credential check. Set this field to enforce identity-based access control per request.
        // +optional
        OIDCConfigRef *MCPOIDCConfigReference `json:"oidcConfigRef,omitempty"`

        // AuthzConfig defines authorization policy configuration for the MCP server.
        // AuthzConfig and AuthzConfigRef are mutually exclusive.
        // +optional
        AuthzConfig *AuthzConfigRef `json:"authzConfig,omitempty"`

        // AuthzConfigRef references a shared MCPAuthzConfig resource for authorization.
        // The referenced MCPAuthzConfig must exist in the same namespace as this MCPServer.
        // Mutually exclusive with authzConfig.
        // +optional
        AuthzConfigRef *MCPAuthzConfigReference `json:"authzConfigRef,omitempty"`

        // Audit defines audit logging configuration for the MCP server
        // +optional
        Audit *AuditConfig `json:"audit,omitempty"`

        // ToolConfigRef references a MCPToolConfig resource for tool filtering and renaming.
        // The referenced MCPToolConfig must exist in the same namespace as this MCPServer.
        // Cross-namespace references are not supported for security and isolation reasons.
        // +optional
        ToolConfigRef *ToolConfigRef `json:"toolConfigRef,omitempty"`

        // ExternalAuthConfigRef references a MCPExternalAuthConfig resource for external authentication.
        // The referenced MCPExternalAuthConfig must exist in the same namespace as this MCPServer.
        // +optional
        ExternalAuthConfigRef *ExternalAuthConfigRef `json:"externalAuthConfigRef,omitempty"`

        // WebhookConfigRef references a MCPWebhookConfig resource for webhook middleware configuration.
        // The referenced MCPWebhookConfig must exist in the same namespace as this MCPServer.
        // +optional
        WebhookConfigRef *WebhookConfigRef `json:"webhookConfigRef,omitempty"`

        // AuthServerRef optionally references a resource that configures an embedded
        // OAuth 2.0/OIDC authorization server to authenticate MCP clients.
        // Currently the only supported kind is MCPExternalAuthConfig (type: embeddedAuthServer).
        // +optional
        AuthServerRef *AuthServerRef `json:"authServerRef,omitempty"`

        // TelemetryConfigRef references an MCPTelemetryConfig resource for shared telemetry configuration.
        // The referenced MCPTelemetryConfig must exist in the same namespace as this MCPServer.
        // Cross-namespace references are not supported for security and isolation reasons.
        // +optional
        TelemetryConfigRef *MCPTelemetryConfigReference `json:"telemetryConfigRef,omitempty"`

        // TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies
        // When enabled, the proxy will use X-Forwarded-Proto, X-Forwarded-Host, X-Forwarded-Port,
        // and X-Forwarded-Prefix headers to construct endpoint URLs
        // +kubebuilder:default=false
        // +optional
        TrustProxyHeaders bool `json:"trustProxyHeaders,omitempty"`

        // EndpointPrefix is the path prefix to prepend to SSE endpoint URLs.
        // This is used to handle path-based ingress routing scenarios where the ingress
        // strips a path prefix before forwarding to the backend.
        // +optional
        EndpointPrefix string `json:"endpointPrefix,omitempty"`

        // MaxRequestBodySize is the maximum inbound MCP proxy request body size in bytes.
        // Zero uses the default limit of 8 MiB.
        // +kubebuilder:validation:Minimum=0
        // +optional
        MaxRequestBodySize int64 `json:"maxRequestBodySize,omitempty"`

        // GroupRef references the MCPGroup this server belongs to.
        // The referenced MCPGroup must be in the same namespace.
        // +optional
        GroupRef *MCPGroupRef `json:"groupRef,omitempty"`

        // SessionAffinity controls whether the Service routes repeated client connections to the same pod.
        // MCP protocols (SSE, streamable-http) are stateful, so ClientIP is the default.
        // Set to "None" for stateless servers or when using an external load balancer with its own affinity.
        // +kubebuilder:validation:Enum=ClientIP;None
        // +kubebuilder:default=ClientIP
        // +optional
        SessionAffinity string `json:"sessionAffinity,omitempty"`

        // Replicas is the desired number of proxy runner (thv run) pod replicas.
        // MCPServer creates two separate Deployments: one for the proxy runner and one
        // for the MCP server backend. This field controls the proxy runner Deployment.
        // When nil, the operator does not set Deployment.Spec.Replicas, leaving replica
        // management to an HPA or other external controller.
        // +kubebuilder:validation:Minimum=0
        // +optional
        Replicas *int32 `json:"replicas,omitempty"`

        // BackendReplicas is the desired number of MCP server backend pod replicas.
        // This controls the backend Deployment (the MCP server container itself),
        // independent of the proxy runner controlled by Replicas.
        // When nil, the operator does not set Deployment.Spec.Replicas, leaving replica
        // management to an HPA or other external controller.
        // +kubebuilder:validation:Minimum=0
        // +optional
        BackendReplicas *int32 `json:"backendReplicas,omitempty"`

        // SessionStorage configures session storage for stateful horizontal scaling.
        // When nil, no session storage is configured.
        // +optional
        SessionStorage *SessionStorageConfig `json:"sessionStorage,omitempty"`

        // RateLimiting defines rate limiting configuration for the MCP server.
        // Requires Redis session storage to be configured for distributed rate limiting.
        // +optional
        RateLimiting *ratelimittypes.RateLimitConfig `json:"rateLimiting,omitempty"`

        // ProxyReadTimeout bounds how long the proxy spends reading a full request
        // (headers + body), mitigating slow-upload connection exhaustion. Applies to
        // all transports. Defaults to 30s if not specified. Example: "1m".
        // +kubebuilder:validation:Type=string
        // +kubebuilder:validation:Format=duration
        // +kubebuilder:validation:XValidation:rule="duration(self) >= duration('0s')",message="proxyReadTimeout must be non-negative"
        // +optional
        ProxyReadTimeout *metav1.Duration `json:"proxyReadTimeout,omitempty"`
}

// ResourceOverrides defines overrides for annotations and labels on created resources
type ResourceOverrides struct {
        // ProxyDeployment defines overrides for the Proxy Deployment resource (toolhive proxy)
        // +optional
        ProxyDeployment *ProxyDeploymentOverrides `json:"proxyDeployment,omitempty"`

        // ProxyService defines overrides for the Proxy Service resource (points to the proxy deployment)
        // +optional
        ProxyService *ResourceMetadataOverrides `json:"proxyService,omitempty"`
}

// ProxyDeploymentOverrides defines overrides specific to the proxy deployment
type ProxyDeploymentOverrides struct {
        // ResourceMetadataOverrides is embedded to inherit annotations and labels fields
        ResourceMetadataOverrides `json:",inline"` // nolint:revive

        PodTemplateMetadataOverrides *ResourceMetadataOverrides `json:"podTemplateMetadataOverrides,omitempty"`

        // Env are environment variables to set in the proxy container (thv run process)
        // These affect the toolhive proxy itself, not the MCP server it manages
        // Use TOOLHIVE_DEBUG=true to enable debug logging in the proxy
        // +listType=map
        // +listMapKey=name
        // +optional
        Env []EnvVar `json:"env,omitempty"`

        // ImagePullSecrets allows specifying image pull secrets for the proxy runner
        // These are applied to both the Deployment and the ServiceAccount
        // +listType=atomic
        // +optional
        ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,omitempty"`

        // NodeSelector constrains the proxy pod to nodes with matching labels.
        // Mirrors the scheduling control podTemplateSpec gives the MCP server pod, so
        // the proxy can be steered onto the same nodes (e.g. a pre-warmed pool).
        // On MCPRemoteProxy, spec.podTemplateSpec also reaches the proxy pod: the two
        // maps are merged, and podTemplateSpec wins on keys set in both.
        // +optional
        NodeSelector map[string]string `json:"nodeSelector,omitempty"`

        // Tolerations allow the proxy pod to schedule onto tainted nodes, such as a
        // dedicated pre-warmed pool.
        // On MCPRemoteProxy, spec.podTemplateSpec also reaches the proxy pod, and this
        // list is atomic: a podTemplateSpec that sets tolerations replaces this field
        // rather than adding to it.
        // +listType=atomic
        // +optional
        Tolerations []corev1.Toleration `json:"tolerations,omitempty"`

        // Affinity sets node/pod affinity and anti-affinity for the proxy pod.
        // On MCPRemoteProxy, spec.podTemplateSpec also reaches the proxy pod: the two
        // are merged per sub-field, and podTemplateSpec wins on sub-fields set in both.
        // +optional
        Affinity *corev1.Affinity `json:"affinity,omitempty"`
}

// ResourceMetadataOverrides defines metadata overrides for a resource
type ResourceMetadataOverrides struct {
        // Annotations to add or override on the resource
        // +optional
        Annotations map[string]string `json:"annotations,omitempty"`

        // Labels to add or override on the resource
        // +optional
        Labels map[string]string `json:"labels,omitempty"`
}

// EnvVar represents an environment variable in a container
type EnvVar struct {
        // Name of the environment variable
        // +kubebuilder:validation:Required
        Name string `json:"name"`

        // Value of the environment variable
        // +kubebuilder:validation:Required
        Value string `json:"value"`
}

// Volume represents a volume to mount in a container
type Volume struct {
        // Name is the name of the volume
        // +kubebuilder:validation:Required
        Name string `json:"name"`

        // HostPath is the path on the host to mount
        // +kubebuilder:validation:Required
        HostPath string `json:"hostPath"`

        // MountPath is the path in the container to mount to
        // +kubebuilder:validation:Required
        MountPath string `json:"mountPath"`

        // ReadOnly specifies whether the volume should be mounted read-only
        // +kubebuilder:default=false
        // +optional
        ReadOnly bool `json:"readOnly,omitempty"`
}

// ResourceRequirements describes the compute resource requirements
type ResourceRequirements struct {
        // Limits describes the maximum amount of compute resources allowed
        // +optional
        Limits ResourceList `json:"limits,omitempty"`

        // Requests describes the minimum amount of compute resources required
        // +optional
        Requests ResourceList `json:"requests,omitempty"`
}

// ResourceList is a set of (resource name, quantity) pairs
type ResourceList struct {
        // CPU is the CPU limit in cores (e.g., "500m" for 0.5 cores)
        // +optional
        CPU string `json:"cpu,omitempty"`

        // Memory is the memory limit in bytes (e.g., "64Mi" for 64 megabytes)
        // +optional
        Memory string `json:"memory,omitempty"`
}

// SecretRef is a reference to a secret
type SecretRef struct {
        // Name is the name of the secret
        // +kubebuilder:validation:Required
        Name string `json:"name"`

        // Key is the key in the secret itself
        // +kubebuilder:validation:Required
        Key string `json:"key"`

        // TargetEnvName is the environment variable to be used when setting up the secret in the MCP server
        // If left unspecified, it defaults to the key
        // +optional
        TargetEnvName string `json:"targetEnvName,omitempty"`
}

// SessionStorageConfig defines session storage configuration for horizontal scaling.
//
// This is the CRD/K8s-aware surface: it uses SecretKeyRef for secret resolution.
// The reconciler resolves PasswordRef to a plain string and builds a
// session.RedisConfig (pkg/transport/session) for the actual storage backend.
// The operator also populates pkg/vmcp/config.SessionStorageConfig (without PasswordRef)
// into the vMCP ConfigMap so the vMCP process receives connection parameters at startup.
//
// +kubebuilder:validation:XValidation:rule="self.provider == 'redis' ? has(self.address) : true",message="address is required"
// +kubebuilder:validation:XValidation:rule="!has(self.tls) || self.provider == 'redis'",message="tls is only supported when provider is redis"
// +kubebuilder:validation:XValidation:rule="!has(self.tls) || !has(self.tls.caCertSecretRef) || !has(self.tls.insecureSkipVerify) || !self.tls.insecureSkipVerify",message="tls.insecureSkipVerify disables certificate verification, so tls.caCertSecretRef would be ignored; set only one"
// +kubebuilder:validation:XValidation:rule="!has(self.tls) || !has(self.tls.caCertSecretRef) || (size(self.tls.caCertSecretRef.name) > 0 && size(self.tls.caCertSecretRef.key) > 0)",message="tls.caCertSecretRef requires a non-empty name and key"
//
//nolint:lll // CEL validation rules exceed line length limit
type SessionStorageConfig struct {
        // Provider is the session storage backend type
        // +kubebuilder:validation:Enum=memory;redis
        // +kubebuilder:validation:Required
        Provider string `json:"provider"`

        // Address is the Redis server address (required when provider is redis)
        // +kubebuilder:validation:MinLength=1
        // +optional
        Address string `json:"address,omitempty"`

        // DB is the Redis database number
        // +kubebuilder:validation:Minimum=0
        // +kubebuilder:default=0
        // +optional
        DB int32 `json:"db,omitempty"`

        // KeyPrefix is an optional prefix for all Redis keys used by ToolHive
        // +optional
        KeyPrefix string `json:"keyPrefix,omitempty"`

        // PasswordRef is a reference to a Secret key containing the Redis password
        // +optional
        PasswordRef *SecretKeyRef `json:"passwordRef,omitempty"`

        // TLS enables TLS for the Redis session storage connection and, when rate
        // limiting is configured, for the rate-limit connection to the same Redis.
        // When omitted, the connections are plaintext and the Redis password and
        // session data cross the network unencrypted. An empty object enables TLS
        // with server certificate verification against the system root CAs; set
        // caCertSecretRef for a private CA. A failed handshake never falls back to
        // plaintext. Only used when provider is redis.
        // +optional
        TLS *RedisTLSConfig `json:"tls,omitempty"`
}

// RateLimitConfig defines rate limiting configuration for an MCP server.
// +gendoc
type RateLimitConfig = ratelimittypes.RateLimitConfig

// RateLimitBucket defines a token bucket configuration with a maximum capacity
// and a refill period. Used by both shared and per-user rate limits.
// +gendoc
type RateLimitBucket = ratelimittypes.RateLimitBucket

// ToolRateLimitConfig defines rate limits for a specific tool.
// +gendoc
type ToolRateLimitConfig = ratelimittypes.ToolRateLimitConfig

// Permission profile types
const (
        // PermissionProfileTypeBuiltin is the deprecated type for built-in permission
        // profiles. Permission profiles are not supported for MCPServer workloads.
        PermissionProfileTypeBuiltin = "builtin"

        // PermissionProfileTypeConfigMap is the deprecated type for permission profiles
        // stored in ConfigMaps. Permission profiles are not supported for MCPServer workloads.
        PermissionProfileTypeConfigMap = "configmap"
)

// Authorization configuration types
const (
        // AuthzConfigTypeConfigMap is the type for authorization configuration stored in ConfigMaps
        AuthzConfigTypeConfigMap = "configMap"

        // AuthzConfigTypeInline is the type for inline authorization configuration
        AuthzConfigTypeInline = "inline"
)

// PermissionProfileRef defines a deprecated reference to a permission profile.
// Permission profiles are not supported for MCPServer workloads.
type PermissionProfileRef struct {
        // Type is the type of permission profile reference
        // +kubebuilder:validation:Enum=builtin;configmap
        // +kubebuilder:default=builtin
        Type string `json:"type"`

        // Name is the name of the permission profile
        // If Type is "builtin", Name must be one of: "none", "network"
        // If Type is "configmap", Name is the name of the ConfigMap
        // +kubebuilder:validation:Required
        Name string `json:"name"`

        // Key is the key in the ConfigMap that contains the permission profile
        // Only used when Type is "configmap"
        // +optional
        Key string `json:"key,omitempty"`
}

// PermissionProfileSpec defines deprecated permissions for an MCP server.
// Permission profiles are not supported for MCPServer workloads.
// +gendoc
type PermissionProfileSpec struct {
        // Read is a list of paths that the MCP server can read from
        // +listType=atomic
        // +optional
        Read []string `json:"read,omitempty"`

        // Write is a list of paths that the MCP server can write to
        // +listType=atomic
        // +optional
        Write []string `json:"write,omitempty"`

        // Network defines the network permissions for the MCP server
        // +optional
        Network *NetworkPermissions `json:"network,omitempty"`
}

// NetworkPermissions defines deprecated network permissions for an MCP server.
// Permission profiles are not supported for MCPServer workloads.
type NetworkPermissions struct {
        // Mode specifies the network mode for the container (e.g., "host", "bridge", "none")
        // When empty, the default container runtime network mode is used
        // +optional
        Mode string `json:"mode,omitempty"`

        // Outbound defines the outbound network permissions
        // +optional
        Outbound *OutboundNetworkPermissions `json:"outbound,omitempty"`
}

// OutboundNetworkPermissions defines the outbound network permissions
type OutboundNetworkPermissions struct {
        // InsecureAllowAll allows all outbound network connections (not recommended)
        // +kubebuilder:default=false
        // +optional
        InsecureAllowAll bool `json:"insecureAllowAll,omitempty"`

        // AllowHost is a list of hosts to allow connections to
        // +listType=set
        // +optional
        AllowHost []string `json:"allowHost,omitempty"`

        // AllowPort is a list of ports to allow connections to
        // +listType=set
        // +optional
        AllowPort []int32 `json:"allowPort,omitempty"`
}

// CABundleSource defines a source for CA certificate bundles.
//
// +kubebuilder:validation:XValidation:rule="has(self.configMapRef) && size(self.configMapRef.name) > 0",message="configMapRef.name is required and must be non-empty"
//
//nolint:lll // CEL validation rule exceeds line length limit
type CABundleSource struct {
        // ConfigMapRef references a ConfigMap containing the CA certificate bundle.
        // The ConfigMap key is required by the API. If omitted in a stored object, it
        // defaults to "ca.crt" for backwards compatibility.
        // +optional
        ConfigMapRef *corev1.ConfigMapKeySelector `json:"configMapRef,omitempty"`
}

// AuthzConfigRef defines a reference to authorization configuration
//
// +kubebuilder:validation:XValidation:rule="self.type == 'configMap' ? has(self.configMap) : !has(self.configMap)",message="configMap must be set when type is 'configMap', and must not be set otherwise"
// +kubebuilder:validation:XValidation:rule="self.type == 'inline' ? has(self.inline) : !has(self.inline)",message="inline must be set when type is 'inline', and must not be set otherwise"
//
//nolint:lll // CEL validation rules exceed line length limit
type AuthzConfigRef struct {
        // Type is the type of authorization configuration
        // +kubebuilder:validation:Enum=configMap;inline
        // +kubebuilder:default=configMap
        Type string `json:"type"`

        // ConfigMap references a ConfigMap containing authorization configuration
        // Only used when Type is "configMap"
        // +optional
        ConfigMap *ConfigMapAuthzRef `json:"configMap,omitempty"`

        // Inline contains direct authorization configuration
        // Only used when Type is "inline"
        // +optional
        Inline *InlineAuthzConfig `json:"inline,omitempty"`

        // GroupClaimName is the JWT claim key that contains group membership for the
        // principal. When set, takes priority over the well-known defaults
        // ("groups", "roles", "cognito:groups"). Use this for IDPs that place
        // groups under a URI-style claim (e.g. "https://example.com/groups"). When
        // Type is "configMap", a group_claim_name entry in the referenced ConfigMap
        // is overridden by this field if both are set.
        // +optional
        // +kubebuilder:validation:MaxLength=253
        GroupClaimName string `json:"groupClaimName,omitempty"`

        // RoleClaimName is the JWT claim key that contains role membership for the
        // principal. When set, the claim is extracted separately from GroupClaimName
        // and both are mapped to the configured GroupEntityType. When Type is
        // "configMap", a role_claim_name entry in the referenced ConfigMap is
        // overridden by this field if both are set.
        // +optional
        // +kubebuilder:validation:MaxLength=253
        RoleClaimName string `json:"roleClaimName,omitempty"`

        // GroupEntityType is the Cedar entity type name used for principal parent
        // UIDs synthesised from JWT group/role claims. Defaults to "THVGroup" when
        // empty. Must match the entity type used in the static entity store for
        // transitive `in` checks (e.g. `ClaimGroup → PlatformRole`) to resolve.
        // Namespaced names (`Foo::Bar`) are not yet supported. When Type is
        // "configMap", a group_entity_type entry in the referenced ConfigMap is
        // overridden by this field if both are set.
        // +optional
        // +kubebuilder:validation:MaxLength=63
        // +kubebuilder:validation:Pattern=`^[A-Za-z_][A-Za-z0-9_]*$`
        GroupEntityType string `json:"groupEntityType,omitempty"`
}

// DeprecatedInlinePrimaryUpstreamProvider returns the legacy inline
// PrimaryUpstreamProvider value, or "" when the field or the AuthzConfigRef
// is nil. The field has moved to spec.authServerConfig.primaryUpstreamProvider
// on VirtualMCPServer; this accessor is the single read point for the
// deprecated location so callers can emit a deprecation warning when it
// returns a non-empty value.
func (r *AuthzConfigRef) DeprecatedInlinePrimaryUpstreamProvider() string {
        if r == nil || r.Inline == nil {
                return ""
        }
        return r.Inline.PrimaryUpstreamProvider
}

// ConfigMapAuthzRef references a ConfigMap containing authorization configuration
type ConfigMapAuthzRef struct {
        // Name is the name of the ConfigMap
        // +kubebuilder:validation:Required
        Name string `json:"name"`

        // Key is the key in the ConfigMap that contains the authorization configuration
        // +kubebuilder:default=authz.json
        // +optional
        Key string `json:"key,omitempty"`
}

// ExternalAuthConfigRef defines a reference to a MCPExternalAuthConfig resource.
// The referenced MCPExternalAuthConfig must be in the same namespace as the MCPServer.
type ExternalAuthConfigRef struct {
        // Name is the name of the MCPExternalAuthConfig resource
        // +kubebuilder:validation:Required
        Name string `json:"name"`
}

// AuthServerRef defines a reference to a resource that configures an embedded
// OAuth 2.0/OIDC authorization server. Currently only MCPExternalAuthConfig is supported;
// the enum will be extended when a dedicated auth server CRD is introduced.
type AuthServerRef struct {
        // Kind identifies the type of the referenced resource.
        // +kubebuilder:validation:Enum=MCPExternalAuthConfig
        // +kubebuilder:default=MCPExternalAuthConfig
        Kind string `json:"kind"`

        // Name is the name of the referenced resource in the same namespace.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Name string `json:"name"`
}

// WebhookConfigRef defines a reference to a MCPWebhookConfig resource.
// The referenced MCPWebhookConfig must be in the same namespace as the MCPServer.
type WebhookConfigRef struct {
        // Name is the name of the MCPWebhookConfig resource
        // +kubebuilder:validation:Required
        Name string `json:"name"`
}

// ToolConfigRef defines a reference to a MCPToolConfig resource.
// The referenced MCPToolConfig must be in the same namespace as the MCPServer.
type ToolConfigRef struct {
        // Name is the name of the MCPToolConfig resource in the same namespace
        // +kubebuilder:validation:Required
        Name string `json:"name"`
}

// MCPGroupRef defines a reference to an MCPGroup resource.
// The referenced MCPGroup must be in the same namespace.
type MCPGroupRef struct {
        // Name is the name of the MCPGroup resource in the same namespace
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Name string `json:"name"`
}

// GetName returns the name, or empty string if the receiver is nil.
func (r *MCPGroupRef) GetName() string {
        if r == nil {
                return ""
        }
        return r.Name
}

// InlineAuthzConfig contains direct authorization configuration.
//
// Source-agnostic Cedar JWT-claim mapping settings (GroupClaimName,
// RoleClaimName, GroupEntityType) live on the parent AuthzConfigRef so they
// work the same way for inline and configMap-sourced authz.
type InlineAuthzConfig struct {
        // Policies is a list of Cedar policy strings
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +listType=atomic
        Policies []string `json:"policies"`

        // EntitiesJSON is a JSON string representing Cedar entities. Required when
        // transitive policies (e.g. `ClaimGroup → PlatformRole`) need a static
        // entity store; defaults to "[]".
        // +kubebuilder:default="[]"
        // +optional
        EntitiesJSON string `json:"entitiesJson,omitempty"`

        // PrimaryUpstreamProvider names the upstream IDP whose access token's
        // claims Cedar should evaluate.
        //
        // Deprecated: on VirtualMCPServer this field has moved to
        // spec.authServerConfig.primaryUpstreamProvider. The old location is
        // still read for one release for backward compatibility; the
        // VirtualMCPServer controller emits an AuthzPrimaryUpstreamProviderDeprecated
        // Warning event whenever it is consumed, and removal is planned for the
        // release after the deprecation cycle.
        //
        // On MCPServer and MCPRemoteProxy this field has always been a structural
        // no-op (those CRDs do not run an embedded auth server). Setting it
        // continues to surface the AuthzPrimaryUpstreamProviderIgnored advisory
        // condition; the deprecation does not change that behaviour.
        // +optional
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=63
        // +kubebuilder:validation:Pattern=`^[a-z0-9]([a-z0-9-]*[a-z0-9])?$`
        PrimaryUpstreamProvider string `json:"primaryUpstreamProvider,omitempty"`
}

// AuditConfig defines audit logging configuration for the MCP server
type AuditConfig struct {
        // Enabled controls whether audit logging is enabled
        // When true, enables audit logging with default configuration
        // +kubebuilder:default=false
        // +optional
        Enabled bool `json:"enabled,omitempty"`
}

// PrometheusConfig defines Prometheus-specific configuration
type PrometheusConfig struct {
        // Enabled controls whether Prometheus metrics endpoint is exposed
        // +kubebuilder:default=false
        // +optional
        Enabled bool `json:"enabled,omitempty"`

        // MetricsOnTransportPort controls whether /metrics is ALSO served on the main
        // transport port, in addition to the dedicated diagnostics port. It exists to
        // give deployments a migration window: while true, an existing scrape
        // configuration aimed at the transport port keeps working, so a scraper can be
        // moved to the diagnostics port and verified before the old location goes away.
        //
        // Leave unset to follow the current default, which is true during the migration
        // window and becomes false when it closes. Setting it explicitly opts out of
        // that change: an explicit value is honoured before and after the cutover.
        //
        // See https://github.com/stacklok/toolhive/issues/6384 for the removal timeline.
        //
        // +optional
        MetricsOnTransportPort *bool `json:"metricsOnTransportPort,omitempty"`
}

// OpenTelemetryTracingConfig defines OpenTelemetry tracing configuration
type OpenTelemetryTracingConfig struct {
        // Enabled controls whether OTLP tracing is sent
        // +kubebuilder:default=false
        // +optional
        Enabled bool `json:"enabled,omitempty"`

        // SamplingRate is the trace sampling rate (0.0-1.0)
        // +kubebuilder:default="0.05"
        // +kubebuilder:validation:Pattern=`^(0(\.\d+)?|1(\.0+)?)$`
        // +optional
        SamplingRate string `json:"samplingRate,omitempty"`
}

// OpenTelemetryMetricsConfig defines OpenTelemetry metrics configuration
type OpenTelemetryMetricsConfig struct {
        // Enabled controls whether OTLP metrics are sent
        // +kubebuilder:default=false
        // +optional
        Enabled bool `json:"enabled,omitempty"`
}

// MCPServerStatus defines the observed state of MCPServer
type MCPServerStatus struct {
        // Conditions represent the latest available observations of the MCPServer's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration reflects the generation most recently observed by the controller
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // ToolConfigHash stores the hash of the referenced ToolConfig for change detection
        // +optional
        ToolConfigHash string `json:"toolConfigHash,omitempty"`

        // ExternalAuthConfigHash is the hash of the referenced MCPExternalAuthConfig spec
        // +optional
        ExternalAuthConfigHash string `json:"externalAuthConfigHash,omitempty"`

        // AuthServerConfigHash is the hash of the referenced authServerRef spec,
        // used to detect configuration changes and trigger reconciliation.
        // +optional
        AuthServerConfigHash string `json:"authServerConfigHash,omitempty"`

        // AuthzConfigHash is the hash of the referenced MCPAuthzConfig spec for change detection
        // +optional
        AuthzConfigHash string `json:"authzConfigHash,omitempty"`

        // OIDCConfigHash is the hash of the referenced MCPOIDCConfig spec for change detection
        // +optional
        OIDCConfigHash string `json:"oidcConfigHash,omitempty"`

        // TelemetryConfigHash is the hash of the referenced MCPTelemetryConfig spec for change detection
        // +optional
        TelemetryConfigHash string `json:"telemetryConfigHash,omitempty"`

        // WebhookConfigHash is the hash of the referenced MCPWebhookConfig spec
        // +optional
        WebhookConfigHash string `json:"webhookConfigHash,omitempty"`

        // URL is the URL where the MCP server can be accessed
        // +optional
        URL string `json:"url,omitempty"`

        // Phase is the current phase of the MCPServer
        // +optional
        Phase MCPServerPhase `json:"phase,omitempty"`

        // Message provides additional information about the current phase
        // +optional
        Message string `json:"message,omitempty"`

        // ReadyReplicas is the number of ready proxy replicas
        // +optional
        ReadyReplicas int32 `json:"readyReplicas,omitempty"`
}

// MCPServerPhase is the phase of the MCPServer
// +kubebuilder:validation:Enum=Pending;Ready;Failed;Terminating;Stopped
type MCPServerPhase string

const (
        // MCPServerPhasePending means the MCPServer is being created
        MCPServerPhasePending MCPServerPhase = "Pending"

        // MCPServerPhaseReady means the MCPServer is ready
        MCPServerPhaseReady MCPServerPhase = "Ready"

        // MCPServerPhaseFailed means the MCPServer failed to start
        MCPServerPhaseFailed MCPServerPhase = "Failed"

        // MCPServerPhaseTerminating means the MCPServer is being deleted
        MCPServerPhaseTerminating MCPServerPhase = "Terminating"

        // MCPServerPhaseStopped means the MCPServer is scaled to zero
        MCPServerPhaseStopped MCPServerPhase = "Stopped"
)

//+kubebuilder:object:root=true
//+kubebuilder:storageversion
//+kubebuilder:subresource:status
//+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
//+kubebuilder:resource:categories=toolhive
//+kubebuilder:printcolumn:name="Status",type="string",JSONPath=".status.phase"
//+kubebuilder:printcolumn:name="Ready",type="string",JSONPath=".status.conditions[?(@.type=='Ready')].status"
//+kubebuilder:printcolumn:name="Replicas",type="integer",JSONPath=".status.readyReplicas"
//+kubebuilder:printcolumn:name="URL",type="string",JSONPath=".status.url"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"

// MCPServer is the Schema for the mcpservers API
type MCPServer struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPServerSpec   `json:"spec,omitempty"`
        Status MCPServerStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// MCPServerList contains a list of MCPServer
type MCPServerList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPServer `json:"items"`
}

// GetName returns the name of the MCPServer
func (m *MCPServer) GetName() string {
        return m.Name
}

// GetNamespace returns the namespace of the MCPServer
func (m *MCPServer) GetNamespace() string {
        return m.Namespace
}

// GetProxyPort returns the proxy port of the MCPServer
func (m *MCPServer) GetProxyPort() int32 {
        if m.Spec.ProxyPort > 0 {
                return m.Spec.ProxyPort
        }
        return 8080
}

// GetMCPPort returns the MCP port of the MCPServer
func (m *MCPServer) GetMCPPort() int32 {
        if m.Spec.MCPPort > 0 {
                return m.Spec.MCPPort
        }
        return 8080
}

func init() {
        SchemeBuilder.Register(&MCPServer{}, &MCPServerList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// MCPServerEntrySpec defines the desired state of MCPServerEntry.
// MCPServerEntry is a zero-infrastructure catalog entry that declares a remote MCP
// server endpoint. Unlike MCPRemoteProxy, it creates no pods, services, or deployments.
type MCPServerEntrySpec struct {
        // RemoteURL is the URL of the remote MCP server.
        // Both HTTP and HTTPS schemes are accepted at admission time.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https?://`
        RemoteURL string `json:"remoteUrl"`

        // AllowPrivateEndpoint permits RemoteURL to point at a private or
        // in-cluster endpoint: RFC 1918 and IPv6 unique-local addresses, and
        // hostnames ending in "cluster.local". (Bare ".svc" hostnames are never
        // blocked, since no DNS resolution is performed.) Enable this to reach a
        // co-located in-cluster backend in-mesh so the backend's workload-identity
        // authorization policy still applies. Loopback, link-local, cloud-metadata,
        // and kubernetes.default endpoints remain blocked regardless of this
        // setting.
        // +kubebuilder:default=false
        // +optional
        AllowPrivateEndpoint bool `json:"allowPrivateEndpoint,omitempty"`

        // Transport is the transport method for the remote server (sse or streamable-http).
        // No default is set (unlike MCPRemoteProxy) because MCPServerEntry points at external
        // servers the user doesn't control — requiring explicit transport avoids silent mismatches.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Enum=sse;streamable-http
        Transport string `json:"transport"`

        // GroupRef references the MCPGroup this entry belongs to.
        // Required — every MCPServerEntry must be part of a group for vMCP discovery.
        // +kubebuilder:validation:Required
        GroupRef *MCPGroupRef `json:"groupRef"`

        // ExternalAuthConfigRef references a MCPExternalAuthConfig resource for token exchange
        // when connecting to the remote MCP server. The referenced MCPExternalAuthConfig must
        // exist in the same namespace as this MCPServerEntry.
        // +optional
        ExternalAuthConfigRef *ExternalAuthConfigRef `json:"externalAuthConfigRef,omitempty"`

        // HeaderForward configures headers to inject into requests to the remote MCP server.
        // Use this to add custom headers like API keys or correlation IDs.
        // +optional
        HeaderForward *HeaderForwardConfig `json:"headerForward,omitempty"`

        // CABundleRef references a ConfigMap containing CA certificates for TLS verification
        // when connecting to the remote MCP server.
        // +optional
        CABundleRef *CABundleSource `json:"caBundleRef,omitempty"`
}

// MCPServerEntryStatus defines the observed state of MCPServerEntry.
type MCPServerEntryStatus struct {
        // ObservedGeneration reflects the generation most recently observed by the controller.
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // Phase indicates the current lifecycle phase of the MCPServerEntry.
        // +optional
        // +kubebuilder:default=Pending
        Phase MCPServerEntryPhase `json:"phase,omitempty"`

        // Conditions represent the latest available observations of the MCPServerEntry's state.
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`
}

// MCPServerEntryPhase represents the lifecycle phase of an MCPServerEntry.
// +kubebuilder:validation:Enum=Valid;Pending;Failed
type MCPServerEntryPhase string

const (
        // MCPServerEntryPhaseValid indicates all validations passed and the entry is usable.
        MCPServerEntryPhaseValid MCPServerEntryPhase = "Valid"

        // MCPServerEntryPhasePending is the initial state before the first reconciliation.
        MCPServerEntryPhasePending MCPServerEntryPhase = "Pending"

        // MCPServerEntryPhaseFailed indicates one or more referenced resources are missing or invalid.
        MCPServerEntryPhaseFailed MCPServerEntryPhase = "Failed"
)

// Condition types for MCPServerEntry.
// Reuses shared condition type constants from mcpserver_types.go where the string
// values match (GroupRefValidated, ExternalAuthConfigValidated, CABundleRefValidated).
const (
        // ConditionTypeMCPServerEntryValid indicates overall validation status of the MCPServerEntry.
        // Uses the shared "Valid" condition type since this is a configuration resource, not a workload.
        ConditionTypeMCPServerEntryValid = ConditionTypeValid

        // ConditionTypeMCPServerEntryGroupRefValidated indicates whether the referenced MCPGroup exists.
        ConditionTypeMCPServerEntryGroupRefValidated = ConditionGroupRefValidated

        // ConditionTypeMCPServerEntryAuthConfigValidated indicates whether the referenced
        // MCPExternalAuthConfig exists (when configured).
        ConditionTypeMCPServerEntryAuthConfigValidated = ConditionTypeExternalAuthConfigValidated

        // ConditionTypeMCPServerEntryCABundleRefValidated indicates whether the referenced
        // CA bundle ConfigMap exists (when configured).
        ConditionTypeMCPServerEntryCABundleRefValidated = ConditionCABundleRefValidated

        // ConditionTypeMCPServerEntryRemoteURLValidated indicates whether the RemoteURL passes
        // format and SSRF safety checks.
        ConditionTypeMCPServerEntryRemoteURLValidated = "RemoteURLValidated"

        // ConditionTypeMCPServerEntryHeaderSecretRefsValidated indicates whether every Kubernetes
        // Secret referenced by spec.headerForward.addHeadersFromSecret exists. Absent when the
        // entry declares no header Secret refs.
        ConditionTypeMCPServerEntryHeaderSecretRefsValidated = "HeaderSecretRefsValidated"
)

// Condition reasons for MCPServerEntry.
// GroupRef reasons reuse shared constants from mcpserver_types.go.
// CABundle reasons reuse shared constants from mcpserver_types.go.
const (
        // ConditionReasonMCPServerEntryValid indicates the entry passed all validations.
        ConditionReasonMCPServerEntryValid = "ConfigValid"

        // ConditionReasonMCPServerEntryInvalid indicates one or more validations failed.
        ConditionReasonMCPServerEntryInvalid = "ConfigInvalid"

        // ConditionReasonMCPServerEntryGroupRefValidated reuses the shared GroupRef reason.
        ConditionReasonMCPServerEntryGroupRefValidated = ConditionReasonGroupRefValidated

        // ConditionReasonMCPServerEntryGroupRefNotFound reuses the shared GroupRef reason.
        ConditionReasonMCPServerEntryGroupRefNotFound = ConditionReasonGroupRefNotFound

        // ConditionReasonMCPServerEntryGroupRefNotReady reuses the shared GroupRef reason.
        ConditionReasonMCPServerEntryGroupRefNotReady = ConditionReasonGroupRefNotReady

        // ConditionReasonMCPServerEntryAuthConfigValid indicates the referenced auth config exists.
        ConditionReasonMCPServerEntryAuthConfigValid = "AuthConfigValid"

        // ConditionReasonMCPServerEntryAuthConfigNotFound indicates the referenced auth config was not found.
        ConditionReasonMCPServerEntryAuthConfigNotFound = "AuthConfigNotFound"

        // ConditionReasonMCPServerEntryAuthConfigNotConfigured indicates no auth config ref is set.
        ConditionReasonMCPServerEntryAuthConfigNotConfigured = "AuthConfigNotConfigured"

        // ConditionReasonMCPServerEntryCABundleRefValid reuses the shared CABundle reason.
        ConditionReasonMCPServerEntryCABundleRefValid = ConditionReasonCABundleRefValid

        // ConditionReasonMCPServerEntryCABundleRefNotFound reuses the shared CABundle reason.
        ConditionReasonMCPServerEntryCABundleRefNotFound = ConditionReasonCABundleRefNotFound

        // ConditionReasonMCPServerEntryCABundleRefNotConfigured indicates no CA bundle ref is set.
        ConditionReasonMCPServerEntryCABundleRefNotConfigured = "CABundleRefNotConfigured"

        // ConditionReasonMCPServerEntryRemoteURLValid indicates the RemoteURL passed all checks.
        ConditionReasonMCPServerEntryRemoteURLValid = "RemoteURLValid"

        // ConditionReasonMCPServerEntryRemoteURLInvalid indicates the RemoteURL is malformed or
        // targets a blocked internal/metadata endpoint.
        ConditionReasonMCPServerEntryRemoteURLInvalid = ConditionReasonRemoteURLInvalid

        // ConditionReasonMCPServerEntryHeaderSecretsValid indicates every referenced header Secret exists.
        ConditionReasonMCPServerEntryHeaderSecretsValid = "HeaderSecretsValid"

        // ConditionReasonMCPServerEntryHeaderSecretNotFound indicates a Secret referenced by
        // spec.headerForward.addHeadersFromSecret was not found in the entry's namespace.
        // Reuses the string value used by MCPRemoteProxy for parity.
        ConditionReasonMCPServerEntryHeaderSecretNotFound = ConditionReasonHeaderSecretNotFound
)

//+kubebuilder:object:root=true
//+kubebuilder:storageversion
//+kubebuilder:subresource:status
//+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
//+kubebuilder:resource:shortName=mcpentry,categories=toolhive
//+kubebuilder:printcolumn:name="Phase",type="string",JSONPath=".status.phase"
//+kubebuilder:printcolumn:name="Transport",type="string",JSONPath=".spec.transport"
//+kubebuilder:printcolumn:name="Remote URL",type="string",JSONPath=".spec.remoteUrl"
//+kubebuilder:printcolumn:name="Group",type="string",JSONPath=".spec.groupRef.name"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"

// MCPServerEntry is the Schema for the mcpserverentries API.
// It declares a remote MCP server endpoint for vMCP discovery and routing
// without deploying any infrastructure.
type MCPServerEntry struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPServerEntrySpec   `json:"spec,omitempty"`
        Status MCPServerEntryStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// MCPServerEntryList contains a list of MCPServerEntry.
type MCPServerEntryList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPServerEntry `json:"items"`
}

func init() {
        SchemeBuilder.Register(&MCPServerEntry{}, &MCPServerEntryList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        "fmt"

        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
        // maxK8sVolumeName is the maximum length for a Kubernetes volume name (RFC 1123 label).
        maxK8sVolumeName = 63
        // telemetryCABundleVolumePrefix must match validation.TelemetryCABundleVolumePrefix.
        telemetryCABundleVolumePrefix = "otel-ca-bundle-"
        // maxTelemetryCABundleConfigMapName is the maximum ConfigMap name length that fits in a volume name.
        maxTelemetryCABundleConfigMapName = maxK8sVolumeName - len(telemetryCABundleVolumePrefix)
)

// SensitiveHeader represents a header whose value is stored in a Kubernetes Secret.
// This allows credential headers (e.g., API keys, bearer tokens) to be securely
// referenced without embedding secrets inline in the MCPTelemetryConfig resource.
type SensitiveHeader struct {
        // Name is the header name (e.g., "Authorization", "X-API-Key")
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Name string `json:"name"`

        // SecretKeyRef is a reference to a Kubernetes Secret key containing the header value
        // +kubebuilder:validation:Required
        SecretKeyRef SecretKeyRef `json:"secretKeyRef"`
}

// MCPTelemetryOTelConfig defines OpenTelemetry configuration for shared MCPTelemetryConfig resources.
// Unlike OpenTelemetryConfig (used by inline MCPServer telemetry), this type:
//   - Omits ServiceName (per-server field set via MCPTelemetryConfigReference)
//   - Uses map[string]string for Headers (not []string)
//   - Adds SensitiveHeaders for Kubernetes Secret-backed credentials
//   - Adds ResourceAttributes for shared OTel resource attributes
//
// +kubebuilder:validation:XValidation:rule="!has(self.headers) || !has(self.sensitiveHeaders) || self.sensitiveHeaders.all(sh, !(sh.name in self.headers))",message="a header name cannot appear in both headers and sensitiveHeaders"
//
//nolint:lll // CEL validation rules exceed line length limit
type MCPTelemetryOTelConfig struct {
        // Enabled controls whether OpenTelemetry is enabled
        // +kubebuilder:default=false
        // +optional
        Enabled bool `json:"enabled,omitempty"`

        // Endpoint is the OTLP endpoint URL for tracing and metrics
        // +optional
        Endpoint string `json:"endpoint,omitempty"`

        // Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint
        // +kubebuilder:default=false
        // +optional
        Insecure bool `json:"insecure,omitempty"`

        // Headers contains authentication headers for the OTLP endpoint.
        // For secret-backed credentials, use sensitiveHeaders instead.
        // +optional
        Headers map[string]string `json:"headers,omitempty"`

        // SensitiveHeaders contains headers whose values are stored in Kubernetes Secrets.
        // Use this for credential headers (e.g., API keys, bearer tokens) instead of
        // embedding secrets in the headers field.
        // +listType=map
        // +listMapKey=name
        // +optional
        SensitiveHeaders []SensitiveHeader `json:"sensitiveHeaders,omitempty"`

        // ResourceAttributes contains custom resource attributes to be added to all telemetry signals.
        // These become OTel resource attributes (e.g., deployment.environment, service.namespace).
        // Note: service.name is intentionally excluded — it is set per-server via
        // MCPTelemetryConfigReference.ServiceName.
        // +optional
        ResourceAttributes map[string]string `json:"resourceAttributes,omitempty"`

        // Metrics defines OpenTelemetry metrics-specific configuration
        // +optional
        Metrics *OpenTelemetryMetricsConfig `json:"metrics,omitempty"`

        // Tracing defines OpenTelemetry tracing configuration
        // +optional
        Tracing *OpenTelemetryTracingConfig `json:"tracing,omitempty"`

        // UseLegacyAttributes controls whether legacy attribute names are emitted alongside
        // the new MCP OTEL semantic convention names. Defaults to true for backward compatibility.
        // This will change to false in a future release and eventually be removed.
        // +kubebuilder:default=true
        // +optional
        UseLegacyAttributes bool `json:"useLegacyAttributes"`

        // CABundleRef references a ConfigMap containing a CA certificate bundle for the OTLP endpoint.
        // When specified, the operator mounts the ConfigMap into the proxyrunner pod and configures
        // the OTLP exporters to trust the custom CA. This is useful when the OTLP collector uses
        // TLS with certificates signed by an internal or private CA.
        // +optional
        CABundleRef *CABundleSource `json:"caBundleRef,omitempty"`
}

// MCPTelemetryConfigSpec defines the desired state of MCPTelemetryConfig.
// The spec uses a nested structure with openTelemetry and prometheus sub-objects
// for clear separation of concerns.
type MCPTelemetryConfigSpec struct {
        // OpenTelemetry defines OpenTelemetry configuration (OTLP endpoint, tracing, metrics)
        // +optional
        OpenTelemetry *MCPTelemetryOTelConfig `json:"openTelemetry,omitempty"`

        // Prometheus defines Prometheus-specific configuration
        // +optional
        Prometheus *PrometheusConfig `json:"prometheus,omitempty"`
}

// MCPTelemetryConfigStatus defines the observed state of MCPTelemetryConfig
type MCPTelemetryConfigStatus struct {
        // Conditions represent the latest available observations of the MCPTelemetryConfig's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration is the most recent generation observed for this MCPTelemetryConfig.
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // ConfigHash is a hash of the current configuration for change detection
        // +optional
        ConfigHash string `json:"configHash,omitempty"`
}

// +kubebuilder:object:root=true
// +kubebuilder:storageversion
// +kubebuilder:subresource:status
// +kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
// +kubebuilder:resource:shortName=mcpotel,categories=toolhive
// +kubebuilder:printcolumn:name="Endpoint",type=string,JSONPath=`.spec.openTelemetry.endpoint`
// +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status`
// +kubebuilder:printcolumn:name="Tracing",type=boolean,JSONPath=`.spec.openTelemetry.tracing.enabled`
// +kubebuilder:printcolumn:name="Metrics",type=boolean,JSONPath=`.spec.openTelemetry.metrics.enabled`
// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp`

// MCPTelemetryConfig is the Schema for the mcptelemetryconfigs API.
// MCPTelemetryConfig resources are namespace-scoped and can only be referenced by
// MCPServer resources within the same namespace. Cross-namespace references
// are not supported for security and isolation reasons.
type MCPTelemetryConfig struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPTelemetryConfigSpec   `json:"spec,omitempty"`
        Status MCPTelemetryConfigStatus `json:"status,omitempty"`
}

// +kubebuilder:object:root=true

// MCPTelemetryConfigList contains a list of MCPTelemetryConfig
type MCPTelemetryConfigList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPTelemetryConfig `json:"items"`
}

// MCPTelemetryConfigReference is a reference to an MCPTelemetryConfig resource
// with per-server overrides. The referenced MCPTelemetryConfig must be in the
// same namespace as the MCPServer.
type MCPTelemetryConfigReference struct {
        // Name is the name of the MCPTelemetryConfig resource
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinLength=1
        Name string `json:"name"`

        // ServiceName overrides the telemetry service name for this specific server.
        // This MUST be unique per server for proper observability (e.g., distinguishing
        // traces and metrics from different servers sharing the same collector).
        // If empty, defaults to the server name with "thv-" prefix at runtime.
        // +optional
        ServiceName string `json:"serviceName,omitempty"`
}

// Validate performs validation on the MCPTelemetryConfig spec.
// This provides defense-in-depth alongside CEL validation rules.
// CEL catches issues at API admission time, but this method also validates
// stored objects to catch any that bypassed CEL or were stored before CEL rules were added.
func (r *MCPTelemetryConfig) Validate() error {
        if err := r.validateEndpointRequiresSignals(); err != nil {
                return err
        }
        if err := r.validateSensitiveHeaders(); err != nil {
                return err
        }
        return r.validateCABundle()
}

// validateEndpointRequiresSignals rejects an endpoint when neither tracing nor metrics is enabled.
// Without this check the config would pass CRD validation but fail at runtime in telemetry.NewProvider.
func (r *MCPTelemetryConfig) validateEndpointRequiresSignals() error {
        if r.Spec.OpenTelemetry == nil {
                return nil
        }
        otel := r.Spec.OpenTelemetry
        if otel.Endpoint == "" {
                return nil
        }
        tracingEnabled := otel.Tracing != nil && otel.Tracing.Enabled
        metricsEnabled := otel.Metrics != nil && otel.Metrics.Enabled
        if !tracingEnabled && !metricsEnabled {
                return fmt.Errorf("endpoint requires at least one of tracing or metrics to be enabled")
        }
        return nil
}

// validateSensitiveHeaders validates sensitive header entries and checks for overlap with plaintext headers.
func (r *MCPTelemetryConfig) validateSensitiveHeaders() error {
        if r.Spec.OpenTelemetry == nil {
                return nil
        }
        otel := r.Spec.OpenTelemetry
        for i, sh := range otel.SensitiveHeaders {
                if sh.Name == "" {
                        return fmt.Errorf("openTelemetry.sensitiveHeaders[%d].name must not be empty", i)
                }
                if sh.SecretKeyRef.Name == "" {
                        return fmt.Errorf("openTelemetry.sensitiveHeaders[%d].secretKeyRef.name must not be empty", i)
                }
                if sh.SecretKeyRef.Key == "" {
                        return fmt.Errorf("openTelemetry.sensitiveHeaders[%d].secretKeyRef.key must not be empty", i)
                }
                if _, exists := otel.Headers[sh.Name]; exists {
                        return fmt.Errorf("header %q appears in both headers and sensitiveHeaders", sh.Name)
                }
        }
        return nil
}

// validateCABundle validates the CA bundle configuration if present.
func (r *MCPTelemetryConfig) validateCABundle() error {
        if r.Spec.OpenTelemetry == nil || r.Spec.OpenTelemetry.CABundleRef == nil {
                return nil
        }
        otel := r.Spec.OpenTelemetry
        if otel.Insecure {
                return fmt.Errorf("openTelemetry.caBundleRef cannot be specified when insecure is true; they are mutually exclusive")
        }
        ref := otel.CABundleRef
        if ref.ConfigMapRef == nil {
                return fmt.Errorf("openTelemetry.caBundleRef.configMapRef must be specified")
        }
        if ref.ConfigMapRef.Name == "" {
                return fmt.Errorf("openTelemetry.caBundleRef.configMapRef.name must not be empty")
        }
        if len(ref.ConfigMapRef.Name) > maxTelemetryCABundleConfigMapName {
                //nolint:lll // error message clarity requires full context
                return fmt.Errorf(
                        "openTelemetry.caBundleRef.configMapRef.name %q is too long (%d chars); maximum is %d",
                        ref.ConfigMapRef.Name, len(ref.ConfigMapRef.Name), maxTelemetryCABundleConfigMapName,
                )
        }
        return nil
}

func init() {
        SchemeBuilder.Register(&MCPTelemetryConfig{}, &MCPTelemetryConfigList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// Condition types for MCPToolConfig
const (
        // ConditionToolConfigValid indicates whether the MCPToolConfig spec is valid.
        ConditionToolConfigValid = ConditionTypeValid
)

const (
        // ConditionReasonToolConfigValidationSucceeded indicates validation passed.
        ConditionReasonToolConfigValidationSucceeded = "ValidationSucceeded"
        // ConditionReasonToolConfigValidationFailed indicates validation failed.
        ConditionReasonToolConfigValidationFailed = "ValidationFailed"
)

// MCPToolConfigSpec defines the desired state of MCPToolConfig.
// MCPToolConfig resources are namespace-scoped and can only be referenced by
// MCPServer resources in the same namespace.
type MCPToolConfigSpec struct {
        // ToolsFilter is a list of tool names to filter (allow list).
        // Only tools in this list will be exposed by the MCP server.
        // If empty, all tools are exposed.
        // +listType=set
        // +optional
        ToolsFilter []string `json:"toolsFilter,omitempty"`

        // ToolsOverride is a map from actual tool names to their overridden configuration.
        // This allows renaming tools and/or changing their descriptions.
        // +optional
        ToolsOverride map[string]ToolOverride `json:"toolsOverride,omitempty"`
}

// ToolAnnotationsOverride defines overrides for tool annotation fields.
// All fields use pointers so nil means "don't override" while zero values
// (empty string, false) mean "explicitly set to this value."
type ToolAnnotationsOverride struct {
        // Title overrides the human-readable title annotation.
        // +optional
        Title *string `json:"title,omitempty"`

        // ReadOnlyHint overrides the read-only hint annotation.
        // +optional
        ReadOnlyHint *bool `json:"readOnlyHint,omitempty"`

        // DestructiveHint overrides the destructive hint annotation.
        // +optional
        DestructiveHint *bool `json:"destructiveHint,omitempty"`

        // IdempotentHint overrides the idempotent hint annotation.
        // +optional
        IdempotentHint *bool `json:"idempotentHint,omitempty"`

        // OpenWorldHint overrides the open-world hint annotation.
        // +optional
        OpenWorldHint *bool `json:"openWorldHint,omitempty"`
}

// ToolOverride represents a tool override configuration.
// Both Name and Description can be overridden independently, but
// they can't be both empty.
type ToolOverride struct {
        // Name is the redefined name of the tool
        // +optional
        Name string `json:"name,omitempty"`

        // Description is the redefined description of the tool
        // +optional
        Description string `json:"description,omitempty"`

        // Annotations overrides specific tool annotation fields.
        // Only specified fields are overridden; others pass through from the backend.
        // +optional
        Annotations *ToolAnnotationsOverride `json:"annotations,omitempty"`
}

// MCPToolConfigStatus defines the observed state of MCPToolConfig
type MCPToolConfigStatus struct {
        // Conditions represent the latest available observations of the MCPToolConfig's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration is the most recent generation observed for this MCPToolConfig.
        // It corresponds to the MCPToolConfig's generation, which is updated on mutation by the API Server.
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // ConfigHash is a hash of the current configuration for change detection
        // +optional
        ConfigHash string `json:"configHash,omitempty"`
}

// +kubebuilder:object:root=true
// +kubebuilder:storageversion
// +kubebuilder:subresource:status
// +kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
// +kubebuilder:resource:shortName=tc;toolconfig,categories=toolhive
// +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status`
// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp`

// MCPToolConfig is the Schema for the mcptoolconfigs API.
// MCPToolConfig resources are namespace-scoped and can only be referenced by
// MCPServer resources within the same namespace. Cross-namespace references
// are not supported for security and isolation reasons.
type MCPToolConfig struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   MCPToolConfigSpec   `json:"spec,omitempty"`
        Status MCPToolConfigStatus `json:"status,omitempty"`
}

// +kubebuilder:object:root=true

// MCPToolConfigList contains a list of MCPToolConfig
type MCPToolConfigList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []MCPToolConfig `json:"items"`
}

func init() {
        SchemeBuilder.Register(&MCPToolConfig{}, &MCPToolConfigList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

        "github.com/stacklok/toolhive/pkg/vmcp/config"
)

// VirtualMCPCompositeToolDefinitionSpec defines the desired state of VirtualMCPCompositeToolDefinition.
// This embeds the CompositeToolConfig from pkg/vmcp/config to share the configuration model
// between CLI and operator usage.
type VirtualMCPCompositeToolDefinitionSpec struct {
        config.CompositeToolConfig `json:",inline"` // nolint:revive // inline is valid
}

// VirtualMCPCompositeToolDefinitionStatus defines the observed state of VirtualMCPCompositeToolDefinition
type VirtualMCPCompositeToolDefinitionStatus struct {
        // ValidationStatus indicates the validation state of the workflow
        // - Valid: Workflow structure is valid
        // - Invalid: Workflow has validation errors
        // +optional
        ValidationStatus ValidationStatus `json:"validationStatus,omitempty"`

        // ValidationErrors contains validation error messages if ValidationStatus is Invalid
        // +listType=atomic
        // +optional
        ValidationErrors []string `json:"validationErrors,omitempty"`

        // ReferencingVirtualServers lists VirtualMCPServer resources that reference this workflow
        // This helps track which servers need to be reconciled when this workflow changes
        // +listType=set
        // +optional
        ReferencingVirtualServers []string `json:"referencingVirtualServers,omitempty"`

        // ObservedGeneration is the most recent generation observed for this VirtualMCPCompositeToolDefinition
        // It corresponds to the resource's generation, which is updated on mutation by the API Server
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // Conditions represent the latest available observations of the workflow's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`
}

// ValidationStatus represents the validation state of a workflow
// +kubebuilder:validation:Enum=Valid;Invalid;Unknown
type ValidationStatus string

const (
        // ValidationStatusValid indicates the workflow is valid
        ValidationStatusValid ValidationStatus = "Valid"

        // ValidationStatusInvalid indicates the workflow has validation errors
        ValidationStatusInvalid ValidationStatus = "Invalid"

        // ValidationStatusUnknown indicates validation hasn't been performed yet
        ValidationStatusUnknown ValidationStatus = "Unknown"
)

// Condition types for VirtualMCPCompositeToolDefinition
const (
        // ConditionTypeWorkflowValidated indicates whether the workflow has been validated
        ConditionTypeWorkflowValidated = "WorkflowValidated"

        // Note: ConditionTypeReady is shared across multiple resources and defined in mcpremoteproxy_types.go
)

// Condition reasons for VirtualMCPCompositeToolDefinition
const (
        // ConditionReasonValidationSuccess indicates workflow validation succeeded
        ConditionReasonValidationSuccess = "ValidationSuccess"

        // ConditionReasonValidationFailed indicates workflow validation failed
        ConditionReasonValidationFailed = "ValidationFailed"

        // ConditionReasonSchemaInvalid indicates parameter or step schema is invalid
        ConditionReasonSchemaInvalid = "SchemaInvalid"

        // ConditionReasonTemplateInvalid indicates template syntax is invalid
        ConditionReasonTemplateInvalid = "TemplateInvalid"

        // ConditionReasonDependencyCycle indicates step dependencies contain cycles
        ConditionReasonDependencyCycle = "DependencyCycle"

        // ConditionReasonToolNotFound indicates a referenced tool doesn't exist
        ConditionReasonToolNotFound = "ToolNotFound"

        // ConditionReasonWorkflowReady indicates the workflow is ready to use
        ConditionReasonWorkflowReady = "WorkflowReady"

        // ConditionReasonWorkflowNotReady indicates the workflow is not ready
        ConditionReasonWorkflowNotReady = "WorkflowNotReady"
)

//+kubebuilder:object:root=true
//+kubebuilder:storageversion
//+kubebuilder:subresource:status
//+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
//+kubebuilder:resource:shortName=vmcpctd;compositetool,categories=toolhive
//+kubebuilder:printcolumn:name="Workflow",type="string",JSONPath=".spec.name",description="Workflow name"
//+kubebuilder:printcolumn:name="Status",type="string",JSONPath=".status.validationStatus",description="Validation status"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp",description="Age"
//+kubebuilder:printcolumn:name="Ready",type="string",JSONPath=".status.conditions[?(@.type=='Ready')].status"

// VirtualMCPCompositeToolDefinition is the Schema for the virtualmcpcompositetooldefinitions API
// VirtualMCPCompositeToolDefinition defines reusable composite workflows that can be referenced
// by multiple VirtualMCPServer instances
type VirtualMCPCompositeToolDefinition struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   VirtualMCPCompositeToolDefinitionSpec   `json:"spec,omitempty"`
        Status VirtualMCPCompositeToolDefinitionStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// VirtualMCPCompositeToolDefinitionList contains a list of VirtualMCPCompositeToolDefinition
type VirtualMCPCompositeToolDefinitionList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []VirtualMCPCompositeToolDefinition `json:"items"`
}

// Validate performs validation for VirtualMCPCompositeToolDefinition
// This method is called by the controller during reconciliation
// It delegates to the shared ValidateCompositeToolConfig in pkg/vmcp/config
func (r *VirtualMCPCompositeToolDefinition) Validate() error {
        return config.ValidateCompositeToolConfig("spec", &r.Spec.CompositeToolConfig)
}

// GetValidationErrors returns a list of validation errors
// This is a helper method for the controller to populate status.validationErrors
func (r *VirtualMCPCompositeToolDefinition) GetValidationErrors() []string {
        if err := r.Validate(); err != nil {
                return []string{err.Error()}
        }
        return nil
}

func init() {
        SchemeBuilder.Register(&VirtualMCPCompositeToolDefinition{}, &VirtualMCPCompositeToolDefinitionList{})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package v1beta1

import (
        "fmt"

        corev1 "k8s.io/api/core/v1"
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/runtime"

        vmcptypes "github.com/stacklok/toolhive/pkg/vmcp"
        "github.com/stacklok/toolhive/pkg/vmcp/config"
)

// VirtualMCPServerSpec defines the desired state of VirtualMCPServer
//
// +kubebuilder:validation:XValidation:rule="!has(self.config) || !has(self.config.rateLimiting) || (has(self.sessionStorage) && self.sessionStorage.provider == 'redis')",message="config.rateLimiting requires sessionStorage with provider 'redis'"
// +kubebuilder:validation:XValidation:rule="!(has(self.config) && has(self.config.rateLimiting) && has(self.config.rateLimiting.perUser)) || (has(self.incomingAuth) && self.incomingAuth.type == 'oidc')",message="config.rateLimiting.perUser requires incomingAuth.type oidc"
// +kubebuilder:validation:XValidation:rule="!has(self.config) || !has(self.config.rateLimiting) || !has(self.config.rateLimiting.tools) || self.config.rateLimiting.tools.all(t, !has(t.perUser)) || (has(self.incomingAuth) && self.incomingAuth.type == 'oidc')",message="per-tool perUser rate limiting requires incomingAuth.type oidc"
// +kubebuilder:validation:XValidation:rule="!has(self.config) || !has(self.config.sessionStorage) || !has(self.config.sessionStorage.tls)",message="config.sessionStorage.tls is ignored by the operator; set sessionStorage.tls instead"
// +kubebuilder:validation:XValidation:rule="!(has(self.embeddingServerRef) && has(self.config) && has(self.config.optimizer) && has(self.config.optimizer.embeddingProvider) && self.config.optimizer.embeddingProvider == 'openai')",message="embeddingServerRef provisions a managed TEI server and cannot be combined with optimizer.embeddingProvider 'openai'; openai mode uses embeddingService directly"
//
// +kubebuilder:validation:XValidation:rule="!has(self.authServerUpstreamFilter) || has(self.authServerConfig)",message="authServerUpstreamFilter requires authServerConfig to be set"
// +kubebuilder:validation:XValidation:rule="!has(self.authServerUpstreamFilter) || (has(self.authServerConfig) && has(self.authServerConfig.upstreamProviders) && size(self.authServerConfig.upstreamProviders) >= 2)",message="authServerUpstreamFilter requires at least two configured upstream providers"
// Deeper cross-reference checks (every referenced upstream provider must
// exist, and no reference to the mandatory first provider) were previously
// expressed as CEL on this struct but were removed because they pushed the
// cluster's x-kubernetes-validations cost budget over the per-schema limit
// in realistic deployments where upstreamProviders is unbounded. The same
// invariants are re-checked at reconciliation time by
// validateAuthServerUpstreamFilter in the VirtualMCPServer controller, which
// surfaces a terminal AuthServerConfigValidated=False condition — only the
// kubectl-apply-time feedback is lost, not the safety.
//
//nolint:lll // CEL validation rules exceed line length limit
type VirtualMCPServerSpec struct {
        // IncomingAuth configures authentication for clients connecting to the Virtual MCP server.
        // Must be explicitly set - use "anonymous" type when no authentication is required.
        // This field takes precedence over config.IncomingAuth and should be preferred because it
        // supports Kubernetes-native secret references (SecretKeyRef, ConfigMapRef) for secure
        // dynamic discovery of credentials, rather than requiring secrets to be embedded in config.
        // +kubebuilder:validation:Required
        IncomingAuth *IncomingAuthConfig `json:"incomingAuth"`

        // OutgoingAuth configures authentication from Virtual MCP to backend MCPServers.
        // This field takes precedence over config.OutgoingAuth and should be preferred because it
        // supports Kubernetes-native secret references (SecretKeyRef, ConfigMapRef) for secure
        // dynamic discovery of credentials, rather than requiring secrets to be embedded in config.
        // +optional
        OutgoingAuth *OutgoingAuthConfig `json:"outgoingAuth,omitempty"`

        // PassthroughHeaders is an allowlist of incoming client request header names
        // forwarded verbatim to all backends (e.g. an API key the backend resolves to
        // a user). Takes precedence over config.PassthroughHeaders. Names must not be
        // restricted headers (Host, hop-by-hop, X-Forwarded-*). Forwarded headers are
        // attacker-influenceable unless a trusted upstream sets them.
        // +optional
        // +listType=atomic
        PassthroughHeaders []string `json:"passthroughHeaders,omitempty"`

        // ServiceType specifies the Kubernetes service type for the Virtual MCP server
        // +kubebuilder:validation:Enum=ClusterIP;NodePort;LoadBalancer
        // +kubebuilder:default=ClusterIP
        // +optional
        ServiceType string `json:"serviceType,omitempty"`

        // SessionAffinity controls whether the Service routes repeated client connections to the same pod.
        // MCP protocols (SSE, streamable-http) are stateful, so ClientIP is the default.
        // Set to "None" for stateless servers or when using an external load balancer with its own affinity.
        // +kubebuilder:validation:Enum=ClientIP;None
        // +kubebuilder:default=ClientIP
        // +optional
        SessionAffinity string `json:"sessionAffinity,omitempty"`

        // ServiceAccount is the name of an already existing service account to use by the Virtual MCP server.
        // If not specified, a ServiceAccount will be created automatically and used by the Virtual MCP server.
        // +optional
        ServiceAccount *string `json:"serviceAccount,omitempty"`

        // PodTemplateSpec defines the pod template to use for the Virtual MCP server
        // This allows for customizing the pod configuration beyond what is provided by the other fields.
        // Note that to modify the specific container the Virtual MCP server runs in, you must specify
        // the 'vmcp' container name in the PodTemplateSpec.
        // This field accepts a PodTemplateSpec object as JSON/YAML.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        PodTemplateSpec *runtime.RawExtension `json:"podTemplateSpec,omitempty"`

        // GroupRef references the MCPGroup that defines backend workloads.
        // The referenced MCPGroup must exist in the same namespace.
        // +kubebuilder:validation:Required
        GroupRef *MCPGroupRef `json:"groupRef"`

        // Config is the Virtual MCP server configuration.
        // The audit config from here is also supported, but not required.
        // +optional
        Config config.Config `json:"config,omitempty"`

        // TelemetryConfigRef references an MCPTelemetryConfig resource for shared telemetry configuration.
        // The referenced MCPTelemetryConfig must exist in the same namespace as this VirtualMCPServer.
        // Cross-namespace references are not supported for security and isolation reasons.
        // +optional
        TelemetryConfigRef *MCPTelemetryConfigReference `json:"telemetryConfigRef,omitempty"`

        // EmbeddingServerRef references an existing EmbeddingServer resource by name.
        // When the optimizer is enabled, this field is required to point to a ready EmbeddingServer
        // that provides embedding capabilities.
        // The referenced EmbeddingServer must exist in the same namespace and be ready.
        // +optional
        EmbeddingServerRef *EmbeddingServerRef `json:"embeddingServerRef,omitempty"`

        // AuthServerConfig configures an embedded OAuth authorization server.
        // When set, the vMCP server acts as an OIDC issuer, drives users through
        // upstream IDPs, and issues ToolHive JWTs. The embedded AS becomes the
        // IncomingAuth OIDC provider — its issuer must match IncomingAuth.OIDCConfigRef
        // so that tokens it issues are accepted by the vMCP's incoming auth middleware.
        // When nil, IncomingAuth uses an external IDP and behavior is unchanged.
        // +optional
        AuthServerConfig *EmbeddedAuthServerConfig `json:"authServerConfig,omitempty"`

        // AuthServerUpstreamFilter optionally narrows the multi-upstream
        // authorization chain per-request based on the first upstream's resolved
        // identity, so principals only walk the upstream providers they are
        // entitled to. Lives beside AuthServerConfig (rather than nested in it)
        // because the shared EmbeddedAuthServerConfig type is also used by
        // MCPExternalAuthConfig, where multi-upstream filtering has no meaning.
        //
        // When nil, every request walks the full upstream chain in configured
        // order — the behavior from before the filter hook existed. The filter
        // only controls which OAuth credentials are collected during login;
        // backend and tool authorization remain Cedar's responsibility.
        // +optional
        AuthServerUpstreamFilter *AuthServerUpstreamFilterConfig `json:"authServerUpstreamFilter,omitempty"`

        // Replicas is the desired number of vMCP pod replicas.
        // VirtualMCPServer creates a single Deployment for the vMCP aggregator process,
        // so there is only one replicas field (unlike MCPServer which has separate
        // Replicas and BackendReplicas for its two Deployments).
        // When nil, the operator does not set Deployment.Spec.Replicas, leaving replica
        // management to an HPA or other external controller.
        // +kubebuilder:validation:Minimum=0
        // +optional
        Replicas *int32 `json:"replicas,omitempty"`

        // SessionStorage configures session storage for stateful horizontal scaling.
        // When nil, no session storage is configured.
        // +optional
        SessionStorage *SessionStorageConfig `json:"sessionStorage,omitempty"`

        // ImagePullSecrets allows specifying image pull secrets for the vMCP workload.
        // These are applied to both the vMCP Deployment's PodSpec.ImagePullSecrets
        // and to the operator-managed ServiceAccount the vMCP server runs as, so private
        // images are pullable through either path.
        //
        // Merge semantics with PodTemplateSpec:
        // The deployed PodSpec.ImagePullSecrets is the Kubernetes-native strategic-merge
        // union of this field and spec.podTemplateSpec.spec.imagePullSecrets, merged by
        // the patchStrategy:"merge" / patchMergeKey:"name" tags on corev1.PodSpec.
        //   - This field is rendered first as the controller-generated default.
        //   - spec.podTemplateSpec.spec.imagePullSecrets is then strategic-merge-patched
        //     on top, keyed by Name. Distinct names from the two sources are unioned in
        //     the resulting list; entries with the same Name are deduplicated and the
        //     PodTemplateSpec entry wins on overlap (user override).
        //   - Order in the resulting list is not guaranteed and should not be relied on:
        //     strategic merge by name is order-insensitive.
        //   - The operator-managed ServiceAccount's imagePullSecrets list is populated
        //     ONLY from this field. spec.podTemplateSpec.spec.imagePullSecrets does not
        //     reach the ServiceAccount because PodTemplateSpec has no notion of a
        //     ServiceAccount. To make a secret usable via the ServiceAccount path
        //     (e.g. for sidecars or init containers that pull images independently),
        //     list it here rather than under spec.podTemplateSpec.
        //
        // Note on cross-CRD consistency:
        // MCPRegistry currently uses an atomic-replace strategy for its imagePullSecrets
        // (the user-provided value replaces the controller-generated list rather than
        // being merged on top). VirtualMCPServer follows the Kubernetes-native
        // strategic-merge-by-name behavior described above. Aligning the two is tracked
        // as a separate follow-up; until then, manifests that set imagePullSecrets on
        // both CRDs will see different override behavior between them.
        //
        // +listType=atomic
        // +optional
        ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,omitempty"`
}

// EmbeddingServerRef references an existing EmbeddingServer resource by name.
// This follows the same pattern as ExternalAuthConfigRef and ToolConfigRef.
type EmbeddingServerRef struct {
        // Name is the name of the EmbeddingServer resource
        // +kubebuilder:validation:Required
        Name string `json:"name"`
}

// IncomingAuthConfig configures authentication for clients connecting to the Virtual MCP server
//
// +kubebuilder:validation:XValidation:rule="self.type == 'oidc' ? has(self.oidcConfigRef) : true",message="spec.incomingAuth.oidcConfigRef is required when type is oidc"
// +kubebuilder:validation:XValidation:rule="!(has(self.authzConfig) && has(self.authzConfigRef))",message="authzConfig and authzConfigRef are mutually exclusive; use authzConfigRef to reference a shared MCPAuthzConfig"
//
//nolint:lll // CEL validation rules exceed line length limit
type IncomingAuthConfig struct {
        // Type defines the authentication type: anonymous or oidc
        // When no authentication is required, explicitly set this to "anonymous"
        // +kubebuilder:validation:Enum=anonymous;oidc
        // +kubebuilder:validation:Required
        Type string `json:"type"`

        // OIDCConfigRef references a shared MCPOIDCConfig resource for OIDC authentication.
        // The referenced MCPOIDCConfig must exist in the same namespace as this VirtualMCPServer.
        // Per-server overrides (audience, scopes) are specified here; shared provider config
        // lives in the MCPOIDCConfig resource.
        // +optional
        OIDCConfigRef *MCPOIDCConfigReference `json:"oidcConfigRef,omitempty"`

        // AuthzConfig defines authorization policy configuration.
        // Reuses MCPServer authz patterns.
        // AuthzConfig and AuthzConfigRef are mutually exclusive.
        // +optional
        AuthzConfig *AuthzConfigRef `json:"authzConfig,omitempty"`

        // AuthzConfigRef references a shared MCPAuthzConfig resource for authorization.
        // The referenced MCPAuthzConfig must exist in the same namespace as this VirtualMCPServer.
        // Mutually exclusive with authzConfig.
        //
        // Only cedarv1 MCPAuthzConfig resources are supported for VirtualMCPServer
        // today; referencing a non-Cedar config fails reconciliation with a clear
        // error because the vMCP runtime authz middleware is Cedar-only.
        // +optional
        AuthzConfigRef *MCPAuthzConfigReference `json:"authzConfigRef,omitempty"`
}

// OutgoingAuthConfig configures authentication from Virtual MCP to backend MCPServers
type OutgoingAuthConfig struct {
        // Source defines how backend authentication configurations are determined
        // - discovered: Automatically discover from backend's MCPServer.spec.externalAuthConfigRef
        // - inline: Explicit per-backend configuration in VirtualMCPServer
        // +kubebuilder:validation:Enum=discovered;inline
        // +kubebuilder:default=discovered
        // +optional
        Source string `json:"source,omitempty"`

        // Default defines default behavior for backends without explicit auth config
        // +optional
        Default *BackendAuthConfig `json:"default,omitempty"`

        // Backends defines per-backend authentication overrides
        // Works in all modes (discovered, inline)
        // +optional
        Backends map[string]BackendAuthConfig `json:"backends,omitempty"`
}

// BackendAuthConfig defines authentication configuration for a backend MCPServer
type BackendAuthConfig struct {
        // Type defines the authentication type
        // +kubebuilder:validation:Enum=discovered;externalAuthConfigRef
        // +kubebuilder:validation:Required
        Type string `json:"type"`

        // ExternalAuthConfigRef references an MCPExternalAuthConfig resource
        // Only used when Type is "externalAuthConfigRef"
        // +optional
        ExternalAuthConfigRef *ExternalAuthConfigRef `json:"externalAuthConfigRef,omitempty"`
}

// OperationalConfig defines operational settings

// Backend status constants for DiscoveredBackend.Status
// These are the user-facing values stored in VirtualMCPServer.Status.DiscoveredBackends.
// Use BackendHealthStatus.ToCRDStatus() to convert from internal health status.
const (
        BackendStatusReady           = "ready"
        BackendStatusUnavailable     = "unavailable"
        BackendStatusDegraded        = "degraded"
        BackendStatusUnknown         = "unknown"
        BackendStatusUnauthenticated = "unauthenticated"
)

// DiscoveredBackend is an alias to the canonical definition in pkg/vmcp/types.go
// This provides a local name for use in the CRD status.
// +gendoc
type DiscoveredBackend = vmcptypes.DiscoveredBackend

// AuthServerUpstreamFilterConfig configures the group-based filter that
// narrows which of authServerConfig.upstreamProviders a given principal must
// walk during multi-upstream authorization. The filter runs exactly once per
// authorization, on the first upstream's callback leg, and may only narrow
// the chain — it cannot reorder it, add unknown providers, or remove the
// mandatory first provider. See the pkg/authserver/upstreamfilter package
// for the full semantic contract.
type AuthServerUpstreamFilterConfig struct {
        // Claim is the exact top-level JWT/OIDC claim name to read group
        // membership from on the first upstream's resolved principal. Defaults
        // to "groups" when empty. Nested claims (e.g. Keycloak's
        // realm_access.roles) are not supported and must be flattened to a
        // top-level string array by the identity provider.
        //
        // The value is read from the identity the first upstream resolves: an
        // OIDC upstream fills it from the validated ID token, a pure OAuth2
        // upstream from the userinfo response. Groups carried only in an access
        // token never reach it. Before enabling the filter, configure the first
        // upstream to include the claim in its ID token or userinfo response,
        // and request any required scopes. A missing claim fails the
        // authorization (defaultUpstreams covers a valid-but-empty claim and
        // unmatched groups, not an absent one), so enabling the filter before
        // the provider emits the claim locks every principal out.
        // +optional
        // +kubebuilder:validation:MinLength=1
        // +kubebuilder:validation:MaxLength=253
        Claim string `json:"claim,omitempty"`

        // Rules assigns groups to upstream provider subsets. A principal matches
        // a rule if any of its groups appears in the rule's groups list
        // (any-of / OR semantics). Multiple matching rules are combined by
        // union — the effective chain contains every upstream from every
        // matching rule, deduplicated and in configured order.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=100
        // +listType=atomic
        Rules []UpstreamFilterGroupRule `json:"rules"`

        // DefaultUpstreams is used when the claim is a valid empty array or
        // when no rule matches. When omitted, no optional upstreams are added —
        // the mandatory first upstream remains but the authorization walks no
        // further. Every entry must name a configured
        // authServerConfig.upstreamProviders provider other than the mandatory
        // first provider; those cross-reference checks are enforced at
        // reconciliation time rather than by CEL (CEL cost-budget constraints
        // moved the check into the Go filter constructor; see the top-level CEL
        // rules on VirtualMCPServerSpec for the shape invariants that remain at
        // admission time).
        // +optional
        // +kubebuilder:validation:MaxItems=100
        // +listType=set
        DefaultUpstreams []string `json:"defaultUpstreams,omitempty"`
}

// UpstreamFilterGroupRule pairs a set of principal group names with the set
// of non-first upstream provider names those principals are entitled to walk.
type UpstreamFilterGroupRule struct {
        // Groups is the set of principal group names this rule matches. Matching
        // is case-sensitive and exact. Per-element emptiness and duplicate
        // detection are enforced at reconciliation time by the Go filter
        // constructor, which runs again before the auth server starts.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=100
        // +listType=set
        Groups []string `json:"groups"`

        // UpstreamProviders lists the non-first upstream provider names to
        // include in the chain when this rule matches. Every entry must name a
        // configured authServerConfig.upstreamProviders provider other than the
        // mandatory first provider; these cross-reference checks are enforced at
        // reconciliation time rather than by CEL (CEL cost-budget constraints
        // moved the check into the Go filter constructor). Per-element emptiness
        // and duplicates are also enforced at reconciliation time.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:MinItems=1
        // +kubebuilder:validation:MaxItems=100
        // +listType=set
        UpstreamProviders []string `json:"upstreamProviders"`
}

// VirtualMCPServerRuntimeStatus is the runtime-owned status snapshot. The
// operator projects this snapshot into the top-level compatibility fields and
// remains the sole writer of the top-level Conditions array.
type VirtualMCPServerRuntimeStatus struct {
        // Phase is the lifecycle phase observed by the running vMCP process.
        // +optional
        Phase VirtualMCPServerPhase `json:"phase,omitempty"`

        // Message provides detail about the runtime phase.
        // +optional
        Message string `json:"message,omitempty"`

        // Conditions contains runtime health observations.
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // DiscoveredBackends contains the runtime's latest backend observations.
        // +listType=map
        // +listMapKey=name
        // +optional
        DiscoveredBackends []DiscoveredBackend `json:"discoveredBackends,omitempty"`

        // BackendCount is the number of routable backends observed by the runtime.
        // +optional
        BackendCount int32 `json:"backendCount,omitempty"`
}

// VirtualMCPServerStatus defines the observed state of VirtualMCPServer
type VirtualMCPServerStatus struct {
        // Runtime is the status snapshot written exclusively by the vMCP process.
        // The operator projects it into the top-level compatibility fields.
        // +optional
        Runtime *VirtualMCPServerRuntimeStatus `json:"runtime,omitempty"`

        // Conditions represent the latest available observations of the VirtualMCPServer's state
        // +listType=map
        // +listMapKey=type
        // +optional
        Conditions []metav1.Condition `json:"conditions,omitempty"`

        // ObservedGeneration is the most recent generation observed for this VirtualMCPServer
        // +optional
        ObservedGeneration int64 `json:"observedGeneration,omitempty"`

        // Phase is the current phase of the VirtualMCPServer
        // +optional
        // +kubebuilder:default=Pending
        Phase VirtualMCPServerPhase `json:"phase,omitempty"`

        // Message provides additional information about the current phase
        // +optional
        Message string `json:"message,omitempty"`

        // URL is the URL where the Virtual MCP server can be accessed
        // +optional
        URL string `json:"url,omitempty"`

        // DiscoveredBackends lists discovered backend configurations from the MCPGroup
        // +listType=map
        // +listMapKey=name
        // +optional
        DiscoveredBackends []DiscoveredBackend `json:"discoveredBackends,omitempty"`

        // BackendCount is the number of routable backends (ready + unauthenticated).
        // Excludes unavailable, degraded, and unknown backends.
        // +optional
        BackendCount int32 `json:"backendCount,omitempty"`

        // AuthzConfigHash is the hash of the referenced MCPAuthzConfig spec for change detection.
        // Only populated when IncomingAuth.AuthzConfigRef is set.
        // +optional
        AuthzConfigHash string `json:"authzConfigHash,omitempty"`

        // OIDCConfigHash is the hash of the referenced MCPOIDCConfig spec for change detection.
        // Only populated when IncomingAuth.OIDCConfigRef is set.
        // +optional
        OIDCConfigHash string `json:"oidcConfigHash,omitempty"`

        // TelemetryConfigHash is the hash of the referenced MCPTelemetryConfig spec for change detection.
        // Only populated when TelemetryConfigRef is set.
        // +optional
        TelemetryConfigHash string `json:"telemetryConfigHash,omitempty"`
}

// VirtualMCPServerPhase represents the lifecycle phase of a VirtualMCPServer
// +kubebuilder:validation:Enum=Pending;Ready;Degraded;Failed
type VirtualMCPServerPhase string

const (
        // VirtualMCPServerPhasePending indicates the VirtualMCPServer is being initialized
        VirtualMCPServerPhasePending VirtualMCPServerPhase = "Pending"

        // VirtualMCPServerPhaseReady indicates the VirtualMCPServer is ready and serving requests
        VirtualMCPServerPhaseReady VirtualMCPServerPhase = "Ready"

        // VirtualMCPServerPhaseDegraded indicates the VirtualMCPServer is running but some backends are unavailable
        VirtualMCPServerPhaseDegraded VirtualMCPServerPhase = "Degraded"

        // VirtualMCPServerPhaseFailed indicates the VirtualMCPServer has failed
        VirtualMCPServerPhaseFailed VirtualMCPServerPhase = "Failed"
)

// Condition types for VirtualMCPServer
// Note: ConditionTypeAuthConfigured is shared with MCPRemoteProxy and defined in mcpremoteproxy_types.go
const (
        // ConditionTypeVirtualMCPServerReady indicates whether the VirtualMCPServer is ready
        ConditionTypeVirtualMCPServerReady = "Ready"

        // ConditionTypeVirtualMCPServerGroupRefValidated indicates whether the GroupRef is valid
        ConditionTypeVirtualMCPServerGroupRefValidated = "GroupRefValidated"

        // ConditionTypeCompositeToolRefsValidated indicates whether the CompositeToolRefs are valid
        ConditionTypeCompositeToolRefsValidated = "CompositeToolRefsValidated"
        // ConditionTypeVirtualMCPServerPodTemplateSpecValid indicates whether the PodTemplateSpec is valid
        ConditionTypeVirtualMCPServerPodTemplateSpecValid = "PodTemplateSpecValid"

        // ConditionTypeVirtualMCPServerBackendsDiscovered indicates whether backends have been discovered
        ConditionTypeVirtualMCPServerBackendsDiscovered = "BackendsDiscovered"

        // ConditionTypeEmbeddingServerReady indicates whether the EmbeddingServer is ready
        ConditionTypeEmbeddingServerReady = "EmbeddingServerReady"

        // ConditionTypeAuthServerConfigValidated indicates whether the AuthServerConfig has been validated
        ConditionTypeAuthServerConfigValidated = "AuthServerConfigValidated"

        // ConditionTypeAuthzUpstreamSelectionWarning is an advisory condition set to True when
        // multiple AuthServerConfig.UpstreamProviders are configured alongside AuthzConfig.
        // Only the first upstream is authoritative for Cedar claim resolution; this warns the
        // operator that the auto-selection has taken effect and names the selected upstream.
        ConditionTypeAuthzUpstreamSelectionWarning = "AuthzUpstreamSelectionWarning"

        // ConditionTypeVirtualMCPServerTelemetryConfigRefValidated indicates whether the TelemetryConfigRef is valid
        ConditionTypeVirtualMCPServerTelemetryConfigRefValidated = "TelemetryConfigRefValidated"

        // ConditionTypeVirtualMCPServerDeprecatedInboundGrantConfiguration reports
        // whether inline auth uses released legacy inbound grant fields.
        ConditionTypeVirtualMCPServerDeprecatedInboundGrantConfiguration = "DeprecatedInboundGrantConfiguration"
)

// Condition reasons for VirtualMCPServer
const (
        // ConditionReasonIncomingAuthValid indicates incoming auth is valid
        ConditionReasonIncomingAuthValid = "IncomingAuthValid"

        // ConditionReasonIncomingAuthInvalid indicates incoming auth is invalid
        ConditionReasonIncomingAuthInvalid = "IncomingAuthInvalid"

        // Note: ConditionReasonAuthzConfigMapNotFound and ConditionReasonAuthzConfigMapInvalid
        // are shared with MCPRemoteProxy and are declared in mcpremoteproxy_types.go.

        // ConditionReasonGroupRefValid indicates the GroupRef is valid
        ConditionReasonVirtualMCPServerGroupRefValid = "GroupRefValid"

        // ConditionReasonGroupRefNotFound indicates the referenced MCPGroup was not found
        ConditionReasonVirtualMCPServerGroupRefNotFound = "GroupRefNotFound"

        // ConditionReasonGroupRefNotReady indicates the referenced MCPGroup is not ready
        ConditionReasonVirtualMCPServerGroupRefNotReady = "GroupRefNotReady"

        // ConditionReasonCompositeToolRefsValid indicates the CompositeToolRefs are valid
        ConditionReasonCompositeToolRefsValid = "CompositeToolRefsValid"

        // ConditionReasonCompositeToolRefNotFound indicates a referenced VirtualMCPCompositeToolDefinition was not found
        ConditionReasonCompositeToolRefNotFound = "CompositeToolRefNotFound"

        // ConditionReasonCompositeToolRefInvalid indicates a referenced VirtualMCPCompositeToolDefinition is invalid
        ConditionReasonCompositeToolRefInvalid = "CompositeToolRefInvalid"

        // ConditionReasonVirtualMCPServerPodTemplateSpecValid indicates PodTemplateSpec validation succeeded
        ConditionReasonVirtualMCPServerPodTemplateSpecValid = "PodTemplateSpecValid"

        // ConditionReasonVirtualMCPServerPodTemplateSpecInvalid indicates PodTemplateSpec validation failed
        ConditionReasonVirtualMCPServerPodTemplateSpecInvalid = "InvalidPodTemplateSpec"

        // ConditionReasonVirtualMCPServerBackendsDiscoveredSuccessfully indicates backends were discovered successfully
        ConditionReasonVirtualMCPServerBackendsDiscoveredSuccessfully = "BackendsDiscoveredSuccessfully"

        // ConditionReasonVirtualMCPServerBackendDiscoveryFailed indicates backend discovery failed
        ConditionReasonVirtualMCPServerBackendDiscoveryFailed = "BackendDiscoveryFailed"

        // ConditionReasonVirtualMCPServerDeploymentFailed indicates the deployment failed
        ConditionReasonVirtualMCPServerDeploymentFailed = "DeploymentFailed"

        // ConditionReasonVirtualMCPServerDeploymentReady indicates the deployment is ready
        ConditionReasonVirtualMCPServerDeploymentReady = "DeploymentReady"

        // ConditionReasonVirtualMCPServerDeploymentNotReady indicates the deployment is not ready
        ConditionReasonVirtualMCPServerDeploymentNotReady = "DeploymentNotReady"

        // ConditionReasonEmbeddingServerReady indicates the EmbeddingServer is ready
        ConditionReasonEmbeddingServerReady = "EmbeddingServerReady"

        // ConditionReasonEmbeddingServerNotFound indicates the referenced EmbeddingServer was not found
        ConditionReasonEmbeddingServerNotFound = "EmbeddingServerNotFound"

        // ConditionReasonEmbeddingServerNotReady indicates the referenced EmbeddingServer is not ready
        ConditionReasonEmbeddingServerNotReady = "EmbeddingServerNotReady"

        // ConditionReasonAuthServerConfigValid indicates the AuthServerConfig is valid
        ConditionReasonAuthServerConfigValid = "AuthServerConfigValid"

        // ConditionReasonAuthServerConfigInvalid indicates the AuthServerConfig is invalid
        ConditionReasonAuthServerConfigInvalid = "AuthServerConfigInvalid"

        // ConditionReasonVirtualMCPServerLegacyInboundGrantFields indicates that
        // inline auth uses released legacy inbound grant fields.
        ConditionReasonVirtualMCPServerLegacyInboundGrantFields = "LegacyInboundGrantFields"

        // ConditionReasonVirtualMCPServerCanonicalInboundGrantConfiguration indicates
        // that inline auth uses only canonical inbound grant configuration.
        ConditionReasonVirtualMCPServerCanonicalInboundGrantConfiguration = "CanonicalInboundGrantConfiguration"

        // ConditionReasonAuthzRequiresUpstream indicates that authorization policies are
        // configured but no upstream IDP is available to source claims from. Without an
        // upstream, Cedar evaluates against the ToolHive-issued AS token, whose claim
        // namespace (sub, aud, tsid) can overlap upstream claims and silently authorize
        // against the wrong identity.
        ConditionReasonAuthzRequiresUpstream = "AuthzRequiresUpstream"

        // ConditionReasonAuthzUpstreamAutoSelected is set when authorization is configured
        // alongside multiple upstream providers and the first upstream has been chosen as
        // the Cedar claim source. The advisory message names the selected upstream.
        ConditionReasonAuthzUpstreamAutoSelected = "AuthzUpstreamAutoSelected"

        // ConditionReasonAuthzUpstreamUnknown indicates that
        // spec.incomingAuth.authzConfig.inline.primaryUpstreamProvider names an upstream
        // IDP that is not declared on spec.authServerConfig.upstreamProviders. Cedar
        // would otherwise deny every request at runtime; reject at admission instead.
        ConditionReasonAuthzUpstreamUnknown = "AuthzUpstreamUnknown"

        // ConditionReasonAuthzPrimaryProviderRequiresAuthServer indicates that
        // spec.incomingAuth.authzConfig.inline.primaryUpstreamProvider is set but
        // spec.authServerConfig is not configured. The field names an upstream IDP
        // on the embedded auth server, which is required for it to take effect.
        // Distinct from AuthzUpstreamUnknown so tooling (alertmanager rules,
        // dashboards) can route the two misconfigurations separately.
        ConditionReasonAuthzPrimaryProviderRequiresAuthServer = "AuthzPrimaryProviderRequiresAuthServer"

        // ConditionReasonVirtualMCPServerTelemetryConfigRefValid indicates the referenced MCPTelemetryConfig is valid
        ConditionReasonVirtualMCPServerTelemetryConfigRefValid = "TelemetryConfigRefValid"

        // ConditionReasonVirtualMCPServerTelemetryConfigRefNotFound indicates the referenced MCPTelemetryConfig was not found
        ConditionReasonVirtualMCPServerTelemetryConfigRefNotFound = "TelemetryConfigRefNotFound"

        // ConditionReasonVirtualMCPServerTelemetryConfigRefInvalid indicates the referenced MCPTelemetryConfig is not valid
        ConditionReasonVirtualMCPServerTelemetryConfigRefInvalid = "TelemetryConfigRefInvalid"

        // ConditionReasonVirtualMCPServerTelemetryConfigRefFetchError indicates a transient error occurred fetching the config
        ConditionReasonVirtualMCPServerTelemetryConfigRefFetchError = "TelemetryConfigRefFetchError"
)

// Backend authentication types
const (
        // BackendAuthTypeDiscovered automatically discovers from backend's externalAuthConfigRef
        BackendAuthTypeDiscovered = "discovered"

        // BackendAuthTypeExternalAuthConfigRef references an MCPExternalAuthConfig resource
        BackendAuthTypeExternalAuthConfigRef = "externalAuthConfigRef"
)

// Workflow step types
const (
        // WorkflowStepTypeToolCall calls a backend tool
        WorkflowStepTypeToolCall = "tool"

        // WorkflowStepTypeElicitation requests user input
        WorkflowStepTypeElicitation = "elicitation"
)

// Error handling actions
const (
        // ErrorActionAbort aborts the workflow on error
        ErrorActionAbort = "abort"

        // ErrorActionContinue continues the workflow on error
        ErrorActionContinue = "continue"

        // ErrorActionRetry retries the step on error
        ErrorActionRetry = "retry"
)

//+kubebuilder:object:root=true
//+kubebuilder:storageversion
//+kubebuilder:subresource:status
//+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true
//+kubebuilder:resource:shortName=vmcp;virtualmcp,categories=toolhive
//+kubebuilder:printcolumn:name="Phase",type="string",JSONPath=".status.phase",description="The phase of the VirtualMCPServer"
//+kubebuilder:printcolumn:name="URL",type="string",JSONPath=".status.url",description="Virtual MCP server URL"
//+kubebuilder:printcolumn:name="Backends",type="integer",JSONPath=".status.backendCount",description="Discovered backends count"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp",description="Age"
//+kubebuilder:printcolumn:name="Ready",type="string",JSONPath=".status.conditions[?(@.type=='Ready')].status"

// VirtualMCPServer is the Schema for the virtualmcpservers API
// VirtualMCPServer aggregates multiple backend MCPServers into a unified endpoint
type VirtualMCPServer struct {
        metav1.TypeMeta   `json:",inline"` // nolint:revive
        metav1.ObjectMeta `json:"metadata,omitempty"`

        Spec   VirtualMCPServerSpec   `json:"spec,omitempty"`
        Status VirtualMCPServerStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// VirtualMCPServerList contains a list of VirtualMCPServer
type VirtualMCPServerList struct {
        metav1.TypeMeta `json:",inline"` // nolint:revive
        metav1.ListMeta `json:"metadata,omitempty"`
        Items           []VirtualMCPServer `json:"items"`
}

// GetProxyPort returns the proxy port for the VirtualMCPServer.
// vMCP uses port 4483 by default.
func (*VirtualMCPServer) GetProxyPort() int32 {
        return 4483
}

// ResolveGroupName returns the group name from spec.groupRef.
func (r *VirtualMCPServer) ResolveGroupName() string {
        return r.Spec.GroupRef.GetName()
}

// ExplicitPrimaryUpstreamProvider returns the user-configured primary upstream
// provider name and a flag indicating whether the value came from the
// deprecated spec.incomingAuth.authzConfig.inline.primaryUpstreamProvider
// location (fromDeprecated=true) or the canonical
// spec.authServerConfig.primaryUpstreamProvider location (fromDeprecated=false).
// Returns ("", false) when neither location is set.
//
// Precedence: the canonical location wins if set; the deprecated location is
// read only as a backward-compatibility fallback. Callers should emit a
// Warning event when fromDeprecated is true.
func (r *VirtualMCPServer) ExplicitPrimaryUpstreamProvider() (name string, fromDeprecated bool) {
        if r.Spec.AuthServerConfig != nil && r.Spec.AuthServerConfig.PrimaryUpstreamProvider != "" {
                return r.Spec.AuthServerConfig.PrimaryUpstreamProvider, false
        }
        if r.Spec.IncomingAuth != nil {
                if dep := r.Spec.IncomingAuth.AuthzConfig.DeprecatedInlinePrimaryUpstreamProvider(); dep != "" {
                        return dep, true
                }
        }
        return "", false
}

// Validate performs validation for VirtualMCPServer
// This method is called by the controller during reconciliation
func (r *VirtualMCPServer) Validate() error {
        // Validate Group is set — spec.groupRef.name is required
        // Note: CEL cannot validate embedded types from other packages
        if r.Spec.GroupRef.GetName() == "" {
                return fmt.Errorf("spec.groupRef.name is required")
        }

        // Note: IncomingAuth validation is handled by kubebuilder markers and CEL rules

        // Validate OutgoingAuth backend configurations
        if r.Spec.OutgoingAuth != nil {
                for backendName, backendAuth := range r.Spec.OutgoingAuth.Backends {
                        if err := r.validateBackendAuth(backendName, backendAuth); err != nil {
                                return err
                        }
                }
        }

        // Validate Aggregation configuration
        if r.Spec.Config.Aggregation != nil {
                if err := r.validateAggregation(); err != nil {
                        return err
                }
        }

        // Validate CompositeTools
        if len(r.Spec.Config.CompositeTools) > 0 {
                if err := r.validateCompositeTools(); err != nil {
                        return err
                }
        }

        // Note: AuthServerConfig validation is handled by the reconciler (validateAuthServerConfig)
        // so it can set the AuthServerConfigValidated condition on failure.

        // Validate EmbeddingServer / EmbeddingServerRef
        return r.validateEmbeddingServer()
}

// validateEmbeddingServer validates EmbeddingServerRef and Optimizer configuration.
// Rules:
// - embeddingServerRef.name must be non-empty when ref is provided
// - optimizer requires either embeddingServerRef or a manually set embeddingService
// - if embeddingServerRef is set without optimizer, auto-populate optimizer with defaults
//
// The controller handles the remaining cases at runtime (event emission, URL population).
func (r *VirtualMCPServer) validateEmbeddingServer() error {
        // Validate ref name is non-empty
        if r.Spec.EmbeddingServerRef != nil && r.Spec.EmbeddingServerRef.Name == "" {
                return fmt.Errorf("spec.embeddingServerRef.name is required")
        }

        hasOptimizer := r.Spec.Config.Optimizer != nil
        hasRef := r.Spec.EmbeddingServerRef != nil
        hasManualService := hasOptimizer && r.Spec.Config.Optimizer.EmbeddingService != ""

        // Optimizer configured without any embedding source is an error.
        // The user must either set embeddingServerRef or manually set optimizer.embeddingService.
        if hasOptimizer && !hasRef && !hasManualService {
                return fmt.Errorf(
                        "spec.config.optimizer requires an embedding service: " +
                                "set spec.embeddingServerRef (recommended) or spec.config.optimizer.embeddingService")
        }

        // EmbeddingServerRef is set but optimizer is not configured: auto-populate
        // optimizer with default values so the embedding server is actually used.
        // The controller emits a Kubernetes event for this case.
        if hasRef && !hasOptimizer {
                r.Spec.Config.Optimizer = &config.OptimizerConfig{}
        }

        return nil
}

// validateBackendAuth validates a single backend auth configuration
func (*VirtualMCPServer) validateBackendAuth(backendName string, auth BackendAuthConfig) error {
        // Validate type is set
        if auth.Type == "" {
                return fmt.Errorf("spec.outgoingAuth.backends[%s].type is required", backendName)
        }

        // Validate type-specific configurations
        switch auth.Type {
        case BackendAuthTypeExternalAuthConfigRef:
                if auth.ExternalAuthConfigRef == nil {
                        return fmt.Errorf(
                                "spec.outgoingAuth.backends[%s].externalAuthConfigRef is required when type is externalAuthConfigRef",
                                backendName)
                }
                if auth.ExternalAuthConfigRef.Name == "" {
                        return fmt.Errorf("spec.outgoingAuth.backends[%s].externalAuthConfigRef.name is required", backendName)
                }

        case BackendAuthTypeDiscovered:
                // No additional validation needed

        default:
                return fmt.Errorf(
                        "spec.outgoingAuth.backends[%s].type must be one of: discovered, externalAuthConfigRef",
                        backendName)
        }

        return nil
}

// validateAggregation validates Aggregation configuration
func (r *VirtualMCPServer) validateAggregation() error {
        agg := r.Spec.Config.Aggregation

        // Validate conflict resolution strategy
        if agg.ConflictResolution != "" {
                validStrategies := map[vmcptypes.ConflictResolutionStrategy]bool{
                        vmcptypes.ConflictStrategyPrefix:   true,
                        vmcptypes.ConflictStrategyPriority: true,
                        vmcptypes.ConflictStrategyManual:   true,
                }
                if !validStrategies[agg.ConflictResolution] {
                        return fmt.Errorf("config.aggregation.conflictResolution must be one of: prefix, priority, manual")
                }
        }

        // Validate conflict resolution config based on strategy
        if agg.ConflictResolutionConfig != nil {
                resConfig := agg.ConflictResolutionConfig

                switch agg.ConflictResolution {
                case vmcptypes.ConflictStrategyPrefix:
                        // Prefix strategy uses PrefixFormat if specified, otherwise defaults
                        // No additional validation required

                case vmcptypes.ConflictStrategyPriority:
                        if len(resConfig.PriorityOrder) == 0 {
                                return fmt.Errorf("config.aggregation.conflictResolutionConfig.priorityOrder is required when conflictResolution is priority")
                        }

                case vmcptypes.ConflictStrategyManual:
                        // For manual resolution, tools must define explicit overrides
                        // This will be validated at runtime when conflicts are detected
                }
        }

        // Validate per-workload tool configurations
        for i, toolConfig := range agg.Tools {
                if toolConfig.Workload == "" {
                        return fmt.Errorf("config.aggregation.tools[%d].workload is required", i)
                }

                // If ToolConfigRef is specified, ensure it has a name
                if toolConfig.ToolConfigRef != nil && toolConfig.ToolConfigRef.Name == "" {
                        return fmt.Errorf("config.aggregation.tools[%d].toolConfigRef.name is required when toolConfigRef is specified", i)
                }
        }

        return nil
}

// validateCompositeTools validates composite tool definitions in spec.config.compositeTools.
// Uses shared validation from pkg/vmcp/config/composite_validation.go.
func (r *VirtualMCPServer) validateCompositeTools() error {
        toolNames := make(map[string]bool)

        for i := range r.Spec.Config.CompositeTools {
                tool := &r.Spec.Config.CompositeTools[i]

                // Check for duplicate tool names
                if toolNames[tool.Name] {
                        return fmt.Errorf("spec.config.compositeTools[%d].name %q is duplicated", i, tool.Name)
                }
                toolNames[tool.Name] = true

                // Use shared validation
                if err := config.ValidateCompositeToolConfig(
                        fmt.Sprintf("spec.config.compositeTools[%d]", i), tool,
                ); err != nil {
                        return err
                }
        }

        return nil
}

func init() {
        SchemeBuilder.Register(&VirtualMCPServer{}, &VirtualMCPServerList{})
}

//go:build !ignore_autogenerated

/*
Copyright 2025 Stacklok

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Code generated by controller-gen. DO NOT EDIT.

package v1beta1

import (
        "github.com/stacklok/toolhive/pkg/ratelimit/types"
        corev1 "k8s.io/api/core/v1"
        apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
        "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/runtime"
)

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AWSStsConfig) DeepCopyInto(out *AWSStsConfig) {
        *out = *in
        if in.RoleMappings != nil {
                in, out := &in.RoleMappings, &out.RoleMappings
                *out = make([]RoleMapping, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.SessionDuration != nil {
                in, out := &in.SessionDuration, &out.SessionDuration
                *out = new(int32)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AWSStsConfig.
func (in *AWSStsConfig) DeepCopy() *AWSStsConfig {
        if in == nil {
                return nil
        }
        out := new(AWSStsConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AuditConfig) DeepCopyInto(out *AuditConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AuditConfig.
func (in *AuditConfig) DeepCopy() *AuditConfig {
        if in == nil {
                return nil
        }
        out := new(AuditConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AuthServerRef) DeepCopyInto(out *AuthServerRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AuthServerRef.
func (in *AuthServerRef) DeepCopy() *AuthServerRef {
        if in == nil {
                return nil
        }
        out := new(AuthServerRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AuthServerStorageConfig) DeepCopyInto(out *AuthServerStorageConfig) {
        *out = *in
        if in.Redis != nil {
                in, out := &in.Redis, &out.Redis
                *out = new(RedisStorageConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.MaxDCRClients != nil {
                in, out := &in.MaxDCRClients, &out.MaxDCRClients
                *out = new(int32)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AuthServerStorageConfig.
func (in *AuthServerStorageConfig) DeepCopy() *AuthServerStorageConfig {
        if in == nil {
                return nil
        }
        out := new(AuthServerStorageConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AuthServerUpstreamFilterConfig) DeepCopyInto(out *AuthServerUpstreamFilterConfig) {
        *out = *in
        if in.Rules != nil {
                in, out := &in.Rules, &out.Rules
                *out = make([]UpstreamFilterGroupRule, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.DefaultUpstreams != nil {
                in, out := &in.DefaultUpstreams, &out.DefaultUpstreams
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AuthServerUpstreamFilterConfig.
func (in *AuthServerUpstreamFilterConfig) DeepCopy() *AuthServerUpstreamFilterConfig {
        if in == nil {
                return nil
        }
        out := new(AuthServerUpstreamFilterConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AuthzConfigRef) DeepCopyInto(out *AuthzConfigRef) {
        *out = *in
        if in.ConfigMap != nil {
                in, out := &in.ConfigMap, &out.ConfigMap
                *out = new(ConfigMapAuthzRef)
                **out = **in
        }
        if in.Inline != nil {
                in, out := &in.Inline, &out.Inline
                *out = new(InlineAuthzConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AuthzConfigRef.
func (in *AuthzConfigRef) DeepCopy() *AuthzConfigRef {
        if in == nil {
                return nil
        }
        out := new(AuthzConfigRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *BackendAuthConfig) DeepCopyInto(out *BackendAuthConfig) {
        *out = *in
        if in.ExternalAuthConfigRef != nil {
                in, out := &in.ExternalAuthConfigRef, &out.ExternalAuthConfigRef
                *out = new(ExternalAuthConfigRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new BackendAuthConfig.
func (in *BackendAuthConfig) DeepCopy() *BackendAuthConfig {
        if in == nil {
                return nil
        }
        out := new(BackendAuthConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *BearerTokenConfig) DeepCopyInto(out *BearerTokenConfig) {
        *out = *in
        if in.TokenSecretRef != nil {
                in, out := &in.TokenSecretRef, &out.TokenSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new BearerTokenConfig.
func (in *BearerTokenConfig) DeepCopy() *BearerTokenConfig {
        if in == nil {
                return nil
        }
        out := new(BearerTokenConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *CABundleSource) DeepCopyInto(out *CABundleSource) {
        *out = *in
        if in.ConfigMapRef != nil {
                in, out := &in.ConfigMapRef, &out.ConfigMapRef
                *out = new(corev1.ConfigMapKeySelector)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new CABundleSource.
func (in *CABundleSource) DeepCopy() *CABundleSource {
        if in == nil {
                return nil
        }
        out := new(CABundleSource)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ConfigMapAuthzRef) DeepCopyInto(out *ConfigMapAuthzRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ConfigMapAuthzRef.
func (in *ConfigMapAuthzRef) DeepCopy() *ConfigMapAuthzRef {
        if in == nil {
                return nil
        }
        out := new(ConfigMapAuthzRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *DCRUpstreamConfig) DeepCopyInto(out *DCRUpstreamConfig) {
        *out = *in
        if in.InitialAccessTokenRef != nil {
                in, out := &in.InitialAccessTokenRef, &out.InitialAccessTokenRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new DCRUpstreamConfig.
func (in *DCRUpstreamConfig) DeepCopy() *DCRUpstreamConfig {
        if in == nil {
                return nil
        }
        out := new(DCRUpstreamConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *DelegateClientConfig) DeepCopyInto(out *DelegateClientConfig) {
        *out = *in
        if in.ClientSecretRef != nil {
                in, out := &in.ClientSecretRef, &out.ClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Audiences != nil {
                in, out := &in.Audiences, &out.Audiences
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new DelegateClientConfig.
func (in *DelegateClientConfig) DeepCopy() *DelegateClientConfig {
        if in == nil {
                return nil
        }
        out := new(DelegateClientConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddedAuthServerCIMDConfig) DeepCopyInto(out *EmbeddedAuthServerCIMDConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddedAuthServerCIMDConfig.
func (in *EmbeddedAuthServerCIMDConfig) DeepCopy() *EmbeddedAuthServerCIMDConfig {
        if in == nil {
                return nil
        }
        out := new(EmbeddedAuthServerCIMDConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddedAuthServerConfig) DeepCopyInto(out *EmbeddedAuthServerConfig) {
        *out = *in
        if in.SigningKeySecretRefs != nil {
                in, out := &in.SigningKeySecretRefs, &out.SigningKeySecretRefs
                *out = make([]SecretKeyRef, len(*in))
                copy(*out, *in)
        }
        if in.HMACSecretRefs != nil {
                in, out := &in.HMACSecretRefs, &out.HMACSecretRefs
                *out = make([]SecretKeyRef, len(*in))
                copy(*out, *in)
        }
        if in.TokenLifespans != nil {
                in, out := &in.TokenLifespans, &out.TokenLifespans
                *out = new(TokenLifespanConfig)
                **out = **in
        }
        if in.SPIFFETrustDomains != nil {
                in, out := &in.SPIFFETrustDomains, &out.SPIFFETrustDomains
                *out = make([]SPIFFETrustDomainConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.InboundGrants != nil {
                in, out := &in.InboundGrants, &out.InboundGrants
                *out = new(InboundGrantsConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.UpstreamProviders != nil {
                in, out := &in.UpstreamProviders, &out.UpstreamProviders
                *out = make([]UpstreamProviderConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.Storage != nil {
                in, out := &in.Storage, &out.Storage
                *out = new(AuthServerStorageConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.BaselineClientScopes != nil {
                in, out := &in.BaselineClientScopes, &out.BaselineClientScopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.DelegateClients != nil {
                in, out := &in.DelegateClients, &out.DelegateClients
                *out = make([]DelegateClientConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.TrustedIssuers != nil {
                in, out := &in.TrustedIssuers, &out.TrustedIssuers
                *out = make([]TrustedIssuerConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.ForceConfidentialRedirectURIs != nil {
                in, out := &in.ForceConfidentialRedirectURIs, &out.ForceConfidentialRedirectURIs
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.CIMD != nil {
                in, out := &in.CIMD, &out.CIMD
                *out = new(EmbeddedAuthServerCIMDConfig)
                **out = **in
        }
        if in.RegistrationRateLimit != nil {
                in, out := &in.RegistrationRateLimit, &out.RegistrationRateLimit
                *out = new(RateLimitBucket)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddedAuthServerConfig.
func (in *EmbeddedAuthServerConfig) DeepCopy() *EmbeddedAuthServerConfig {
        if in == nil {
                return nil
        }
        out := new(EmbeddedAuthServerConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddingResourceOverrides) DeepCopyInto(out *EmbeddingResourceOverrides) {
        *out = *in
        if in.StatefulSet != nil {
                in, out := &in.StatefulSet, &out.StatefulSet
                *out = new(EmbeddingStatefulSetOverrides)
                (*in).DeepCopyInto(*out)
        }
        if in.Service != nil {
                in, out := &in.Service, &out.Service
                *out = new(ResourceMetadataOverrides)
                (*in).DeepCopyInto(*out)
        }
        if in.PersistentVolumeClaim != nil {
                in, out := &in.PersistentVolumeClaim, &out.PersistentVolumeClaim
                *out = new(ResourceMetadataOverrides)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddingResourceOverrides.
func (in *EmbeddingResourceOverrides) DeepCopy() *EmbeddingResourceOverrides {
        if in == nil {
                return nil
        }
        out := new(EmbeddingResourceOverrides)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddingServer) DeepCopyInto(out *EmbeddingServer) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddingServer.
func (in *EmbeddingServer) DeepCopy() *EmbeddingServer {
        if in == nil {
                return nil
        }
        out := new(EmbeddingServer)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *EmbeddingServer) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddingServerList) DeepCopyInto(out *EmbeddingServerList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]EmbeddingServer, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddingServerList.
func (in *EmbeddingServerList) DeepCopy() *EmbeddingServerList {
        if in == nil {
                return nil
        }
        out := new(EmbeddingServerList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *EmbeddingServerList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddingServerRef) DeepCopyInto(out *EmbeddingServerRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddingServerRef.
func (in *EmbeddingServerRef) DeepCopy() *EmbeddingServerRef {
        if in == nil {
                return nil
        }
        out := new(EmbeddingServerRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddingServerSpec) DeepCopyInto(out *EmbeddingServerSpec) {
        *out = *in
        if in.HFTokenSecretRef != nil {
                in, out := &in.HFTokenSecretRef, &out.HFTokenSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.Args != nil {
                in, out := &in.Args, &out.Args
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Env != nil {
                in, out := &in.Env, &out.Env
                *out = make([]EnvVar, len(*in))
                copy(*out, *in)
        }
        out.Resources = in.Resources
        if in.ModelCache != nil {
                in, out := &in.ModelCache, &out.ModelCache
                *out = new(ModelCacheConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.PodTemplateSpec != nil {
                in, out := &in.PodTemplateSpec, &out.PodTemplateSpec
                *out = new(runtime.RawExtension)
                (*in).DeepCopyInto(*out)
        }
        if in.ResourceOverrides != nil {
                in, out := &in.ResourceOverrides, &out.ResourceOverrides
                *out = new(EmbeddingResourceOverrides)
                (*in).DeepCopyInto(*out)
        }
        if in.Replicas != nil {
                in, out := &in.Replicas, &out.Replicas
                *out = new(int32)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddingServerSpec.
func (in *EmbeddingServerSpec) DeepCopy() *EmbeddingServerSpec {
        if in == nil {
                return nil
        }
        out := new(EmbeddingServerSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddingServerStatus) DeepCopyInto(out *EmbeddingServerStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddingServerStatus.
func (in *EmbeddingServerStatus) DeepCopy() *EmbeddingServerStatus {
        if in == nil {
                return nil
        }
        out := new(EmbeddingServerStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EmbeddingStatefulSetOverrides) DeepCopyInto(out *EmbeddingStatefulSetOverrides) {
        *out = *in
        in.ResourceMetadataOverrides.DeepCopyInto(&out.ResourceMetadataOverrides)
        if in.PodTemplateMetadataOverrides != nil {
                in, out := &in.PodTemplateMetadataOverrides, &out.PodTemplateMetadataOverrides
                *out = new(ResourceMetadataOverrides)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbeddingStatefulSetOverrides.
func (in *EmbeddingStatefulSetOverrides) DeepCopy() *EmbeddingStatefulSetOverrides {
        if in == nil {
                return nil
        }
        out := new(EmbeddingStatefulSetOverrides)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *EnvVar) DeepCopyInto(out *EnvVar) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EnvVar.
func (in *EnvVar) DeepCopy() *EnvVar {
        if in == nil {
                return nil
        }
        out := new(EnvVar)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ExternalAuthConfigRef) DeepCopyInto(out *ExternalAuthConfigRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ExternalAuthConfigRef.
func (in *ExternalAuthConfigRef) DeepCopy() *ExternalAuthConfigRef {
        if in == nil {
                return nil
        }
        out := new(ExternalAuthConfigRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *HeaderForwardConfig) DeepCopyInto(out *HeaderForwardConfig) {
        *out = *in
        if in.AddPlaintextHeaders != nil {
                in, out := &in.AddPlaintextHeaders, &out.AddPlaintextHeaders
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.AddHeadersFromSecret != nil {
                in, out := &in.AddHeadersFromSecret, &out.AddHeadersFromSecret
                *out = make([]HeaderFromSecret, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new HeaderForwardConfig.
func (in *HeaderForwardConfig) DeepCopy() *HeaderForwardConfig {
        if in == nil {
                return nil
        }
        out := new(HeaderForwardConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *HeaderFromSecret) DeepCopyInto(out *HeaderFromSecret) {
        *out = *in
        if in.ValueSecretRef != nil {
                in, out := &in.ValueSecretRef, &out.ValueSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new HeaderFromSecret.
func (in *HeaderFromSecret) DeepCopy() *HeaderFromSecret {
        if in == nil {
                return nil
        }
        out := new(HeaderFromSecret)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *HeaderInjectionConfig) DeepCopyInto(out *HeaderInjectionConfig) {
        *out = *in
        if in.ValueSecretRef != nil {
                in, out := &in.ValueSecretRef, &out.ValueSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new HeaderInjectionConfig.
func (in *HeaderInjectionConfig) DeepCopy() *HeaderInjectionConfig {
        if in == nil {
                return nil
        }
        out := new(HeaderInjectionConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *IdentityFromTokenConfig) DeepCopyInto(out *IdentityFromTokenConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new IdentityFromTokenConfig.
func (in *IdentityFromTokenConfig) DeepCopy() *IdentityFromTokenConfig {
        if in == nil {
                return nil
        }
        out := new(IdentityFromTokenConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *InboundGrantsConfig) DeepCopyInto(out *InboundGrantsConfig) {
        *out = *in
        if in.SPIFFEClientAuth != nil {
                in, out := &in.SPIFFEClientAuth, &out.SPIFFEClientAuth
                *out = make([]SPIFFEClientConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.TokenExchange != nil {
                in, out := &in.TokenExchange, &out.TokenExchange
                *out = new(TokenExchangeInboundGrantConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.JWTBearer != nil {
                in, out := &in.JWTBearer, &out.JWTBearer
                *out = new(JWTBearerInboundGrantConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InboundGrantsConfig.
func (in *InboundGrantsConfig) DeepCopy() *InboundGrantsConfig {
        if in == nil {
                return nil
        }
        out := new(InboundGrantsConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *IncomingAuthConfig) DeepCopyInto(out *IncomingAuthConfig) {
        *out = *in
        if in.OIDCConfigRef != nil {
                in, out := &in.OIDCConfigRef, &out.OIDCConfigRef
                *out = new(MCPOIDCConfigReference)
                (*in).DeepCopyInto(*out)
        }
        if in.AuthzConfig != nil {
                in, out := &in.AuthzConfig, &out.AuthzConfig
                *out = new(AuthzConfigRef)
                (*in).DeepCopyInto(*out)
        }
        if in.AuthzConfigRef != nil {
                in, out := &in.AuthzConfigRef, &out.AuthzConfigRef
                *out = new(MCPAuthzConfigReference)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new IncomingAuthConfig.
func (in *IncomingAuthConfig) DeepCopy() *IncomingAuthConfig {
        if in == nil {
                return nil
        }
        out := new(IncomingAuthConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *InlineAuthzConfig) DeepCopyInto(out *InlineAuthzConfig) {
        *out = *in
        if in.Policies != nil {
                in, out := &in.Policies, &out.Policies
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InlineAuthzConfig.
func (in *InlineAuthzConfig) DeepCopy() *InlineAuthzConfig {
        if in == nil {
                return nil
        }
        out := new(InlineAuthzConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *InlineOIDCSharedConfig) DeepCopyInto(out *InlineOIDCSharedConfig) {
        *out = *in
        if in.ClientSecretRef != nil {
                in, out := &in.ClientSecretRef, &out.ClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.CABundleRef != nil {
                in, out := &in.CABundleRef, &out.CABundleRef
                *out = new(CABundleSource)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InlineOIDCSharedConfig.
func (in *InlineOIDCSharedConfig) DeepCopy() *InlineOIDCSharedConfig {
        if in == nil {
                return nil
        }
        out := new(InlineOIDCSharedConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *JWTBearerGrantConfig) DeepCopyInto(out *JWTBearerGrantConfig) {
        *out = *in
        if in.MaxAssertionAge != nil {
                in, out := &in.MaxAssertionAge, &out.MaxAssertionAge
                *out = new(v1.Duration)
                **out = **in
        }
        if in.SubjectBindings != nil {
                in, out := &in.SubjectBindings, &out.SubjectBindings
                *out = make([]JWTBearerSubjectBinding, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.AcceptedAudiences != nil {
                in, out := &in.AcceptedAudiences, &out.AcceptedAudiences
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new JWTBearerGrantConfig.
func (in *JWTBearerGrantConfig) DeepCopy() *JWTBearerGrantConfig {
        if in == nil {
                return nil
        }
        out := new(JWTBearerGrantConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *JWTBearerInboundGrantConfig) DeepCopyInto(out *JWTBearerInboundGrantConfig) {
        *out = *in
        if in.IssuerPolicies != nil {
                in, out := &in.IssuerPolicies, &out.IssuerPolicies
                *out = make([]JWTBearerIssuerPolicyConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new JWTBearerInboundGrantConfig.
func (in *JWTBearerInboundGrantConfig) DeepCopy() *JWTBearerInboundGrantConfig {
        if in == nil {
                return nil
        }
        out := new(JWTBearerInboundGrantConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *JWTBearerIssuerPolicyConfig) DeepCopyInto(out *JWTBearerIssuerPolicyConfig) {
        *out = *in
        in.JWTBearerGrantConfig.DeepCopyInto(&out.JWTBearerGrantConfig)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new JWTBearerIssuerPolicyConfig.
func (in *JWTBearerIssuerPolicyConfig) DeepCopy() *JWTBearerIssuerPolicyConfig {
        if in == nil {
                return nil
        }
        out := new(JWTBearerIssuerPolicyConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *JWTBearerSubjectBinding) DeepCopyInto(out *JWTBearerSubjectBinding) {
        *out = *in
        if in.AllowedResources != nil {
                in, out := &in.AllowedResources, &out.AllowedResources
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new JWTBearerSubjectBinding.
func (in *JWTBearerSubjectBinding) DeepCopy() *JWTBearerSubjectBinding {
        if in == nil {
                return nil
        }
        out := new(JWTBearerSubjectBinding)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *KubernetesServiceAccountOIDCConfig) DeepCopyInto(out *KubernetesServiceAccountOIDCConfig) {
        *out = *in
        if in.UseClusterAuth != nil {
                in, out := &in.UseClusterAuth, &out.UseClusterAuth
                *out = new(bool)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new KubernetesServiceAccountOIDCConfig.
func (in *KubernetesServiceAccountOIDCConfig) DeepCopy() *KubernetesServiceAccountOIDCConfig {
        if in == nil {
                return nil
        }
        out := new(KubernetesServiceAccountOIDCConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPAuthzConfig) DeepCopyInto(out *MCPAuthzConfig) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPAuthzConfig.
func (in *MCPAuthzConfig) DeepCopy() *MCPAuthzConfig {
        if in == nil {
                return nil
        }
        out := new(MCPAuthzConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPAuthzConfig) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPAuthzConfigList) DeepCopyInto(out *MCPAuthzConfigList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPAuthzConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPAuthzConfigList.
func (in *MCPAuthzConfigList) DeepCopy() *MCPAuthzConfigList {
        if in == nil {
                return nil
        }
        out := new(MCPAuthzConfigList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPAuthzConfigList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPAuthzConfigReference) DeepCopyInto(out *MCPAuthzConfigReference) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPAuthzConfigReference.
func (in *MCPAuthzConfigReference) DeepCopy() *MCPAuthzConfigReference {
        if in == nil {
                return nil
        }
        out := new(MCPAuthzConfigReference)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPAuthzConfigSpec) DeepCopyInto(out *MCPAuthzConfigSpec) {
        *out = *in
        in.Config.DeepCopyInto(&out.Config)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPAuthzConfigSpec.
func (in *MCPAuthzConfigSpec) DeepCopy() *MCPAuthzConfigSpec {
        if in == nil {
                return nil
        }
        out := new(MCPAuthzConfigSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPAuthzConfigStatus) DeepCopyInto(out *MCPAuthzConfigStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPAuthzConfigStatus.
func (in *MCPAuthzConfigStatus) DeepCopy() *MCPAuthzConfigStatus {
        if in == nil {
                return nil
        }
        out := new(MCPAuthzConfigStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPExternalAuthConfig) DeepCopyInto(out *MCPExternalAuthConfig) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPExternalAuthConfig.
func (in *MCPExternalAuthConfig) DeepCopy() *MCPExternalAuthConfig {
        if in == nil {
                return nil
        }
        out := new(MCPExternalAuthConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPExternalAuthConfig) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPExternalAuthConfigList) DeepCopyInto(out *MCPExternalAuthConfigList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPExternalAuthConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPExternalAuthConfigList.
func (in *MCPExternalAuthConfigList) DeepCopy() *MCPExternalAuthConfigList {
        if in == nil {
                return nil
        }
        out := new(MCPExternalAuthConfigList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPExternalAuthConfigList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPExternalAuthConfigSpec) DeepCopyInto(out *MCPExternalAuthConfigSpec) {
        *out = *in
        if in.TokenExchange != nil {
                in, out := &in.TokenExchange, &out.TokenExchange
                *out = new(TokenExchangeConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.HeaderInjection != nil {
                in, out := &in.HeaderInjection, &out.HeaderInjection
                *out = new(HeaderInjectionConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.BearerToken != nil {
                in, out := &in.BearerToken, &out.BearerToken
                *out = new(BearerTokenConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.EmbeddedAuthServer != nil {
                in, out := &in.EmbeddedAuthServer, &out.EmbeddedAuthServer
                *out = new(EmbeddedAuthServerConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.AWSSts != nil {
                in, out := &in.AWSSts, &out.AWSSts
                *out = new(AWSStsConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.UpstreamInject != nil {
                in, out := &in.UpstreamInject, &out.UpstreamInject
                *out = new(UpstreamInjectSpec)
                **out = **in
        }
        if in.OBO != nil {
                in, out := &in.OBO, &out.OBO
                *out = new(OBOConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.XAA != nil {
                in, out := &in.XAA, &out.XAA
                *out = new(XAASpec)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPExternalAuthConfigSpec.
func (in *MCPExternalAuthConfigSpec) DeepCopy() *MCPExternalAuthConfigSpec {
        if in == nil {
                return nil
        }
        out := new(MCPExternalAuthConfigSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPExternalAuthConfigStatus) DeepCopyInto(out *MCPExternalAuthConfigStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPExternalAuthConfigStatus.
func (in *MCPExternalAuthConfigStatus) DeepCopy() *MCPExternalAuthConfigStatus {
        if in == nil {
                return nil
        }
        out := new(MCPExternalAuthConfigStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPGroup) DeepCopyInto(out *MCPGroup) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        out.Spec = in.Spec
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPGroup.
func (in *MCPGroup) DeepCopy() *MCPGroup {
        if in == nil {
                return nil
        }
        out := new(MCPGroup)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPGroup) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPGroupList) DeepCopyInto(out *MCPGroupList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPGroup, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPGroupList.
func (in *MCPGroupList) DeepCopy() *MCPGroupList {
        if in == nil {
                return nil
        }
        out := new(MCPGroupList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPGroupList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPGroupRef) DeepCopyInto(out *MCPGroupRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPGroupRef.
func (in *MCPGroupRef) DeepCopy() *MCPGroupRef {
        if in == nil {
                return nil
        }
        out := new(MCPGroupRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPGroupSpec) DeepCopyInto(out *MCPGroupSpec) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPGroupSpec.
func (in *MCPGroupSpec) DeepCopy() *MCPGroupSpec {
        if in == nil {
                return nil
        }
        out := new(MCPGroupSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPGroupStatus) DeepCopyInto(out *MCPGroupStatus) {
        *out = *in
        if in.Servers != nil {
                in, out := &in.Servers, &out.Servers
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.RemoteProxies != nil {
                in, out := &in.RemoteProxies, &out.RemoteProxies
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Entries != nil {
                in, out := &in.Entries, &out.Entries
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPGroupStatus.
func (in *MCPGroupStatus) DeepCopy() *MCPGroupStatus {
        if in == nil {
                return nil
        }
        out := new(MCPGroupStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPOIDCConfig) DeepCopyInto(out *MCPOIDCConfig) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPOIDCConfig.
func (in *MCPOIDCConfig) DeepCopy() *MCPOIDCConfig {
        if in == nil {
                return nil
        }
        out := new(MCPOIDCConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPOIDCConfig) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPOIDCConfigList) DeepCopyInto(out *MCPOIDCConfigList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPOIDCConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPOIDCConfigList.
func (in *MCPOIDCConfigList) DeepCopy() *MCPOIDCConfigList {
        if in == nil {
                return nil
        }
        out := new(MCPOIDCConfigList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPOIDCConfigList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPOIDCConfigReference) DeepCopyInto(out *MCPOIDCConfigReference) {
        *out = *in
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPOIDCConfigReference.
func (in *MCPOIDCConfigReference) DeepCopy() *MCPOIDCConfigReference {
        if in == nil {
                return nil
        }
        out := new(MCPOIDCConfigReference)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPOIDCConfigSpec) DeepCopyInto(out *MCPOIDCConfigSpec) {
        *out = *in
        if in.KubernetesServiceAccount != nil {
                in, out := &in.KubernetesServiceAccount, &out.KubernetesServiceAccount
                *out = new(KubernetesServiceAccountOIDCConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Inline != nil {
                in, out := &in.Inline, &out.Inline
                *out = new(InlineOIDCSharedConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPOIDCConfigSpec.
func (in *MCPOIDCConfigSpec) DeepCopy() *MCPOIDCConfigSpec {
        if in == nil {
                return nil
        }
        out := new(MCPOIDCConfigSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPOIDCConfigStatus) DeepCopyInto(out *MCPOIDCConfigStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPOIDCConfigStatus.
func (in *MCPOIDCConfigStatus) DeepCopy() *MCPOIDCConfigStatus {
        if in == nil {
                return nil
        }
        out := new(MCPOIDCConfigStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRegistry) DeepCopyInto(out *MCPRegistry) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRegistry.
func (in *MCPRegistry) DeepCopy() *MCPRegistry {
        if in == nil {
                return nil
        }
        out := new(MCPRegistry)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPRegistry) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRegistryList) DeepCopyInto(out *MCPRegistryList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPRegistry, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRegistryList.
func (in *MCPRegistryList) DeepCopy() *MCPRegistryList {
        if in == nil {
                return nil
        }
        out := new(MCPRegistryList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPRegistryList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRegistrySpec) DeepCopyInto(out *MCPRegistrySpec) {
        *out = *in
        if in.Volumes != nil {
                in, out := &in.Volumes, &out.Volumes
                *out = make([]apiextensionsv1.JSON, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.VolumeMounts != nil {
                in, out := &in.VolumeMounts, &out.VolumeMounts
                *out = make([]apiextensionsv1.JSON, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.PGPassSecretRef != nil {
                in, out := &in.PGPassSecretRef, &out.PGPassSecretRef
                *out = new(corev1.SecretKeySelector)
                (*in).DeepCopyInto(*out)
        }
        if in.PodTemplateSpec != nil {
                in, out := &in.PodTemplateSpec, &out.PodTemplateSpec
                *out = new(runtime.RawExtension)
                (*in).DeepCopyInto(*out)
        }
        if in.ImagePullSecrets != nil {
                in, out := &in.ImagePullSecrets, &out.ImagePullSecrets
                *out = make([]corev1.LocalObjectReference, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRegistrySpec.
func (in *MCPRegistrySpec) DeepCopy() *MCPRegistrySpec {
        if in == nil {
                return nil
        }
        out := new(MCPRegistrySpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRegistryStatus) DeepCopyInto(out *MCPRegistryStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRegistryStatus.
func (in *MCPRegistryStatus) DeepCopy() *MCPRegistryStatus {
        if in == nil {
                return nil
        }
        out := new(MCPRegistryStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRemoteProxy) DeepCopyInto(out *MCPRemoteProxy) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRemoteProxy.
func (in *MCPRemoteProxy) DeepCopy() *MCPRemoteProxy {
        if in == nil {
                return nil
        }
        out := new(MCPRemoteProxy)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPRemoteProxy) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRemoteProxyList) DeepCopyInto(out *MCPRemoteProxyList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPRemoteProxy, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRemoteProxyList.
func (in *MCPRemoteProxyList) DeepCopy() *MCPRemoteProxyList {
        if in == nil {
                return nil
        }
        out := new(MCPRemoteProxyList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPRemoteProxyList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRemoteProxySpec) DeepCopyInto(out *MCPRemoteProxySpec) {
        *out = *in
        if in.OIDCConfigRef != nil {
                in, out := &in.OIDCConfigRef, &out.OIDCConfigRef
                *out = new(MCPOIDCConfigReference)
                (*in).DeepCopyInto(*out)
        }
        if in.ExternalAuthConfigRef != nil {
                in, out := &in.ExternalAuthConfigRef, &out.ExternalAuthConfigRef
                *out = new(ExternalAuthConfigRef)
                **out = **in
        }
        if in.AuthServerRef != nil {
                in, out := &in.AuthServerRef, &out.AuthServerRef
                *out = new(AuthServerRef)
                **out = **in
        }
        if in.HeaderForward != nil {
                in, out := &in.HeaderForward, &out.HeaderForward
                *out = new(HeaderForwardConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.AuthzConfig != nil {
                in, out := &in.AuthzConfig, &out.AuthzConfig
                *out = new(AuthzConfigRef)
                (*in).DeepCopyInto(*out)
        }
        if in.AuthzConfigRef != nil {
                in, out := &in.AuthzConfigRef, &out.AuthzConfigRef
                *out = new(MCPAuthzConfigReference)
                **out = **in
        }
        if in.Audit != nil {
                in, out := &in.Audit, &out.Audit
                *out = new(AuditConfig)
                **out = **in
        }
        if in.ToolConfigRef != nil {
                in, out := &in.ToolConfigRef, &out.ToolConfigRef
                *out = new(ToolConfigRef)
                **out = **in
        }
        if in.TelemetryConfigRef != nil {
                in, out := &in.TelemetryConfigRef, &out.TelemetryConfigRef
                *out = new(MCPTelemetryConfigReference)
                **out = **in
        }
        out.Resources = in.Resources
        if in.ServiceAccount != nil {
                in, out := &in.ServiceAccount, &out.ServiceAccount
                *out = new(string)
                **out = **in
        }
        if in.PodTemplateSpec != nil {
                in, out := &in.PodTemplateSpec, &out.PodTemplateSpec
                *out = new(runtime.RawExtension)
                (*in).DeepCopyInto(*out)
        }
        if in.ResourceOverrides != nil {
                in, out := &in.ResourceOverrides, &out.ResourceOverrides
                *out = new(ResourceOverrides)
                (*in).DeepCopyInto(*out)
        }
        if in.GroupRef != nil {
                in, out := &in.GroupRef, &out.GroupRef
                *out = new(MCPGroupRef)
                **out = **in
        }
        if in.Replicas != nil {
                in, out := &in.Replicas, &out.Replicas
                *out = new(int32)
                **out = **in
        }
        if in.SessionStorage != nil {
                in, out := &in.SessionStorage, &out.SessionStorage
                *out = new(SessionStorageConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRemoteProxySpec.
func (in *MCPRemoteProxySpec) DeepCopy() *MCPRemoteProxySpec {
        if in == nil {
                return nil
        }
        out := new(MCPRemoteProxySpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPRemoteProxyStatus) DeepCopyInto(out *MCPRemoteProxyStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPRemoteProxyStatus.
func (in *MCPRemoteProxyStatus) DeepCopy() *MCPRemoteProxyStatus {
        if in == nil {
                return nil
        }
        out := new(MCPRemoteProxyStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServer) DeepCopyInto(out *MCPServer) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServer.
func (in *MCPServer) DeepCopy() *MCPServer {
        if in == nil {
                return nil
        }
        out := new(MCPServer)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPServer) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServerEntry) DeepCopyInto(out *MCPServerEntry) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServerEntry.
func (in *MCPServerEntry) DeepCopy() *MCPServerEntry {
        if in == nil {
                return nil
        }
        out := new(MCPServerEntry)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPServerEntry) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServerEntryList) DeepCopyInto(out *MCPServerEntryList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPServerEntry, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServerEntryList.
func (in *MCPServerEntryList) DeepCopy() *MCPServerEntryList {
        if in == nil {
                return nil
        }
        out := new(MCPServerEntryList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPServerEntryList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServerEntrySpec) DeepCopyInto(out *MCPServerEntrySpec) {
        *out = *in
        if in.GroupRef != nil {
                in, out := &in.GroupRef, &out.GroupRef
                *out = new(MCPGroupRef)
                **out = **in
        }
        if in.ExternalAuthConfigRef != nil {
                in, out := &in.ExternalAuthConfigRef, &out.ExternalAuthConfigRef
                *out = new(ExternalAuthConfigRef)
                **out = **in
        }
        if in.HeaderForward != nil {
                in, out := &in.HeaderForward, &out.HeaderForward
                *out = new(HeaderForwardConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.CABundleRef != nil {
                in, out := &in.CABundleRef, &out.CABundleRef
                *out = new(CABundleSource)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServerEntrySpec.
func (in *MCPServerEntrySpec) DeepCopy() *MCPServerEntrySpec {
        if in == nil {
                return nil
        }
        out := new(MCPServerEntrySpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServerEntryStatus) DeepCopyInto(out *MCPServerEntryStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServerEntryStatus.
func (in *MCPServerEntryStatus) DeepCopy() *MCPServerEntryStatus {
        if in == nil {
                return nil
        }
        out := new(MCPServerEntryStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServerList) DeepCopyInto(out *MCPServerList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPServer, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServerList.
func (in *MCPServerList) DeepCopy() *MCPServerList {
        if in == nil {
                return nil
        }
        out := new(MCPServerList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPServerList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServerSpec) DeepCopyInto(out *MCPServerSpec) {
        *out = *in
        if in.Args != nil {
                in, out := &in.Args, &out.Args
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Env != nil {
                in, out := &in.Env, &out.Env
                *out = make([]EnvVar, len(*in))
                copy(*out, *in)
        }
        if in.Volumes != nil {
                in, out := &in.Volumes, &out.Volumes
                *out = make([]Volume, len(*in))
                copy(*out, *in)
        }
        out.Resources = in.Resources
        if in.Secrets != nil {
                in, out := &in.Secrets, &out.Secrets
                *out = make([]SecretRef, len(*in))
                copy(*out, *in)
        }
        if in.ServiceAccount != nil {
                in, out := &in.ServiceAccount, &out.ServiceAccount
                *out = new(string)
                **out = **in
        }
        if in.PermissionProfile != nil {
                in, out := &in.PermissionProfile, &out.PermissionProfile
                *out = new(PermissionProfileRef)
                **out = **in
        }
        if in.PodTemplateSpec != nil {
                in, out := &in.PodTemplateSpec, &out.PodTemplateSpec
                *out = new(runtime.RawExtension)
                (*in).DeepCopyInto(*out)
        }
        if in.ResourceOverrides != nil {
                in, out := &in.ResourceOverrides, &out.ResourceOverrides
                *out = new(ResourceOverrides)
                (*in).DeepCopyInto(*out)
        }
        if in.OIDCConfigRef != nil {
                in, out := &in.OIDCConfigRef, &out.OIDCConfigRef
                *out = new(MCPOIDCConfigReference)
                (*in).DeepCopyInto(*out)
        }
        if in.AuthzConfig != nil {
                in, out := &in.AuthzConfig, &out.AuthzConfig
                *out = new(AuthzConfigRef)
                (*in).DeepCopyInto(*out)
        }
        if in.AuthzConfigRef != nil {
                in, out := &in.AuthzConfigRef, &out.AuthzConfigRef
                *out = new(MCPAuthzConfigReference)
                **out = **in
        }
        if in.Audit != nil {
                in, out := &in.Audit, &out.Audit
                *out = new(AuditConfig)
                **out = **in
        }
        if in.ToolConfigRef != nil {
                in, out := &in.ToolConfigRef, &out.ToolConfigRef
                *out = new(ToolConfigRef)
                **out = **in
        }
        if in.ExternalAuthConfigRef != nil {
                in, out := &in.ExternalAuthConfigRef, &out.ExternalAuthConfigRef
                *out = new(ExternalAuthConfigRef)
                **out = **in
        }
        if in.WebhookConfigRef != nil {
                in, out := &in.WebhookConfigRef, &out.WebhookConfigRef
                *out = new(WebhookConfigRef)
                **out = **in
        }
        if in.AuthServerRef != nil {
                in, out := &in.AuthServerRef, &out.AuthServerRef
                *out = new(AuthServerRef)
                **out = **in
        }
        if in.TelemetryConfigRef != nil {
                in, out := &in.TelemetryConfigRef, &out.TelemetryConfigRef
                *out = new(MCPTelemetryConfigReference)
                **out = **in
        }
        if in.GroupRef != nil {
                in, out := &in.GroupRef, &out.GroupRef
                *out = new(MCPGroupRef)
                **out = **in
        }
        if in.Replicas != nil {
                in, out := &in.Replicas, &out.Replicas
                *out = new(int32)
                **out = **in
        }
        if in.BackendReplicas != nil {
                in, out := &in.BackendReplicas, &out.BackendReplicas
                *out = new(int32)
                **out = **in
        }
        if in.SessionStorage != nil {
                in, out := &in.SessionStorage, &out.SessionStorage
                *out = new(SessionStorageConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.RateLimiting != nil {
                in, out := &in.RateLimiting, &out.RateLimiting
                *out = new(types.RateLimitConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.ProxyReadTimeout != nil {
                in, out := &in.ProxyReadTimeout, &out.ProxyReadTimeout
                *out = new(v1.Duration)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServerSpec.
func (in *MCPServerSpec) DeepCopy() *MCPServerSpec {
        if in == nil {
                return nil
        }
        out := new(MCPServerSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPServerStatus) DeepCopyInto(out *MCPServerStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPServerStatus.
func (in *MCPServerStatus) DeepCopy() *MCPServerStatus {
        if in == nil {
                return nil
        }
        out := new(MCPServerStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPTelemetryConfig) DeepCopyInto(out *MCPTelemetryConfig) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPTelemetryConfig.
func (in *MCPTelemetryConfig) DeepCopy() *MCPTelemetryConfig {
        if in == nil {
                return nil
        }
        out := new(MCPTelemetryConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPTelemetryConfig) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPTelemetryConfigList) DeepCopyInto(out *MCPTelemetryConfigList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPTelemetryConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPTelemetryConfigList.
func (in *MCPTelemetryConfigList) DeepCopy() *MCPTelemetryConfigList {
        if in == nil {
                return nil
        }
        out := new(MCPTelemetryConfigList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPTelemetryConfigList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPTelemetryConfigReference) DeepCopyInto(out *MCPTelemetryConfigReference) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPTelemetryConfigReference.
func (in *MCPTelemetryConfigReference) DeepCopy() *MCPTelemetryConfigReference {
        if in == nil {
                return nil
        }
        out := new(MCPTelemetryConfigReference)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPTelemetryConfigSpec) DeepCopyInto(out *MCPTelemetryConfigSpec) {
        *out = *in
        if in.OpenTelemetry != nil {
                in, out := &in.OpenTelemetry, &out.OpenTelemetry
                *out = new(MCPTelemetryOTelConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Prometheus != nil {
                in, out := &in.Prometheus, &out.Prometheus
                *out = new(PrometheusConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPTelemetryConfigSpec.
func (in *MCPTelemetryConfigSpec) DeepCopy() *MCPTelemetryConfigSpec {
        if in == nil {
                return nil
        }
        out := new(MCPTelemetryConfigSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPTelemetryConfigStatus) DeepCopyInto(out *MCPTelemetryConfigStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPTelemetryConfigStatus.
func (in *MCPTelemetryConfigStatus) DeepCopy() *MCPTelemetryConfigStatus {
        if in == nil {
                return nil
        }
        out := new(MCPTelemetryConfigStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPTelemetryOTelConfig) DeepCopyInto(out *MCPTelemetryOTelConfig) {
        *out = *in
        if in.Headers != nil {
                in, out := &in.Headers, &out.Headers
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.SensitiveHeaders != nil {
                in, out := &in.SensitiveHeaders, &out.SensitiveHeaders
                *out = make([]SensitiveHeader, len(*in))
                copy(*out, *in)
        }
        if in.ResourceAttributes != nil {
                in, out := &in.ResourceAttributes, &out.ResourceAttributes
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.Metrics != nil {
                in, out := &in.Metrics, &out.Metrics
                *out = new(OpenTelemetryMetricsConfig)
                **out = **in
        }
        if in.Tracing != nil {
                in, out := &in.Tracing, &out.Tracing
                *out = new(OpenTelemetryTracingConfig)
                **out = **in
        }
        if in.CABundleRef != nil {
                in, out := &in.CABundleRef, &out.CABundleRef
                *out = new(CABundleSource)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPTelemetryOTelConfig.
func (in *MCPTelemetryOTelConfig) DeepCopy() *MCPTelemetryOTelConfig {
        if in == nil {
                return nil
        }
        out := new(MCPTelemetryOTelConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPToolConfig) DeepCopyInto(out *MCPToolConfig) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPToolConfig.
func (in *MCPToolConfig) DeepCopy() *MCPToolConfig {
        if in == nil {
                return nil
        }
        out := new(MCPToolConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPToolConfig) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPToolConfigList) DeepCopyInto(out *MCPToolConfigList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]MCPToolConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPToolConfigList.
func (in *MCPToolConfigList) DeepCopy() *MCPToolConfigList {
        if in == nil {
                return nil
        }
        out := new(MCPToolConfigList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *MCPToolConfigList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPToolConfigSpec) DeepCopyInto(out *MCPToolConfigSpec) {
        *out = *in
        if in.ToolsFilter != nil {
                in, out := &in.ToolsFilter, &out.ToolsFilter
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.ToolsOverride != nil {
                in, out := &in.ToolsOverride, &out.ToolsOverride
                *out = make(map[string]ToolOverride, len(*in))
                for key, val := range *in {
                        (*out)[key] = *val.DeepCopy()
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPToolConfigSpec.
func (in *MCPToolConfigSpec) DeepCopy() *MCPToolConfigSpec {
        if in == nil {
                return nil
        }
        out := new(MCPToolConfigSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPToolConfigStatus) DeepCopyInto(out *MCPToolConfigStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPToolConfigStatus.
func (in *MCPToolConfigStatus) DeepCopy() *MCPToolConfigStatus {
        if in == nil {
                return nil
        }
        out := new(MCPToolConfigStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPWebhookConfigSpec) DeepCopyInto(out *MCPWebhookConfigSpec) {
        *out = *in
        if in.Validating != nil {
                in, out := &in.Validating, &out.Validating
                *out = make([]WebhookSpec, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.Mutating != nil {
                in, out := &in.Mutating, &out.Mutating
                *out = make([]WebhookSpec, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPWebhookConfigSpec.
func (in *MCPWebhookConfigSpec) DeepCopy() *MCPWebhookConfigSpec {
        if in == nil {
                return nil
        }
        out := new(MCPWebhookConfigSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *MCPWebhookConfigStatus) DeepCopyInto(out *MCPWebhookConfigStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPWebhookConfigStatus.
func (in *MCPWebhookConfigStatus) DeepCopy() *MCPWebhookConfigStatus {
        if in == nil {
                return nil
        }
        out := new(MCPWebhookConfigStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ModelCacheConfig) DeepCopyInto(out *ModelCacheConfig) {
        *out = *in
        if in.StorageClassName != nil {
                in, out := &in.StorageClassName, &out.StorageClassName
                *out = new(string)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ModelCacheConfig.
func (in *ModelCacheConfig) DeepCopy() *ModelCacheConfig {
        if in == nil {
                return nil
        }
        out := new(ModelCacheConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *NetworkPermissions) DeepCopyInto(out *NetworkPermissions) {
        *out = *in
        if in.Outbound != nil {
                in, out := &in.Outbound, &out.Outbound
                *out = new(OutboundNetworkPermissions)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NetworkPermissions.
func (in *NetworkPermissions) DeepCopy() *NetworkPermissions {
        if in == nil {
                return nil
        }
        out := new(NetworkPermissions)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OAuth2UpstreamConfig) DeepCopyInto(out *OAuth2UpstreamConfig) {
        *out = *in
        if in.UserInfo != nil {
                in, out := &in.UserInfo, &out.UserInfo
                *out = new(UserInfoConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.ClientSecretRef != nil {
                in, out := &in.ClientSecretRef, &out.ClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.TokenResponseMapping != nil {
                in, out := &in.TokenResponseMapping, &out.TokenResponseMapping
                *out = new(TokenResponseMapping)
                **out = **in
        }
        if in.IdentityFromToken != nil {
                in, out := &in.IdentityFromToken, &out.IdentityFromToken
                *out = new(IdentityFromTokenConfig)
                **out = **in
        }
        if in.AdditionalAuthorizationParams != nil {
                in, out := &in.AdditionalAuthorizationParams, &out.AdditionalAuthorizationParams
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.AdditionalTokenParams != nil {
                in, out := &in.AdditionalTokenParams, &out.AdditionalTokenParams
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.CABundleRef != nil {
                in, out := &in.CABundleRef, &out.CABundleRef
                *out = new(CABundleSource)
                (*in).DeepCopyInto(*out)
        }
        if in.DCRConfig != nil {
                in, out := &in.DCRConfig, &out.DCRConfig
                *out = new(DCRUpstreamConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OAuth2UpstreamConfig.
func (in *OAuth2UpstreamConfig) DeepCopy() *OAuth2UpstreamConfig {
        if in == nil {
                return nil
        }
        out := new(OAuth2UpstreamConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OBOConfig) DeepCopyInto(out *OBOConfig) {
        *out = *in
        if in.ClientSecretRef != nil {
                in, out := &in.ClientSecretRef, &out.ClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.CacheSkew != nil {
                in, out := &in.CacheSkew, &out.CacheSkew
                *out = new(v1.Duration)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OBOConfig.
func (in *OBOConfig) DeepCopy() *OBOConfig {
        if in == nil {
                return nil
        }
        out := new(OBOConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OIDCUpstreamConfig) DeepCopyInto(out *OIDCUpstreamConfig) {
        *out = *in
        if in.DCRConfig != nil {
                in, out := &in.DCRConfig, &out.DCRConfig
                *out = new(DCRUpstreamConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.ClientSecretRef != nil {
                in, out := &in.ClientSecretRef, &out.ClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.UserInfoOverride != nil {
                in, out := &in.UserInfoOverride, &out.UserInfoOverride
                *out = new(UserInfoConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.AdditionalAuthorizationParams != nil {
                in, out := &in.AdditionalAuthorizationParams, &out.AdditionalAuthorizationParams
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.AdditionalTokenParams != nil {
                in, out := &in.AdditionalTokenParams, &out.AdditionalTokenParams
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.CABundleRef != nil {
                in, out := &in.CABundleRef, &out.CABundleRef
                *out = new(CABundleSource)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OIDCUpstreamConfig.
func (in *OIDCUpstreamConfig) DeepCopy() *OIDCUpstreamConfig {
        if in == nil {
                return nil
        }
        out := new(OIDCUpstreamConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OpenTelemetryMetricsConfig) DeepCopyInto(out *OpenTelemetryMetricsConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OpenTelemetryMetricsConfig.
func (in *OpenTelemetryMetricsConfig) DeepCopy() *OpenTelemetryMetricsConfig {
        if in == nil {
                return nil
        }
        out := new(OpenTelemetryMetricsConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OpenTelemetryTracingConfig) DeepCopyInto(out *OpenTelemetryTracingConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OpenTelemetryTracingConfig.
func (in *OpenTelemetryTracingConfig) DeepCopy() *OpenTelemetryTracingConfig {
        if in == nil {
                return nil
        }
        out := new(OpenTelemetryTracingConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OutboundNetworkPermissions) DeepCopyInto(out *OutboundNetworkPermissions) {
        *out = *in
        if in.AllowHost != nil {
                in, out := &in.AllowHost, &out.AllowHost
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.AllowPort != nil {
                in, out := &in.AllowPort, &out.AllowPort
                *out = make([]int32, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OutboundNetworkPermissions.
func (in *OutboundNetworkPermissions) DeepCopy() *OutboundNetworkPermissions {
        if in == nil {
                return nil
        }
        out := new(OutboundNetworkPermissions)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OutgoingAuthConfig) DeepCopyInto(out *OutgoingAuthConfig) {
        *out = *in
        if in.Default != nil {
                in, out := &in.Default, &out.Default
                *out = new(BackendAuthConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Backends != nil {
                in, out := &in.Backends, &out.Backends
                *out = make(map[string]BackendAuthConfig, len(*in))
                for key, val := range *in {
                        (*out)[key] = *val.DeepCopy()
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OutgoingAuthConfig.
func (in *OutgoingAuthConfig) DeepCopy() *OutgoingAuthConfig {
        if in == nil {
                return nil
        }
        out := new(OutgoingAuthConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *PermissionProfileRef) DeepCopyInto(out *PermissionProfileRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PermissionProfileRef.
func (in *PermissionProfileRef) DeepCopy() *PermissionProfileRef {
        if in == nil {
                return nil
        }
        out := new(PermissionProfileRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *PermissionProfileSpec) DeepCopyInto(out *PermissionProfileSpec) {
        *out = *in
        if in.Read != nil {
                in, out := &in.Read, &out.Read
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Write != nil {
                in, out := &in.Write, &out.Write
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Network != nil {
                in, out := &in.Network, &out.Network
                *out = new(NetworkPermissions)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PermissionProfileSpec.
func (in *PermissionProfileSpec) DeepCopy() *PermissionProfileSpec {
        if in == nil {
                return nil
        }
        out := new(PermissionProfileSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *PrometheusConfig) DeepCopyInto(out *PrometheusConfig) {
        *out = *in
        if in.MetricsOnTransportPort != nil {
                in, out := &in.MetricsOnTransportPort, &out.MetricsOnTransportPort
                *out = new(bool)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PrometheusConfig.
func (in *PrometheusConfig) DeepCopy() *PrometheusConfig {
        if in == nil {
                return nil
        }
        out := new(PrometheusConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ProxyDeploymentOverrides) DeepCopyInto(out *ProxyDeploymentOverrides) {
        *out = *in
        in.ResourceMetadataOverrides.DeepCopyInto(&out.ResourceMetadataOverrides)
        if in.PodTemplateMetadataOverrides != nil {
                in, out := &in.PodTemplateMetadataOverrides, &out.PodTemplateMetadataOverrides
                *out = new(ResourceMetadataOverrides)
                (*in).DeepCopyInto(*out)
        }
        if in.Env != nil {
                in, out := &in.Env, &out.Env
                *out = make([]EnvVar, len(*in))
                copy(*out, *in)
        }
        if in.ImagePullSecrets != nil {
                in, out := &in.ImagePullSecrets, &out.ImagePullSecrets
                *out = make([]corev1.LocalObjectReference, len(*in))
                copy(*out, *in)
        }
        if in.NodeSelector != nil {
                in, out := &in.NodeSelector, &out.NodeSelector
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.Tolerations != nil {
                in, out := &in.Tolerations, &out.Tolerations
                *out = make([]corev1.Toleration, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.Affinity != nil {
                in, out := &in.Affinity, &out.Affinity
                *out = new(corev1.Affinity)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ProxyDeploymentOverrides.
func (in *ProxyDeploymentOverrides) DeepCopy() *ProxyDeploymentOverrides {
        if in == nil {
                return nil
        }
        out := new(ProxyDeploymentOverrides)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RedisACLUserConfig) DeepCopyInto(out *RedisACLUserConfig) {
        *out = *in
        if in.UsernameSecretRef != nil {
                in, out := &in.UsernameSecretRef, &out.UsernameSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.PasswordSecretRef != nil {
                in, out := &in.PasswordSecretRef, &out.PasswordSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RedisACLUserConfig.
func (in *RedisACLUserConfig) DeepCopy() *RedisACLUserConfig {
        if in == nil {
                return nil
        }
        out := new(RedisACLUserConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RedisSentinelConfig) DeepCopyInto(out *RedisSentinelConfig) {
        *out = *in
        if in.SentinelAddrs != nil {
                in, out := &in.SentinelAddrs, &out.SentinelAddrs
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.SentinelService != nil {
                in, out := &in.SentinelService, &out.SentinelService
                *out = new(SentinelServiceRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RedisSentinelConfig.
func (in *RedisSentinelConfig) DeepCopy() *RedisSentinelConfig {
        if in == nil {
                return nil
        }
        out := new(RedisSentinelConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RedisStorageConfig) DeepCopyInto(out *RedisStorageConfig) {
        *out = *in
        if in.SentinelConfig != nil {
                in, out := &in.SentinelConfig, &out.SentinelConfig
                *out = new(RedisSentinelConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.ACLUserConfig != nil {
                in, out := &in.ACLUserConfig, &out.ACLUserConfig
                *out = new(RedisACLUserConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.TLS != nil {
                in, out := &in.TLS, &out.TLS
                *out = new(RedisTLSConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.SentinelTLS != nil {
                in, out := &in.SentinelTLS, &out.SentinelTLS
                *out = new(RedisTLSConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RedisStorageConfig.
func (in *RedisStorageConfig) DeepCopy() *RedisStorageConfig {
        if in == nil {
                return nil
        }
        out := new(RedisStorageConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RedisTLSConfig) DeepCopyInto(out *RedisTLSConfig) {
        *out = *in
        if in.CACertSecretRef != nil {
                in, out := &in.CACertSecretRef, &out.CACertSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RedisTLSConfig.
func (in *RedisTLSConfig) DeepCopy() *RedisTLSConfig {
        if in == nil {
                return nil
        }
        out := new(RedisTLSConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ResourceList) DeepCopyInto(out *ResourceList) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ResourceList.
func (in *ResourceList) DeepCopy() *ResourceList {
        if in == nil {
                return nil
        }
        out := new(ResourceList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ResourceMetadataOverrides) DeepCopyInto(out *ResourceMetadataOverrides) {
        *out = *in
        if in.Annotations != nil {
                in, out := &in.Annotations, &out.Annotations
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.Labels != nil {
                in, out := &in.Labels, &out.Labels
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ResourceMetadataOverrides.
func (in *ResourceMetadataOverrides) DeepCopy() *ResourceMetadataOverrides {
        if in == nil {
                return nil
        }
        out := new(ResourceMetadataOverrides)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ResourceOverrides) DeepCopyInto(out *ResourceOverrides) {
        *out = *in
        if in.ProxyDeployment != nil {
                in, out := &in.ProxyDeployment, &out.ProxyDeployment
                *out = new(ProxyDeploymentOverrides)
                (*in).DeepCopyInto(*out)
        }
        if in.ProxyService != nil {
                in, out := &in.ProxyService, &out.ProxyService
                *out = new(ResourceMetadataOverrides)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ResourceOverrides.
func (in *ResourceOverrides) DeepCopy() *ResourceOverrides {
        if in == nil {
                return nil
        }
        out := new(ResourceOverrides)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ResourceRequirements) DeepCopyInto(out *ResourceRequirements) {
        *out = *in
        out.Limits = in.Limits
        out.Requests = in.Requests
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ResourceRequirements.
func (in *ResourceRequirements) DeepCopy() *ResourceRequirements {
        if in == nil {
                return nil
        }
        out := new(ResourceRequirements)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RoleMapping) DeepCopyInto(out *RoleMapping) {
        *out = *in
        if in.Priority != nil {
                in, out := &in.Priority, &out.Priority
                *out = new(int32)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RoleMapping.
func (in *RoleMapping) DeepCopy() *RoleMapping {
        if in == nil {
                return nil
        }
        out := new(RoleMapping)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SPIFFEBundleEndpointSourceConfig) DeepCopyInto(out *SPIFFEBundleEndpointSourceConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SPIFFEBundleEndpointSourceConfig.
func (in *SPIFFEBundleEndpointSourceConfig) DeepCopy() *SPIFFEBundleEndpointSourceConfig {
        if in == nil {
                return nil
        }
        out := new(SPIFFEBundleEndpointSourceConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SPIFFEBundleSourceConfig) DeepCopyInto(out *SPIFFEBundleSourceConfig) {
        *out = *in
        if in.Endpoint != nil {
                in, out := &in.Endpoint, &out.Endpoint
                *out = new(SPIFFEBundleEndpointSourceConfig)
                **out = **in
        }
        if in.File != nil {
                in, out := &in.File, &out.File
                *out = new(SPIFFEFileBundleSourceConfig)
                **out = **in
        }
        if in.WorkloadAPI != nil {
                in, out := &in.WorkloadAPI, &out.WorkloadAPI
                *out = new(SPIFFEWorkloadAPIBundleSourceConfig)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SPIFFEBundleSourceConfig.
func (in *SPIFFEBundleSourceConfig) DeepCopy() *SPIFFEBundleSourceConfig {
        if in == nil {
                return nil
        }
        out := new(SPIFFEBundleSourceConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SPIFFEClientConfig) DeepCopyInto(out *SPIFFEClientConfig) {
        *out = *in
        if in.Methods != nil {
                in, out := &in.Methods, &out.Methods
                *out = make([]SPIFFEAuthenticationMethod, len(*in))
                copy(*out, *in)
        }
        if in.Resources != nil {
                in, out := &in.Resources, &out.Resources
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Audiences != nil {
                in, out := &in.Audiences, &out.Audiences
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SPIFFEClientConfig.
func (in *SPIFFEClientConfig) DeepCopy() *SPIFFEClientConfig {
        if in == nil {
                return nil
        }
        out := new(SPIFFEClientConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SPIFFEFileBundleSourceConfig) DeepCopyInto(out *SPIFFEFileBundleSourceConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SPIFFEFileBundleSourceConfig.
func (in *SPIFFEFileBundleSourceConfig) DeepCopy() *SPIFFEFileBundleSourceConfig {
        if in == nil {
                return nil
        }
        out := new(SPIFFEFileBundleSourceConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SPIFFETrustDomainConfig) DeepCopyInto(out *SPIFFETrustDomainConfig) {
        *out = *in
        if in.Methods != nil {
                in, out := &in.Methods, &out.Methods
                *out = make([]SPIFFEAuthenticationMethod, len(*in))
                copy(*out, *in)
        }
        in.BundleSource.DeepCopyInto(&out.BundleSource)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SPIFFETrustDomainConfig.
func (in *SPIFFETrustDomainConfig) DeepCopy() *SPIFFETrustDomainConfig {
        if in == nil {
                return nil
        }
        out := new(SPIFFETrustDomainConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SPIFFEWorkloadAPIBundleSourceConfig) DeepCopyInto(out *SPIFFEWorkloadAPIBundleSourceConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SPIFFEWorkloadAPIBundleSourceConfig.
func (in *SPIFFEWorkloadAPIBundleSourceConfig) DeepCopy() *SPIFFEWorkloadAPIBundleSourceConfig {
        if in == nil {
                return nil
        }
        out := new(SPIFFEWorkloadAPIBundleSourceConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SecretKeyRef) DeepCopyInto(out *SecretKeyRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SecretKeyRef.
func (in *SecretKeyRef) DeepCopy() *SecretKeyRef {
        if in == nil {
                return nil
        }
        out := new(SecretKeyRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SecretRef) DeepCopyInto(out *SecretRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SecretRef.
func (in *SecretRef) DeepCopy() *SecretRef {
        if in == nil {
                return nil
        }
        out := new(SecretRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SensitiveHeader) DeepCopyInto(out *SensitiveHeader) {
        *out = *in
        out.SecretKeyRef = in.SecretKeyRef
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SensitiveHeader.
func (in *SensitiveHeader) DeepCopy() *SensitiveHeader {
        if in == nil {
                return nil
        }
        out := new(SensitiveHeader)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SentinelServiceRef) DeepCopyInto(out *SentinelServiceRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SentinelServiceRef.
func (in *SentinelServiceRef) DeepCopy() *SentinelServiceRef {
        if in == nil {
                return nil
        }
        out := new(SentinelServiceRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SessionStorageConfig) DeepCopyInto(out *SessionStorageConfig) {
        *out = *in
        if in.PasswordRef != nil {
                in, out := &in.PasswordRef, &out.PasswordRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.TLS != nil {
                in, out := &in.TLS, &out.TLS
                *out = new(RedisTLSConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SessionStorageConfig.
func (in *SessionStorageConfig) DeepCopy() *SessionStorageConfig {
        if in == nil {
                return nil
        }
        out := new(SessionStorageConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TokenExchangeConfig) DeepCopyInto(out *TokenExchangeConfig) {
        *out = *in
        if in.ClientSecretRef != nil {
                in, out := &in.ClientSecretRef, &out.ClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TokenExchangeConfig.
func (in *TokenExchangeConfig) DeepCopy() *TokenExchangeConfig {
        if in == nil {
                return nil
        }
        out := new(TokenExchangeConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TokenExchangeInboundGrantConfig) DeepCopyInto(out *TokenExchangeInboundGrantConfig) {
        *out = *in
        if in.DelegateClients != nil {
                in, out := &in.DelegateClients, &out.DelegateClients
                *out = make([]DelegateClientConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.IssuerPolicies != nil {
                in, out := &in.IssuerPolicies, &out.IssuerPolicies
                *out = make([]TokenExchangeIssuerPolicyConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TokenExchangeInboundGrantConfig.
func (in *TokenExchangeInboundGrantConfig) DeepCopy() *TokenExchangeInboundGrantConfig {
        if in == nil {
                return nil
        }
        out := new(TokenExchangeInboundGrantConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TokenExchangeIssuerPolicyConfig) DeepCopyInto(out *TokenExchangeIssuerPolicyConfig) {
        *out = *in
        if in.AllowedActors != nil {
                in, out := &in.AllowedActors, &out.AllowedActors
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.AllowedDelegateClients != nil {
                in, out := &in.AllowedDelegateClients, &out.AllowedDelegateClients
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TokenExchangeIssuerPolicyConfig.
func (in *TokenExchangeIssuerPolicyConfig) DeepCopy() *TokenExchangeIssuerPolicyConfig {
        if in == nil {
                return nil
        }
        out := new(TokenExchangeIssuerPolicyConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TokenLifespanConfig) DeepCopyInto(out *TokenLifespanConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TokenLifespanConfig.
func (in *TokenLifespanConfig) DeepCopy() *TokenLifespanConfig {
        if in == nil {
                return nil
        }
        out := new(TokenLifespanConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TokenResponseMapping) DeepCopyInto(out *TokenResponseMapping) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TokenResponseMapping.
func (in *TokenResponseMapping) DeepCopy() *TokenResponseMapping {
        if in == nil {
                return nil
        }
        out := new(TokenResponseMapping)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ToolAnnotationsOverride) DeepCopyInto(out *ToolAnnotationsOverride) {
        *out = *in
        if in.Title != nil {
                in, out := &in.Title, &out.Title
                *out = new(string)
                **out = **in
        }
        if in.ReadOnlyHint != nil {
                in, out := &in.ReadOnlyHint, &out.ReadOnlyHint
                *out = new(bool)
                **out = **in
        }
        if in.DestructiveHint != nil {
                in, out := &in.DestructiveHint, &out.DestructiveHint
                *out = new(bool)
                **out = **in
        }
        if in.IdempotentHint != nil {
                in, out := &in.IdempotentHint, &out.IdempotentHint
                *out = new(bool)
                **out = **in
        }
        if in.OpenWorldHint != nil {
                in, out := &in.OpenWorldHint, &out.OpenWorldHint
                *out = new(bool)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ToolAnnotationsOverride.
func (in *ToolAnnotationsOverride) DeepCopy() *ToolAnnotationsOverride {
        if in == nil {
                return nil
        }
        out := new(ToolAnnotationsOverride)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ToolConfigRef) DeepCopyInto(out *ToolConfigRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ToolConfigRef.
func (in *ToolConfigRef) DeepCopy() *ToolConfigRef {
        if in == nil {
                return nil
        }
        out := new(ToolConfigRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ToolOverride) DeepCopyInto(out *ToolOverride) {
        *out = *in
        if in.Annotations != nil {
                in, out := &in.Annotations, &out.Annotations
                *out = new(ToolAnnotationsOverride)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ToolOverride.
func (in *ToolOverride) DeepCopy() *ToolOverride {
        if in == nil {
                return nil
        }
        out := new(ToolOverride)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TrustedIssuerConfig) DeepCopyInto(out *TrustedIssuerConfig) {
        *out = *in
        if in.CABundleRef != nil {
                in, out := &in.CABundleRef, &out.CABundleRef
                *out = new(CABundleSource)
                (*in).DeepCopyInto(*out)
        }
        if in.AllowedActors != nil {
                in, out := &in.AllowedActors, &out.AllowedActors
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.AllowedDelegateClients != nil {
                in, out := &in.AllowedDelegateClients, &out.AllowedDelegateClients
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.JWTBearerGrant != nil {
                in, out := &in.JWTBearerGrant, &out.JWTBearerGrant
                *out = new(JWTBearerGrantConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TrustedIssuerConfig.
func (in *TrustedIssuerConfig) DeepCopy() *TrustedIssuerConfig {
        if in == nil {
                return nil
        }
        out := new(TrustedIssuerConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *UpstreamFilterGroupRule) DeepCopyInto(out *UpstreamFilterGroupRule) {
        *out = *in
        if in.Groups != nil {
                in, out := &in.Groups, &out.Groups
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.UpstreamProviders != nil {
                in, out := &in.UpstreamProviders, &out.UpstreamProviders
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new UpstreamFilterGroupRule.
func (in *UpstreamFilterGroupRule) DeepCopy() *UpstreamFilterGroupRule {
        if in == nil {
                return nil
        }
        out := new(UpstreamFilterGroupRule)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *UpstreamInjectSpec) DeepCopyInto(out *UpstreamInjectSpec) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new UpstreamInjectSpec.
func (in *UpstreamInjectSpec) DeepCopy() *UpstreamInjectSpec {
        if in == nil {
                return nil
        }
        out := new(UpstreamInjectSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *UpstreamProviderConfig) DeepCopyInto(out *UpstreamProviderConfig) {
        *out = *in
        if in.OIDCConfig != nil {
                in, out := &in.OIDCConfig, &out.OIDCConfig
                *out = new(OIDCUpstreamConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.OAuth2Config != nil {
                in, out := &in.OAuth2Config, &out.OAuth2Config
                *out = new(OAuth2UpstreamConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new UpstreamProviderConfig.
func (in *UpstreamProviderConfig) DeepCopy() *UpstreamProviderConfig {
        if in == nil {
                return nil
        }
        out := new(UpstreamProviderConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *UserInfoConfig) DeepCopyInto(out *UserInfoConfig) {
        *out = *in
        if in.AdditionalHeaders != nil {
                in, out := &in.AdditionalHeaders, &out.AdditionalHeaders
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.FieldMapping != nil {
                in, out := &in.FieldMapping, &out.FieldMapping
                *out = new(UserInfoFieldMapping)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new UserInfoConfig.
func (in *UserInfoConfig) DeepCopy() *UserInfoConfig {
        if in == nil {
                return nil
        }
        out := new(UserInfoConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *UserInfoFieldMapping) DeepCopyInto(out *UserInfoFieldMapping) {
        *out = *in
        if in.SubjectFields != nil {
                in, out := &in.SubjectFields, &out.SubjectFields
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.NameFields != nil {
                in, out := &in.NameFields, &out.NameFields
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.EmailFields != nil {
                in, out := &in.EmailFields, &out.EmailFields
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new UserInfoFieldMapping.
func (in *UserInfoFieldMapping) DeepCopy() *UserInfoFieldMapping {
        if in == nil {
                return nil
        }
        out := new(UserInfoFieldMapping)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPCompositeToolDefinition) DeepCopyInto(out *VirtualMCPCompositeToolDefinition) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPCompositeToolDefinition.
func (in *VirtualMCPCompositeToolDefinition) DeepCopy() *VirtualMCPCompositeToolDefinition {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPCompositeToolDefinition)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *VirtualMCPCompositeToolDefinition) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPCompositeToolDefinitionList) DeepCopyInto(out *VirtualMCPCompositeToolDefinitionList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]VirtualMCPCompositeToolDefinition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPCompositeToolDefinitionList.
func (in *VirtualMCPCompositeToolDefinitionList) DeepCopy() *VirtualMCPCompositeToolDefinitionList {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPCompositeToolDefinitionList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *VirtualMCPCompositeToolDefinitionList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPCompositeToolDefinitionSpec) DeepCopyInto(out *VirtualMCPCompositeToolDefinitionSpec) {
        *out = *in
        in.CompositeToolConfig.DeepCopyInto(&out.CompositeToolConfig)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPCompositeToolDefinitionSpec.
func (in *VirtualMCPCompositeToolDefinitionSpec) DeepCopy() *VirtualMCPCompositeToolDefinitionSpec {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPCompositeToolDefinitionSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPCompositeToolDefinitionStatus) DeepCopyInto(out *VirtualMCPCompositeToolDefinitionStatus) {
        *out = *in
        if in.ValidationErrors != nil {
                in, out := &in.ValidationErrors, &out.ValidationErrors
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.ReferencingVirtualServers != nil {
                in, out := &in.ReferencingVirtualServers, &out.ReferencingVirtualServers
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPCompositeToolDefinitionStatus.
func (in *VirtualMCPCompositeToolDefinitionStatus) DeepCopy() *VirtualMCPCompositeToolDefinitionStatus {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPCompositeToolDefinitionStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPServer) DeepCopyInto(out *VirtualMCPServer) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ObjectMeta.DeepCopyInto(&out.ObjectMeta)
        in.Spec.DeepCopyInto(&out.Spec)
        in.Status.DeepCopyInto(&out.Status)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPServer.
func (in *VirtualMCPServer) DeepCopy() *VirtualMCPServer {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPServer)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *VirtualMCPServer) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPServerList) DeepCopyInto(out *VirtualMCPServerList) {
        *out = *in
        out.TypeMeta = in.TypeMeta
        in.ListMeta.DeepCopyInto(&out.ListMeta)
        if in.Items != nil {
                in, out := &in.Items, &out.Items
                *out = make([]VirtualMCPServer, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPServerList.
func (in *VirtualMCPServerList) DeepCopy() *VirtualMCPServerList {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPServerList)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object.
func (in *VirtualMCPServerList) DeepCopyObject() runtime.Object {
        if c := in.DeepCopy(); c != nil {
                return c
        }
        return nil
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPServerRuntimeStatus) DeepCopyInto(out *VirtualMCPServerRuntimeStatus) {
        *out = *in
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.DiscoveredBackends != nil {
                in, out := &in.DiscoveredBackends, &out.DiscoveredBackends
                *out = make([]DiscoveredBackend, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPServerRuntimeStatus.
func (in *VirtualMCPServerRuntimeStatus) DeepCopy() *VirtualMCPServerRuntimeStatus {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPServerRuntimeStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPServerSpec) DeepCopyInto(out *VirtualMCPServerSpec) {
        *out = *in
        if in.IncomingAuth != nil {
                in, out := &in.IncomingAuth, &out.IncomingAuth
                *out = new(IncomingAuthConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.OutgoingAuth != nil {
                in, out := &in.OutgoingAuth, &out.OutgoingAuth
                *out = new(OutgoingAuthConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.PassthroughHeaders != nil {
                in, out := &in.PassthroughHeaders, &out.PassthroughHeaders
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.ServiceAccount != nil {
                in, out := &in.ServiceAccount, &out.ServiceAccount
                *out = new(string)
                **out = **in
        }
        if in.PodTemplateSpec != nil {
                in, out := &in.PodTemplateSpec, &out.PodTemplateSpec
                *out = new(runtime.RawExtension)
                (*in).DeepCopyInto(*out)
        }
        if in.GroupRef != nil {
                in, out := &in.GroupRef, &out.GroupRef
                *out = new(MCPGroupRef)
                **out = **in
        }
        in.Config.DeepCopyInto(&out.Config)
        if in.TelemetryConfigRef != nil {
                in, out := &in.TelemetryConfigRef, &out.TelemetryConfigRef
                *out = new(MCPTelemetryConfigReference)
                **out = **in
        }
        if in.EmbeddingServerRef != nil {
                in, out := &in.EmbeddingServerRef, &out.EmbeddingServerRef
                *out = new(EmbeddingServerRef)
                **out = **in
        }
        if in.AuthServerConfig != nil {
                in, out := &in.AuthServerConfig, &out.AuthServerConfig
                *out = new(EmbeddedAuthServerConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.AuthServerUpstreamFilter != nil {
                in, out := &in.AuthServerUpstreamFilter, &out.AuthServerUpstreamFilter
                *out = new(AuthServerUpstreamFilterConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Replicas != nil {
                in, out := &in.Replicas, &out.Replicas
                *out = new(int32)
                **out = **in
        }
        if in.SessionStorage != nil {
                in, out := &in.SessionStorage, &out.SessionStorage
                *out = new(SessionStorageConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.ImagePullSecrets != nil {
                in, out := &in.ImagePullSecrets, &out.ImagePullSecrets
                *out = make([]corev1.LocalObjectReference, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPServerSpec.
func (in *VirtualMCPServerSpec) DeepCopy() *VirtualMCPServerSpec {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPServerSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *VirtualMCPServerStatus) DeepCopyInto(out *VirtualMCPServerStatus) {
        *out = *in
        if in.Runtime != nil {
                in, out := &in.Runtime, &out.Runtime
                *out = new(VirtualMCPServerRuntimeStatus)
                (*in).DeepCopyInto(*out)
        }
        if in.Conditions != nil {
                in, out := &in.Conditions, &out.Conditions
                *out = make([]v1.Condition, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.DiscoveredBackends != nil {
                in, out := &in.DiscoveredBackends, &out.DiscoveredBackends
                *out = make([]DiscoveredBackend, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VirtualMCPServerStatus.
func (in *VirtualMCPServerStatus) DeepCopy() *VirtualMCPServerStatus {
        if in == nil {
                return nil
        }
        out := new(VirtualMCPServerStatus)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *Volume) DeepCopyInto(out *Volume) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Volume.
func (in *Volume) DeepCopy() *Volume {
        if in == nil {
                return nil
        }
        out := new(Volume)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *WebhookConfigRef) DeepCopyInto(out *WebhookConfigRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new WebhookConfigRef.
func (in *WebhookConfigRef) DeepCopy() *WebhookConfigRef {
        if in == nil {
                return nil
        }
        out := new(WebhookConfigRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *WebhookSpec) DeepCopyInto(out *WebhookSpec) {
        *out = *in
        if in.Timeout != nil {
                in, out := &in.Timeout, &out.Timeout
                *out = new(v1.Duration)
                **out = **in
        }
        if in.TLSConfig != nil {
                in, out := &in.TLSConfig, &out.TLSConfig
                *out = new(WebhookTLSConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.HMACSecretRef != nil {
                in, out := &in.HMACSecretRef, &out.HMACSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new WebhookSpec.
func (in *WebhookSpec) DeepCopy() *WebhookSpec {
        if in == nil {
                return nil
        }
        out := new(WebhookSpec)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *WebhookTLSConfig) DeepCopyInto(out *WebhookTLSConfig) {
        *out = *in
        if in.CASecretRef != nil {
                in, out := &in.CASecretRef, &out.CASecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.ClientCertSecretRef != nil {
                in, out := &in.ClientCertSecretRef, &out.ClientCertSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.ClientKeySecretRef != nil {
                in, out := &in.ClientKeySecretRef, &out.ClientKeySecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new WebhookTLSConfig.
func (in *WebhookTLSConfig) DeepCopy() *WebhookTLSConfig {
        if in == nil {
                return nil
        }
        out := new(WebhookTLSConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *WorkloadReference) DeepCopyInto(out *WorkloadReference) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new WorkloadReference.
func (in *WorkloadReference) DeepCopy() *WorkloadReference {
        if in == nil {
                return nil
        }
        out := new(WorkloadReference)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *XAASpec) DeepCopyInto(out *XAASpec) {
        *out = *in
        if in.IDPClientSecretRef != nil {
                in, out := &in.IDPClientSecretRef, &out.IDPClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.TargetClientSecretRef != nil {
                in, out := &in.TargetClientSecretRef, &out.TargetClientSecretRef
                *out = new(SecretKeyRef)
                **out = **in
        }
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new XAASpec.
func (in *XAASpec) DeepCopy() *XAASpec {
        if in == nil {
                return nil
        }
        out := new(XAASpec)
        in.DeepCopyInto(out)
        return out
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package audit provides audit logging functionality for ToolHive.
package audit

import (
        "bytes"
        "context"
        "encoding/json"
        "io"
        "log/slog"
        "net"
        "net/http"
        "os"
        "strings"
        "time"

        coreaudit "github.com/stacklok/toolhive-core/audit"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// LevelAudit is a custom audit log level - between Info and Warn
const LevelAudit = coreaudit.LevelAudit

// contextKey is an unexported type for context keys to avoid collisions
type contextKey struct{}

// backendInfoKey is the context key for storing backend routing information
var backendInfoKey = contextKey{}

// BackendInfo stores backend routing information that can be mutated by handlers.
// This allows handlers deep in the call stack to provide backend info to the audit middleware.
type BackendInfo struct {
        BackendName string
}

// WithBackendInfo returns a new context with BackendInfo attached.
func WithBackendInfo(ctx context.Context, info *BackendInfo) context.Context {
        return context.WithValue(ctx, backendInfoKey, info)
}

// BackendInfoFromContext retrieves BackendInfo from the context.
// Returns (nil, false) if BackendInfo is not found in the context.
func BackendInfoFromContext(ctx context.Context) (*BackendInfo, bool) {
        info, ok := ctx.Value(backendInfoKey).(*BackendInfo)
        return info, ok
}

// NewAuditLogger creates a new structured audit logger that writes to the specified writer.
func NewAuditLogger(w io.Writer) *slog.Logger {
        return coreaudit.NewAuditLogger(w)
}

// Auditor handles audit logging for HTTP requests.
type Auditor struct {
        config        *Config
        auditLogger   *slog.Logger
        transportType string // e.g., "sse", "streamable-http"
        logWriter     io.Writer
}

// NewAuditorWithTransport creates a new Auditor with the given configuration and transport information.
func NewAuditorWithTransport(config *Config, transportType string) (*Auditor, error) {
        var logWriter io.Writer = os.Stdout // default to stdout

        if config != nil {
                w, err := config.GetLogWriter()
                if err != nil {
                        // Log error and fall back to stdout
                        slog.Error("failed to open audit log file, falling back to stdout",
                                "error", err)
                        return nil, err
                }
                logWriter = w
        }

        return &Auditor{
                config:        config,
                auditLogger:   NewAuditLogger(logWriter),
                transportType: transportType,
                logWriter:     logWriter,
        }, nil
}

// Close closes the underlying log writer if it implements io.Closer.
// This should be called when the auditor is no longer needed to properly release resources.
func (a *Auditor) Close() error {
        return closeLogWriter(a.logWriter)
}

func closeLogWriter(logWriter io.Writer) error {
        if logWriter == os.Stdout || logWriter == os.Stderr {
                return nil
        }

        if closer, ok := logWriter.(io.Closer); ok {
                return closer.Close()
        }
        return nil
}

// isSSETransport checks if the current transport is SSE
func (a *Auditor) isSSETransport() bool {
        return a.transportType == types.TransportTypeSSE.String()
}

// errorDetectionBufferSize is the maximum number of bytes buffered from the
// response body for JSON-RPC error detection. JSON-RPC error responses have
// the "error" field near the top of the object, so a small prefix is
// sufficient. This buffer is allocated independently of IncludeResponseData.
const errorDetectionBufferSize = 512

// maxAuditErrorMessageLength caps the JSON-RPC error message length stored
// in audit event metadata to keep log entries compact.
const maxAuditErrorMessageLength = 256

// responseWriter wraps http.ResponseWriter to capture response data and status.
type responseWriter struct {
        http.ResponseWriter
        statusCode int
        body       *bytes.Buffer
        // errorDetectionBody is a small prefix buffer used exclusively for
        // JSON-RPC error detection. It is allocated when DetectApplicationErrors
        // is true, independent of IncludeResponseData.
        errorDetectionBody *bytes.Buffer
        auditor            *Auditor
}

func (rw *responseWriter) WriteHeader(statusCode int) {
        rw.statusCode = statusCode
        rw.ResponseWriter.WriteHeader(statusCode)
}

func (rw *responseWriter) Write(data []byte) (int, error) {
        // Capture response data if configured
        if rw.auditor.config.IncludeResponseData && rw.body != nil {
                // Limit the size of captured data
                if rw.body.Len()+len(data) <= rw.auditor.config.MaxDataSize {
                        rw.body.Write(data)
                }
        }
        // Capture a small prefix for JSON-RPC error detection
        if rw.errorDetectionBody != nil && rw.errorDetectionBody.Len() < errorDetectionBufferSize {
                remaining := errorDetectionBufferSize - rw.errorDetectionBody.Len()
                if len(data) <= remaining {
                        rw.errorDetectionBody.Write(data)
                } else {
                        rw.errorDetectionBody.Write(data[:remaining])
                }
        }
        return rw.ResponseWriter.Write(data)
}

// Flush implements http.Flusher if the underlying ResponseWriter supports it.
func (rw *responseWriter) Flush() {
        if flusher, ok := rw.ResponseWriter.(http.Flusher); ok {
                flusher.Flush()
        }
}

// Unwrap exposes the underlying ResponseWriter so http.ResponseController
// can reach interfaces this wrapper does not re-implement (e.g. SetWriteDeadline).
func (rw *responseWriter) Unwrap() http.ResponseWriter {
        return rw.ResponseWriter
}

// isMCPStreamOpenRequest returns true only for MCP "stream" opens:
// - SSE transport's SSE endpoint (GET + Accept: text/event-stream)
// - Streamable HTTP's GET stream (same header pattern)
// Everything else (including POST message sends) is non-sticky.
func (*Auditor) isMCPStreamOpenRequest(r *http.Request) bool {
        // Optional hardening: limit to your MCP base path(s)
        // if !strings.HasPrefix(r.URL.Path, a.config.MCPBasePath) { return false }

        if r.Method != http.MethodGet {
                return false
        }
        accept := r.Header.Get("Accept")
        return strings.Contains(strings.ToLower(accept), "text/event-stream")
}

// ensureAuditContext injects the mutable carriers the auditor reads after the
// inner chain returns: BackendInfo (backend routing), an auth.IdentityHolder
// (identity attached by an auth middleware running INSIDE audit), an
// mcp.ParsedRequestHolder (parsed MCP data from a parser running INSIDE
// audit), and an mcp.AuthzDenialMarker (pre-parse refusals flagged by the
// authz middleware running INSIDE audit). Each is only injected when absent
// so nested auditors share carriers.
func ensureAuditContext(r *http.Request) *http.Request {
        ctx := r.Context()
        changed := false
        if _, ok := BackendInfoFromContext(ctx); !ok {
                ctx = WithBackendInfo(ctx, &BackendInfo{})
                changed = true
        }
        if _, ok := auth.IdentityHolderFromContext(ctx); !ok {
                ctx = auth.WithIdentityHolder(ctx, &auth.IdentityHolder{})
                changed = true
        }
        if _, ok := mcp.ParsedRequestHolderFromContext(ctx); !ok {
                ctx = mcp.WithParsedRequestHolder(ctx, &mcp.ParsedRequestHolder{})
                changed = true
        }
        if _, ok := mcp.AuthzDenialMarkerFromContext(ctx); !ok {
                ctx = mcp.WithAuthzDenialMarker(ctx, &mcp.AuthzDenialMarker{})
                changed = true
        }
        if !changed {
                return r
        }
        return r.WithContext(ctx)
}

// Middleware creates an HTTP middleware that logs audit events.
func (a *Auditor) Middleware(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                r = ensureAuditContext(r)

                // Handle MCP stream opens (SSE endpoint, streamable GET) specially:
                // these connections are long-lived, so instead of waiting for the
                // response to complete, the connection event is logged on the FIRST
                // write. By then any inner auth middleware has run, so the event
                // carries the authenticated identity (or records the 401/403 denial).
                if a.isMCPStreamOpenRequest(r) {
                        sw := &streamOpenWriter{ResponseWriter: w, auditor: a, req: r}
                        // Deferred BEFORE ServeHTTP so a panic in an inner handler still
                        // produces the connection event during unwinding — some chains run
                        // the recovery middleware OUTSIDE audit (e.g. the vMCP Serve path),
                        // which would otherwise swallow the event entirely. If the stream
                        // already logged on first write this is a no-op.
                        defer sw.logOnce(http.StatusOK)
                        next.ServeHTTP(sw, r)
                        return
                }

                startTime := time.Now()

                // Capture request data if configured
                var requestData []byte
                if a.config.IncludeRequestData && r.Body != nil {
                        body, err := io.ReadAll(r.Body)
                        if err == nil {
                                // Always restore the body for the next handler
                                r.Body = io.NopCloser(bytes.NewReader(body))
                                // Only capture for auditing if within size limit
                                if len(body) <= a.config.MaxDataSize {
                                        requestData = body
                                }
                        }
                }

                // Wrap the response writer to capture response data and status
                rw := &responseWriter{
                        ResponseWriter: w,
                        statusCode:     http.StatusOK, // Default status
                        auditor:        a,
                }

                if a.config.IncludeResponseData {
                        rw.body = &bytes.Buffer{}
                }

                // Allocate a small prefix buffer for JSON-RPC error detection,
                // independent of IncludeResponseData. When IncludeResponseData
                // is already true, we reuse rw.body instead of double-buffering.
                if a.config.ShouldDetectApplicationErrors() && !a.config.IncludeResponseData {
                        rw.errorDetectionBody = &bytes.Buffer{}
                }

                // Process the request
                next.ServeHTTP(rw, r)

                // Calculate duration
                duration := time.Since(startTime)

                // Create and log the audit event
                a.logAuditEvent(r, rw, requestData, duration)
        })
}

// logAuditEvent creates and logs an audit event for the HTTP request.
func (a *Auditor) logAuditEvent(r *http.Request, rw *responseWriter, requestData []byte, duration time.Duration) {
        // Determine event type based on the request
        eventType := a.determineEventType(r)

        // Determine outcome based on status code
        outcome := a.determineOutcome(rw.statusCode)

        // A refusal by the authz middleware before message-level authorization
        // could run (e.g. a client response or missing parsing middleware)
        // writes a 400, which determineOutcome maps to a generic failure. The
        // marker reclassifies it as a denial so blocked sweeps are alertable.
        if marker, ok := mcp.AuthzDenialMarkerFromContext(r.Context()); ok && marker.Denied {
                outcome = OutcomeDenied
        }

        // When HTTP status indicates success, check for JSON-RPC errors
        // hidden inside HTTP 200 responses.
        var mcpResponse *mcp.ParsedMCPResponse
        if outcome == OutcomeSuccess && a.config.ShouldDetectApplicationErrors() {
                mcpResponse = a.detectApplicationError(rw)
                if mcpResponse != nil && mcpResponse.HasError {
                        outcome = OutcomeApplicationError
                }
        }

        // Check if we should audit this event
        if !a.config.ShouldAuditEvent(eventType) {
                return
        }

        // Extract source information
        source := a.extractSource(r)

        // Extract subject information
        subjects := a.extractSubjects(r)

        // Determine component name
        component := a.determineComponent(r)

        // Create the audit event
        event := NewAuditEvent(eventType, source, outcome, subjects, component)

        // Attach the RFC 8693 delegation chain, if the caller used a delegated token
        event.WithDelegationChain(a.extractDelegationChain(r))

        // Add target information
        target := a.extractTarget(r, eventType)
        if len(target) > 0 {
                event.WithTarget(target)
        }

        // Add metadata
        a.addMetadata(event, r, duration, rw)

        // Attach JSON-RPC error details so operators can see the error code
        // and message without enabling full response data capture.
        if outcome == OutcomeApplicationError {
                if event.Metadata.Extra == nil {
                        event.Metadata.Extra = make(map[string]any)
                }
                event.Metadata.Extra["jsonrpc_error_code"] = mcpResponse.ErrorCode
                msg := mcpResponse.ErrorMessage
                if len(msg) > maxAuditErrorMessageLength {
                        msg = msg[:maxAuditErrorMessageLength]
                }
                event.Metadata.Extra["jsonrpc_error_message"] = msg
        }

        // Add request/response data if configured
        a.addEventData(event, r, rw, requestData)

        // Log the audit event
        event.LogTo(r.Context(), a.auditLogger, LevelAudit)
}

// mcpMethodFor returns the parsed MCP method for the request, whether the
// parser ran outside audit (context value) or inside it (holder filled by
// the parser and read back after the inner chain returns).
func mcpMethodFor(r *http.Request) string {
        if m := mcp.GetMCPMethod(r.Context()); m != "" {
                return m
        }
        if holder, ok := mcp.ParsedRequestHolderFromContext(r.Context()); ok && holder.Parsed != nil {
                return holder.Parsed.Method
        }
        return ""
}

// mcpResourceIDFor returns the parsed MCP resource ID for the request, with
// the same context-then-holder fallback as mcpMethodFor.
func mcpResourceIDFor(r *http.Request) string {
        if id := mcp.GetMCPResourceID(r.Context()); id != "" {
                return id
        }
        if holder, ok := mcp.ParsedRequestHolderFromContext(r.Context()); ok && holder.Parsed != nil {
                return holder.Parsed.ResourceID
        }
        return ""
}

// determineEventType determines the event type based on the HTTP request.
func (a *Auditor) determineEventType(r *http.Request) string {
        // First, try to get the parsed MCP method
        if mcpMethod := mcpMethodFor(r); mcpMethod != "" {
                return a.mapMCPMethodToEventType(mcpMethod)
        }

        // Handle SSE connection establishment
        if a.isSSETransport() && r.Method == http.MethodGet {
                return EventTypeSSEConnection
        }
        // Handle MCP message endpoints that weren't parsed (malformed requests)
        if a.isSSETransport() && r.Method == http.MethodPost {
                return EventTypeMCPRequest
        }

        // Default for non-MCP requests
        return EventTypeHTTPRequest
}

// mapMCPMethodToEventType maps MCP method names to event types.
func (*Auditor) mapMCPMethodToEventType(mcpMethod string) string {
        switch mcpMethod {
        case "initialize":
                return EventTypeMCPInitialize
        case "tools/call":
                return EventTypeMCPToolCall
        case "tools/list":
                return EventTypeMCPToolsList
        case "resources/read":
                return EventTypeMCPResourceRead
        case "resources/list":
                return EventTypeMCPResourcesList
        case "prompts/get":
                return EventTypeMCPPromptGet
        case "prompts/list":
                return EventTypeMCPPromptsList
        case "notifications/message":
                return EventTypeMCPNotification
        case "ping":
                return EventTypeMCPPing
        case "logging/setLevel":
                return EventTypeMCPLogging
        case "completion/complete":
                return EventTypeMCPCompletion
        case "notifications/roots/list_changed":
                return EventTypeMCPRootsListChanged
        default:
                return EventTypeMCPRequest
        }
}

// determineOutcome determines the outcome based on the HTTP status code.
func (*Auditor) determineOutcome(statusCode int) string {
        switch {
        case statusCode >= 200 && statusCode < 300:
                return OutcomeSuccess
        case statusCode == 401 || statusCode == 403:
                return OutcomeDenied
        case statusCode >= 400 && statusCode < 500:
                return OutcomeFailure
        case statusCode >= 500:
                return OutcomeError
        default:
                return OutcomeSuccess
        }
}

// detectApplicationError inspects the captured response body prefix for a
// JSON-RPC error field. It reuses rw.body when IncludeResponseData is
// enabled to avoid double-buffering.
func (*Auditor) detectApplicationError(rw *responseWriter) *mcp.ParsedMCPResponse {
        var prefix []byte
        if rw.body != nil && rw.body.Len() > 0 {
                prefix = rw.body.Bytes()
                if len(prefix) > errorDetectionBufferSize {
                        prefix = prefix[:errorDetectionBufferSize]
                }
        } else if rw.errorDetectionBody != nil && rw.errorDetectionBody.Len() > 0 {
                prefix = rw.errorDetectionBody.Bytes()
        }
        if len(prefix) > 0 && prefix[0] == '{' {
                return mcp.ParseMCPResponse(prefix)
        }
        return nil
}

// extractSource extracts source information from the HTTP request.
func (a *Auditor) extractSource(r *http.Request) EventSource {
        // Get the client IP address
        clientIP := a.getClientIP(r)

        source := EventSource{
                Type:  SourceTypeNetwork,
                Value: clientIP,
                Extra: make(map[string]any),
        }

        // Add user agent if available
        if userAgent := r.Header.Get("User-Agent"); userAgent != "" {
                source.Extra[SourceExtraKeyUserAgent] = userAgent
        }

        // Add request ID if available
        if requestID := r.Header.Get("X-Request-ID"); requestID != "" {
                source.Extra[SourceExtraKeyRequestID] = requestID
        }

        return source
}

// getClientIP extracts the client IP address from the request.
func (*Auditor) getClientIP(r *http.Request) string {
        // Check X-Forwarded-For header first
        if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
                // Take the first IP in the list
                if ips := strings.Split(xff, ","); len(ips) > 0 {
                        return strings.TrimSpace(ips[0])
                }
        }

        // Check X-Real-IP header
        if xri := r.Header.Get("X-Real-IP"); xri != "" {
                return xri
        }

        // Fall back to RemoteAddr
        if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
                return host
        }

        return r.RemoteAddr
}

// extractSubjectsFromIdentity extracts subject information from an Identity.
// This helper ensures consistent fallback order and validation across all auditors.
// Fallback order for user: Name → PreferredUsername → Email
func extractSubjectsFromIdentity(identity *auth.Identity) map[string]string {
        subjects := make(map[string]string)

        // Extract user ID (subject)
        if identity.Subject != "" {
                subjects[SubjectKeyUserID] = identity.Subject
        }

        // Extract user name with fallback order: Name → PreferredUsername → Email
        if identity.Name != "" {
                subjects[SubjectKeyUser] = identity.Name
        } else if preferredUsername, ok := identity.Claims["preferred_username"].(string); ok && preferredUsername != "" {
                subjects[SubjectKeyUser] = preferredUsername
        } else if identity.Email != "" {
                subjects[SubjectKeyUser] = identity.Email
        }

        // Add client information if available
        if clientName, ok := identity.Claims["client_name"].(string); ok && clientName != "" {
                subjects[SubjectKeyClientName] = clientName
        }

        if clientVersion, ok := identity.Claims["client_version"].(string); ok && clientVersion != "" {
                subjects[SubjectKeyClientVersion] = clientVersion
        }

        return subjects
}

// extractDelegationChainFromIdentity returns the RFC 8693 delegation chain
// carried by the identity, bound to maxDepth, or nil when there is none.
// The identity layer (auth.claimsToIdentity) always parses at
// auth.DefaultMaxDelegationDepth, so a larger maxDepth requires re-parsing the
// raw "act" claim — rebinding an already-parsed chain can only shrink it. The
// re-parse is only trusted when it yields at least as many hops as are
// already parsed, so it can widen the chain but never narrow it.
func extractDelegationChainFromIdentity(identity *auth.Identity, maxDepth int) *coreaudit.DelegationChain {
        if identity == nil {
                return nil
        }
        rawAct := identity.Claims["act"]
        chain := identity.DelegationChain
        existingHops := 0
        if chain != nil {
                existingHops = len(chain.Chain)
        }
        if rawAct != nil && (chain.IsZero() ||
                (chain.Truncated && maxDepth > existingHops)) {
                parsed := coreaudit.ParseDelegationChain(rawAct, maxDepth)
                if !parsed.IsZero() && len(parsed.Chain) >= existingHops {
                        return parsed
                }
        }
        if !chain.IsZero() {
                return rebindDelegationChain(chain, maxDepth)
        }
        return nil
}

// rebindDelegationChain applies maxDepth to an already-parsed chain: hops
// beyond the cap are dropped (keeping the outermost hops, so Chain[0] — the
// current actor — is preserved), the chain is marked truncated, and the
// omitted count accounts for both hops truncated at parse time and hops
// dropped by this re-binding.
func rebindDelegationChain(chain *coreaudit.DelegationChain, maxDepth int) *coreaudit.DelegationChain {
        if maxDepth <= 0 {
                maxDepth = auth.DefaultMaxDelegationDepth
        }
        bound := *chain
        if len(bound.Chain) > maxDepth {
                bound.Omitted += len(bound.Chain) - maxDepth
                bound.Chain = bound.Chain[:maxDepth]
                bound.Truncated = true
        }
        return &bound
}

// identityFromRequest resolves the authenticated identity for the request.
// The context value is present when an auth middleware runs OUTSIDE audit;
// the holder covers the audit-wraps-auth arrangement, where the identity
// attached for inner handlers is published back up via auth.WithIdentity.
func identityFromRequest(r *http.Request) (*auth.Identity, bool) {
        if identity, ok := auth.IdentityFromContext(r.Context()); ok {
                return identity, true
        }
        if holder, hok := auth.IdentityHolderFromContext(r.Context()); hok && holder.Identity != nil {
                return holder.Identity, true
        }
        return nil, false
}

// extractSubjects extracts subject information from the HTTP request.
func (*Auditor) extractSubjects(r *http.Request) map[string]string {
        subjects := make(map[string]string)

        // Extract user information from Identity.
        if identity, ok := identityFromRequest(r); ok {
                subjects = extractSubjectsFromIdentity(identity)
        }

        // If no user found in claims, set anonymous
        if subjects[SubjectKeyUser] == "" {
                subjects[SubjectKeyUser] = "anonymous"
        }

        return subjects
}

// extractDelegationChain extracts the RFC 8693 delegation chain from the
// HTTP request's identity, or nil when there is no identity or no chain.
func (a *Auditor) extractDelegationChain(r *http.Request) *coreaudit.DelegationChain {
        identity, ok := identityFromRequest(r)
        if !ok {
                return nil
        }
        return extractDelegationChainFromIdentity(identity, a.config.MaxDelegationDepthOrDefault())
}

// determineComponent determines the component name based on the request.
func (a *Auditor) determineComponent(_ *http.Request) string {
        // Use the component from configuration if set
        if a.config.Component != "" {
                return a.config.Component
        }
        // For MCP requests, we could extract the server name from the path or headers
        // For now, we'll use a default component name
        return ComponentToolHive
}

// extractTarget extracts target information from the HTTP request.
func (*Auditor) extractTarget(r *http.Request, eventType string) map[string]string {
        target := make(map[string]string)

        target[TargetKeyEndpoint] = r.URL.Path
        target[TargetKeyMethod] = r.Method

        // Add MCP method if available from parsed data
        if mcpMethod := mcpMethodFor(r); mcpMethod != "" {
                target[TargetKeyMethod] = mcpMethod
        }

        // Add resource ID if available from parsed data
        if resourceID := mcpResourceIDFor(r); resourceID != "" {
                target[TargetKeyName] = resourceID
        }

        // Add event-specific target information
        switch eventType {
        case EventTypeMCPToolCall:
                target[TargetKeyType] = TargetTypeTool
        case EventTypeMCPResourceRead:
                target[TargetKeyType] = TargetTypeResource
        case EventTypeMCPPromptGet:
                target[TargetKeyType] = TargetTypePrompt
        default:
                target[TargetKeyType] = "endpoint"
        }

        return target
}

// addMetadata adds metadata to the audit event.
func (a *Auditor) addMetadata(event *AuditEvent, r *http.Request, duration time.Duration, rw *responseWriter) {
        if event.Metadata.Extra == nil {
                event.Metadata.Extra = make(map[string]any)
        }

        // Add duration
        event.Metadata.Extra[MetadataExtraKeyDuration] = duration.Milliseconds()

        // Add transport information
        if a.isSSETransport() {
                event.Metadata.Extra[MetadataExtraKeyTransport] = "sse"
        } else {
                event.Metadata.Extra[MetadataExtraKeyTransport] = "http"
        }

        // Add response size if available
        if rw.body != nil {
                event.Metadata.Extra[MetadataExtraKeyResponseSize] = rw.body.Len()
        }

        // Add backend routing information from context if available.
        // BackendInfo is populated upstream: by the backend-enrichment middleware on the
        // vMCP legacy path, or by the vMCP Serve-path session handlers (which label the
        // request directly). Either way the value is written before this runs.
        if backendInfo, ok := BackendInfoFromContext(r.Context()); ok && backendInfo != nil && backendInfo.BackendName != "" {
                event.Metadata.Extra["backend_name"] = backendInfo.BackendName
        }
}

// addEventData adds request/response data to the audit event if configured.
func (a *Auditor) addEventData(event *AuditEvent, _ *http.Request, rw *responseWriter, requestData []byte) {
        if !a.config.IncludeRequestData && !a.config.IncludeResponseData {
                return
        }

        data := make(map[string]any)

        if a.config.IncludeRequestData && len(requestData) > 0 {
                // Try to parse as JSON, otherwise store as string
                var requestJSON any
                if err := json.Unmarshal(requestData, &requestJSON); err == nil {
                        data["request"] = requestJSON
                } else {
                        data["request"] = string(requestData)
                }
        }

        if a.config.IncludeResponseData && rw.body != nil && rw.body.Len() > 0 {
                responseData := rw.body.Bytes()
                // Try to parse as JSON, otherwise store as string
                var responseJSON any
                if err := json.Unmarshal(responseData, &responseJSON); err == nil {
                        data["response"] = responseJSON
                } else {
                        data["response"] = string(responseData)
                }
        }

        if len(data) > 0 {
                if dataBytes, err := json.Marshal(data); err == nil {
                        rawMsg := json.RawMessage(dataBytes)
                        event.WithData(&rawMsg)
                }
        }
}

// streamOpenWriter wraps the ResponseWriter for MCP stream-open requests
// (SSE endpoint, streamable GET). It logs the connection audit event exactly
// once, on the first WriteHeader/Write, so the event reflects the actual
// outcome (200 stream established, 401/403 denied by inner middleware) and
// carries the identity the inner auth middleware attached by that point.
type streamOpenWriter struct {
        http.ResponseWriter
        auditor *Auditor
        req     *http.Request
        logged  bool
}

func (sw *streamOpenWriter) WriteHeader(statusCode int) {
        // Informational (1xx) responses are not the final status — don't consume
        // the one-shot connection event on them.
        if statusCode >= http.StatusOK {
                sw.logOnce(statusCode)
        }
        sw.ResponseWriter.WriteHeader(statusCode)
}

func (sw *streamOpenWriter) Write(data []byte) (int, error) {
        // An implicit WriteHeader(200) happens on first Write.
        sw.logOnce(http.StatusOK)
        return sw.ResponseWriter.Write(data)
}

// Flush implements http.Flusher if the underlying ResponseWriter supports it.
// A flush commits the response headers with an implicit 200, so it counts as
// the stream being established — log the connection event here too, otherwise
// a handler that flushes before its first write would delay (or, on a stream
// that only ever flushes, lose) the event.
func (sw *streamOpenWriter) Flush() {
        sw.logOnce(http.StatusOK)
        if flusher, ok := sw.ResponseWriter.(http.Flusher); ok {
                flusher.Flush()
        }
}

// Unwrap exposes the underlying ResponseWriter so http.ResponseController
// can reach interfaces this wrapper does not re-implement (e.g. SetWriteDeadline).
func (sw *streamOpenWriter) Unwrap() http.ResponseWriter {
        return sw.ResponseWriter
}

// logOnce logs the stream connection event with the given status on the first
// call; subsequent calls are no-ops.
func (sw *streamOpenWriter) logOnce(statusCode int) {
        if sw.logged {
                return
        }
        sw.logged = true
        sw.auditor.logSSEConnectionEvent(sw.req, statusCode)
}

// logSSEConnectionEvent logs an audit event for SSE connection initiation.
func (a *Auditor) logSSEConnectionEvent(r *http.Request, statusCode int) {
        // Honor the configured event-type filter, like logAuditEvent does.
        if !a.config.ShouldAuditEvent(EventTypeSSEConnection) {
                return
        }

        // Extract source information
        source := a.extractSource(r)

        // Extract subject information
        subjects := a.extractSubjects(r)

        // Determine component name
        component := a.determineComponent(r)

        // Create the audit event for SSE connection
        event := NewAuditEvent(EventTypeSSEConnection, source, a.determineOutcome(statusCode), subjects, component)

        // Attach the RFC 8693 delegation chain, if the caller used a delegated token
        event.WithDelegationChain(a.extractDelegationChain(r))

        // Add target information
        target := map[string]string{
                "endpoint": r.URL.Path,
                "method":   r.Method,
                "type":     "sse_endpoint",
        }
        event.WithTarget(target)

        // Add metadata
        event.Metadata.Extra = map[string]any{
                "transport":  a.transportType,
                "user_agent": r.Header.Get("User-Agent"),
        }

        // Log the event
        event.LogTo(r.Context(), a.auditLogger, LevelAudit)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package audit provides audit logging configuration for ToolHive.
package audit

import (
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "os"
        "path/filepath"

        "github.com/stacklok/toolhive/pkg/auth"
)

// Config represents the audit logging configuration.
// +kubebuilder:object:generate=true
// +gendoc
type Config struct {
        // Enabled controls whether audit logging is enabled.
        // When true, enables audit logging with the configured options.
        // +kubebuilder:default=false
        // +optional
        Enabled bool `json:"enabled,omitempty" yaml:"enabled,omitempty"`
        // Component is the component name to use in audit events.
        // +optional
        Component string `json:"component,omitempty" yaml:"component,omitempty"`
        // EventTypes specifies which event types to audit. If empty, all events are audited.
        // +optional
        EventTypes []string `json:"eventTypes,omitempty" yaml:"eventTypes,omitempty"`
        // ExcludeEventTypes specifies which event types to exclude from auditing.
        // This takes precedence over EventTypes.
        // +optional
        ExcludeEventTypes []string `json:"excludeEventTypes,omitempty" yaml:"excludeEventTypes,omitempty"`
        // IncludeRequestData determines whether to include request data in audit logs.
        // +kubebuilder:default=false
        // +optional
        IncludeRequestData bool `json:"includeRequestData,omitempty" yaml:"includeRequestData,omitempty"`
        // IncludeResponseData determines whether to include response data in audit logs.
        // +kubebuilder:default=false
        // +optional
        IncludeResponseData bool `json:"includeResponseData,omitempty" yaml:"includeResponseData,omitempty"`
        // DetectApplicationErrors controls whether the audit middleware inspects
        // JSON-RPC response bodies for application-level errors when the HTTP
        // status code indicates success (2xx). When enabled, a small prefix of
        // the response body is buffered to detect JSON-RPC error fields,
        // independent of the IncludeResponseData setting.
        // +kubebuilder:default=true
        // +optional
        DetectApplicationErrors *bool `json:"detectApplicationErrors,omitempty" yaml:"detectApplicationErrors,omitempty"`
        // MaxDataSize limits the size of request/response data included in audit logs (in bytes).
        // +kubebuilder:default=1024
        // +optional
        MaxDataSize int `json:"maxDataSize,omitempty" yaml:"maxDataSize,omitempty"`
        // MaxDelegationDepth caps how many nested RFC 8693 "act" entries are
        // recorded in an audit event's delegation chain. Deeper chains are
        // truncated (marked with truncated=true). Defaults to 10 when unset.
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:default=10
        // +optional
        MaxDelegationDepth *int `json:"maxDelegationDepth,omitempty" yaml:"maxDelegationDepth,omitempty"`
        // LogFile specifies the file path for audit logs. If empty, logs to stdout.
        // +optional
        LogFile string `json:"logFile,omitempty" yaml:"logFile,omitempty"`
}

// GetLogWriter creates and returns the appropriate io.Writer based on the configuration.
func (c *Config) GetLogWriter() (io.Writer, error) {
        if c == nil || c.LogFile == "" {
                return os.Stdout, nil
        }

        // Clean the path to prevent directory traversal
        file, err := os.OpenFile(filepath.Clean(c.LogFile), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
        if err != nil {
                return nil, fmt.Errorf("failed to open audit log file %s: %w", c.LogFile, err)
        }

        return file, nil
}

// DefaultConfig returns a default audit configuration.
func DefaultConfig() *Config {
        detectErrors := true
        return &Config{
                // Note, these defaults are also present on the kubebuilder annotations above.
                // If you change these defaults, you must also change the kubebuilder annotations.
                IncludeRequestData:      false,         // Disabled by default for privacy
                IncludeResponseData:     false,         // Disabled by default for privacy
                MaxDataSize:             1024,          // 1KB default limit
                DetectApplicationErrors: &detectErrors, // Enabled by default to surface JSON-RPC errors
        }
}

// ShouldDetectApplicationErrors returns whether JSON-RPC error detection is enabled.
// Defaults to true when DetectApplicationErrors is nil.
func (c *Config) ShouldDetectApplicationErrors() bool {
        if c.DetectApplicationErrors == nil {
                return true
        }
        return *c.DetectApplicationErrors
}

// MaxDelegationDepthOrDefault returns the configured cap on RFC 8693
// delegation chain depth, or auth.DefaultMaxDelegationDepth (10) when unset
// or non-positive.
//
// The default is deliberately ToolHive's minting cap (10), NOT toolhive-core's
// audit.DefaultMaxDelegationDepth (16): core's value is a library-wide parse
// ceiling, while passing the mint cap here preserves the signal that a chain
// with truncated=true could not have come from ToolHive's own issuance path.
func (c *Config) MaxDelegationDepthOrDefault() int {
        if c == nil || c.MaxDelegationDepth == nil || *c.MaxDelegationDepth <= 0 {
                return auth.DefaultMaxDelegationDepth
        }
        return *c.MaxDelegationDepth
}

// LoadFromFile loads audit configuration from a file.
func LoadFromFile(path string) (*Config, error) {
        // Clean the path to prevent directory traversal
        file, err := os.Open(filepath.Clean(path))
        if err != nil {
                return nil, fmt.Errorf("failed to open audit config file: %w", err)
        }
        defer func() {
                if err := file.Close(); err != nil {
                        slog.Warn("failed to close audit config file", "error", err)
                }
        }()

        return LoadFromReader(file)
}

// LoadFromReader loads audit configuration from an io.Reader.
func LoadFromReader(r io.Reader) (*Config, error) {
        var config Config
        decoder := json.NewDecoder(r)
        if err := decoder.Decode(&config); err != nil {
                return nil, fmt.Errorf("failed to decode audit config: %w", err)
        }

        return &config, nil
}

// ShouldAuditEvent determines whether an event should be audited based on the configuration.
func (c *Config) ShouldAuditEvent(eventType string) bool {
        // Check if event type is excluded
        for _, excludeType := range c.ExcludeEventTypes {
                if excludeType == eventType {
                        return false
                }
        }

        // If specific event types are configured, check if this event type is included
        if len(c.EventTypes) > 0 {
                found := false
                for _, allowedType := range c.EventTypes {
                        if allowedType == eventType {
                                found = true
                                break
                        }
                }
                if !found {
                        return false
                }
        }

        return true
}

// Validate validates the audit configuration.
func (c *Config) Validate() error {
        // Apply default for MaxDataSize if not set (0 means use default)
        if c.MaxDataSize == 0 {
                c.MaxDataSize = DefaultConfig().MaxDataSize
        }

        if c.MaxDataSize < 0 {
                return fmt.Errorf("maxDataSize cannot be negative")
        }

        if c.MaxDelegationDepth != nil && *c.MaxDelegationDepth <= 0 {
                return fmt.Errorf("maxDelegationDepth must be positive")
        }

        // Validate event types (basic validation - could be extended)
        validEventTypes := map[string]bool{
                EventTypeMCPInitialize:       true,
                EventTypeMCPToolCall:         true,
                EventTypeMCPToolsList:        true,
                EventTypeMCPResourceRead:     true,
                EventTypeMCPResourcesList:    true,
                EventTypeMCPPromptGet:        true,
                EventTypeMCPPromptsList:      true,
                EventTypeMCPNotification:     true,
                EventTypeMCPPing:             true,
                EventTypeMCPLogging:          true,
                EventTypeMCPCompletion:       true,
                EventTypeMCPRootsListChanged: true,
                // Workflow event types for vMCP composite workflows
                EventTypeWorkflowStarted:       true,
                EventTypeWorkflowCompleted:     true,
                EventTypeWorkflowFailed:        true,
                EventTypeWorkflowTimedOut:      true,
                EventTypeWorkflowStepStarted:   true,
                EventTypeWorkflowStepCompleted: true,
                EventTypeWorkflowStepFailed:    true,
                EventTypeWorkflowStepSkipped:   true,
                // Fallback event types that can also be emitted by the middleware
                EventTypeMCPRequest:  true,
                EventTypeHTTPRequest: true,
        }

        for _, eventType := range c.EventTypes {
                if !validEventTypes[eventType] {
                        return fmt.Errorf("unknown event type: %s", eventType)
                }
        }

        for _, eventType := range c.ExcludeEventTypes {
                if !validEventTypes[eventType] {
                        return fmt.Errorf("unknown exclude event type: %s", eventType)
                }
        }

        return nil
}

// Package audit provides audit logging functionality for ToolHive.
// This package includes audit event structures and utilities based on
// the auditevent library from metal-toolbox/auditevent to ensure
// NIST SP 800-53 compliance.
package audit

import (
        "context"
        "encoding/json"
        "log/slog"
        "time"

        "github.com/google/uuid"

        coreaudit "github.com/stacklok/toolhive-core/audit"
)

// The following code is adapted from github.com/metal-toolbox/auditevent
// Original copyright notice:
/*
Copyright 2022 Equinix, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// AuditEvent represents an audit event.
// It provides the minimal information needed to audit an event, as well as
// a uniform format to persist the events in audit logs.
//
// It is highly recommended to use the NewAuditEvent function to create
// audit events and set the required fields.
//
//nolint:revive // AuditEvent name is intentional for compatibility with auditevent library
type AuditEvent struct {
        Metadata EventMetadata `json:"metadata"`
        // Type: Defines the type of event that occurred
        // This is a small identifier to quickly determine what happened.
        // e.g. UserLogin, UserLogout, UserCreate, UserDelete, etc.
        Type string `json:"type"`
        // LoggedAt: determines when the event occurred.
        // Note that this should have sufficient information to authoritatively
        // determine the exact time the event was logged at. The output must be in
        // Coordinated Universal Time (UTC) format, a modern continuation of
        // Greenwich Mean Time (GMT), or local time with an offset from UTC to satisfy
        // NIST SP 800-53 requirement AU-8.
        LoggedAt time.Time `json:"loggedAt"`
        // Source: determines the source of the event.
        // Normally, using the IP address of the client, or pod name is sufficient.
        // One must be careful of the data that's added here as we don't want to
        // leak Personally Identifiable Information.
        Source EventSource `json:"source"`
        // Outcome: determines whether the event was successful or not, e.g. successful login
        // It may also determine if the event was approved or denied.
        Outcome string `json:"outcome"`
        // Subject: is the identity of the subject of the event.
        // e.g. who triggered the event? Additional information
        // may be added, such as group membership and/or role
        Subjects map[string]string `json:"subjects"`
        // DelegationChain: when the caller authenticated with an RFC 8693
        // delegated token, this records the full chain of acting parties
        // (outermost/most recent first) extracted from the token's "act" claim.
        // The type and the "delegation" wire key are toolhive-core's canonical
        // delegation-chain schema.
        DelegationChain *coreaudit.DelegationChain `json:"delegation,omitempty"`
        // Component: allows to determine in which component the event occurred
        // (Answering the "Where" question of section c in the NIST SP 800-53
        // Revision 5.1 Control AU-3).
        Component string `json:"component"`
        // Target: Defines where the target of the operation. e.g. the path of
        // the REST resource
        // (Answering the "Where" question of section c in the NIST SP 800-53
        // Revision 5.1 Control AU-3 as well as indicating an entity
        // associated for section f).
        Target map[string]string `json:"target,omitempty"`
        // Data: enhances the audit event with extra information that may be
        // useful for forensic analysis.
        Data *json.RawMessage `json:"data,omitempty"`
}

// EventMetadata contains metadata about the audit event.
type EventMetadata struct {
        // AuditID: is a unique identifier for the audit event.
        AuditID string `json:"auditId"`
        // Extra allows for including additional information about the event
        // that aids in tracking, parsing or auditing
        Extra map[string]any `json:"extra,omitempty"`
}

// EventSource represents the source of an audit event.
type EventSource struct {
        // Type indicates the source type. e.g. Network, File, local, etc.
        // The intent is to determine where a request came from.
        Type string `json:"type"`
        // Value aims to indicate the source of the event. e.g. IP address,
        // hostname, etc.
        Value string `json:"value"`
        // Extra allows for including additional information about the event
        // source that aids in tracking, parsing or auditing
        Extra map[string]any `json:"extra,omitempty"`
}

// NewAuditEvent returns a new AuditEvent with an appropriately set AuditID and logging time.
func NewAuditEvent(
        eventType string,
        source EventSource,
        outcome string,
        subjects map[string]string,
        component string,
) *AuditEvent {
        return &AuditEvent{
                Metadata: EventMetadata{
                        AuditID: uuid.New().String(),
                },
                Type:      eventType,
                LoggedAt:  time.Now().UTC(),
                Source:    source,
                Outcome:   outcome,
                Subjects:  subjects,
                Component: component,
        }
}

// NewAuditEventWithID returns a new AuditEvent with the passed AuditID.
func NewAuditEventWithID(
        auditID string,
        eventType string,
        source EventSource,
        outcome string,
        subjects map[string]string,
        component string,
) *AuditEvent {
        return &AuditEvent{
                Metadata: EventMetadata{
                        AuditID: auditID,
                },
                Type:      eventType,
                LoggedAt:  time.Now().UTC(),
                Source:    source,
                Outcome:   outcome,
                Subjects:  subjects,
                Component: component,
        }
}

// WithDelegationChain sets the RFC 8693 delegation chain of the event.
// A nil or zero chain (no hops, no truncation, no malformation — see
// [coreaudit.DelegationChain.IsZero]) clears any previously set chain. A
// hopless but malformed chain is kept: "the caller asserted delegation we
// could not read" is audit-relevant and distinct from "no delegation".
func (e *AuditEvent) WithDelegationChain(chain *coreaudit.DelegationChain) *AuditEvent {
        if chain.IsZero() {
                e.DelegationChain = nil
                return e
        }
        e.DelegationChain = chain
        return e
}

// WithTarget sets the target of the event.
func (e *AuditEvent) WithTarget(target map[string]string) *AuditEvent {
        e.Target = target
        return e
}

// WithData sets the data of the event.
func (e *AuditEvent) WithData(data *json.RawMessage) *AuditEvent {
        e.Data = data
        return e
}

// WithDataFromString sets the data of the event from a string.
// Note that validating that this is properly JSON-formatted
// is the responsibility of the caller.
func (e *AuditEvent) WithDataFromString(data string) *AuditEvent {
        rawMsg := json.RawMessage(data)
        return e.WithData(&rawMsg)
}

// LogTo logs the audit event to the provided slog.Logger using the custom audit level.
func (e *AuditEvent) LogTo(ctx context.Context, logger *slog.Logger, level slog.Level) {
        // Create slog attributes for the audit event
        attrs := []slog.Attr{
                slog.String("audit_id", e.Metadata.AuditID),
                slog.String("type", e.Type),
                slog.Time("logged_at", e.LoggedAt),
                slog.String("outcome", e.Outcome),
                slog.String("component", e.Component),
                slog.Group("source",
                        slog.String("type", e.Source.Type),
                        slog.String("value", e.Source.Value),
                        slog.Any("extra", e.Source.Extra),
                ),
                slog.Any("subjects", e.Subjects),
        }

        // Add delegation chain if present and non-zero. The chain resolves through
        // coreaudit's [log/slog.LogValuer] implementation, which mirrors its JSON
        // shape (so the slog and JSON wire forms are structurally identical) and
        // never exposes the per-hop in-memory extra claims (PII guard).
        if !e.DelegationChain.IsZero() {
                attrs = append(attrs, slog.Any("delegation", e.DelegationChain))
        }

        // Add target if present
        if e.Target != nil {
                attrs = append(attrs, slog.Any("target", e.Target))
        }

        // Add metadata extra if present
        if e.Metadata.Extra != nil {
                attrs = append(attrs, slog.Group("metadata", slog.Any("extra", e.Metadata.Extra)))
        }

        // Add data if present
        if e.Data != nil {
                attrs = append(attrs, slog.Any("data", e.Data))
        }

        // Log with the specified level
        logger.LogAttrs(ctx, level, "audit_event", attrs...)
}

// Common event outcomes
const (
        // OutcomeSuccess indicates the event was successful
        OutcomeSuccess = "success"
        // OutcomeFailure indicates the event failed
        OutcomeFailure = "failure"
        // OutcomeError indicates the event resulted in an error
        OutcomeError = "error"
        // OutcomeDenied indicates the event was denied (e.g., by authorization)
        OutcomeDenied = "denied"
        // OutcomeApplicationError indicates the HTTP transport succeeded but the
        // JSON-RPC response body contained an application-level error (e.g.,
        // expired tokens, backend failures, invalid parameters).
        OutcomeApplicationError = "application_error"
)

// Common source types
const (
        // SourceTypeNetwork indicates the event came from a network request
        SourceTypeNetwork = "network"
        // SourceTypeLocal indicates the event came from a local source
        SourceTypeLocal = "local"
)

// Component name for ToolHive
const (
        // ComponentToolHive is the component name for ToolHive API audit events.
        // Note that events directed for an MCP server will have the name of the
        // MCP server as the component instead.
        ComponentToolHive = "toolhive-api"
)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package audit

import (
        "encoding/json"
        "fmt"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// Middleware type constant
const (
        MiddlewareType = "audit"
)

// MiddlewareParams represents the parameters for audit middleware
type MiddlewareParams struct {
        ConfigPath string  `json:"config_path,omitempty"` // Kept for backwards compatibility
        ConfigData *Config `json:"config_data,omitempty"` // New field for config contents
        Component  string  `json:"component,omitempty"`
        // Transport information for dynamic transport detection
        TransportType string `json:"transport_type,omitempty"` // e.g., "sse", "streamable-http"
}

// Middleware wraps audit middleware functionality
type Middleware struct {
        middleware types.MiddlewareFunction
        auditor    *Auditor
}

// Handler returns the middleware function used by the proxy.
func (m *Middleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (m *Middleware) Close() error {
        if m.auditor != nil {
                return m.auditor.Close()
        }
        return nil
}

// CreateMiddleware factory function for audit middleware
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {

        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal audit middleware parameters: %w", err)
        }

        var auditConfig *Config
        var err error

        if params.ConfigData != nil {
                // Use provided config data (preferred method)
                auditConfig = params.ConfigData
        } else if params.ConfigPath != "" {
                // Load config from file (backwards compatibility)
                auditConfig, err = LoadFromFile(params.ConfigPath)
                if err != nil {
                        return fmt.Errorf("failed to load audit configuration: %w", err)
                }
        } else {
                // Use default config
                auditConfig = DefaultConfig()
        }

        // Set component name if provided and config doesn't already have one
        if params.Component != "" && auditConfig.Component == "" {
                auditConfig.Component = params.Component
        }

        // Validate and apply defaults to the config
        if err := auditConfig.Validate(); err != nil {
                return fmt.Errorf("invalid audit configuration: %w", err)
        }

        // Create the auditor directly so we can store a reference for cleanup
        auditor, err := NewAuditorWithTransport(auditConfig, params.TransportType)
        if err != nil {
                return fmt.Errorf("failed to create audit middleware: %w", err)
        }

        auditMw := &Middleware{
                middleware: auditor.Middleware,
                auditor:    auditor,
        }
        runner.AddMiddleware(config.Type, auditMw)
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package audit provides audit logging functionality for ToolHive.
package audit

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "time"

        "github.com/stacklok/toolhive/pkg/auth"
)

// WorkflowAuditor provides audit logging for workflow execution.
// This struct abstracts workflow-specific audit operations from the
// HTTP middleware-based Auditor.
type WorkflowAuditor struct {
        auditLogger *slog.Logger
        config      *Config
        component   string
        logWriter   io.Writer
}

// NewWorkflowAuditor creates a new workflow auditor.
// If config is nil, creates a default configuration with stdout logging.
func NewWorkflowAuditor(config *Config) (*WorkflowAuditor, error) {
        if config == nil {
                config = DefaultConfig()
        }

        logWriter, err := config.GetLogWriter()
        if err != nil {
                return nil, fmt.Errorf("failed to create log writer: %w", err)
        }

        // Use configured component or default to vmcp-composer
        component := config.Component
        if component == "" {
                component = "vmcp-composer"
        }

        return &WorkflowAuditor{
                auditLogger: NewAuditLogger(logWriter),
                config:      config,
                component:   component,
                logWriter:   logWriter,
        }, nil
}

// Close closes the underlying log writer if it owns a closeable resource.
// This should be called when the workflow auditor is no longer needed.
func (w *WorkflowAuditor) Close() error {
        return closeLogWriter(w.logWriter)
}

// LogWorkflowStarted logs the start of workflow execution.
func (w *WorkflowAuditor) LogWorkflowStarted(
        ctx context.Context,
        workflowID string,
        workflowName string,
        parameters map[string]any,
        timeout time.Duration,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowStarted) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowStarted, OutcomeSuccess)

        target := map[string]string{
                TargetKeyWorkflowID:   workflowID,
                TargetKeyWorkflowName: workflowName,
                TargetKeyType:         TargetTypeWorkflow,
        }
        event.WithTarget(target)

        // Add timeout to metadata
        event.Metadata.Extra = map[string]any{
                MetadataExtraKeyTimeout: timeout.Milliseconds(),
        }

        // Add workflow parameters as data (if configured)
        // Using same structure as HTTP auditor for consistency
        if w.config.IncludeRequestData && parameters != nil {
                data := map[string]any{
                        "request": parameters,
                }
                if dataBytes, err := json.Marshal(data); err == nil {
                        rawMsg := json.RawMessage(dataBytes)
                        event.WithData(&rawMsg)
                }
        }

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// LogWorkflowCompleted logs successful workflow completion.
func (w *WorkflowAuditor) LogWorkflowCompleted(
        ctx context.Context,
        workflowID string,
        workflowName string,
        duration time.Duration,
        stepCount int,
        output map[string]any,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowCompleted) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowCompleted, OutcomeSuccess)

        target := map[string]string{
                TargetKeyWorkflowID:   workflowID,
                TargetKeyWorkflowName: workflowName,
                TargetKeyType:         TargetTypeWorkflow,
        }
        event.WithTarget(target)

        // Add metadata
        event.Metadata.Extra = map[string]any{
                MetadataExtraKeyDuration:  duration.Milliseconds(),
                MetadataExtraKeyStepCount: stepCount,
        }

        // Add output data (if configured)
        // Using same structure as HTTP auditor for consistency
        if w.config.IncludeResponseData && output != nil {
                data := map[string]any{
                        "response": output,
                }
                if dataBytes, err := json.Marshal(data); err == nil {
                        rawMsg := json.RawMessage(dataBytes)
                        event.WithData(&rawMsg)
                }
        }

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// LogWorkflowFailed logs workflow failure.
func (w *WorkflowAuditor) LogWorkflowFailed(
        ctx context.Context,
        workflowID string,
        workflowName string,
        duration time.Duration,
        stepCount int,
        _ error,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowFailed) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowFailed, OutcomeFailure)

        target := map[string]string{
                TargetKeyWorkflowID:   workflowID,
                TargetKeyWorkflowName: workflowName,
                TargetKeyType:         TargetTypeWorkflow,
        }
        event.WithTarget(target)

        // Add metadata
        event.Metadata.Extra = map[string]any{
                MetadataExtraKeyDuration:  duration.Milliseconds(),
                MetadataExtraKeyStepCount: stepCount,
        }

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// LogWorkflowTimedOut logs workflow timeout.
func (w *WorkflowAuditor) LogWorkflowTimedOut(
        ctx context.Context,
        workflowID string,
        workflowName string,
        duration time.Duration,
        stepCount int,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowTimedOut) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowTimedOut, OutcomeFailure)

        target := map[string]string{
                TargetKeyWorkflowID:   workflowID,
                TargetKeyWorkflowName: workflowName,
                TargetKeyType:         TargetTypeWorkflow,
        }
        event.WithTarget(target)

        // Add metadata
        event.Metadata.Extra = map[string]any{
                MetadataExtraKeyDuration:  duration.Milliseconds(),
                MetadataExtraKeyStepCount: stepCount,
        }

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// LogStepStarted logs the start of step execution.
func (w *WorkflowAuditor) LogStepStarted(
        ctx context.Context,
        workflowID string,
        stepID string,
        stepType string,
        toolName string,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowStepStarted) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowStepStarted, OutcomeSuccess)

        target := map[string]string{
                TargetKeyWorkflowID: workflowID,
                TargetKeyStepID:     stepID,
                TargetKeyStepType:   stepType,
                TargetKeyType:       TargetTypeWorkflowStep,
        }
        if toolName != "" {
                target[TargetKeyToolName] = toolName
        }
        event.WithTarget(target)

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// LogStepCompleted logs successful step completion.
func (w *WorkflowAuditor) LogStepCompleted(
        ctx context.Context,
        workflowID string,
        stepID string,
        duration time.Duration,
        retryCount int,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowStepCompleted) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowStepCompleted, OutcomeSuccess)

        target := map[string]string{
                TargetKeyWorkflowID: workflowID,
                TargetKeyStepID:     stepID,
                TargetKeyType:       TargetTypeWorkflowStep,
        }
        event.WithTarget(target)

        event.Metadata.Extra = map[string]any{
                MetadataExtraKeyDuration:   duration.Milliseconds(),
                MetadataExtraKeyRetryCount: retryCount,
        }

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// LogStepFailed logs step failure.
func (w *WorkflowAuditor) LogStepFailed(
        ctx context.Context,
        workflowID string,
        stepID string,
        duration time.Duration,
        retryCount int,
        _ error,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowStepFailed) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowStepFailed, OutcomeFailure)

        target := map[string]string{
                TargetKeyWorkflowID: workflowID,
                TargetKeyStepID:     stepID,
                TargetKeyType:       TargetTypeWorkflowStep,
        }
        event.WithTarget(target)

        event.Metadata.Extra = map[string]any{
                MetadataExtraKeyDuration:   duration.Milliseconds(),
                MetadataExtraKeyRetryCount: retryCount,
        }

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// LogStepSkipped logs conditional step skip.
func (w *WorkflowAuditor) LogStepSkipped(
        ctx context.Context,
        workflowID string,
        stepID string,
        condition string,
) {
        if !w.config.ShouldAuditEvent(EventTypeWorkflowStepSkipped) {
                return
        }

        event := w.newEvent(ctx, EventTypeWorkflowStepSkipped, OutcomeSuccess)

        target := map[string]string{
                TargetKeyWorkflowID: workflowID,
                TargetKeyStepID:     stepID,
                TargetKeyType:       TargetTypeWorkflowStep,
        }
        event.WithTarget(target)

        // Add condition as metadata
        if condition != "" {
                event.Metadata.Extra = map[string]any{
                        "condition": condition,
                }
        }

        event.LogTo(ctx, w.auditLogger, LevelAudit)
}

// newEvent creates an audit event of the given type and outcome with the
// source, subjects, and RFC 8693 delegation chain extracted from the context.
// Every Log* method MUST build its event through this helper so that the
// delegation chain cannot be forgotten at any individual call site.
func (w *WorkflowAuditor) newEvent(ctx context.Context, eventType, outcome string) *AuditEvent {
        event := NewAuditEvent(
                eventType,
                w.extractSource(ctx),
                outcome,
                w.extractSubjects(ctx),
                w.component,
        )
        w.attachDelegation(ctx, event)
        return event
}

// extractSource extracts source information from context.
// For workflows, source is always local since they're internal orchestration.
func (*WorkflowAuditor) extractSource(_ context.Context) EventSource {
        return EventSource{
                Type:  SourceTypeLocal,
                Value: "vmcp-composer",
                Extra: map[string]any{},
        }
}

// extractSubjects extracts subject information from context.
func (*WorkflowAuditor) extractSubjects(ctx context.Context) map[string]string {
        subjects := make(map[string]string)

        // Extract user information from Identity
        if identity, ok := auth.IdentityFromContext(ctx); ok {
                subjects = extractSubjectsFromIdentity(identity)
        }

        // If no user found, set anonymous
        if subjects[SubjectKeyUser] == "" {
                subjects[SubjectKeyUser] = "anonymous"
        }

        return subjects
}

// attachDelegation attaches the RFC 8693 delegation chain from the context's
// identity to the event. It is a no-op when there is no identity or the
// identity carries no delegation chain.
func (w *WorkflowAuditor) attachDelegation(ctx context.Context, event *AuditEvent) {
        identity, ok := auth.IdentityFromContext(ctx)
        if !ok {
                return
        }
        event.WithDelegationChain(
                extractDelegationChainFromIdentity(identity, w.config.MaxDelegationDepthOrDefault()))
}

//go:build !ignore_autogenerated

/*
Copyright 2025 Stacklok

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Code generated by controller-gen. DO NOT EDIT.

package audit

import ()

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *Config) DeepCopyInto(out *Config) {
        *out = *in
        if in.EventTypes != nil {
                in, out := &in.EventTypes, &out.EventTypes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.ExcludeEventTypes != nil {
                in, out := &in.ExcludeEventTypes, &out.ExcludeEventTypes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.DetectApplicationErrors != nil {
                in, out := &in.DetectApplicationErrors, &out.DetectApplicationErrors
                *out = new(bool)
                **out = **in
        }
        if in.MaxDelegationDepth != nil {
                in, out := &in.MaxDelegationDepth, &out.MaxDelegationDepth
                *out = new(int)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Config.
func (in *Config) DeepCopy() *Config {
        if in == nil {
                return nil
        }
        out := new(Config)
        in.DeepCopyInto(out)
        return out
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

import (
        "net/http"
        "time"

        "github.com/golang-jwt/jwt/v5"
)

// AnonymousMiddleware creates an HTTP middleware that sets up anonymous identity.
// This is useful for testing and local environments where authorization policies
// need to work without requiring actual authentication.
//
// The middleware sets up basic anonymous identity that can be used by authorization
// policies, allowing them to function even when authentication is disabled.
// This is heavily discouraged in production settings but is handy for testing
// and local development environments.
func AnonymousMiddleware(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                // Create anonymous claims with basic information
                claims := jwt.MapClaims{
                        "sub":   "anonymous",
                        "iss":   "toolhive-local",
                        "aud":   "toolhive",
                        "exp":   time.Now().Add(24 * time.Hour).Unix(), // Valid for 24 hours
                        "iat":   time.Now().Unix(),
                        "nbf":   time.Now().Unix(),
                        "email": "anonymous@localhost",
                        "name":  "Anonymous User",
                }

                // Create Identity from claims
                identity := &Identity{
                        PrincipalInfo: PrincipalInfo{
                                Subject: "anonymous",
                                Name:    "Anonymous User",
                                Email:   "anonymous@localhost",
                                Claims:  claims,
                        },
                        Token:     "", // No token for anonymous auth
                        TokenType: "Bearer",
                }

                // Add the Identity to the request context
                ctx := WithIdentity(r.Context(), identity)
                next.ServeHTTP(w, r.WithContext(ctx))
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package awssts provides AWS STS token exchange with SigV4 signing support.
package awssts

// MinSessionDuration is the minimum allowed session duration (AWS limit).
const MinSessionDuration int32 = 900

// MaxSessionDuration is the maximum allowed session duration (12 hours).
const MaxSessionDuration int32 = 43200

// defaultRoleClaim is the default JWT claim to use for role mapping.
const defaultRoleClaim = "groups"

// Config holds configuration for AWS STS token exchange.
type Config struct {
        // Region is the AWS region for STS and SigV4 signing.
        Region string `json:"region" yaml:"region"`

        // Service is the AWS service name for SigV4 signing (default: "aws-mcp").
        Service string `json:"service" yaml:"service"`

        // FallbackRoleArn is the IAM role ARN to assume when no role mapping matches.
        FallbackRoleArn string `json:"fallback_role_arn,omitempty" yaml:"fallback_role_arn,omitempty"`

        // RoleMappings maps JWT claim values to IAM roles with priority.
        RoleMappings []RoleMapping `json:"role_mappings,omitempty" yaml:"role_mappings,omitempty"`

        // RoleClaim is the JWT claim to use for role mapping (default: "groups").
        RoleClaim string `json:"role_claim,omitempty" yaml:"role_claim,omitempty"`

        // SessionDuration is the duration in seconds for assumed role credentials (default: 3600).
        SessionDuration int32 `json:"session_duration,omitempty" yaml:"session_duration,omitempty"`

        // SessionNameClaim is the JWT claim to use for role session name (default: "sub").
        SessionNameClaim string `json:"session_name_claim,omitempty" yaml:"session_name_claim,omitempty"`

        // SubjectProviderName identifies which upstream provider's access token to use
        // for STS AssumeRoleWithWebIdentity. Used by vMCP only. When empty, the bearer
        // token from the incoming HTTP request is used.
        SubjectProviderName string `json:"subject_provider_name,omitempty" yaml:"subject_provider_name,omitempty"`
}

// defaultSessionDuration is the default session duration in seconds (1 hour).
const defaultSessionDuration int32 = 3600

// GetRoleClaim returns the configured role claim or the default.
func (c *Config) GetRoleClaim() string {
        if c.RoleClaim != "" {
                return c.RoleClaim
        }
        return defaultRoleClaim
}

// GetService returns the configured service name or the default ("aws-mcp").
func (c *Config) GetService() string {
        if c.Service != "" {
                return c.Service
        }
        return defaultService
}

// GetSessionDuration returns the configured session duration or the default (3600s).
func (c *Config) GetSessionDuration() int32 {
        if c.SessionDuration != 0 {
                return c.SessionDuration
        }
        return defaultSessionDuration
}

// RoleMapping maps a JWT claim value or CEL expression to an IAM role with explicit priority.
type RoleMapping struct {
        // Claim is the simple claim value to match (e.g., group name).
        // Internally compiles to a CEL expression: "<claim_value>" in claims["<role_claim>"]
        // Mutually exclusive with Matcher.
        Claim string `json:"claim,omitempty" yaml:"claim,omitempty"`

        // Matcher is a CEL expression for complex matching against JWT claims.
        // The expression has access to a "claims" variable containing all JWT claims.
        // Examples:
        //   - "admins" in claims["groups"]
        //   - claims["sub"] == "user123" && !("act" in claims)
        // Mutually exclusive with Claim.
        Matcher string `json:"matcher,omitempty" yaml:"matcher,omitempty"`

        // RoleArn is the IAM role ARN to assume when this mapping matches.
        RoleArn string `json:"role_arn" yaml:"role_arn"`

        // Priority determines selection order (lower number = higher priority).
        // When multiple mappings match, the one with the lowest priority is selected.
        // When nil (omitted), the mapping has the lowest possible priority, and
        // configuration order acts as tie-breaker via stable sort.
        Priority *int `json:"priority,omitempty" yaml:"priority,omitempty"`
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package awssts

import (
        "context"
        "fmt"
        "log/slog"
        "regexp"

        "github.com/aws/aws-sdk-go-v2/aws"
        "github.com/aws/aws-sdk-go-v2/config"
        "github.com/aws/aws-sdk-go-v2/service/sts"
)

// STSClient defines the interface for STS operations, enabling mock injection for testing.
type STSClient interface {
        AssumeRoleWithWebIdentity(
                ctx context.Context,
                params *sts.AssumeRoleWithWebIdentityInput,
                optFns ...func(*sts.Options),
        ) (*sts.AssumeRoleWithWebIdentityOutput, error)
}

// Exchanger handles STS token exchange operations.
type Exchanger struct {
        client STSClient
}

// NewExchanger creates a new Exchanger with a regional STS client.
func NewExchanger(ctx context.Context, region string) (*Exchanger, error) {
        if region == "" {
                return nil, ErrMissingRegion
        }

        client, err := newRegionalSTSClient(ctx, region)
        if err != nil {
                return nil, err
        }

        return &Exchanger{client: client}, nil
}

// newRegionalSTSClient creates an STS client configured for the specified region.
// The SDK automatically resolves regional STS endpoints for lower latency.
func newRegionalSTSClient(ctx context.Context, region string) (STSClient, error) {
        cfg, err := config.LoadDefaultConfig(ctx,
                config.WithRegion(region),
                config.WithCredentialsProvider(aws.AnonymousCredentials{}),
        )
        if err != nil {
                return nil, fmt.Errorf("failed to load AWS config: %w", err)
        }

        return sts.NewFromConfig(cfg), nil
}

// ExchangeToken performs AssumeRoleWithWebIdentity to exchange an identity token
// for temporary AWS credentials.
func (e *Exchanger) ExchangeToken(
        ctx context.Context,
        token, roleArn, sessionName string,
        durationSeconds int32,
) (*aws.Credentials, error) {
        if err := validateInputs(token, roleArn, sessionName, durationSeconds); err != nil {
                return nil, err
        }

        input := &sts.AssumeRoleWithWebIdentityInput{
                RoleArn:          aws.String(roleArn),
                RoleSessionName:  aws.String(sessionName),
                WebIdentityToken: aws.String(token),
                DurationSeconds:  aws.Int32(durationSeconds),
        }

        output, err := e.client.AssumeRoleWithWebIdentity(ctx, input)
        if err != nil {
                slog.Debug("STS AssumeRoleWithWebIdentity failed", "error", err)
                return nil, ErrSTSExchangeFailed
        }

        if output == nil || output.Credentials == nil {
                return nil, ErrSTSNilCredentials
        }

        return &aws.Credentials{
                AccessKeyID:     aws.ToString(output.Credentials.AccessKeyId),
                SecretAccessKey: aws.ToString(output.Credentials.SecretAccessKey),
                SessionToken:    aws.ToString(output.Credentials.SessionToken),
                Expires:         aws.ToTime(output.Credentials.Expiration),
                CanExpire:       true,
        }, nil
}

// sessionNamePattern validates AWS RoleSessionName values.
// AWS allows: letters (a-z, A-Z), digits (0-9), and the characters _+=,.@-
// See: https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
var sessionNamePattern = regexp.MustCompile(`^[a-zA-Z0-9_+=,.@-]+$`)

const (
        // minSessionNameLen is the minimum length for an AWS RoleSessionName.
        minSessionNameLen = 2
        // maxSessionNameLen is the maximum length for an AWS RoleSessionName.
        maxSessionNameLen = 64
)

// ValidateSessionName checks that a session name meets AWS RoleSessionName constraints:
// 2-64 characters, only letters, digits, and _+=,.@- are allowed.
func ValidateSessionName(name string) error {
        if len(name) < minSessionNameLen {
                return fmt.Errorf("%w: must be at least %d characters", ErrInvalidSessionName, minSessionNameLen)
        }
        if len(name) > maxSessionNameLen {
                return fmt.Errorf("%w: must be at most %d characters", ErrInvalidSessionName, maxSessionNameLen)
        }
        if !sessionNamePattern.MatchString(name) {
                return fmt.Errorf("%w: contains invalid characters (allowed: letters, digits, _+=,.@-)", ErrInvalidSessionName)
        }
        return nil
}

// validateInputs validates the exchange inputs.
func validateInputs(token, roleArn, sessionName string, durationSeconds int32) error {
        if token == "" {
                return ErrMissingToken
        }

        if err := ValidateRoleArn(roleArn); err != nil {
                return err
        }

        if err := ValidateSessionName(sessionName); err != nil {
                return err
        }

        if durationSeconds < MinSessionDuration {
                return fmt.Errorf("%w: %d is below minimum %d seconds", ErrInvalidSessionDuration, durationSeconds, MinSessionDuration)
        }

        if durationSeconds > MaxSessionDuration {
                return fmt.Errorf("%w: %d exceeds maximum %d seconds", ErrInvalidSessionDuration, durationSeconds, MaxSessionDuration)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package awssts provides AWS STS token exchange with SigV4 signing support.
package awssts

import (
        "bytes"
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "net/textproto"
        "net/url"
        "strings"

        "github.com/aws/aws-sdk-go-v2/aws"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// Middleware type constant
const (
        MiddlewareType = "awssts"
)

// Default session name claim when not specified in config.
const defaultSessionNameClaim = "sub"

// MiddlewareParams represents the parameters for AWS STS middleware.
type MiddlewareParams struct {
        AWSStsConfig *Config `json:"aws_sts_config,omitempty"`
        // TargetURL is the remote MCP server URL for SigV4 signing.
        // The request must be signed with the target host, not the proxy host.
        TargetURL string `json:"target_url,omitempty"`
}

// Middleware wraps AWS STS middleware functionality.
type Middleware struct {
        middleware types.MiddlewareFunction
        exchanger  *Exchanger
}

// Handler returns the middleware function used by the proxy.
func (m *Middleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (*Middleware) Close() error {
        return nil
}

// CreateMiddleware is the factory function for AWS STS middleware.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal AWS STS middleware parameters: %w", err)
        }

        // AWS STS config is required when this middleware type is specified
        if params.AWSStsConfig == nil {
                return fmt.Errorf("AWS STS configuration is required but not provided")
        }

        // Validate configuration at startup
        if err := ValidateConfig(params.AWSStsConfig); err != nil {
                return fmt.Errorf("invalid AWS STS configuration: %w", err)
        }

        // Parse and validate target URL if provided
        var targetURL *url.URL
        if params.TargetURL != "" {
                var err error
                targetURL, err = url.Parse(params.TargetURL)
                if err != nil {
                        return fmt.Errorf("invalid target URL: %w", err)
                }
                if targetURL.Scheme == "" || targetURL.Host == "" {
                        return fmt.Errorf("target URL must include scheme and host (e.g., https://example.com)")
                }
        }

        // Create the middleware
        // TODO(jakub): MiddlewareFactory interface does not accept a context; pass context.TODO
        // because we don't really have a better option here.
        mw, err := newAWSStsMiddleware(context.TODO(), params.AWSStsConfig, targetURL)
        if err != nil {
                return fmt.Errorf("failed to create AWS STS middleware: %w", err)
        }

        // Add middleware to runner
        runner.AddMiddleware(config.Type, mw)

        return nil
}

// newAWSStsMiddleware creates a new AWS STS middleware with all required components.
// targetURL is the remote MCP server URL used for SigV4 signing (can be nil if not proxying).
func newAWSStsMiddleware(ctx context.Context, cfg *Config, targetURL *url.URL) (*Middleware, error) {
        // Create the STS exchanger with regional endpoint
        exchanger, err := NewExchanger(ctx, cfg.Region)
        if err != nil {
                return nil, fmt.Errorf("failed to create STS exchanger: %w", err)
        }

        // Create the role mapper
        roleMapper, err := NewRoleMapper(cfg)
        if err != nil {
                return nil, fmt.Errorf("failed to create role mapper: %w", err)
        }

        // Create the SigV4 signer
        signer, err := newRequestSigner(cfg.Region, withService(cfg.GetService()))
        if err != nil {
                return nil, fmt.Errorf("failed to create SigV4 signer: %w", err)
        }

        // Determine session name claim
        sessionNameClaim := cfg.SessionNameClaim
        if sessionNameClaim == "" {
                sessionNameClaim = defaultSessionNameClaim
        }

        // Get session duration
        sessionDuration := cfg.GetSessionDuration()

        // Create the middleware function
        middlewareFunc := createAWSStsMiddlewareFunc(exchanger, roleMapper, signer, sessionNameClaim, sessionDuration, targetURL)

        return &Middleware{
                middleware: middlewareFunc,
                exchanger:  exchanger,
        }, nil
}

// createAWSStsMiddlewareFunc creates the HTTP middleware function.
// targetURL is the remote MCP server URL used for SigV4 signing.
// SigV4 requires signing with the actual target host, not the proxy host.
func createAWSStsMiddlewareFunc(
        exchanger *Exchanger,
        roleMapper *RoleMapper,
        signer *RequestSigner,
        sessionNameClaim string,
        sessionDuration int32,
        targetURL *url.URL,
) types.MiddlewareFunction {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // Get identity from the auth middleware.
                        // Unlike token exchange/upstream swap middleware, AWS STS requires valid
                        // credentials and cannot fall through — every request must be signed.
                        identity, ok := auth.IdentityFromContext(r.Context())
                        if !ok {
                                slog.Warn("No identity found in context, rejecting request")
                                http.Error(w, "Authentication required", http.StatusUnauthorized)
                                return
                        }

                        // Extract JWT claims from identity
                        claims := identity.Claims
                        if claims == nil {
                                slog.Warn("No claims in identity, rejecting request")
                                http.Error(w, "Authentication required", http.StatusUnauthorized)
                                return
                        }

                        // Use RoleMapper to select the appropriate IAM role based on claims
                        roleArn, err := roleMapper.SelectRole(claims)
                        if err != nil {
                                slog.Warn("Failed to select IAM role", "error", err)
                                http.Error(w, "Failed to determine IAM role", http.StatusForbidden)
                                return
                        }

                        //nolint:gosec // G706: roleArn is from server config, not user input
                        slog.Debug("Selected IAM role", "role_arn", roleArn)

                        // Extract bearer token from request
                        bearerToken, err := auth.ExtractBearerToken(r)
                        if err != nil {
                                slog.Warn("No valid Bearer token found", "error", err)
                                http.Error(w, "Bearer token required", http.StatusUnauthorized)
                                return
                        }

                        // Extract and validate session name from claims
                        sessionName, err := ExtractSessionName(claims, sessionNameClaim)
                        if err != nil {
                                slog.Warn("Failed to extract session name", "error", err)
                                http.Error(w, "Missing session name claim", http.StatusUnauthorized)
                                return
                        }
                        if err := ValidateSessionName(sessionName); err != nil {
                                slog.Warn("Invalid session name from claim", "claim", sessionNameClaim, "error", err)
                                //nolint:gosec // G706: logged for debugging invalid input
                                slog.Debug("Invalid session name value", "session_name", sessionName)
                                http.Error(w, "Invalid session name", http.StatusUnauthorized)
                                return
                        }

                        //nolint:gosec // G706: session name is from validated JWT claims
                        slog.Debug("Exchanging token for AWS credentials", "session", sessionName)

                        // Exchange token for AWS credentials via STS
                        creds, err := exchanger.ExchangeToken(r.Context(), bearerToken, roleArn, sessionName, sessionDuration)
                        if err != nil {
                                slog.Warn("STS token exchange failed", "error", err)
                                http.Error(w, "AWS credential exchange failed", http.StatusUnauthorized)
                                return
                        }

                        // Sign the request with SigV4 using a clone so we don't permanently
                        // overwrite r.Host / r.URL.Host — that rewriting is the reverse
                        // proxy's responsibility, not ours. We only add the SigV4 headers.
                        if err := signRequestForTarget(r, signer, creds, targetURL); err != nil {
                                slog.Warn("Failed to sign request with SigV4", "error", err)
                                http.Error(w, "Request signing failed", http.StatusInternalServerError)
                                return
                        }

                        slog.Debug("Request signed with AWS SigV4")

                        next.ServeHTTP(w, r)
                })
        }
}

// hopByHopHeaders lists hop-by-hop headers that httputil.ReverseProxy
// removes before forwarding. This matches the standard library's
// hopHeaders list in net/http/httputil/reverseproxy.go.
var hopByHopHeaders = []string{
        "Connection",
        "Proxy-Connection",
        "Keep-Alive",
        "Proxy-Authenticate",
        "Proxy-Authorization",
        "Te",
        "Trailer",
        "Transfer-Encoding",
        "Upgrade",
}

// stripHopByHopHeaders removes hop-by-hop headers from a header set,
// including any header named in the Connection value (RFC 7230 §6.1,
// RFC 9110 §7.6.1).
func stripHopByHopHeaders(header http.Header) {
        for _, f := range header.Values("Connection") {
                for _, name := range strings.Split(f, ",") {
                        if name = textproto.TrimString(name); name != "" {
                                header.Del(name)
                        }
                }
        }
        for _, h := range hopByHopHeaders {
                header.Del(h)
        }
}

// signRequestForTarget signs the request with SigV4 for the given target host
// without permanently modifying r.Host or r.URL. When targetURL is non-nil, a
// clone is used for signing so that only the SigV4 headers are copied back to
// the original request; the reverse proxy's Director is left to handle host
// rewriting. When targetURL is nil the request is signed in-place.
func signRequestForTarget(r *http.Request, signer *RequestSigner, creds *aws.Credentials, targetURL *url.URL) error {
        if targetURL == nil {
                return signer.SignRequest(r.Context(), r, creds)
        }

        // Buffer the body so both the signing clone and the original request
        // can read it. The SigV4 signer consumes the body to compute the
        // payload hash and then replaces it on the request it receives.
        // Because Clone() shares the same Body reader, we must buffer once
        // and provide fresh readers to each side.
        var bodyBytes []byte
        if r.Body != nil && r.Body != http.NoBody {
                var err error
                bodyBytes, err = io.ReadAll(io.LimitReader(r.Body, maxPayloadSize+1))
                if err != nil {
                        return fmt.Errorf("failed to read request body for signing: %w", err)
                }
                if len(bodyBytes) > maxPayloadSize {
                        return fmt.Errorf("request body exceeds maximum size of %d bytes", maxPayloadSize)
                }
                _ = r.Body.Close()
        }

        // Strip hop-by-hop headers from the actual request before cloning.
        // httputil.ReverseProxy removes these before forwarding; if they
        // remain on r they can trigger removal of newly-signed headers
        // (e.g. Connection: X-Amz-Date would delete the SigV4 X-Amz-Date).
        stripHopByHopHeaders(r.Header)

        // Build a signing-only clone with the target host.
        signingReq := r.Clone(r.Context())
        signingReq.URL.Scheme = targetURL.Scheme
        signingReq.URL.Host = targetURL.Host
        signingReq.Host = targetURL.Host

        // Strip headers that upstream gateways inject and that
        // httputil.ReverseProxy.SetXForwarded() rewrites after signing.
        // Including them in the SigV4 canonical headers produces a
        // signature mismatch because the values change in flight.
        signingReq.Header.Del("X-Forwarded-For")
        signingReq.Header.Del("X-Forwarded-Host")
        signingReq.Header.Del("X-Forwarded-Proto")
        signingReq.Header.Del("X-Real-Ip")
        signingReq.Header.Del("Forwarded") // RFC 7239

        if bodyBytes != nil {
                signingReq.Body = io.NopCloser(bytes.NewReader(bodyBytes))
                signingReq.ContentLength = int64(len(bodyBytes))
        }

        //nolint:gosec // G706: target host is from server configuration
        slog.Debug("Signing request for target host", "host", targetURL.Host)

        if err := signer.SignRequest(r.Context(), signingReq, creds); err != nil {
                return err
        }

        // Copy only the SigV4 headers back — these are the only headers the
        // AWS SDK v4 signer sets during SignHTTP.
        r.Header.Set("Authorization", signingReq.Header.Get("Authorization"))
        r.Header.Set("X-Amz-Date", signingReq.Header.Get("X-Amz-Date"))
        if tok := signingReq.Header.Get("X-Amz-Security-Token"); tok != "" {
                r.Header.Set("X-Amz-Security-Token", tok)
        }

        // Restore the body on the original request for downstream handlers
        // (the reverse proxy and tracingTransport both read it again).
        if bodyBytes != nil {
                r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
                r.ContentLength = int64(len(bodyBytes))
        }

        return nil
}

// ExtractSessionName extracts the session name from JWT claims.
// Returns an error if the configured claim is missing or empty, since a missing
// claim likely indicates a misconfiguration and would produce untraceable
// CloudTrail entries.
//
// The returned value is passed directly to AWS STS as RoleSessionName.
// and returns a clear error if the value doesn't conform.
//
// See: https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
func ExtractSessionName(claims map[string]interface{}, claimName string) (string, error) {
        value, ok := claims[claimName]
        if !ok {
                return "", fmt.Errorf("claim %q not found in token", claimName)
        }
        strValue, ok := value.(string)
        if !ok || strValue == "" {
                return "", fmt.Errorf("claim %q is not a non-empty string", claimName)
        }
        return strValue, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package awssts

import (
        "cmp"
        "fmt"
        "log/slog"
        "math"
        "slices"
        "strings"

        celgo "cel.dev/cel-go/cel"
        "github.com/aws/aws-sdk-go-v2/aws/arn"

        "github.com/stacklok/toolhive-core/cel"
)

// claimBindingExpression is the generic CEL expression used for claim-based role mappings.
// Instead of interpolating user-supplied claim values into CEL expression strings,
// we bind them as variables at evaluation time — making CEL injection impossible by design.
const claimBindingExpression = `claim_value in claims[role_claim_key]`

// newMatcherEngine creates a CEL engine for admin-authored matcher expressions.
// The only available variable is "claims" as a map[string]any.
func newMatcherEngine() *cel.Engine {
        return cel.NewEngine(
                celgo.Variable("claims", celgo.MapType(celgo.StringType, celgo.DynType)),
        )
}

// newClaimBindingEngine creates a CEL engine for claim-based mappings that uses
// variable binding instead of string interpolation. Three variables are available:
//   - claims: the JWT claims map
//   - claim_value: the claim value to match (e.g. "admins")
//   - role_claim_key: the claims map key to look up (e.g. "groups")
func newClaimBindingEngine() *cel.Engine {
        return cel.NewEngine(
                celgo.Variable("claims", celgo.MapType(celgo.StringType, celgo.DynType)),
                celgo.Variable("claim_value", celgo.StringType),
                celgo.Variable("role_claim_key", celgo.StringType),
        )
}

// ValidateRoleArn validates that the given string is a valid IAM role ARN.
// It accepts ARNs from all AWS partitions (aws, aws-cn, aws-us-gov) and
// supports role paths (e.g., arn:aws:iam::123456789012:role/service-role/MyRole).
func ValidateRoleArn(roleArn string) error {
        if roleArn == "" {
                return fmt.Errorf("%w: ARN is empty", ErrInvalidRoleArn)
        }

        // Use AWS SDK to parse the ARN
        parsed, err := arn.Parse(roleArn)
        if err != nil {
                return fmt.Errorf("%w: %s", ErrInvalidRoleArn, roleArn)
        }

        // Verify it's an IAM role
        if parsed.Service != "iam" {
                return fmt.Errorf("%w: not an IAM ARN: %s", ErrInvalidRoleArn, roleArn)
        }

        // Resource should start with "role/"
        if !strings.HasPrefix(parsed.Resource, "role/") {
                return fmt.Errorf("%w: not a role ARN: %s", ErrInvalidRoleArn, roleArn)
        }

        // Verify account ID is present and valid (12 digits)
        if len(parsed.AccountID) != 12 {
                return fmt.Errorf("%w: invalid account ID: %s", ErrInvalidRoleArn, roleArn)
        }
        for _, c := range parsed.AccountID {
                if c < '0' || c > '9' {
                        return fmt.Errorf("%w: invalid account ID: %s", ErrInvalidRoleArn, roleArn)
                }
        }

        return nil
}

// compiledMapping holds a role mapping with its compiled CEL expression.
type compiledMapping struct {
        roleArn    string
        priority   int
        expr       *cel.CompiledExpression
        claimValue string // non-empty for claim-based mappings; empty for matcher-based
}

// evalContext builds the CEL variable bindings for evaluating this mapping.
// Claim-based mappings bind claim_value and role_claim_key as variables so that
// user-supplied values are never interpolated into CEL expression strings,
// eliminating CEL injection by design. Matcher-based mappings only need claims.
func (cm *compiledMapping) evalContext(
        claims map[string]any,
        normalizedClaims map[string]any,
        roleClaim string,
) map[string]any {
        if cm.claimValue != "" {
                return map[string]any{
                        "claims":         normalizedClaims,
                        "claim_value":    cm.claimValue,
                        "role_claim_key": roleClaim,
                }
        }
        return map[string]any{"claims": claims}
}

// normalizeRoleClaim returns a copy of claims whose role claim value is
// normalized to a list so the claim binding expression
// `claim_value in claims[role_claim_key]` performs exact element membership
// regardless of how the IdP serializes the claim: a single string is wrapped
// into a one-element list, and a list passes through unchanged.
//
// The documented contract (config.go) treats the role claim as a value list, so
// any other shape is an unsupported deviation and fails closed: without this
// guard, a string-typed claim made CEL `in` raise a "no such overload" error
// that SelectRole swallowed as a non-match (silently granting FallbackRoleArn),
// and an object-typed claim made `in` test map-key membership (spuriously
// matching when the configured value was a key).
func normalizeRoleClaim(claims map[string]any, roleClaim string) (map[string]any, error) {
        v, ok := claims[roleClaim]
        if !ok {
                // A missing role claim is a normal "no match", not an error: index the
                // expression against an empty list so it evaluates false and SelectRole
                // falls back as it always did.
                return cloneClaimsWithRoleClaim(claims, roleClaim, []any{}), nil
        }
        switch t := v.(type) {
        case string:
                return cloneClaimsWithRoleClaim(claims, roleClaim, []any{t}), nil
        case []any, []string:
                return claims, nil
        default:
                return nil, fmt.Errorf("role claim %q has unsupported value type %T (want string or list of strings)", roleClaim, v)
        }
}

func cloneClaimsWithRoleClaim(claims map[string]any, roleClaim string, value any) map[string]any {
        clone := make(map[string]any, len(claims)+1)
        for key, claim := range claims {
                clone[key] = claim
        }
        clone[roleClaim] = value
        return clone
}

// RoleMapper handles mapping JWT claims to IAM roles with priority-based selection.
// It uses CEL expressions for flexible claim matching.
type RoleMapper struct {
        config   *Config
        mappings []compiledMapping
}

// NewRoleMapper creates a new RoleMapper with the provided configuration.
// It validates the configuration and compiles all CEL expressions during construction.
// Returns an error if the configuration is invalid or any expression fails to compile.
//
// ValidateConfig is called internally, so callers do not need to call both.
func NewRoleMapper(cfg *Config) (*RoleMapper, error) {
        if err := ValidateConfig(cfg); err != nil {
                return nil, fmt.Errorf("invalid config: %w", err)
        }

        claimEngine := newClaimBindingEngine()
        matcherEngine := newMatcherEngine()

        claimExpr, err := claimEngine.Compile(claimBindingExpression)
        if err != nil {
                return nil, fmt.Errorf("compiling claim binding expression: %w", err)
        }

        rm := &RoleMapper{
                config:   cfg,
                mappings: make([]compiledMapping, 0, len(cfg.RoleMappings)),
        }

        for i, mapping := range cfg.RoleMappings {
                if mapping.Claim != "" {
                        rm.mappings = append(rm.mappings, compiledMapping{
                                roleArn:    mapping.RoleArn,
                                priority:   effectivePriority(mapping.Priority),
                                expr:       claimExpr,
                                claimValue: mapping.Claim,
                        })
                        continue
                }

                expr, err := matcherEngine.Compile(mapping.Matcher)
                if err != nil {
                        return nil, fmt.Errorf("role mapping at index %d: %w: %w", i, ErrInvalidMatcher, err)
                }
                rm.mappings = append(rm.mappings, compiledMapping{
                        roleArn:  mapping.RoleArn,
                        priority: effectivePriority(mapping.Priority),
                        expr:     expr,
                })
        }

        return rm, nil
}

// SelectRole selects the appropriate IAM role based on JWT claims.
// It returns the role ARN to assume based on the following logic:
//  1. If no role mappings are configured, return the FallbackRoleArn
//  2. Evaluate each mapping's CEL expression against the claims
//  3. Collect all matching mappings
//  4. Sort matches by priority (lower number = higher priority)
//  5. Return the highest priority match
//  6. If no matches found, fall back to the FallbackRoleArn
func (rm *RoleMapper) SelectRole(claims map[string]any) (string, error) {
        // If no role mappings configured, use default role
        if len(rm.mappings) == 0 {
                return rm.fallbackRole(ErrMissingRoleConfig)
        }

        // Find all matching mappings
        roleClaim := rm.config.GetRoleClaim()
        normalizedClaims := claims
        var normalizationErr error
        for _, mapping := range rm.mappings {
                if mapping.claimValue != "" {
                        normalizedClaims, normalizationErr = normalizeRoleClaim(claims, roleClaim)
                        break
                }
        }

        var matches []compiledMapping
        var claimMappingErr error
        for _, mapping := range rm.mappings {
                if mapping.claimValue != "" && normalizationErr != nil {
                        // Keep evaluating other mappings: a valid matcher mapping may have
                        // already matched. If none do, fail closed below rather than granting
                        // the fallback role for an unsupported claim shape.
                        slog.Warn("role claim has unsupported shape, failing closed",
                                "role_arn", mapping.roleArn, "error", normalizationErr)
                        if claimMappingErr == nil {
                                claimMappingErr = normalizationErr
                        }
                        continue
                }

                ctx := mapping.evalContext(claims, normalizedClaims, roleClaim)
                match, err := mapping.expr.EvaluateBool(ctx)
                if err != nil {
                        if mapping.claimValue != "" {
                                // Claim-based mappings are normalized before evaluation, so an
                                // error here is unexpected. Keep evaluating other mappings; if no
                                // valid mapping matches, fail closed below instead of falling back.
                                slog.Warn("claim-based role mapping evaluation failed, failing closed",
                                        "role_arn", mapping.roleArn, "error", err)
                                if claimMappingErr == nil {
                                        claimMappingErr = err
                                }
                                continue
                        }
                        // Matcher expressions are admin-authored; keep the historical
                        // skip-and-fall-back behavior but surface the failure at Warn so
                        // operators can see it.
                        //nolint:gosec // G706: role ARN is from server configuration
                        slog.Warn("CEL expression evaluation failed, skipping mapping",
                                "role_arn", mapping.roleArn, "error", err)
                        continue
                }

                if match {
                        matches = append(matches, mapping)
                }
        }

        // A malformed claim-based mapping must not fall through to the fallback
        // role, but a valid mapping match always takes precedence over that error.
        if len(matches) == 0 && claimMappingErr != nil {
                return "", fmt.Errorf("%w: %w", ErrNoRoleMapping, claimMappingErr)
        }

        // If no matches, fall back to default role
        if len(matches) == 0 {
                return rm.fallbackRole(fmt.Errorf("%w: no mapping matched for the provided claims", ErrNoRoleMapping))
        }

        // Sort by priority (lower number = higher priority).
        // SortStableFunc preserves configuration order as a tie-breaker
        // when priorities are equal.
        slices.SortStableFunc(matches, func(a, b compiledMapping) int {
                return cmp.Compare(a.priority, b.priority)
        })

        // Return the highest priority match (lowest priority number)
        return matches[0].roleArn, nil
}

// fallbackRole returns the configured fallback role, or missingErr when none
// is configured. Callers provide their context-specific error to preserve the
// distinction between missing configuration and unmatched mappings.
func (rm *RoleMapper) fallbackRole(missingErr error) (string, error) {
        if rm.config.FallbackRoleArn == "" {
                return "", missingErr
        }
        return rm.config.FallbackRoleArn, nil
}

// ValidateConfig validates the AWS STS configuration structure.
// It checks that required fields are present, ARNs are well-formed,
// and session duration is within bounds.
//
// This performs structural validation only — CEL expression compilation is handled
// by NewRoleMapper. It is safe to call standalone for early validation at config
// load time. NewRoleMapper calls this internally, so callers do not need to call both.
func ValidateConfig(cfg *Config) error {
        if cfg == nil {
                return fmt.Errorf("config is nil")
        }

        // Region is required
        if cfg.Region == "" {
                return ErrMissingRegion
        }

        // Either FallbackRoleArn or RoleMappings must be configured
        if cfg.FallbackRoleArn == "" && len(cfg.RoleMappings) == 0 {
                return ErrMissingRoleConfig
        }

        // Validate FallbackRoleArn if provided
        if cfg.FallbackRoleArn != "" {
                if err := ValidateRoleArn(cfg.FallbackRoleArn); err != nil {
                        return err
                }
        }

        // Validate all role mappings (structural checks only)
        for i, mapping := range cfg.RoleMappings {
                if err := validateRoleMapping(i, mapping); err != nil {
                        return err
                }
        }

        // Validate session duration if specified
        if cfg.SessionDuration != 0 {
                if cfg.SessionDuration < MinSessionDuration {
                        return fmt.Errorf("session duration %d is below minimum %d seconds", cfg.SessionDuration, MinSessionDuration)
                }
                if cfg.SessionDuration > MaxSessionDuration {
                        return fmt.Errorf("session duration %d exceeds maximum %d seconds", cfg.SessionDuration, MaxSessionDuration)
                }
        }

        return nil
}

// validateRoleMapping validates the structural properties of a single role mapping.
func validateRoleMapping(index int, mapping RoleMapping) error {
        // Exactly one of Claim or Matcher must be set
        if mapping.Claim == "" && mapping.Matcher == "" {
                return fmt.Errorf("%w at index %d: either claim or matcher must be set", ErrInvalidRoleMapping, index)
        }
        if mapping.Claim != "" && mapping.Matcher != "" {
                return fmt.Errorf("%w at index %d: claim and matcher are mutually exclusive", ErrInvalidRoleMapping, index)
        }

        // RoleArn is required
        if mapping.RoleArn == "" {
                return fmt.Errorf("role mapping at index %d has empty role ARN", index)
        }

        // Validate the role ARN
        if err := ValidateRoleArn(mapping.RoleArn); err != nil {
                return fmt.Errorf("role mapping at index %d: %w", index, err)
        }

        return nil
}

// effectivePriority returns the priority value from the pointer, or math.MaxInt
// if nil. This makes omitted priority act as lowest-possible priority so that
// config order (via stable sort) is the natural tie-breaker.
func effectivePriority(p *int) int {
        if p != nil {
                return *p
        }
        return math.MaxInt
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package awssts provides AWS STS token exchange and SigV4 signing functionality.
package awssts

import (
        "bytes"
        "context"
        "crypto/sha256"
        "encoding/hex"
        "fmt"
        "io"
        "net/http"
        "time"

        "github.com/aws/aws-sdk-go-v2/aws"
        v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
)

// maxPayloadSize is the maximum request body size (10 MB) for SigV4 signing.
const maxPayloadSize = 10 * 1024 * 1024

// emptySHA256 is the well-known SHA-256 hash of an empty string, used for
// SigV4 signing of requests with no body.
const emptySHA256 = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"

// defaultService is the AWS service name used in SigV4 signing for AWS MCP Server.
// This value appears in the credential scope of the Authorization header:
//
//        Credential=AKIAEXAMPLE/20260206/us-east-1/aws-mcp/aws4_request
//
// The service name must match what AWS expects. For AWS MCP Server, this is "aws-mcp",
// as documented in the IAM actions (aws-mcp:InvokeMcp, aws-mcp:CallReadOnlyTool, etc.)
// and the endpoint URL pattern (aws-mcp.{region}.api.aws).
//
// See: https://docs.aws.amazon.com/aws-mcp/latest/userguide/getting-started-aws-mcp-server.html
const defaultService = "aws-mcp"

// RequestSigner signs HTTP requests using AWS Signature Version 4.
//
// SigV4 signing should run as close to the backend as possible.
// Modifying signed headers or the request body after signing will
// invalidate the signature.
type RequestSigner struct {
        signer  *v4.Signer
        region  string
        service string
}

type signerOption func(*RequestSigner)

// withService sets a custom service name for SigV4 signing.
func withService(service string) signerOption {
        return func(s *RequestSigner) {
                s.service = service
        }
}

// newRequestSigner creates a new SigV4 request signer for the specified region.
//
// By default, it uses "aws-mcp" as the service name for AWS MCP Server.
// Use withService to override for other AWS services.
func newRequestSigner(region string, opts ...signerOption) (*RequestSigner, error) {
        if region == "" {
                return nil, ErrMissingRegion
        }

        s := &RequestSigner{
                signer:  v4.NewSigner(),
                region:  region,
                service: defaultService,
        }

        for _, opt := range opts {
                opt(s)
        }

        return s, nil
}

// NewRequestSigner creates a new SigV4 request signer for the specified region
// and service name. An empty service string defaults to "aws-mcp".
//
// Exported so that pkg/vmcp/auth/strategies and other packages can sign requests
// outside the HTTP middleware flow.
func NewRequestSigner(region, service string) (*RequestSigner, error) {
        opts := []signerOption{}
        if service != "" {
                opts = append(opts, withService(service))
        }
        return newRequestSigner(region, opts...)
}

// SignRequest signs an HTTP request using AWS SigV4.
//
// This method:
//  1. Reads and hashes the request body with SHA-256
//  2. Signs the request with the provided credentials
//  3. Adds required headers: Authorization, X-Amz-Date, X-Amz-Security-Token
//
// The request body is consumed and replaced with a new reader containing
// the same content, allowing the request to be sent after signing.
//
// Parameters:
//   - ctx: Context for the signing operation
//   - req: HTTP request to sign (will be modified in place)
//   - creds: AWS credentials from STS token exchange
//
// Returns an error if:
//   - The request body cannot be read
//   - Signing fails
func (s *RequestSigner) SignRequest(ctx context.Context, req *http.Request, creds *aws.Credentials) error {
        if creds == nil {
                return fmt.Errorf("credentials are required for signing")
        }

        // Read and hash the request body
        payloadHash, bodyBytes, err := s.hashPayload(req)
        if err != nil {
                return fmt.Errorf("failed to hash request payload: %w", err)
        }

        // Replace the body with a new reader (the original was consumed)
        if bodyBytes != nil {
                req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
                req.ContentLength = int64(len(bodyBytes))
        }

        // Sign the request
        err = s.signer.SignHTTP(ctx, *creds, req, payloadHash, s.service, s.region, time.Now())
        if err != nil {
                return fmt.Errorf("failed to sign request: %w", err)
        }

        return nil
}

// hashPayload reads and hashes the request body with SHA-256.
//
// Returns:
//   - payloadHash: Hex-encoded SHA-256 hash of the body
//   - bodyBytes: The body content (for replacing the consumed reader)
//   - error: Any error reading the body
func (*RequestSigner) hashPayload(req *http.Request) (string, []byte, error) {
        // Handle empty body
        if req.Body == nil || req.Body == http.NoBody {
                return emptySHA256, nil, nil
        }

        defer func() { _ = req.Body.Close() }()

        // Read the body with a size limit to prevent memory exhaustion
        bodyBytes, err := io.ReadAll(io.LimitReader(req.Body, maxPayloadSize+1))
        if err != nil {
                return "", nil, err
        }
        if len(bodyBytes) > maxPayloadSize {
                return "", nil, fmt.Errorf("request body exceeds maximum size of %d bytes", maxPayloadSize)
        }

        // Hash the body
        hash := sha256.Sum256(bodyBytes)
        return hex.EncodeToString(hash[:]), bodyBytes, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package claims provides shared helpers for decoding JWT/OIDC claim values.
package claims

import (
        "slices"

        oauthclaims "github.com/stacklok/toolhive/pkg/oauthproto/claims"
)

// DecodeStringArray preserves the historical host-package API while the strict
// OAuth claim decoder is owned by pkg/oauthproto/claims.
func DecodeStringArray(raw any) ([]string, error) {
        return oauthclaims.DecodeStringArray(raw)
}

// DecodeStringArrayLenient decodes a claim value using the permissive contract
// the Cedar authorizer has always applied to group/role claims, where a claim
// that is present but partially malformed must still yield whatever usable
// names it carries rather than failing the whole authorization decision.
//
// It differs from DecodeStringArray in three ways, each of which Cedar depends
// on today:
//
//   - non-string elements of a []any are dropped instead of raising an error,
//     so an array of all non-strings decodes to a non-nil empty slice;
//   - a typed nil slice is a recognized array, not an error: []any(nil)
//     decodes to a non-nil empty slice and []string(nil) decodes to nil,
//     preserving each shape's existing nil-ness;
//   - a recognized but empty []any stays non-nil, so a caller can tell "the
//     IdP asserted this claim with no usable values" (which must suppress a
//     fallback to other claim names) from "the claim is absent" (nil).
//
// recognized reports whether raw was a supported array shape at all. A false
// value means the claim holds a scalar, object, or other unsupported type;
// the caller decides whether to log and ignore it. As with DecodeStringArray,
// the result never aliases the input.
//
// New security-relevant callers should prefer DecodeStringArray: dropping
// malformed elements silently is only appropriate where the existing
// behavior must be preserved.
func DecodeStringArrayLenient(raw any) (values []string, recognized bool) {
        switch v := raw.(type) {
        case []string:
                // slices.Clone preserves nil-ness: nil stays nil, non-nil empty stays
                // non-nil empty. Cedar distinguishes the two.
                return slices.Clone(v), true
        case []any:
                out := make([]string, 0, len(v))
                for _, item := range v {
                        if s, ok := item.(string); ok {
                                out = append(out, s)
                        }
                }
                return out, true
        default:
                return nil, false
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

import (
        "context"
        "errors"

        "github.com/golang-jwt/jwt/v5"

        "github.com/stacklok/toolhive-core/audit"
        serveridentity "github.com/stacklok/toolhive/pkg/authserver/identity"
)

// DefaultMaxDelegationDepth is ToolHive's default cap on how many RFC 8693
// delegation hops are retained when parsing an "act" claim. It is aligned
// with the issuance-side cap in pkg/authserver/server/tokenexchange/handler.go
// (maxDelegationDepth) and the claim-nesting cap in
// pkg/authz/authorizers/cedar/entity.go (maxClaimNestingDepth), both 10.
//
// This is deliberately LOWER than toolhive-core's audit.DefaultMaxDelegationDepth
// (16): core's value is a library-wide parse ceiling, while 10 is ToolHive's own
// minting cap. Parsing at the mint cap preserves the signal that a chain with
// Truncated=true could not have been minted by ToolHive's own issuance path.
const DefaultMaxDelegationDepth = 10

// IdentityContextKey is the key used to store a host Identity in the request context.
// Authentication middleware typically attaches it on authenticated request paths;
// its presence in a context alone does not prove that authentication succeeded.
//
// Using an empty struct as the key prevents collisions with other context keys,
// as each empty struct type is distinct even if they have the same name in different packages.
type IdentityContextKey struct{}

// WithIdentity stores an Identity in the context.
// If identity is nil, the original context is returned unchanged.
//
// This function is typically called by authentication middleware after successful
// authentication to make the identity available to downstream handlers.
//
// Side effect: if the context carries an IdentityHolder (injected by the audit
// middleware wrapping this call), the identity is also published into it. Two
// invariants follow for callers on request-derived contexts:
//   - Last write wins: the holder reports the most recently attached identity,
//     so audit events carry that principal. Do not attach a different principal
//     (e.g. a re-minted service identity) on a request-derived context unless
//     that is the principal audit should report.
//   - Call only on the request goroutine: the audit middleware reads the holder
//     when the response is written, so off-goroutine writes would race.
//
// Example:
//
//        identity := &Identity{PrincipalInfo: PrincipalInfo{Subject: "user123", Name: "Alice"}}
//        ctx = WithIdentity(ctx, identity)
func WithIdentity(ctx context.Context, identity *Identity) context.Context {
        if identity == nil {
                return ctx
        }
        // Also publish the identity to an IdentityHolder if one is present, so
        // middleware wrapping the auth middleware (e.g. audit) can observe the
        // identity even though the derived context only flows downstream.
        if holder, ok := IdentityHolderFromContext(ctx); ok {
                holder.Identity = identity
        }
        return context.WithValue(ctx, IdentityContextKey{}, identity)
}

// IdentityHolderContextKey is the key used to store an IdentityHolder in the
// request context.
type IdentityHolderContextKey struct{}

// IdentityHolder is a mutable carrier that lets middleware running OUTSIDE the
// auth middleware observe the authenticated identity. Context values only flow
// downstream, so a wrapper such as the audit middleware cannot read the
// identity that auth attaches for inner handlers. The wrapper injects an empty
// holder via WithIdentityHolder before calling the inner chain; WithIdentity
// fills it when the identity is attached; the wrapper reads it after the inner
// chain returns.
//
// The holder is written and read by the single request goroutine (writes
// happen-before the wrapper's post-ServeHTTP read), mirroring the audit
// package's BackendInfo pattern, so no synchronization is needed.
type IdentityHolder struct {
        Identity *Identity
}

// WithIdentityHolder returns a new context carrying the given IdentityHolder.
func WithIdentityHolder(ctx context.Context, holder *IdentityHolder) context.Context {
        return context.WithValue(ctx, IdentityHolderContextKey{}, holder)
}

// IdentityHolderFromContext retrieves the IdentityHolder from the context.
// Returns (nil, false) if no holder is present.
func IdentityHolderFromContext(ctx context.Context) (*IdentityHolder, bool) {
        holder, ok := ctx.Value(IdentityHolderContextKey{}).(*IdentityHolder)
        return holder, ok && holder != nil
}

// IdentityFromContext retrieves the host Identity from the context.
// Returns the identity and true if a non-nil identity is present, nil and false otherwise.
// A typed-nil *Identity stored directly in the context (bypassing WithIdentity) is
// treated as absent so callers can safely gate on the boolean without nil checks.
// Authentication middleware typically attaches this Identity on authenticated
// request paths, but context lookup alone does not prove that authentication succeeded.
//
// Example:
//
//        identity, ok := IdentityFromContext(ctx)
//        if !ok {
//            return errors.New("no identity in context")
//        }
//        log.Printf("Request from user: %s", identity.Subject)
func IdentityFromContext(ctx context.Context) (*Identity, bool) {
        identity, ok := ctx.Value(IdentityContextKey{}).(*Identity)
        return identity, ok && identity != nil
}

// PlatformUserContextKey is a compatibility alias for the authserver's
// credential-free canonical-user context key. It is deliberately distinct from
// IdentityContextKey.
//
// A value under this key means only "this is the canonical user to key storage
// on" — it is NOT proof that an authenticated principal is present. The two are
// kept structurally separate so a storage-scoped user id cannot be mistaken for a
// host Identity. Consumers that need the authenticated caller may retrieve the host
// Identity with IdentityFromContext on trusted authenticated request paths, but the
// context lookup alone is not proof that the Identity was validated. Storage layers
// that key on the canonical user use CanonicalUserFromContext.
type PlatformUserContextKey = serveridentity.CanonicalUserContextKey

// WithPlatformUser stores the platform's canonical user identifier in the
// context using the authserver's shared canonical-user contract. If userID is
// empty, the original context is returned unchanged.
//
// Use this — not WithIdentity — on paths that have resolved the user but have no
// host Identity to attach (e.g. the OAuth callback, which resolves the user while
// it is still minting the ToolHive bearer, so no validated token/claims exist yet).
// Reusing WithIdentity there would place a credential-free stub under the identity
// key that later readers could mistake for an authenticated principal.
//
// Request-serving paths that already carry the host Identity, typically attached
// after authentication, do NOT need to call this — CanonicalUserFromContext falls
// back to that Identity's PlatformUserID, so the canonical user is read from one
// accessor without storing it twice.
func WithPlatformUser(ctx context.Context, userID string) context.Context {
        return serveridentity.WithCanonicalUser(ctx, userID)
}

// PlatformUserFromContext retrieves the platform's canonical user identifier from
// the shared canonical-user key. Returns the identifier and true if present,
// including when an empty string was stored directly; returns "" and false when
// absent or stored with a non-string type.
//
// Most callers should use CanonicalUserFromContext, which also resolves the user
// from the host Identity, typically attached on authenticated request paths. Context
// lookup alone is not proof of authentication. Use this only when you specifically
// need the dedicated key and must NOT fall back to an Identity.
func PlatformUserFromContext(ctx context.Context) (string, bool) {
        return serveridentity.CanonicalUserFromContext(ctx)
}

// CanonicalUserFromContext returns the platform's canonical user identifier for
// storage keying. It is the single accessor storage layers should use, regardless
// of which path produced the context.
//
// It prefers the dedicated canonical-user key (set on identity-less paths such as
// the OAuth callback) and falls back to the host Identity's PlatformUserID (typically
// attached on authenticated request paths). Context lookup alone does not prove that
// authentication succeeded. Returns "" and false if neither is present. A string
// directly stored under the canonical-user key wins even if empty, so an explicit
// value overrides the identity-derived value.
func CanonicalUserFromContext(ctx context.Context) (string, bool) {
        if userID, ok := PlatformUserFromContext(ctx); ok {
                return userID, true
        }
        if identity, ok := IdentityFromContext(ctx); ok && identity.PlatformUserID != "" {
                return identity.PlatformUserID, true
        }
        return "", false
}

// claimsToIdentity converts JWT claims to Identity struct.
// It requires the 'sub' claim per OIDC Core 1.0 spec § 5.1.
// The original token can be provided for passthrough scenarios.
//
// Note: The Groups field is intentionally NOT populated here.
// Authorization logic MUST extract groups from the Claims map, as group claim
// names vary by provider (e.g., "groups", "roles", "cognito:groups").
func claimsToIdentity(claims jwt.MapClaims, token string) (*Identity, error) {
        // Validate required 'sub' claim per OIDC Core 1.0 spec
        sub, ok := claims["sub"].(string)
        if !ok || sub == "" {
                return nil, errors.New("missing or invalid 'sub' claim (required by OIDC Core 1.0 § 5.1)")
        }

        // Filter internal claims that should not be externalized (e.g., in
        // webhook payloads or audit logs). The tsid is a session identifier
        // used to look up upstream tokens in storage; exposing it widens the
        // attack surface if a webhook receiver is compromised.
        filteredClaims := filterInternalClaims(claims)

        identity := &Identity{
                PrincipalInfo: PrincipalInfo{
                        Subject:        sub,
                        PlatformUserID: sub,
                        Claims:         filteredClaims,
                },
                Token:     token,
                TokenType: "Bearer",
        }

        // Extract optional standard claims
        if name, ok := claims["name"].(string); ok {
                identity.Name = name
        }
        if email, ok := claims["email"].(string); ok {
                identity.Email = email
        }

        // Parse the RFC 8693 "act" claim (if present) into a delegation chain.
        // This is primarily for audit: per RFC 8693 §4.1, an authorizer applying
        // access control MUST consider only the current actor (chain.Chain[0]);
        // any nested actors in chain.Chain[1:] are informational only.
        // The parse cannot fail by design: a malformed act claim must never fail
        // authentication, only be recorded (Malformed/MalformedReason on the chain).
        if act, ok := claims["act"]; ok && act != nil {
                chain := audit.ParseDelegationChain(act, DefaultMaxDelegationDepth)
                if !chain.IsZero() {
                        identity.DelegationChain = chain
                }
        }

        return identity, nil
}

// internalClaims are JWT claim keys used internally by the auth server
// that must not be externalized in webhook payloads, audit logs, etc.
// "tsid" is the token session ID used to look up upstream tokens in storage.
var internalClaims = []string{"tsid"}

// filterInternalClaims returns a copy of claims with internal keys removed.
func filterInternalClaims(claims jwt.MapClaims) jwt.MapClaims {
        filtered := make(jwt.MapClaims, len(claims))
        for k, v := range claims {
                filtered[k] = v
        }
        for _, key := range internalClaims {
                delete(filtered, key)
        }
        return filtered
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package dcr

import "fmt"

// RegistrationProbe selects how ResolveCredentials confirms, on a cache hit,
// that the upstream still recognises the cached client (see
// verifyCachedRegistration). The zero value means ProbeTokenEndpoint.
type RegistrationProbe string

const (
        // ProbeTokenEndpoint redeems a random authorization code at the token
        // endpoint with the cached credentials. invalid_client means the
        // registration is gone; invalid_grant means it is alive. It decides only
        // for upstreams that authenticate the client before validating the grant
        // (a ToolHive auth server does). This is the default.
        ProbeTokenEndpoint RegistrationProbe = "tokenEndpoint"

        // ProbeAuthorizationEndpoint sends an authorization request for the
        // cached client without following redirects and no credentials, so it
        // covers every auth method. RFC 6749 §4.1.2.1 has the server answer an
        // unknown client_id or redirect_uri directly, with no redirect, while a
        // known client is redirected (to login, or back to redirect_uri with an
        // error). A direct 400, 401 or 403 carrying an RFC 6749 JSON error object
        // with error invalid_client, unauthorized_client or invalid_request means
        // the registration is gone; any other direct answer (an HTML page, a
        // gateway's 401/403, another error code) is unverified and keeps the
        // cached client. Use it for upstreams whose token endpoint validates the
        // grant before the client, where ProbeTokenEndpoint cannot decide. A
        // false rejection re-registers at every start and, when registration
        // also fails, fails startup.
        ProbeAuthorizationEndpoint RegistrationProbe = "authorizationEndpoint"

        // ProbeNone never contacts the upstream on a cache hit; the cached
        // registration is used as-is and retired only through RotationKey.
        ProbeNone RegistrationProbe = "none"
)

// Validate rejects a RegistrationProbe value other than the empty default
// and the three named strategies.
func (p RegistrationProbe) Validate() error {
        switch p {
        case "", ProbeTokenEndpoint, ProbeAuthorizationEndpoint, ProbeNone:
                return nil
        default:
                return fmt.Errorf("registration_probe must be one of %q, %q or %q (got %q)",
                        ProbeTokenEndpoint, ProbeAuthorizationEndpoint, ProbeNone, string(p))
        }
}

// effective returns the strategy to run, resolving the empty default.
func (p RegistrationProbe) effective() RegistrationProbe {
        if p == "" {
                return ProbeTokenEndpoint
        }
        return p
}

// Request is the profile-neutral input to ResolveCredentials. Each consumer
// (embedded authserver, CLI OAuth flow) translates its domain types into a
// Request at the call site so the resolver does not import any consumer's
// shapes.
//
// The struct collects exactly the fields the resolver reads:
//
//   - Identity: Issuer (the caller's logical scope — see the Issuer field
//     doc for what each consumer puts there), Scopes, and RedirectURI.
//   - Endpoint discovery: DiscoveryURL or RegistrationEndpoint, with optional
//     explicit AuthorizationEndpoint / TokenEndpoint overrides.
//   - Registration metadata: InitialAccessToken, ClientName. The caller is
//     responsible for resolving any file-or-env reference (e.g. the embedded
//     authserver's InitialAccessTokenFile / InitialAccessTokenEnvVar) into
//     InitialAccessToken before constructing a Request.
//
// Mutually exclusive constraints are enforced at validation time:
//
//   - Exactly one of DiscoveryURL / RegistrationEndpoint must be non-empty.
//   - Issuer must be non-empty.
//
// Constructing a Request is the caller's responsibility; the resolver does
// not clone or mutate it.
type Request struct {
        // Issuer is the caller's logical scope for cache keying and (when
        // RedirectURI is empty) the basis for the default redirect URI
        // ({Issuer}/oauth/callback). Required.
        //
        // What each consumer puts here:
        //
        //   - Embedded authserver: its own issuer identifier. The authserver
        //     owns a distinct logical identity from the upstream IdP, and
        //     defaulting RedirectURI to {Issuer}/oauth/callback lands the
        //     callback on the authserver's origin (correct).
        //   - CLI OAuth flow: the upstream OAuth provider's issuer URL,
        //     because the CLI has no separate logical issuer of its own. The
        //     CLI always supplies an explicit loopback RedirectURI per
        //     RFC 8252 §7.3, so the redirect-URI defaulting is never
        //     exercised on this path.
        //
        // The resolver's cache key is (Issuer, UpstreamID, RedirectURI,
        // ScopesHash) plus RotationKey when set, where UpstreamID is derived by
        // the resolver from DiscoveryURL / RegistrationEndpoint (issue #5823).
        // The two consumers
        // do not collide on the key because the embedded authserver's
        // RedirectURI lives on the authserver's origin and the CLI's lives on a
        // loopback (RFC 8252 §7.3) — even when Issuer happens to match between
        // the two profiles, the RedirectURI component keeps the cache key
        // distinct; UpstreamID additionally keeps two upstreams within a single
        // consumer apart. See the cross-consumer caveat on dcrFlight in
        // resolver.go for the wider invariant.
        //
        // Issuer is *not* used for RFC 8414 §3.3 metadata verification —
        // that uses an issuer recovered from DiscoveryURL inside the
        // resolver.
        Issuer string

        // RedirectURI is the redirect URI to register with the upstream. When
        // empty, the resolver derives {Issuer}/oauth/callback. HTTPS is required
        // except for loopback hosts (RFC 8252 §7.3 — the CLI flow uses
        // http://localhost:{port}/callback).
        RedirectURI string

        // Scopes are the OAuth scopes to request in the registration body.
        // May be empty; the resolver will fall back to discovered scopes if the
        // upstream advertises any.
        Scopes []string

        // DiscoveryURL points at the upstream's RFC 8414 / OIDC Discovery
        // document. The resolver fetches this URL exactly once and reads
        // registration_endpoint, authorization_endpoint, token_endpoint,
        // token_endpoint_auth_methods_supported, scopes_supported, and
        // code_challenge_methods_supported from it. On this branch the resolver
        // also derives the cache key's UpstreamID from the upstream issuer
        // recovered from this URL.
        //
        // Mutually exclusive with RegistrationEndpoint.
        DiscoveryURL string

        // RegistrationEndpoint is used directly when set, bypassing discovery.
        // On this branch the caller is expected to also supply AuthorizationEndpoint
        // and TokenEndpoint explicitly; the resolver's auth-method selection
        // defaults to client_secret_basic because no server-capability fields
        // are available unless CodeChallengeMethodsSupported is also set. This
        // URL is also the cache key's UpstreamID on this branch.
        //
        // Mutually exclusive with DiscoveryURL.
        RegistrationEndpoint string

        // RotationKey is an optional operator-chosen label folded into the cache
        // key (and the in-process singleflight key) when non-empty. It is never
        // sent to the upstream. Changing it forces a fresh RFC 7591 registration
        // on the next resolve, without editing the credential store by hand.
        // ResolveCredentials replaces a registration the upstream rejects on its
        // own when the configured RegistrationProbe can tell (see
        // verifyCachedRegistration); RotationKey is how an operator retires one
        // the probe cannot tell is gone, or rotates one on request. An empty
        // value addresses the same cache row as a Request written before the
        // field existed.
        RotationKey string

        // RegistrationProbe selects how a cached registration is checked against
        // the upstream on a cache hit. Empty means ProbeTokenEndpoint. See the
        // RegistrationProbe constants for when each strategy can decide.
        RegistrationProbe RegistrationProbe

        // CodeChallengeMethodsSupported lists the PKCE challenge methods the
        // upstream advertises. Only consulted when RegistrationEndpoint is set;
        // on the DiscoveryURL branch the resolver reads this value from the
        // fetched metadata document. Callers that have already performed
        // multi-URL AS metadata discovery (e.g. via
        // oauthproto.FetchAuthorizationServerMetadata) should populate this
        // field so the S256 PKCE gate fires correctly without a redundant
        // round-trip.
        CodeChallengeMethodsSupported []string

        // AuthorizationEndpoint, when non-empty, overrides any value discovered
        // via DiscoveryURL. Explicit caller configuration always wins.
        AuthorizationEndpoint string

        // TokenEndpoint, when non-empty, overrides any value discovered via
        // DiscoveryURL. Explicit caller configuration always wins.
        TokenEndpoint string

        // InitialAccessToken is the RFC 7591 initial access token presented to
        // the registration endpoint as a Bearer header. The caller resolves any
        // file-or-env reference before populating this field.
        InitialAccessToken string

        // ClientName is sent as the RFC 7591 "client_name" registration
        // metadata. When empty, the resolver uses
        // oauthproto.ToolHiveMCPClientName.
        ClientName string

        // PublicClient, when true, instructs the resolver to register as a
        // public PKCE client (token_endpoint_auth_method=none) regardless of
        // what other methods the upstream advertises. This matches the CLI
        // flow's intent (RFC 8252 §8.4 native public clients) and the
        // resolver still enforces the RFC 7636 / OAuth 2.1 S256 PKCE gate:
        // when the upstream's discovery metadata does not advertise S256 in
        // code_challenge_methods_supported, the registration is refused with
        // a clear error rather than silently downgrading.
        //
        // When false (the embedded-authserver default), the resolver picks
        // the most preferred auth method the upstream advertises among the
        // ones our OAuth2 runtime can actually use (client_secret_basic >
        // client_secret_post > none, with the same S256 gate on "none").
        // private_key_jwt is excluded from this preference even when
        // advertised: it requires a signed client assertion the runtime has
        // no way to produce.
        //
        // Has no effect on the RegistrationEndpoint-direct branch when neither
        // DiscoveryURL nor CodeChallengeMethodsSupported is set: without
        // code_challenge_methods_supported the S256 gate cannot be evaluated,
        // so the resolver refuses to register as a public client.
        PublicClient bool

        // CAFilePath is the path to a PEM CA bundle added to the system roots.
        CAFilePath string

        // AllowPrivateIPs permits both of the resolver's outbound calls — the
        // discovery fetch to DiscoveryURL and the registration POST to the
        // resolved registration endpoint — to connect to private IP ranges. For
        // ServerSuppliedEndpoints, this is the only way to permit private,
        // loopback, or link-local dialing; localhost and
        // INSECURE_DISABLE_URL_VALIDATION do not widen that gate. Set this to true
        // only when the upstream authorization server is reachable solely over a
        // private address (e.g. an in-cluster IdP with no public endpoint);
        // defaults to false. Mirrors the AllowPrivateIPs posture already carried
        // by the OAuth2 and OIDC upstream configs so a DCR upstream is not the one
        // outbound-facing path without an SSRF guard.
        AllowPrivateIPs bool

        // ServerSuppliedEndpoints reports that DiscoveryURL or RegistrationEndpoint
        // originated with a remote MCP server rather than with the operator. The
        // resolver then uses the strict server-supplied client policy, which does
        // not grant private-IP access merely because an endpoint is localhost or
        // INSECURE_DISABLE_URL_VALIDATION is set. Leave false for trusted,
        // operator-configured embedded-authserver upstreams.
        //
        // A registration endpoint read out of a fetched metadata document is always
        // treated as server-supplied regardless of this field: the upstream controls
        // that document even when the operator chose the DiscoveryURL that returned it.
        ServerSuppliedEndpoints bool
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package dcr is the shared RFC 7591 Dynamic Client Registration client used
// by every consumer in the codebase that needs to register a downstream
// OAuth 2.x client at runtime. The package owns the stateful concerns of the
// flow — credential cache, in-process singleflight deduplication, scope-set
// canonicalisation, token-endpoint auth-method selection (with the RFC 7636 /
// OAuth 2.1 S256 PKCE gate), RFC 7591 §3.2.1 expiry-driven cache invalidation,
// re-registration when the upstream no longer recognises a cached client,
// the bearer-token transport with redirect refusal, and panic recovery around
// the registration body. Stateless RFC 7591 wire-shape primitives live in
// pkg/oauthproto.
//
// # API surface
//
// ResolveCredentials takes a profile-neutral Request and a CredentialStore
// and returns a Resolution. The Request carries only the fields the resolver
// actually reads (issuer, redirect URI, scopes, discovery URL or registration
// endpoint, optional explicit endpoint overrides, optional initial access
// token, optional registration metadata); each consumer translates its
// domain types into a Request at the call site so the resolver does not
// import any consumer's shapes.
//
// Two consumers exist today:
//
//   - pkg/authserver/runner constructs a Request from
//     *authserver.OAuth2UpstreamRunConfig and uses its own adapter helpers
//     (needsDCR / consumeResolution / applyResolutionToOAuth2Config in
//     runner/dcr_adapter.go) to fold the resolution back into its
//     run-config and built upstream.OAuth2Config.
//   - pkg/auth/discovery constructs a Request from *OAuthFlowConfig for the
//     CLI OAuth flow and copies the returned Resolution onto OAuthFlowResult.
//
// # Concurrency
//
// The package maintains a process-global singleflight keyed on the tuple
// (issuer, upstreamID, redirectURI, scopesHash, rotationKey) so concurrent
// ResolveCredentials calls across all consumers in a single process
// coalesce when their cache keys match. Consumers that share all five
// values will share a flight — the deduplication is a feature for the
// embedded authserver but means callers cannot assume per-call-site flight
// isolation. See the dcrFlight doc comment below for the rationale.
//
// See issue #5145 for the design discussion that motivated lifting this out
// of pkg/authserver/runner.
package dcr

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "net/url"
        "regexp"
        "runtime/debug"
        "slices"
        "strings"
        "time"

        "golang.org/x/sync/singleflight"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// dcrFlight coalesces concurrent ResolveCredentials calls that share the
// same Key. Two goroutines hitting the resolver for the same upstream and
// scope set will both miss the cache, so without coalescing they would both
// call RegisterClientDynamically and the loser's registration would become
// orphaned at the upstream IdP — an operator-visible cleanup task and
// possibly a transient startup failure if the upstream rate-limits
// concurrent registrations. Followers wait for the leader's result and
// observe the same Resolution.
//
// Lifetime: process-wide. This intentionally contrasts with the
// CredentialStore the embedded authserver constructs and injects into
// ResolveCredentials, which is per-instance for the memory backend and
// shared across replicas for Redis. The asymmetry is load-bearing: the
// singleflight only deduplicates the in-flight network call, while the
// cache deduplicates the resolution itself across calls. Process-wide
// flight means concurrent EmbeddedAuthServer instances in the same
// process targeting the same upstream still get deduplicated; the
// injected cache decides whether the resolution is fresh enough to
// reuse. A Redis-backed store still wants this in-process gate so a
// single replica does not double-register against itself.
//
// Cross-consumer caveat: because dcrFlight is package-global, two
// consumers that happen to construct identical Keys (same issuer, same
// upstream ID, same redirect URI, same scopes hash) will share a single
// in-flight registration even if they semantically want different client
// profiles.
// The two current call sites do not collide by construction — the embedded
// authserver's redirect URI lives on the AS origin, and the CLI flow's
// redirect URI lives on a loopback (http://localhost:{port}/callback per
// RFC 8252 §7.3), so the disjoint RedirectURI address spaces separate
// the public-client and confidential-client profiles at both the
// singleflight and the persistent-cache layers. A future consumer that
// defaults its redirect URI into either of those spaces would silently
// coalesce; if a third profile is added the flight key (and persisted
// DCRKey) MUST gain a consumer-identifier component so a collision is
// impossible rather than improbable.
var dcrFlight singleflight.Group

// flightKeyOf canonicalises a Key into the singleflight string used by
// dcrFlight. Each field is length-prefixed (the same scheme redisDCRKey
// uses in pkg/authserver/storage) so the key is unambiguous regardless of
// field contents; ScopesHash is a SHA-256 hex digest with no colons, so it
// is appended without a prefix, exactly as redisDCRKey does. Exposed as a
// function so tests and future inspection helpers can compute the exact key
// the resolver would route through dcrFlight without re-implementing the
// encoding.
//
// Length-prefixing rather than a "separator byte that never appears":
// UpstreamID is NOT guaranteed to be a clean URL when this key is built. On
// the CLI/remote path it is req.RegistrationEndpoint copied verbatim from
// the upstream's discovery-document JSON, and validateUpstreamEndpointURL
// does not run until later, inside the singleflight body
// (resolveDCREndpoints) — after this key already exists. A newline (or any
// byte) in a discovery-supplied registration_endpoint would therefore reach
// a naive "\n"-joined key, where a crafted value could byte-collide with a
// concurrently-registering upstream and, as the singleflight follower,
// observe that upstream's Resolution. Length-prefixing removes the
// assumption entirely and aligns the flight key with the persisted DCRKey.
//
// UpstreamID keeps two upstreams that share Issuer, RedirectURI, and
// scopes from coalescing into a single flight (issue #5823) — the same
// distinction the persisted DCRKey draws, so the singleflight and cache
// layers stay aligned. RotationKey, when non-empty, is appended
// length-prefixed exactly as redisDCRKey appends it, so a rotation also
// separates the in-flight registration from one still running under the
// previous key; an empty RotationKey leaves the flight key unchanged.
//
// PublicClient is intentionally NOT part of the flight key — the
// dcrFlight doc above explains why: today's two consumers register on
// disjoint RedirectURI address spaces (AS-origin vs RFC 8252 loopback),
// so the public-client and confidential-client profiles cannot share a
// flight key by construction. The same property protects the persisted
// DCRKey from cross-profile coalescence; the two layers are aligned
// rather than asymmetric.
func flightKeyOf(key Key) string {
        base := fmt.Sprintf("%d:%s:%d:%s:%d:%s:%s",
                len(key.Issuer), key.Issuer,
                len(key.UpstreamID), key.UpstreamID,
                len(key.RedirectURI), key.RedirectURI,
                key.ScopesHash)
        if key.RotationKey == "" {
                return base
        }
        return fmt.Sprintf("%s:%d:%s", base, len(key.RotationKey), key.RotationKey)
}

// defaultUpstreamRedirectPath is the redirect path derived from the issuer
// origin when the caller's run-config does not supply an explicit RedirectURI.
// Matches the authserver's public callback route.
const defaultUpstreamRedirectPath = "/oauth/callback"

// authMethodPreference is the preferred order of token_endpoint_auth_methods,
// most preferred first. The resolver intersects this list with the server's
// advertised methods and picks the first match.
//
// Rationale: client_secret_basic and client_secret_post are equally secure in
// transit but basic is marginally preferred because the credentials do not
// appear in request-body logs. "none" is the fallback for public PKCE clients.
// private_key_jwt is intentionally excluded because the OAuth2 runtime cannot
// produce the required client assertion.
var authMethodPreference = []string{
        "client_secret_basic",
        "client_secret_post",
        "none",
}

// Resolution captures the full RFC 7591 + RFC 7592 response for a
// successful Dynamic Client Registration, together with the endpoints the
// upstream advertises so the caller need not re-discover them.
//
// The struct is the unit of storage in CredentialStore and the unit of
// application that consumers project back into their own domain types
// (e.g. consumeResolution / applyResolutionToOAuth2Config in
// pkg/authserver/runner/dcr_adapter.go, or direct OAuthFlowResult field
// writes in pkg/auth/discovery).
//
// MUST update both converters (resolutionToCredentials and
// credentialsToResolution in store.go) when adding, renaming, or
// removing a field here. The two converters are the seam between this
// dcr-package type and the persisted *storage.DCRCredentials shape; a
// field added here without a paired converter update will silently fail
// to round-trip across an authserver restart, the exact "parallel types
// drift" failure mode .claude/rules/go-style.md warns about. The
// round-trip behaviour is pinned by TestResolutionCredentialsRoundTrip
// in store_test.go.
type Resolution struct {
        // ClientID is the RFC 7591 "client_id" returned by the authorization
        // server.
        ClientID string

        // ClientSecret is the RFC 7591 "client_secret" returned by the
        // authorization server. Empty for public PKCE clients.
        ClientSecret string

        // AuthorizationEndpoint is the discovered (or configured) authorization
        // endpoint for this upstream.
        AuthorizationEndpoint string

        // TokenEndpoint is the discovered (or configured) token endpoint for this
        // upstream.
        TokenEndpoint string

        // RegistrationAccessToken is the RFC 7592 "registration_access_token"
        // required for subsequent registration management operations (update,
        // read, delete).
        RegistrationAccessToken string

        // RegistrationClientURI is the RFC 7592 "registration_client_uri" for
        // registration management operations.
        RegistrationClientURI string

        // TokenEndpointAuthMethod is the authentication method negotiated at the
        // token endpoint for this client.
        TokenEndpointAuthMethod string

        // RedirectURI is the redirect URI presented to the authorization server
        // during registration. When the caller's Request did not specify one,
        // this holds the defaulted value derived from the issuer + /oauth/callback
        // (via resolveUpstreamRedirectURI). Persisting it on the resolution lets
        // consumers (pkg/authserver/runner.consumeResolution) write it back onto
        // the run-config COPY so that downstream code (buildPureOAuth2Config,
        // upstream.OAuth2Config validation) sees a non-empty RedirectURI.
        RedirectURI string

        // ClientIDIssuedAt is the RFC 7591 §3.2.1 "client_id_issued_at" value
        // converted to a Go time.Time. Zero when the upstream omitted the field
        // (the field is OPTIONAL per RFC 7591). Informational; not used to
        // invalidate the cache.
        ClientIDIssuedAt time.Time

        // ClientSecretExpiresAt is the RFC 7591 §3.2.1 "client_secret_expires_at"
        // value converted to a Go time.Time. The wire convention is that 0 means
        // "the secret does not expire"; in this struct that is represented by
        // the zero time.Time so callers can use IsZero() rather than special-
        // casing 0.
        //
        // When non-zero, this field is the authoritative signal that
        // lookupCachedResolution uses to refetch credentials before the upstream
        // rejects them at the token endpoint. The 90-day dcrStaleAgeThreshold
        // is a heuristic for "consider rotating"; this is a hard expiry asserted
        // by the upstream itself.
        ClientSecretExpiresAt time.Time

        // CreatedAt is the wall-clock time at which the resolution was completed.
        // Used by the staleness observability in lookupCachedResolution to
        // compute age against dcrStaleAgeThreshold and emit a warn log when a
        // cached registration exceeds the threshold.
        CreatedAt time.Time
}

// Step identifiers for structured error logs emitted by the caller of
// ResolveCredentials. These values flow through the "step" attribute so
// operators can narrow failures to a specific phase without parsing error
// messages. They are reported only at the boundary log — see
// dcrStepError — so a single failure produces a single slog.Error record.
const (
        dcrStepValidate         = "validate"
        dcrStepResolveRedirect  = "resolve_redirect_uri"
        dcrStepResolveUpstream  = "resolve_upstream_id"
        dcrStepCacheRead        = "cache_read"
        dcrStepMetadata         = "metadata_discovery"
        dcrStepSelectAuthMethod = "select_auth_method"
        dcrStepRegister         = "dcr_call"
        dcrStepCacheWrite       = "cache_write"
        //nolint:gosec // G101: this is a log/error "step" tag, not a credential value.
        dcrStepExpiredCredential = "expired_credential"
)

// dcrStepError annotates a resolver error with the phase it was produced
// in. The boundary caller (buildUpstreamConfigs) emits the single
// slog.Error record for the failure; individual error branches inside
// ResolveCredentials do not log so that each failure surfaces exactly
// once in the combined log stream.
//
// RedirectURI is included when known so that operators can correlate the
// failure with a specific upstream registration without parsing the
// wrapped error string. Stack carries a captured stack trace for the
// dcrStepRegister panic-recovery branch so LogStepError can include
// it in the single boundary record without the in-defer site emitting
// its own duplicate slog.Error. A zero-value dcrStepError is invalid;
// construct via newDCRStepError or the resolver's internal helpers.
type dcrStepError struct {
        Step        string
        Issuer      string
        RedirectURI string
        Stack       string
        Err         error
}

// Error implements error. The "step" tag mirrors the structured-log
// attribute so command-line log scrapers see the same phase identifier.
func (e *dcrStepError) Error() string {
        return fmt.Sprintf("dcr: %s: %s", e.Step, e.Err.Error())
}

// Unwrap lets errors.Is / errors.As reach the wrapped cause.
func (e *dcrStepError) Unwrap() error { return e.Err }

// newDCRStepError builds a dcrStepError. It never returns nil for a
// non-nil cause.
func newDCRStepError(step, issuer, redirectURI string, err error) *dcrStepError {
        return &dcrStepError{
                Step:        step,
                Issuer:      issuer,
                RedirectURI: redirectURI,
                Err:         err,
        }
}

// ResolveCredentials performs Dynamic Client Registration for req against
// the upstream authorization server identified by req.DiscoveryURL or
// req.RegistrationEndpoint, caching the resulting credentials in cache.
// On cache hit, when cache implements Replacer, the resolver sends one
// request to confirm the upstream still recognises the cached client, using
// the strategy req.RegistrationProbe selects (verifyCachedRegistration): the
// token endpoint by default, or the authorization endpoint. When the upstream
// rejects the client, the resolver registers a new client and replaces the
// cached entry; otherwise, and always when cache does not implement Replacer
// or req.RegistrationProbe is ProbeNone, it returns the cached resolution
// without discovery or registration I/O.
//
// Neither probe can tell every dropped registration apart from a live one.
// The token-endpoint probe is blind to an upstream that validates the grant
// before authenticating a public client, and it does not sign for
// private_key_jwt clients, so it skips them. The authorization-endpoint probe
// sends no credentials and probes every auth method, but it is blind to an
// upstream that renders its unknown-client error with 200, a redirect, or a
// page instead of an RFC 6749 JSON error. For those, the operator retires the
// registration with Request.RotationKey.
//
// The check runs only here, so a registration the upstream drops while the
// caller keeps using a returned resolution is recovered on the next
// ResolveCredentials call (for the embedded authserver, its next start), not
// in place.
//
// req.Issuer is *this caller's* logical issuer identifier, NOT the upstream's.
// It is used to key the cache and to default the redirect URI to
// {req.Issuer}/oauth/callback when req.RedirectURI is empty. The upstream's
// issuer is recovered separately from req.DiscoveryURL inside the resolver
// and is used solely for RFC 8414 §3.3 metadata verification. Passing the
// upstream's issuer in req.Issuer would produce a wrong-origin default
// redirect and a cache key that does not identify the caller context.
//
// The resolver does not mutate req or the cache on failure.
//
// Observability: this function never calls slog.Error directly — all
// failures are annotated with a *dcrStepError and returned to the caller,
// which is expected to emit the boundary Error record. This avoids the
// double-logging pattern where both the resolver and the outer frame
// report the same failure. Cache-hit Debug / stale-Warn logs and the
// successful-registration Debug log are emitted here because they have no
// outer-frame equivalent. No secret values (client_secret,
// registration_access_token, initial_access_token) are ever logged — only
// public metadata such as client_id and redirect_uri.
func ResolveCredentials(
        ctx context.Context,
        req *Request,
        cache CredentialStore,
) (*Resolution, error) {
        if err := validateResolveInputs(req, cache); err != nil {
                issuer := ""
                if req != nil {
                        issuer = req.Issuer
                }
                return nil, newDCRStepError(dcrStepValidate, issuer, "", err)
        }

        redirectURI, err := resolveUpstreamRedirectURI(req.RedirectURI, req.Issuer)
        if err != nil {
                return nil, newDCRStepError(dcrStepResolveRedirect, req.Issuer, "",
                        fmt.Errorf("resolve redirect uri: %w", err))
        }

        // Identify the upstream authorization server so two upstreams that share
        // req.Issuer, redirectURI, and scopes — the common case within a single
        // embedded authserver — do not collide on one cache entry (issue #5823).
        upstreamID, err := resolveUpstreamKeyIdentity(req)
        if err != nil {
                return nil, newDCRStepError(dcrStepResolveUpstream, req.Issuer, redirectURI,
                        fmt.Errorf("resolve upstream identity: %w", err))
        }

        scopes := slices.Clone(req.Scopes)
        key := Key{
                Issuer:      req.Issuer,
                UpstreamID:  upstreamID,
                RedirectURI: redirectURI,
                ScopesHash:  storage.ScopesHash(scopes),
                RotationKey: req.RotationKey,
        }

        // Coalesce concurrent cache lookups, cache-hit verifications and
        // registrations for the same Key — see dcrFlight doc comment. The leader
        // runs the closure; followers receive the leader's *Resolution result.
        // The flight key embeds the Key fields with a separator that cannot
        // appear in any of them (newline is not valid in OAuth scope tokens,
        // URLs, or hex digests).
        //
        // The cache is read inside the flight, so a flight always verifies or
        // registers against the row as it stands when the flight starts, never
        // one a caller read before an earlier flight for the same Key replaced
        // it. The lookup still precedes any discovery or registration I/O.
        //
        // A defer/recover inside the closure converts a panic in the lookup,
        // registerAndCache or revalidateCachedResolution (or anything they call)
        // into a normal error. Without this, singleflight re-panics the leader's panic in every
        // follower — N concurrent callers
        // for the same Key would all crash with the same value. The panic is
        // still surfaced: the captured stack trace is attached to the wrapped
        // dcrStepError and surfaces in the single boundary log emitted by
        // LogStepError, so the failure produces exactly one Error record (no
        // in-defer log here) and callers can react to it as a normal failure.
        flightKey := flightKeyOf(key)
        resolutionAny, err, _ := dcrFlight.Do(flightKey, func() (res any, err error) {
                defer func() {
                        if r := recover(); r != nil {
                                stepErr := newDCRStepError(dcrStepRegister, req.Issuer, redirectURI,
                                        fmt.Errorf("registration panicked: %v", r))
                                stepErr.Stack = string(debug.Stack())
                                err = stepErr
                                res = nil
                        }
                }()
                cached, hit, err := lookupCachedResolution(ctx, cache, key, req.Issuer, redirectURI)
                if err != nil {
                        return nil, newDCRStepError(dcrStepCacheRead, req.Issuer, redirectURI, err)
                }
                if hit {
                        return revalidateCachedResolution(ctx, req, redirectURI, scopes, key, cache, cached)
                }
                return registerAndCache(ctx, req, redirectURI, scopes, key, cache)
        })
        if err != nil {
                return nil, err
        }
        return resolutionAny.(*Resolution), nil
}

// revalidateCachedResolution is the cache-hit body of ResolveCredentials
// wrapped by the singleflight. It keeps cached unless the upstream rejects
// the client, in which case it registers a new client and replaces the
// stale entry with Replacer.ReplaceIfUnchanged so concurrent replicas
// converge on a single registration. A failed re-registration is returned
// as an error rather than falling back to cached: the upstream has just said
// that client no longer exists.
//
// A cache that does not implement Replacer could not act on a rejection, so
// the upstream is not contacted at all and cached is returned unchecked.
func revalidateCachedResolution(
        ctx context.Context,
        req *Request,
        redirectURI string,
        scopes []string,
        key Key,
        cache CredentialStore,
        cached *Resolution,
) (*Resolution, error) {
        replacer, ok := cache.(Replacer)
        if !ok {
                return cached, nil
        }

        status, reason := verifyCachedRegistration(ctx, req, cached)
        if status == registrationUnverified {
                logUnverifiedRegistration(ctx, req, key, redirectURI, cached, reason)
        }
        if status != registrationRejected {
                return cached, nil
        }

        //nolint:gosec // G706: client_id is public metadata per RFC 7591.
        slog.Warn("dcr: upstream rejected the cached client registration; registering a new client",
                "local_issuer", req.Issuer,
                "upstream_id", key.UpstreamID,
                "redirect_uri", redirectURI,
                "rotation_key", key.RotationKey,
                "registration_probe", string(req.RegistrationProbe.effective()),
                "client_id", cached.ClientID,
        )
        resolution, err := registerClient(ctx, req, redirectURI, scopes)
        if err != nil {
                return nil, err
        }
        authoritative, err := replacer.ReplaceIfUnchanged(ctx, key, cached, resolution)
        if err != nil {
                return nil, newDCRStepError(dcrStepCacheWrite, req.Issuer, redirectURI,
                        fmt.Errorf("cache replace: %w", err))
        }
        if authoritative.ClientID == cached.ClientID {
                // A concurrent rewrite of the same row won the compare-and-set and
                // left the rejected client in place; using it would only reproduce
                // the upstream's invalid_client.
                return nil, newDCRStepError(dcrStepCacheWrite, req.Issuer, redirectURI,
                        errors.New("cache replace: the rejected registration is still stored"))
        }
        return checkAuthoritativeResolution(req, key, redirectURI, resolution, authoritative)
}

// logUnverifiedRegistration records why a cached registration is used
// without the upstream having confirmed it. The token-endpoint probe skips
// private_key_jwt clients and ProbeNone asks for no probe, so those log at
// Debug on every resolve; any other reason is unexpected and logs at Warn.
func logUnverifiedRegistration(
        ctx context.Context, req *Request, key Key, redirectURI string, cached *Resolution, reason error,
) {
        level := slog.LevelWarn
        attrs := []any{
                "local_issuer", req.Issuer,
                "upstream_id", key.UpstreamID,
                "redirect_uri", redirectURI,
                "rotation_key", key.RotationKey,
                "registration_probe", string(req.RegistrationProbe.effective()),
                "client_id", cached.ClientID,
                "token_endpoint_auth_method", cached.TokenEndpointAuthMethod,
        }
        if errors.Is(reason, errProbeUnsupportedAuthMethod) || errors.Is(reason, errProbeDisabled) {
                level = slog.LevelDebug
        } else {
                attrs = append(attrs, probeErrorAttrs(reason)...)
        }
        //nolint:gosec // G706: client_id is public metadata per RFC 7591.
        slog.Log(ctx, level, "dcr: cached registration not verified; using it as-is", attrs...)
}

// registerAndCache is the cache-miss body of ResolveCredentials wrapped
// by the singleflight. Registration (registerClient) and the durable Put
// live here; the cache was already read inside the same flight.
func registerAndCache(
        ctx context.Context,
        req *Request,
        redirectURI string,
        scopes []string,
        key Key,
        cache CredentialStore,
) (*Resolution, error) {
        resolution, err := registerClient(ctx, req, redirectURI, scopes)
        if err != nil {
                return nil, err
        }

        // Write to durable storage before returning the resolution so a Put
        // failure leaves no in-memory state diverging from the cache: the
        // next call simply re-resolves rather than reading a value the cache
        // never saw.
        //
        // authoritative may differ from resolution: RFC 7591 dynamic
        // registration mints a brand-new, unique client_id/client_secret on
        // every call, so if another replica raced this one to register against
        // the same Key and won the durable claim first, cache.PutIfAbsent
        // returns THAT replica's credentials — the only ones the shared cache
        // (and hence any other replica or a future restart) will agree this
        // caller holds. Returning resolution here instead would leave this
        // process using a client_id the durable store does not recognize.
        authoritative, err := cache.PutIfAbsent(ctx, key, resolution)
        if err != nil {
                return nil, newDCRStepError(dcrStepCacheWrite, req.Issuer, redirectURI,
                        fmt.Errorf("cache put: %w", err))
        }
        return checkAuthoritativeResolution(req, key, redirectURI, resolution, authoritative)
}

// registerClient resolves the upstream endpoints, selects the token endpoint
// auth method, and performs one RFC 7591 registration. It does not touch the
// cache; callers persist the returned resolution.
func registerClient(
        ctx context.Context,
        req *Request,
        redirectURI string,
        scopes []string,
) (*Resolution, error) {
        // Endpoint resolution: discover metadata when configured, otherwise use
        // the caller-supplied RegistrationEndpoint directly. The upstream's
        // expected issuer is recovered from req.DiscoveryURL inside the helper.
        // req.Issuer here is *this caller's* issuer — correct for cache keying
        // and redirect URI defaulting, but it must not be used for RFC 8414
        // §3.3 metadata verification (which is the upstream's concern).
        endpoints, err := resolveDCREndpoints(ctx, req)
        if err != nil {
                return nil, newDCRStepError(dcrStepMetadata, req.Issuer, redirectURI, err)
        }
        if err := applyExplicitEndpointOverrides(endpoints, req); err != nil {
                return nil, newDCRStepError(dcrStepMetadata, req.Issuer, redirectURI, err)
        }

        // Token-endpoint auth method: intersect server support with our
        // preference order; default to client_secret_basic if the server does
        // not advertise the field at all. When the caller declared the
        // registration is for a public PKCE client (CLI flow), force
        // token_endpoint_auth_method=none while still enforcing the S256
        // PKCE gate.
        authMethod, err := selectTokenEndpointAuthMethod(
                endpoints.tokenEndpointAuthMethodsSupported,
                endpoints.codeChallengeMethodsSupported,
                req.PublicClient,
        )
        if err != nil {
                return nil, newDCRStepError(dcrStepSelectAuthMethod, req.Issuer, redirectURI, err)
        }

        registrationScopes := chooseRegistrationScopes(scopes, endpoints.scopesSupported, req.Issuer)

        response, err := performRegistration(ctx, req, endpoints.registrationEndpoint,
                redirectURI, authMethod, registrationScopes, endpoints.registrationEndpointServerSupplied,
                req.CAFilePath)
        if err != nil {
                return nil, newDCRStepError(dcrStepRegister, req.Issuer, redirectURI, err)
        }

        return buildResolution(response, endpoints, authMethod, redirectURI), nil
}

// checkAuthoritativeResolution vets the durable value a cache write returned
// for a freshly registered resolution and logs the outcome.
func checkAuthoritativeResolution(
        req *Request,
        key Key,
        redirectURI string,
        resolution, authoritative *Resolution,
) (*Resolution, error) {
        // The authoritative row can be a concurrent claimant's stable-but-expired
        // registration: when both the existing stored row and this replica's
        // fresh registration are already expired, dcrClaimOrReturnWinner
        // (pkg/authserver/storage/redis.go) deliberately returns the existing
        // row without error rather than retrying forever. That's the right call
        // at the storage layer, but "expired means unusable" is DCR policy, not
        // storage policy, so it belongs here: treat an already-expired
        // authoritative credential as a resolution failure instead of handing
        // callers a dead client_secret disguised as a success.
        if !authoritative.ClientSecretExpiresAt.IsZero() && time.Now().After(authoritative.ClientSecretExpiresAt) {
                return nil, newDCRStepError(dcrStepExpiredCredential, req.Issuer, redirectURI,
                        fmt.Errorf("authoritative credential is already expired (client_secret_expires_at=%s)",
                                authoritative.ClientSecretExpiresAt.UTC().Format(time.RFC3339)))
        }

        if authoritative.ClientID != resolution.ClientID {
                //nolint:gosec // G706: client_id is public metadata per RFC 7591.
                slog.Debug("dcr: registration superseded by concurrent winner",
                        "local_issuer", req.Issuer,
                        "upstream_id", key.UpstreamID,
                        "redirect_uri", redirectURI,
                        "registered_client_id", resolution.ClientID,
                        "authoritative_client_id", authoritative.ClientID,
                )
        }

        //nolint:gosec // G706: client_id is public metadata per RFC 7591.
        slog.Debug("dcr: registered new client",
                "local_issuer", req.Issuer,
                "upstream_id", key.UpstreamID,
                "redirect_uri", redirectURI,
                "rotation_key", key.RotationKey,
                "client_id", authoritative.ClientID,
        )
        return authoritative, nil
}

// LogStepError emits the single boundary slog.Error record for a DCR
// resolver failure, carrying the step / issuer / redirect_uri attributes
// extracted from err. If err is not a *dcrStepError, it is logged with a
// generic "unknown" step — ResolveCredentials always wraps its errors,
// so this branch indicates a programming error in a future caller rather
// than a runtime condition. err == nil is a no-op so this function is
// safe to call without an outer guard.
//
// Every wrapped error is passed through SanitizeErrorForLog to strip URL
// query parameters that could plausibly contain sensitive tokens (defense
// in depth — the current DCR flow sends the initial access token as an
// Authorization header, not a query parameter, but nothing in the type
// system prevents a future refactor from doing otherwise).
func LogStepError(upstreamName string, err error) {
        if err == nil {
                return
        }
        var stepErr *dcrStepError
        if !errors.As(err, &stepErr) {
                slog.Error("dcr: resolve failed",
                        "upstream", upstreamName,
                        "step", "unknown",
                        "error", SanitizeErrorForLog(err),
                )
                return
        }

        attrs := []any{
                "upstream", upstreamName,
                "step", stepErr.Step,
                "issuer", stepErr.Issuer,
                "error", SanitizeErrorForLog(stepErr.Err),
        }
        if stepErr.RedirectURI != "" {
                attrs = append(attrs, "redirect_uri", stepErr.RedirectURI)
        }
        if stepErr.Stack != "" {
                attrs = append(attrs, "stack", stepErr.Stack)
        }
        slog.Error("dcr: resolve failed", attrs...)
}

// SanitizeErrorForLog strips secret-bearing components from any URLs
// embedded in err's message. The Go HTTP client, url.Error, and other
// net/* wrappers embed the full request URL — including userinfo,
// query, and fragment — in their error strings. Any of those can carry
// credentials or tokens (e.g. https://user:pass@host, ?token=…,
// implicit-flow callbacks #access_token=…); the current DCR flow does
// not embed any of them today, but stripping them here is defense in
// depth that protects the log regardless of future changes.
//
// Scheme, host, and path are preserved so operators retain enough
// context to correlate with upstream server logs. Trailing sentence
// punctuation adjacent to the URL (e.g. the comma in "reaching
// URL?q=1, retrying") is preserved — see trimURLTrailingPunctuation
// for the list of characters considered terminators.
//
// Scope: the regex matches http://, https://, redis://, and rediss://
// schemes (case-insensitively per RFC 3986 §3.1). The redis schemes are
// included because the embedded-authserver DCR path persists through
// pkg/authserver/storage/redis.go and a redis-go error chain on the
// Get/Put critical path can embed a sentinel/cluster URL with
// credentials. Other schemes (file://, postgres://, smtp://, raw
// host:port) are NOT sanitised; if a future call site flows error
// strings from one of those into LogStepError, extend the regex here
// rather than re-implementing the sanitiser at the call site.
//
// IMPORTANT — caller responsibility: this function strips credentials
// only from the schemes listed above. Callers that may receive errors
// containing other credential-bearing URL schemes MUST verify those
// URLs are not credential-bearing before logging, or sanitise them
// separately.
func SanitizeErrorForLog(err error) string {
        if err == nil {
                return ""
        }
        msg := err.Error()
        return queryStrippingPattern.ReplaceAllStringFunc(msg, func(match string) string {
                // Split any trailing sentence punctuation off the match before
                // handing it to url.Parse. Without this, a period / comma /
                // closing bracket at the end of the sentence is absorbed into
                // the URL's raw query and dropped along with the rest of the
                // query component, mangling the error text. The trimmed
                // punctuation is re-appended to the replacement so the
                // surrounding prose is preserved verbatim.
                core, tail := trimURLTrailingPunctuation(match)
                u, parseErr := url.Parse(core)
                if parseErr != nil {
                        return match
                }
                if u.User == nil && u.RawQuery == "" && u.Fragment == "" {
                        return match
                }
                u.User = nil
                u.RawQuery = ""
                u.Fragment = ""
                return u.String() + tail
        })
}

// trimURLTrailingPunctuation returns (core, tail) where tail is the run of
// trailing ASCII punctuation that commonly terminates a URL inside prose
// but is never a meaningful part of the URL itself. The characters chosen
// here mirror those used by general-purpose URL extractors (e.g.,
// Chromium's autolinker): sentence-ending punctuation, closing brackets,
// and a few separators that appear between URLs in freeform text.
//
// Note that ')' and ']' are stripped unconditionally — a URL legitimately
// containing a percent-encoded closing bracket will have it as "%29" or
// "%5D", not as a literal, so this cannot truncate a real URL path or
// query. The reverse case (an unescaped ')' inside a path) is
// non-conforming per RFC 3986 and out of scope for a log sanitiser.
func trimURLTrailingPunctuation(s string) (core, tail string) {
        // terminators is intentionally ASCII-only; Unicode terminators (e.g.
        // '」') are out of scope for this log sanitiser, so byte indexing is
        // safe and avoids the rune-decoding overhead of strings.ContainsRune.
        const terminators = ".,;:!?)]}>"
        i := len(s)
        for i > 0 && strings.IndexByte(terminators, s[i-1]) >= 0 {
                i--
        }
        return s[:i], s[i:]
}

// queryStrippingPattern matches URL-shaped substrings inside an error
// message — sufficient to reach the url.Parse path in SanitizeErrorForLog
// and let it decide whether a secret-bearing component exists to strip.
// The regexp covers the schemes that actually flow through the DCR
// resolver's error paths today (http/https from upstream calls, redis/
// rediss from the persistent CredentialStore on the embedded-authserver
// path). It matches schemes case-insensitively per RFC 3986 §3.1 since
// upstream metadata or user input can carry mixed-case schemes.
// Trailing sentence punctuation that the character class happens to
// include (e.g. '.', ',', ')') is stripped by
// trimURLTrailingPunctuation before the match is parsed.
var queryStrippingPattern = regexp.MustCompile(`(?i)(?:https?|rediss?)://[^\s"']+`)

// -----------------------------------------------------------------------------
// Private helpers
// -----------------------------------------------------------------------------

// validateResolveInputs performs the defensive re-check of resolver
// preconditions. Each consumer is expected to validate its own domain
// types upstream of the resolver; this is the final-line check for an
// entry point that programmatic callers can reach with partially-
// constructed Requests.
func validateResolveInputs(
        req *Request,
        cache CredentialStore,
) error {
        if req == nil {
                return fmt.Errorf("dcr: request is required")
        }
        if req.Issuer == "" {
                return fmt.Errorf("dcr: issuer is required")
        }
        if req.DiscoveryURL == "" && req.RegistrationEndpoint == "" {
                return fmt.Errorf("dcr: request must set either discovery_url or registration_endpoint")
        }
        if req.DiscoveryURL != "" && req.RegistrationEndpoint != "" {
                return fmt.Errorf("dcr: discovery_url and registration_endpoint are mutually exclusive")
        }
        if err := req.RegistrationProbe.Validate(); err != nil {
                return fmt.Errorf("dcr: %w", err)
        }
        if cache == nil {
                return fmt.Errorf("dcr: credential store is required")
        }
        return nil
}

// lookupCachedResolution checks the cache and logs the hit. On hit it
// returns (resolution, true, nil). On miss it returns (nil, false, nil). An
// error is returned only on backend failure.
//
// Two distinct staleness signals shape the hit/miss decision and the log:
//
//   - Hard expiry (RFC 7591 §3.2.1 client_secret_expires_at): when the
//     cached resolution's ClientSecretExpiresAt is non-zero and in the
//     past, the entry is treated as a miss so the singleflight body
//     (registerAndCache) re-runs the registration and overwrites the stale
//     entry via cache.Put. Without this check the cache would serve an
//     expired secret indefinitely; the upstream's token endpoint would 401
//     on every use and the resolver would have no signal to refetch. The
//     check is skipped when the field is zero, per the RFC 7591 convention
//     "0 means the secret does not expire". This is the authoritative
//     signal — the upstream said when its credential expires.
//   - Soft staleness (now - CreatedAt vs dcrStaleAgeThreshold): the age in
//     days is logged on every hit, and if it exceeds the threshold an
//     additional slog.Warn points operators at Request.RotationKey, the way
//     to rotate a long-lived registration. This is observability only, not a
//     cache-invalidation trigger.
func lookupCachedResolution(
        ctx context.Context,
        cache CredentialStore,
        key Key,
        localIssuer, redirectURI string,
) (*Resolution, bool, error) {
        cached, ok, err := cache.Get(ctx, key)
        if err != nil {
                return nil, false, fmt.Errorf("dcr: cache lookup: %w", err)
        }
        if !ok {
                return nil, false, nil
        }
        if !cached.ClientSecretExpiresAt.IsZero() && time.Now().After(cached.ClientSecretExpiresAt) {
                //nolint:gosec // G706: client_id is public metadata per RFC 7591.
                slog.Debug("dcr: cache hit ignored; cached secret expired per upstream client_secret_expires_at",
                        "local_issuer", localIssuer,
                        "upstream_id", key.UpstreamID,
                        "redirect_uri", redirectURI,
                        "client_id", cached.ClientID,
                        "client_secret_expires_at", cached.ClientSecretExpiresAt.UTC().Format(time.RFC3339),
                )
                return nil, false, nil
        }

        age := time.Since(cached.CreatedAt)
        ageDays := int(age / (24 * time.Hour))

        //nolint:gosec // G706: client_id is public metadata per RFC 7591.
        slog.Debug("dcr: cache hit",
                "local_issuer", localIssuer,
                "upstream_id", key.UpstreamID,
                "redirect_uri", redirectURI,
                "rotation_key", key.RotationKey,
                "client_id", cached.ClientID,
                "dcr_age_days", ageDays,
        )

        if age > dcrStaleAgeThreshold {
                //nolint:gosec // G706: client_id is public metadata per RFC 7591.
                slog.Warn(
                        "dcr: cached registration exceeds staleness threshold; "+
                                "to rotate it, change dcrConfig.rotationKey and restart",
                        "local_issuer", localIssuer,
                        "upstream_id", key.UpstreamID,
                        "redirect_uri", redirectURI,
                        "rotation_key", key.RotationKey,
                        "client_id", cached.ClientID,
                        "dcr_age_days", ageDays,
                        "stale_threshold_days", int(dcrStaleAgeThreshold/(24*time.Hour)),
                )
        }

        return cached, true, nil
}

// applyExplicitEndpointOverrides overwrites the discovered
// authorizationEndpoint / tokenEndpoint in endpoints with explicit values
// from req when req specifies them. Explicit caller configuration always
// wins over discovery.
//
// Both overrides run through validateUpstreamEndpointURL even though the
// known consumer call sites only ever populate these fields from
// already-validated sources. validateUpstreamEndpointURL is the documented
// gatekeeper "every point where an endpoint URL enters the resolver from
// outside" — applying it here closes the override path's gap rather than
// relying on the caller-side validation to never drift.
func applyExplicitEndpointOverrides(endpoints *dcrEndpoints, req *Request) error {
        if req.AuthorizationEndpoint != "" {
                if err := validateUpstreamEndpointURL(req.AuthorizationEndpoint, "authorization_endpoint"); err != nil {
                        return fmt.Errorf("explicit %w", err)
                }
                endpoints.authorizationEndpoint = req.AuthorizationEndpoint
        }
        if req.TokenEndpoint != "" {
                if err := validateUpstreamEndpointURL(req.TokenEndpoint, "token_endpoint"); err != nil {
                        return fmt.Errorf("explicit %w", err)
                }
                endpoints.tokenEndpoint = req.TokenEndpoint
        }
        return nil
}

// chooseRegistrationScopes selects the scopes to send in the registration
// request: explicit caller scopes > discovered scopes_supported > empty.
// Logs a warning when neither source produces any scopes.
func chooseRegistrationScopes(explicit, discovered []string, localIssuer string) []string {
        if len(explicit) > 0 {
                return explicit
        }
        if len(discovered) > 0 {
                return discovered
        }
        slog.Warn("dcr: no scopes configured or discovered; registering with empty scope",
                "local_issuer", localIssuer,
        )
        return nil
}

// performRegistration executes the HTTP registration request exactly once.
// The initial access token (if configured) is injected as a
// bearer-token Authorization header via a wrapping http.Client.
//
// The caller is responsible for resolving any file-or-env initial access
// token reference into req.InitialAccessToken before calling
// ResolveCredentials; the resolver does not touch the filesystem or
// environment.
func performRegistration(
        ctx context.Context,
        req *Request,
        registrationEndpoint, redirectURI, authMethod string,
        scopes []string,
        registrationEndpointServerSupplied bool,
        caFilePath string,
) (*oauthproto.DynamicClientRegistrationResponse, error) {
        httpClient, err := newDCRHTTPClient(
                req.InitialAccessToken,
                registrationEndpoint,
                req.AllowPrivateIPs,
                req.ServerSuppliedEndpoints || registrationEndpointServerSupplied,
                caFilePath,
        )
        if err != nil {
                return nil, fmt.Errorf("dcr: build registration http client: %w", err)
        }

        clientName := req.ClientName
        if clientName == "" {
                clientName = oauthproto.ToolHiveMCPClientName
        }

        registrationRequest := &oauthproto.DynamicClientRegistrationRequest{
                RedirectURIs:            []string{redirectURI},
                ClientName:              clientName,
                TokenEndpointAuthMethod: authMethod,
                GrantTypes:              []string{oauthproto.GrantTypeAuthorizationCode, oauthproto.GrantTypeRefreshToken},
                ResponseTypes:           []string{oauthproto.ResponseTypeCode},
                Scopes:                  scopes,
        }

        // Call exactly once — no retry loop. If retry/backoff against
        // transient upstream failures becomes useful it belongs at a higher
        // layer (above ResolveCredentials), so the singleflight and cache
        // semantics here stay simple: one attempt per resolution per process.
        response, err := oauthproto.RegisterClientDynamically(ctx, registrationEndpoint, registrationRequest, httpClient)
        if err != nil {
                return nil, fmt.Errorf("dcr: register client: %w", err)
        }
        return response, nil
}

// buildResolution assembles the Resolution from the RFC 7591 response and
// the resolved endpoints. If the server did not echo a
// token_endpoint_auth_method in the response, the method actually sent is
// recorded so downstream consumers see a definite value. redirectURI is the
// value passed to the registration endpoint (caller-supplied or defaulted
// via resolveUpstreamRedirectURI); it is persisted on the resolution so
// consumers (pkg/authserver/runner.consumeResolution) can propagate a
// defaulted value back to the run-config.
//
// RFC 7591 §3.2.1 client_id_issued_at and client_secret_expires_at are
// converted from int64 epoch seconds to time.Time. The wire value 0 means
// "field absent" or "secret does not expire"; both map to the zero time.Time
// so callers can use IsZero() uniformly.
func buildResolution(
        response *oauthproto.DynamicClientRegistrationResponse,
        endpoints *dcrEndpoints,
        sentAuthMethod string,
        redirectURI string,
) *Resolution {
        authMethod := response.TokenEndpointAuthMethod
        if authMethod == "" {
                authMethod = sentAuthMethod
        }
        return &Resolution{
                ClientID:                response.ClientID,
                ClientSecret:            response.ClientSecret,
                AuthorizationEndpoint:   endpoints.authorizationEndpoint,
                TokenEndpoint:           endpoints.tokenEndpoint,
                RegistrationAccessToken: response.RegistrationAccessToken,
                RegistrationClientURI:   response.RegistrationClientURI,
                TokenEndpointAuthMethod: authMethod,
                RedirectURI:             redirectURI,
                ClientIDIssuedAt:        epochSecondsToTime(response.ClientIDIssuedAt),
                ClientSecretExpiresAt:   epochSecondsToTime(derefInt64(response.ClientSecretExpiresAt)),
                CreatedAt:               time.Now(),
        }
}

// derefInt64 unwraps the pointer form of client_secret_expires_at; a nil
// pointer (key absent) and a pointed-to 0 both mean "does not expire" and
// collapse to 0, which epochSecondsToTime maps to the zero time.Time.
func derefInt64(p *int64) int64 {
        if p == nil {
                return 0
        }
        return *p
}

// epochSecondsToTime converts the int64 epoch-seconds form used by RFC 7591
// into a time.Time. Zero passes through to the zero time.Time so callers can
// rely on IsZero() to mean "field absent" / "does not expire".
func epochSecondsToTime(epoch int64) time.Time {
        if epoch == 0 {
                return time.Time{}
        }
        return time.Unix(epoch, 0).UTC()
}

// dcrEndpoints is the internal bundle of endpoints produced by endpoint
// resolution. The separation from Resolution lets the resolver reason
// about discovered vs. overridden values before committing to a resolution.
type dcrEndpoints struct {
        authorizationEndpoint             string
        tokenEndpoint                     string
        registrationEndpoint              string
        tokenEndpointAuthMethodsSupported []string
        scopesSupported                   []string
        // codeChallengeMethodsSupported is consumed by
        // selectTokenEndpointAuthMethod to gate the public-client (none) auth
        // method on S256 PKCE being advertised. RFC 7636 / OAuth 2.1 require
        // PKCE-with-S256 for public clients; registering as none against an
        // upstream that advertises only plain (or omits the field) would be a
        // compliance gap.
        codeChallengeMethodsSupported []string
        // registrationEndpointServerSupplied records that registrationEndpoint came
        // out of a fetched metadata document AND names an authority the operator
        // never configured. Even an operator-configured DiscoveryURL yields a
        // document the upstream controls, so an endpoint it points at some other
        // host must use the strict server-supplied client policy — the registration
        // POST carries the initial access token (CWE-918). An upstream naming its
        // own authority is still covered by the operator's decision to configure it,
        // and keeps the host-scoped policy so an in-cluster or loopback IdP that
        // worked before still works.
        registrationEndpointServerSupplied bool
}

// resolveDCREndpoints produces the endpoint bundle from req.
//
// Two branches, in priority order:
//
//  1. req.RegistrationEndpoint set — use it directly and skip discovery
//     entirely. Server-capability fields (token_endpoint_auth_methods_supported,
//     scopes_supported) are unavailable on this branch; the caller is
//     expected to also supply AuthorizationEndpoint, TokenEndpoint, and an
//     explicit Scopes list. Auth method falls back to the
//     selectTokenEndpointAuthMethod default.
//  2. req.DiscoveryURL set — fetch the exact document the operator
//     configured (bypassing the well-known path fallback). RFC 8414 §3.3
//     requires the metadata's "issuer" field to match the authorization
//     server's issuer identifier; that identifier is the upstream's, not
//     this caller's, so it is recovered from the discovery URL via
//     deriveExpectedIssuerFromDiscoveryURL rather than reusing
//     req.Issuer (which names this caller and is used elsewhere in
//     ResolveCredentials for redirect URI defaulting and cache keying).
//
// validateResolveInputs enforces that exactly one of the two is set, so
// this function does not re-check the both-empty case.
//
// When metadata is returned but omits registration_endpoint, the resolver
// synthesises {origin}/register — a convention used by nanobot and Hydra
// for providers that ship DCR without advertising it in discovery. Origin
// is taken from the upstream issuer, not this caller's issuer, so the
// synthesised endpoint lands at the upstream.
func resolveDCREndpoints(
        ctx context.Context,
        req *Request,
) (*dcrEndpoints, error) {
        if req.RegistrationEndpoint != "" {
                // Validate locally so a non-HTTPS or malformed URL fails before
                // performRegistration constructs a bearer-token transport for it.
                if err := validateUpstreamEndpointURL(req.RegistrationEndpoint, "registration_endpoint"); err != nil {
                        return nil, fmt.Errorf("dcr: %w", err)
                }
                return &dcrEndpoints{
                        registrationEndpoint:          req.RegistrationEndpoint,
                        codeChallengeMethodsSupported: req.CodeChallengeMethodsSupported,
                }, nil
        }

        upstreamIssuer, err := deriveExpectedIssuerFromDiscoveryURL(req.DiscoveryURL)
        if err != nil {
                return nil, err
        }

        // Dial the discovery fetch through the private-IP-guarded client so a
        // DiscoveryURL that resolves to a private/loopback/link-local address —
        // or rebinds to one after the caller validated it — is refused at connect
        // time (CWE-918). Guarded by default; req.AllowPrivateIPs opts in to
        // private ranges for an in-cluster upstream. Server-supplied discovery
        // endpoints use the strict policy, where localhost and the environment
        // validation override cannot bypass this gate.
        discoveryHost, err := hostFromURL(req.DiscoveryURL)
        if err != nil {
                return nil, err
        }
        discoveryClient, err := newGuardedDCRClient(
                discoveryHost,
                req.AllowPrivateIPs,
                req.ServerSuppliedEndpoints,
                req.CAFilePath,
        )
        if err != nil {
                return nil, fmt.Errorf("dcr: build discovery http client: %w", err)
        }
        // The discovery URL is operator-configured, but the document it returns —
        // and any 30x the upstream serves in response — is upstream-controlled.
        // Restrict the fetch to same-host redirects so a malicious redirect cannot
        // walk the request onto another host (CWE-918); this matters even when the
        // dial guard is relaxed for a loopback discovery host. The registration
        // client refuses redirects outright because it carries the bearer token;
        // the discovery GET carries no secret, so same-host redirects are allowed —
        // matching the CLI discovery clients in pkg/auth/discovery.
        discoveryClient.CheckRedirect = networking.SameHostRedirectPolicy()

        metadata, err := oauthproto.FetchAuthorizationServerMetadataFromURL(ctx, req.DiscoveryURL, upstreamIssuer, discoveryClient)
        return endpointsFromMetadata(metadata, err, upstreamIssuer)
}

// resolveUpstreamKeyIdentity returns the stable identifier for the upstream
// authorization server a registration is bound to, used as the
// Key.UpstreamID cache component. It disambiguates upstreams that share the
// caller's Issuer, RedirectURI, and scope set — the common case inside a
// single embedded authserver — so they receive distinct dynamically-
// registered clients instead of colliding on one cache entry (issue #5823).
//
// The identity source mirrors resolveDCREndpoints' own branch order so the
// key names the exact server the client registers against:
//
//   - RegistrationEndpoint set: the registration endpoint URL itself.
//   - DiscoveryURL set: the upstream issuer recovered from the discovery URL
//     via deriveExpectedIssuerFromDiscoveryURL — the same value used for
//     RFC 8414 §3.3 metadata verification.
//
// validateResolveInputs guarantees exactly one of the two is set, so the
// final return is only reached on the DiscoveryURL branch. The derivation is
// pure URL parsing (no network I/O), so computing it before the cache lookup
// keeps the cache-hit path free of I/O; a malformed DiscoveryURL that fails
// here could never have produced a stored entry to hit anyway, because a
// successful registration requires the same derivation to succeed first.
//
// Residual limitation: on the DiscoveryURL branch the identity is the derived
// upstream *issuer*, not the discovery URL. For a custom (non-well-known)
// discovery URL, deriveExpectedIssuerFromDiscoveryURL falls back to the bare
// origin and discards the path, so two upstreams on the same host with
// distinct custom metadata paths, the same scopes, and both declaring that
// origin as their metadata "issuer" derive the same UpstreamID and still
// collide. This is intentional: UpstreamID names the authorization server
// (aligned with the RFC 8414 §3.3 issuer used for metadata verification), so
// two configs that resolve to the same issuer are the same AS and correctly
// share one registration. The collision is therefore confined to the
// nonstandard case where a single issuer serves distinct registration
// endpoints under custom (non-well-known) discovery paths: both upstreams
// resolve to that one issuer, so their UpstreamIDs match. Upstreams that
// resolve to different issuers derive different UpstreamIDs and never collide
// — the key itself keeps them apart, independent of §3.3 verification.
func resolveUpstreamKeyIdentity(req *Request) (string, error) {
        if req.RegistrationEndpoint != "" {
                return req.RegistrationEndpoint, nil
        }
        return deriveExpectedIssuerFromDiscoveryURL(req.DiscoveryURL)
}

// deriveExpectedIssuerFromDiscoveryURL recovers the issuer identifier the
// upstream is expected to claim in its RFC 8414 / OIDC Discovery document,
// given an operator-configured DiscoveryURL.
//
// Three recognised conventions:
//
//  1. Well-known suffix: the URL ends with /.well-known/oauth-authorization-server
//     or /.well-known/openid-configuration. The suffix is stripped to recover
//     the issuer; this covers single-tenant providers (e.g.
//     https://mcp.atlassian.com/.well-known/oauth-authorization-server →
//     https://mcp.atlassian.com) and the OIDC-style suffix-append shape used
//     by some multi-tenant providers (e.g.
//     https://idp.example.com/tenants/acme/.well-known/openid-configuration
//     → https://idp.example.com/tenants/acme).
//  2. RFC 8414 §3.1 path-insertion: the well-known path is inserted BETWEEN
//     the host and the issuer's tenant path. Per RFC 8414 §3 / RFC 8615,
//     this is the canonical form for issuers with a path component, e.g.
//     issuer https://example.com/v1/mcp →
//     discovery URL https://example.com/.well-known/oauth-authorization-server/v1/mcp.
//     The tenant suffix that appears AFTER the well-known segment is
//     re-attached to the origin to recover the issuer.
//  3. Non-well-known path: the URL points at a custom metadata endpoint that
//     contains neither suffix in a recognisable position. Origin
//     (scheme://host) is used as a best-effort fallback; this matches the
//     common shape where the upstream issuer is the host root.
//
// Case (1) and case (2) are disambiguated by where the well-known segment
// sits in the path: at the end ⇒ suffix-append, immediately after the host
// with more path following ⇒ path-insertion.
func deriveExpectedIssuerFromDiscoveryURL(discoveryURL string) (string, error) {
        const (
                oauthSuffix = "/.well-known/oauth-authorization-server"
                oidcSuffix  = "/.well-known/openid-configuration"
        )

        u, err := url.Parse(discoveryURL)
        if err != nil {
                return "", fmt.Errorf("parse discovery url %q: %w", discoveryURL, err)
        }
        if u.Scheme == "" || u.Host == "" {
                return "", fmt.Errorf("discovery url missing scheme or host: %q", discoveryURL)
        }

        switch {
        // Suffix-append form (case 1): well-known segment at end of path.
        case strings.HasSuffix(u.Path, oauthSuffix):
                u.Path = strings.TrimSuffix(u.Path, oauthSuffix)
        case strings.HasSuffix(u.Path, oidcSuffix):
                u.Path = strings.TrimSuffix(u.Path, oidcSuffix)
        // RFC 8414 §3.1 path-insertion form (case 2): well-known segment at the
        // start of the path with tenant path following. Strip just the well-known
        // segment to recover {origin}{tenant-path}.
        //
        // Two shapes hit this branch:
        //   1. A real tenant suffix follows the well-known segment, e.g.
        //      /.well-known/oauth-authorization-server/v1/mcp →
        //      issuer https://host/v1/mcp.
        //   2. A trailing slash with no tenant, e.g.
        //      /.well-known/oauth-authorization-server/ — TrimPrefix leaves
        //      a stray "/", which would yield a spurious issuer
        //      "https://host/" that fails the §3.3 byte-equality check
        //      against the upstream's declared "https://host". Normalise
        //      that stray "/" back to empty so case (2.2) and the bare
        //      suffix case derive the same origin issuer.
        case strings.HasPrefix(u.Path, oauthSuffix+"/"):
                u.Path = strings.TrimPrefix(u.Path, oauthSuffix)
                if u.Path == "/" {
                        u.Path = ""
                }
        case strings.HasPrefix(u.Path, oidcSuffix+"/"):
                u.Path = strings.TrimPrefix(u.Path, oidcSuffix)
                if u.Path == "/" {
                        u.Path = ""
                }
        default:
                // Custom (non-well-known) discovery URL — fall back to origin.
                u.Path = ""
        }
        u.RawQuery = ""
        u.Fragment = ""
        return u.String(), nil
}

// endpointsFromMetadata converts a FetchAuthorizationServerMetadata* result
// into a dcrEndpoints bundle. Handles the ErrRegistrationEndpointMissing
// sentinel by synthesising {origin}/register.
//
// authorization_endpoint and token_endpoint are validated for HTTPS / well-
// formedness before being copied into the bundle. A self-consistent metadata
// document — possible if TLS to the metadata host is compromised, or if the
// upstream is misconfigured — could otherwise plant http:// URLs that flow
// through to the authorization-code and token-exchange call paths.
//
// Contract with oauthproto: FetchAuthorizationServerMetadata* guarantees a
// non-nil *AuthorizationServerMetadata whenever fetchErr is nil OR fetchErr
// is ErrRegistrationEndpointMissing (in the latter case the metadata is
// otherwise valid; only registration_endpoint is missing). The defensive
// nil guard below catches a future cross-package contract regression — e.g.,
// a new oauthproto sentinel that returns nil metadata alongside a non-fatal
// error — and converts it into a clean validation error rather than a
// nil-pointer dereference at the field accesses.
func endpointsFromMetadata(
        metadata *oauthproto.AuthorizationServerMetadata,
        fetchErr error,
        upstreamIssuer string,
) (*dcrEndpoints, error) {
        if fetchErr != nil && !errors.Is(fetchErr, oauthproto.ErrRegistrationEndpointMissing) {
                return nil, fmt.Errorf("discover authorization server metadata: %w", fetchErr)
        }
        if metadata == nil {
                return nil, fmt.Errorf(
                        "dcr: authorization server metadata is nil (oauthproto contract " +
                                "violation: nil metadata returned alongside a non-fatal fetch error)")
        }

        if err := validateUpstreamEndpointURL(metadata.AuthorizationEndpoint, "authorization_endpoint"); err != nil {
                return nil, fmt.Errorf("dcr: discovered %w", err)
        }
        if err := validateUpstreamEndpointURL(metadata.TokenEndpoint, "token_endpoint"); err != nil {
                return nil, fmt.Errorf("dcr: discovered %w", err)
        }

        registrationEndpoint := metadata.RegistrationEndpoint
        if errors.Is(fetchErr, oauthproto.ErrRegistrationEndpointMissing) {
                // Metadata is otherwise valid — synthesise the registration
                // endpoint from the upstream issuer's origin.
                // FetchAuthorizationServerMetadata* deliberately returns
                // ErrRegistrationEndpointMissing alongside a non-nil metadata
                // document, so we still use the returned endpoints/scopes.
                synth, err := synthesiseRegistrationEndpoint(upstreamIssuer)
                if err != nil {
                        return nil, fmt.Errorf("synthesise registration endpoint: %w", err)
                }
                registrationEndpoint = synth
        } else if err := validateUpstreamEndpointURL(registrationEndpoint, "registration_endpoint"); err != nil {
                // Unlike the synthesised branch above, this value came straight from
                // the discovery document — validate it the same as
                // authorization_endpoint/token_endpoint rather than letting it reach
                // hostFromURL unvalidated.
                return nil, fmt.Errorf("dcr: discovered %w", err)
        }

        return &dcrEndpoints{
                authorizationEndpoint:             metadata.AuthorizationEndpoint,
                tokenEndpoint:                     metadata.TokenEndpoint,
                registrationEndpoint:              registrationEndpoint,
                tokenEndpointAuthMethodsSupported: metadata.TokenEndpointAuthMethodsSupported,
                scopesSupported:                   metadata.ScopesSupported,
                codeChallengeMethodsSupported:     metadata.CodeChallengeMethodsSupported,
                // The document is upstream-controlled, so a registration_endpoint on any
                // authority the operator did not name is server-supplied input.
                // upstreamIssuer is derived from the operator's DiscoveryURL (not from the
                // document), so an upstream naming its own authority — or the synthesised
                // {issuer}/register above — stays within the operator's decision to
                // configure it, while an endpoint pointing anywhere else does not.
                registrationEndpointServerSupplied: !networking.AuthorityMatchesAny(registrationEndpoint, upstreamIssuer),
        }, nil
}

// synthesiseRegistrationEndpoint builds {upstreamIssuer}/register, used when
// discovery succeeds but omits registration_endpoint. The argument is the
// upstream's issuer (recovered from the discovery URL), not this auth
// server's local issuer.
//
// The issuer's path is preserved so multi-tenant upstreams that ship DCR
// without advertising it (e.g. https://idp.example.com/tenants/acme) keep
// their tenant prefix in the synthesised URL. Stripping the path would land
// the registration request at a global /register that does not match the
// tenant-aware token/authorize URLs already accepted from metadata.
func synthesiseRegistrationEndpoint(upstreamIssuer string) (string, error) {
        u, err := url.Parse(upstreamIssuer)
        if err != nil {
                return "", fmt.Errorf("parse issuer: %w", err)
        }
        if u.Scheme == "" || u.Host == "" {
                return "", fmt.Errorf("issuer missing scheme or host: %q", upstreamIssuer)
        }
        synth := &url.URL{
                Scheme: u.Scheme,
                Host:   u.Host,
                Path:   strings.TrimRight(u.Path, "/") + "/register",
        }
        return synth.String(), nil
}

// resolveUpstreamRedirectURI returns the redirect URI to present to the
// upstream. The caller-supplied value wins; otherwise a default is derived
// from {localIssuer}/oauth/callback. HTTPS is required except for loopback
// hosts (development).
//
// localIssuer here is *this* auth server's issuer — the redirect URI is
// where the upstream sends the user back to us, so it must live on our
// origin, not the upstream's.
//
// The issuer's path is preserved when defaulting: an issuer with a tenant
// prefix produces a redirect URI under that prefix, not at the host root.
// url.URL.ResolveReference would replace the path entirely because
// defaultUpstreamRedirectPath starts with "/", so we explicitly concatenate
// instead.
func resolveUpstreamRedirectURI(configured, localIssuer string) (string, error) {
        if configured != "" {
                u, err := url.Parse(configured)
                if err != nil {
                        return "", fmt.Errorf("invalid redirect uri %q: %w", configured, err)
                }
                if err := validateRedirectURL(u); err != nil {
                        return "", err
                }
                return configured, nil
        }

        issuerURL, err := url.Parse(localIssuer)
        if err != nil {
                return "", fmt.Errorf("invalid issuer %q: %w", localIssuer, err)
        }
        resolved := &url.URL{
                Scheme: issuerURL.Scheme,
                Host:   issuerURL.Host,
                Path:   strings.TrimRight(issuerURL.Path, "/") + defaultUpstreamRedirectPath,
        }
        if err := validateRedirectURL(resolved); err != nil {
                return "", err
        }
        return resolved.String(), nil
}

// validateRedirectURL enforces the HTTPS-except-loopback rule shared across
// OAuth URLs.
func validateRedirectURL(u *url.URL) error {
        if u.Scheme == "" || u.Host == "" {
                return fmt.Errorf("redirect uri missing scheme or host: %q", u.String())
        }
        if u.Scheme != "https" && !networking.IsLocalhost(u.Host) {
                return fmt.Errorf("redirect uri must use https (got %q) unless host is loopback", u.Scheme)
        }
        return nil
}

// validateUpstreamEndpointURL enforces well-formedness and the
// HTTPS-except-loopback rule for an upstream-supplied OAuth endpoint URL.
//
// Used at every point where an endpoint URL enters the resolver from outside
// — operator-configured RegistrationEndpoint, or authorization_endpoint /
// token_endpoint copied out of an upstream's metadata document. The
// downstream oauthproto.validateRegistrationEndpoint enforces HTTPS for the
// registration URL too, but only after a bearer-token transport has already
// been constructed, so a local fail-fast check keeps the
// "no bearer-token transport for a non-HTTPS endpoint" invariant local.
//
// label is included in the error message ("registration_endpoint",
// "authorization_endpoint", "token_endpoint", …) so failures can be tied
// back to the specific field without an additional wrapper.
func validateUpstreamEndpointURL(rawURL, label string) error {
        if rawURL == "" {
                return fmt.Errorf("%s is required", label)
        }
        u, err := url.Parse(rawURL)
        if err != nil {
                return fmt.Errorf("%s %q is not a valid URL: %w", label, rawURL, err)
        }
        if u.Scheme == "" || u.Host == "" {
                return fmt.Errorf("%s %q missing scheme or host", label, rawURL)
        }
        if u.Scheme != "https" && !networking.IsLocalhost(u.Host) {
                return fmt.Errorf("%s %q must use https unless host is loopback (got scheme %q)",
                        label, rawURL, u.Scheme)
        }
        return nil
}

// selectTokenEndpointAuthMethod returns the preferred token endpoint auth
// method given the server's advertised set, intersected with our preference
// order. When the server does not advertise any methods the caller's default
// of client_secret_basic is used (RFC 6749 §2.3.1 baseline) — unless
// publicClient is true, in which case the absence of advertised methods is
// not enough to safely register as a public client and the function returns
// an error (the S256 PKCE gate cannot be evaluated without
// code_challenge_methods_supported).
//
// PKCE coupling for "none": the public-client method "none" is selected only
// when the upstream also advertises S256 in code_challenge_methods_supported.
// RFC 7636 §4.2 / OAuth 2.1 require S256 PKCE for public clients; registering
// as none against an upstream that advertises only "plain" — or omits the
// field entirely — would be a compliance gap. When S256 is missing, "none"
// is skipped (the iteration continues to the next less-preferred method),
// and if no other method is mutually supported the function returns an error
// so the operator sees a clear failure at boot rather than a silent
// downgrade at runtime.
//
// publicClient: when true (CLI flow), the function returns "none" if and
// only if S256 PKCE is advertised. Any other outcome is an error — the
// caller has declared intent and the resolver must not silently switch to
// a confidential-client method.
func selectTokenEndpointAuthMethod(serverSupported, codeChallengeMethodsSupported []string, publicClient bool) (string, error) {
        pkceS256Advertised := slices.Contains(codeChallengeMethodsSupported, oauthproto.PKCEMethodS256)

        if publicClient {
                if !pkceS256Advertised {
                        return "", fmt.Errorf(
                                "public client requested but upstream does not advertise S256 in "+
                                        "code_challenge_methods_supported (got %v); refusing to register a "+
                                        "public client without S256 PKCE per RFC 7636 / OAuth 2.1",
                                codeChallengeMethodsSupported)
                }
                // When serverSupported is non-empty, require that it includes "none"
                // so we don't claim a method the AS would reject.
                if len(serverSupported) > 0 {
                        if !slices.Contains(serverSupported, "none") {
                                return "", fmt.Errorf(
                                        "public client requested but upstream does not advertise "+
                                                "token_endpoint_auth_method=none (got %v); the CLI flow is a "+
                                                "public PKCE client and cannot register with a confidential method",
                                        serverSupported)
                        }
                }
                return "none", nil
        }

        if len(serverSupported) == 0 {
                return "client_secret_basic", nil
        }

        supported := make(map[string]struct{}, len(serverSupported))
        for _, m := range serverSupported {
                supported[m] = struct{}{}
        }

        for _, m := range authMethodPreference {
                if _, ok := supported[m]; !ok {
                        continue
                }
                if m == "none" && !pkceS256Advertised {
                        // Public-client registration without S256 PKCE is non-compliant
                        // per RFC 7636 / OAuth 2.1. Try the next less-preferred method.
                        continue
                }
                return m, nil
        }
        if _, noneOnly := supported["none"]; noneOnly && !pkceS256Advertised {
                return "", fmt.Errorf(
                        "upstream advertises only token_endpoint_auth_method=none but does not advertise "+
                                "S256 in code_challenge_methods_supported (got %v); refusing to register a public "+
                                "client without S256 PKCE per RFC 7636 / OAuth 2.1", codeChallengeMethodsSupported)
        }
        return "", fmt.Errorf(
                "no supported token_endpoint_auth_method in server advertisement %v; "+
                        "client supports %v", serverSupported, authMethodPreference)
}

// bearerTokenTransport is an http.RoundTripper that adds
// Authorization: Bearer {token} to each outgoing request. Used to supply the
// RFC 7591 initial access token to oauthproto.RegisterClientDynamically
// without leaking the abstraction up into that package.
//
// The wrapping http.Client (see newDCRHTTPClient) refuses to follow HTTP
// redirects via CheckRedirect, so this transport is only ever invoked for
// the original registration request — never for a redirected request whose
// URL the upstream chose. That is what prevents this token from being
// forwarded to a foreign origin.
type bearerTokenTransport struct {
        token string
        next  http.RoundTripper
}

// Compile-time assertion: a rename or typo of CloseIdleConnections would
// otherwise silently make http.Client.CloseIdleConnections a no-op on any
// client using this transport (see networking.IdleConnectionCloser).
var _ networking.IdleConnectionCloser = (*bearerTokenTransport)(nil)

// RoundTrip implements http.RoundTripper.
func (t *bearerTokenTransport) RoundTrip(req *http.Request) (*http.Response, error) {
        // Clone per http.RoundTripper contract: implementations must not modify
        // the input request's headers.
        cp := req.Clone(req.Context())
        cp.Header.Set("Authorization", "Bearer "+t.token)
        return t.next.RoundTrip(cp)
}

// CloseIdleConnections forwards to the wrapped RoundTripper so
// http.Client.CloseIdleConnections reaches the underlying connection pool
// instead of stopping at this wrapper (see networking.IdleConnectionCloser).
func (t *bearerTokenTransport) CloseIdleConnections() {
        networking.ForwardCloseIdle(t.next)
}

// errDCRRedirectRefused is returned when a DCR registration endpoint
// responds with a 30x. Net/http surfaces it via *url.Error so callers
// observe a clear failure mode instead of a confusing JSON decode error.
var errDCRRedirectRefused = errors.New(
        "dcr: registration endpoint returned a redirect; refusing to follow " +
                "to avoid forwarding the RFC 7591 initial access token to a foreign origin")

// newDCRHTTPClient returns the http.Client to pass to
// oauthproto.RegisterClientDynamically for a registration POST to
// registrationEndpoint. The client dials through the private-IP-guarded
// transport built by newGuardedDCRClient (CWE-918 SSRF protection) and always
// blocks HTTP redirects so that an upstream cannot use a 30x to coerce us into
// re-issuing the registration request (and any attached
// Authorization: Bearer header) against a different origin. RFC 7591 §3
// does not require redirect support, so refusing them is safe.
//
// When initialAccessToken is non-empty the client also wraps the guarded
// transport with a bearerTokenTransport that injects the Authorization header.
// The combination of the bearer transport plus the redirect block is what
// prevents the token-leak class of bug.
func newDCRHTTPClient(
        initialAccessToken, registrationEndpoint string,
        allowPrivateIPs, serverSuppliedEndpoints bool,
        caFilePath string,
) (*http.Client, error) {
        host, err := hostFromURL(registrationEndpoint)
        if err != nil {
                return nil, err
        }
        client, err := newGuardedDCRClient(host, allowPrivateIPs, serverSuppliedEndpoints, caFilePath)
        if err != nil {
                return nil, err
        }
        client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
                return errDCRRedirectRefused
        }

        if initialAccessToken == "" {
                return client, nil
        }

        next := client.Transport
        if next == nil {
                next = http.DefaultTransport
        }
        client.Transport = &bearerTokenTransport{
                token: initialAccessToken,
                next:  next,
        }
        return client, nil
}

// newGuardedDCRClient builds the private-IP-guarded *http.Client used for both
// of the resolver's outbound calls — the discovery fetch and the registration
// POST. Trusted operator-configured endpoints use
// networking.NewHostScopedClientBuilder; server-supplied endpoints use the
// strict builder, which permits private/loopback/link-local dialing only when
// allowPrivateIPs is true. Keep-alive is disabled so the dial-time check
// re-runs on every request rather than being bypassed by a pooled connection.
//
// The returned client has no CheckRedirect policy: each caller layers its own
// (the registration client refuses all redirects to protect the bearer token;
// the discovery client restricts them to the same host). The dial guard alone
// does not stop a redirect to a different public host, so the redirect policy
// is a required complement, not an optional one.
func newGuardedDCRClient(host string, allowPrivateIPs, serverSuppliedEndpoints bool, caFilePath string) (*http.Client, error) {
        builder := func() *networking.HttpClientBuilder {
                if serverSuppliedEndpoints {
                        return networking.NewServerSuppliedHostClientBuilder(host, allowPrivateIPs, false)
                }
                return networking.NewHostScopedClientBuilder(host, allowPrivateIPs, false)
        }()
        if caFilePath != "" {
                builder.WithSystemRootsPlusCABundle(caFilePath)
        }
        return builder.WithDisableKeepAlives(true).Build()
}

// hostFromURL extracts the host[:port] component used to scope the guarded
// HTTP client. Every URL reaching this helper has already passed
// scheme-and-host validation at the resolver's entry points
// (validateUpstreamEndpointURL for the registration endpoint — both the
// caller-supplied case and the metadata-discovered case handled in
// endpointsFromMetadata — and FetchAuthorizationServerMetadataFromURL for the
// discovery URL; the synthesised-endpoint case derives its host from an
// upstream issuer that already passed the RFC 8414 §3.3 issuer-match check),
// so a parse failure or empty host here signals an internal inconsistency
// rather than untrusted input.
func hostFromURL(rawURL string) (string, error) {
        u, err := url.Parse(rawURL)
        if err != nil {
                return "", fmt.Errorf("dcr: parse url for http client host: %w", err)
        }
        if u.Host == "" {
                return "", fmt.Errorf("dcr: url missing host: %q", rawURL)
        }
        return u.Host, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package dcr

import (
        "context"
        "errors"
        "fmt"
        "io"
        "sync"
        "time"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

// dcrStaleAgeThreshold is the age beyond which a cached DCR resolution is
// considered stale and logged as such by higher-level wiring. The store itself
// does not expire or evict entries. An upstream can still drop a registration
// on its own (a ToolHive auth server evicts idle DCR clients); the resolver
// can detect that on the next cache hit and replace the entry through
// Replacer.ReplaceIfUnchanged (see verifyCachedRegistration for when it can).
// An operator retires a cached registration by changing Request.RotationKey,
// which addresses a new row; the old row is left in place.
const dcrStaleAgeThreshold = 90 * 24 * time.Hour

// Key is a re-export of storage.DCRKey, kept as a package-local alias so
// callers in this package can reference the canonical cache key without an
// explicit storage. qualifier on every call site, while the canonical
// definition lives in pkg/authserver/storage. The canonical form (and its
// ScopesHash constructor) MUST live in a single place so any future Redis
// backend hashes keys identically to the in-memory backend; see
// storage.DCRKey for the field documentation.
type Key = storage.DCRKey

// CredentialStore caches RFC 7591 Dynamic Client Registration resolutions
// keyed by the (Issuer, UpstreamID, RedirectURI, ScopesHash) tuple plus the
// optional RotationKey. Implementations must be safe for concurrent use.
//
// This is the runner-facing interface used by the DCR resolver. It is a
// narrow re-projection of storage.DCRCredentialStore that exchanges
// *Resolution values (the resolver's working type) instead of
// *storage.DCRCredentials so the resolver internals stay agnostic to the
// persistence layer's exact field shape.
//
// Implementations in this package are thin adapters around a
// storage.DCRCredentialStore — the durable map / Redis hash lives over
// there, and this interface adds a per-call Resolution <-> DCRCredentials
// translation. There is exactly one persistence implementation per backend:
// storage.MemoryStorage and storage.RedisStorage. See NewStorageBackedStore
// for the adapter.
type CredentialStore interface {
        // Get returns the cached resolution for key, or (nil, false, nil) if the
        // key is not present. An error is returned only on backend failure.
        Get(ctx context.Context, key Key) (*Resolution, bool, error)

        // PutIfAbsent claims key for resolution. Returns the authoritative
        // durable value: the caller's own resolution on a successful claim, the
        // concurrent winner's otherwise. Callers MUST use the returned value,
        // not their input resolution — RFC 7591 dynamic registration mints a
        // unique client_id/client_secret on every call, so a caller that lost
        // the race and kept using its own resolution would hold credentials the
        // durable store does not agree it owns.
        //
        // Implementations must reject a nil resolution with an error rather
        // than silently succeeding — a no-op would leave callers with no
        // debug trail for the subsequent Get miss.
        //
        // The returned *Resolution is always non-nil when err is nil.
        PutIfAbsent(ctx context.Context, key Key, resolution *Resolution) (*Resolution, error)
}

// Replacer is implemented by a CredentialStore that can replace a
// registration the upstream has rejected. It is optional: ResolveCredentials
// verifies a cached registration against the upstream only when its store
// implements Replacer, and otherwise returns the cached registration
// unchecked, without contacting the upstream.
//
// It is a separate interface rather than a CredentialStore method so that
// stores implemented outside this package keep compiling, and so that a
// wrapper which embeds CredentialStore does not inherit a replacement that
// skips whatever the wrapper does around writes (key namespacing, a lock
// that lets only one replica register). A wrapper that supports replacement
// must implement and forward ReplaceIfUnchanged explicitly.
type Replacer interface {
        // ReplaceIfUnchanged replaces the resolution stored at key with
        // replacement, iff the stored resolution is still the registration the
        // caller read (it carries expected's ClientID). It is the write path for
        // re-registering after the upstream rejected a cached client, and it
        // never overwrites a newer registration that a concurrent replica already
        // persisted.
        //
        // Returns the authoritative durable value: replacement on success, the
        // concurrent writer's row when the stored registration has already
        // changed, or the result of claiming the key as PutIfAbsent would when no
        // row exists any more. Callers MUST use the returned value, for the same
        // reason as with CredentialStore.PutIfAbsent.
        //
        // Implementations must reject a nil expected or replacement with an
        // error. The returned *Resolution is always non-nil when err is nil.
        ReplaceIfUnchanged(ctx context.Context, key Key, expected, replacement *Resolution) (*Resolution, error)
}

// Both in-package stores support replacement.
var (
        _ Replacer = (*inMemoryStore)(nil)
        _ Replacer = (*storageBackedStore)(nil)
)

// NewStorageBackedStore returns a CredentialStore that delegates to a
// storage.DCRCredentialStore for durable persistence and translates
// Resolution values into DCRCredentials at the boundary. The returned
// store is safe for concurrent use because the underlying
// storage.DCRCredentialStore must be (per its interface contract).
//
// Panics if backend is nil — a nil backend is unambiguously a programming
// error and silent acceptance would only delay the eventual nil-pointer
// dereference to the first Get/Put call, far from the constructor site.
func NewStorageBackedStore(backend storage.DCRCredentialStore) CredentialStore {
        if backend == nil {
                panic("dcr: NewStorageBackedStore: backend must not be nil")
        }
        return &storageBackedStore{backend: backend}
}

// CloseableCredentialStore is a CredentialStore that also releases an
// underlying resource (typically a background goroutine) when closed.
// NewInMemoryStore returns this interface so callers can call Close()
// directly without a runtime type assertion — Close() being part of the
// returned type at compile time prevents the silent-no-op failure mode
// that .claude/rules/go-style.md warns about ("Never define a local
// anonymous interface inside an option and type-assert against it to
// check capability — a silent no-op results if the target doesn't
// implement it.").
//
// NewStorageBackedStore returns the base CredentialStore because its
// backends (storage.MemoryStorage shared across the authserver,
// storage.RedisStorage) are owned by the authserver runner and closed
// through that lifecycle, not by the dcr package.
type CloseableCredentialStore interface {
        CredentialStore
        io.Closer
}

// NewInMemoryStore returns a CloseableCredentialStore whose entries live
// only for the lifetime of the returned value. This is the constructor
// consumers reach for when they have no shared durable backend — most
// notably the CLI OAuth flow, which manages cross-invocation credential
// persistence outside the resolver (in pkg/auth/remote/handler.go's
// CachedClientID / CachedClientSecretRef fields) and only needs the
// resolver's intra-call singleflight + S256 PKCE / expiry-refetch
// behaviour for one PerformOAuthFlow call.
//
// The implementation delegates to storage.NewMemoryStorage to share the
// same Get/Put/scope-hash semantics as the durable backends, including
// the background cleanup goroutine. Callers that need to release that
// goroutine before process exit MUST call Close on the returned value
// when finished — the return type is CloseableCredentialStore precisely
// so the call site can `defer store.Close()` without a runtime
// type-assertion. (CLI flows that complete in a single invocation can
// also rely on process exit.)
func NewInMemoryStore() CloseableCredentialStore {
        return &inMemoryStore{mem: storage.NewMemoryStorage()}
}

// inMemoryStore is the CloseableCredentialStore returned by
// NewInMemoryStore. It holds exactly one *storage.MemoryStorage handle,
// which serves Get / Put / Close. No embedding of storageBackedStore —
// embedding would introduce a second handle to the same backend (typed
// as the broad storage.DCRCredentialStore interface) and require a
// manual "keep these two in sync" invariant the compiler could not
// enforce. The single-field shape gives Get / Put / Close one
// authoritative source.
//
// Get / Put logic is intentionally duplicated with storageBackedStore
// rather than delegated through it: the two methods are three lines
// each, and avoiding the embedding is a strictly stronger structural
// guarantee than DRY here. credentialsToResolution /
// resolutionToCredentials remain shared.
//
// closeOnce makes Close() idempotent at the dcr-package boundary even
// though storage.MemoryStorage.Close() is not — calling
// MemoryStorage.Close() twice panics on `close of closed channel`. The
// CLI's `defer store.Close()` is single-call today, but guarding here
// means future call sites that close-on-shortcircuit-then-defer (or
// test helpers that close-then-reuse) cannot panic the process.
type inMemoryStore struct {
        mem       *storage.MemoryStorage
        closeOnce sync.Once
        closeErr  error
}

// Get implements CredentialStore by delegating to the embedded
// *storage.MemoryStorage. The ErrNotFound translation matches
// storageBackedStore.Get; see that method for the rationale.
func (s *inMemoryStore) Get(ctx context.Context, key Key) (*Resolution, bool, error) {
        creds, err := s.mem.GetDCRCredentials(ctx, key)
        if err != nil {
                if errors.Is(err, storage.ErrNotFound) {
                        return nil, false, nil
                }
                return nil, false, err
        }
        return credentialsToResolution(creds), true, nil
}

// PutIfAbsent implements CredentialStore by delegating to the embedded
// *storage.MemoryStorage. The nil-resolution rejection matches
// storageBackedStore.PutIfAbsent; see that method for the rationale.
func (s *inMemoryStore) PutIfAbsent(ctx context.Context, key Key, resolution *Resolution) (*Resolution, error) {
        if resolution == nil {
                return nil, fmt.Errorf("dcr: resolution must not be nil")
        }
        creds := resolutionToCredentials(key, resolution)
        authoritative, err := s.mem.StoreDCRCredentialsIfAbsent(ctx, creds)
        if err != nil {
                return nil, err
        }
        return credentialsToResolution(authoritative), nil
}

// ReplaceIfUnchanged implements Replacer by delegating to the embedded
// *storage.MemoryStorage; see replaceIfUnchanged.
func (s *inMemoryStore) ReplaceIfUnchanged(
        ctx context.Context, key Key, expected, replacement *Resolution,
) (*Resolution, error) {
        return replaceIfUnchanged(ctx, s.mem, key, expected, replacement)
}

// Close releases the embedded MemoryStorage cleanup goroutine. Safe to
// call multiple times: the underlying storage.MemoryStorage.Close
// closes a channel and is NOT itself idempotent (a second call panics
// on `close of closed channel`); the closeOnce here makes the
// dcr-package boundary contract idempotent regardless. Subsequent
// calls return the same error (if any) the first call produced.
func (s *inMemoryStore) Close() error {
        s.closeOnce.Do(func() {
                s.closeErr = s.mem.Close()
        })
        return s.closeErr
}

// storageBackedStore is the dcr-package CredentialStore wrapping a
// storage.DCRCredentialStore. Its methods are the only place that converts
// between the resolver's *Resolution and the persisted
// *storage.DCRCredentials shapes.
type storageBackedStore struct {
        backend storage.DCRCredentialStore
}

// Get implements CredentialStore.
//
// A storage-level ErrNotFound is translated into the (nil, false, nil)
// miss-tuple advertised by the interface. Other errors propagate as-is.
func (s *storageBackedStore) Get(ctx context.Context, key Key) (*Resolution, bool, error) {
        creds, err := s.backend.GetDCRCredentials(ctx, key)
        if err != nil {
                if errors.Is(err, storage.ErrNotFound) {
                        return nil, false, nil
                }
                return nil, false, err
        }
        return credentialsToResolution(creds), true, nil
}

// PutIfAbsent implements CredentialStore.
//
// A nil resolution is rejected rather than silently no-oped: a caller
// passing nil would otherwise get a successful return, observe a miss on
// the next Get, and have no error trail to debug from. Failing loudly at
// the boundary makes such bugs visible at the first call.
//
// The returned *Resolution is the authoritative durable value —
// s.backend.StoreDCRCredentialsIfAbsent's own contract — so a caller whose
// registration lost a concurrent claim on this key gets back the winner's
// credentials, not its own.
func (s *storageBackedStore) PutIfAbsent(ctx context.Context, key Key, resolution *Resolution) (*Resolution, error) {
        if resolution == nil {
                return nil, fmt.Errorf("dcr: resolution must not be nil")
        }
        creds := resolutionToCredentials(key, resolution)
        authoritative, err := s.backend.StoreDCRCredentialsIfAbsent(ctx, creds)
        if err != nil {
                return nil, err
        }
        return credentialsToResolution(authoritative), nil
}

// ReplaceIfUnchanged implements Replacer; see replaceIfUnchanged.
func (s *storageBackedStore) ReplaceIfUnchanged(
        ctx context.Context, key Key, expected, replacement *Resolution,
) (*Resolution, error) {
        return replaceIfUnchanged(ctx, s.backend, key, expected, replacement)
}

// replaceIfUnchanged is the shared ReplaceIfUnchanged body. It re-reads the
// row and uses that physical row, not the caller's expected resolution, as
// the compare-and-set witness: storage.DCRCredentialStore requires the value
// its own GetDCRCredentials returned, and a Resolution does not carry every
// persisted field. Identity is decided by ClientID because DCR mints a new
// client_id on every registration.
//
// When the row has changed or vanished between the read and the write, the
// key is claimed with StoreDCRCredentialsIfAbsent, which returns whatever a
// concurrent writer left there and claims an empty slot otherwise.
func replaceIfUnchanged(
        ctx context.Context, backend storage.DCRCredentialStore, key Key, expected, replacement *Resolution,
) (*Resolution, error) {
        if expected == nil || replacement == nil {
                return nil, fmt.Errorf("dcr: expected and replacement resolutions must not be nil")
        }
        creds := resolutionToCredentials(key, replacement)
        claim := func() (*Resolution, error) {
                authoritative, err := backend.StoreDCRCredentialsIfAbsent(ctx, creds)
                if err != nil {
                        return nil, err
                }
                return credentialsToResolution(authoritative), nil
        }

        current, err := backend.GetDCRCredentials(ctx, key)
        switch {
        case errors.Is(err, storage.ErrNotFound):
                return claim()
        case err != nil:
                return nil, err
        case current.ClientID != expected.ClientID:
                // Another writer already replaced the stale registration.
                return credentialsToResolution(current), nil
        }

        updated, err := backend.UpdateDCRCredentialsIfUnchanged(ctx, creds, current)
        if errors.Is(err, storage.ErrDCRCredentialsChanged) || errors.Is(err, storage.ErrNotFound) {
                return claim()
        }
        if err != nil {
                return nil, err
        }
        return credentialsToResolution(updated), nil
}

// resolutionToCredentials converts a resolver-side *Resolution into the
// persisted *storage.DCRCredentials shape. The Key is supplied separately
// because storage.DCRCredentials carries the key as a struct field rather
// than implicitly via a map key, so the persistence layer can round-trip it
// across processes and backends.
//
// Fields that exist on Resolution but not on DCRCredentials are dropped:
//   - ClientIDIssuedAt: informational only per RFC 7591 §3.2.1; the resolver
//     does not consult it for cache invalidation, so it does not need to
//     survive a process restart.
//   - RedirectURI: already encoded into key.RedirectURI; storing it twice
//     would risk drift between the canonical key and the persisted value.
//
// CreatedAt and ClientSecretExpiresAt are preserved so cache observers
// (e.g. lookupCachedResolution's staleness Warn) and TTL-aware backends
// (Redis) keep their existing behaviour after a restart.
func resolutionToCredentials(key Key, res *Resolution) *storage.DCRCredentials {
        if res == nil {
                return nil
        }
        return &storage.DCRCredentials{
                Key:                     key,
                ClientID:                res.ClientID,
                ClientSecret:            res.ClientSecret,
                TokenEndpointAuthMethod: res.TokenEndpointAuthMethod,
                RegistrationAccessToken: res.RegistrationAccessToken,
                RegistrationClientURI:   res.RegistrationClientURI,
                AuthorizationEndpoint:   res.AuthorizationEndpoint,
                TokenEndpoint:           res.TokenEndpoint,
                CreatedAt:               res.CreatedAt,
                ClientSecretExpiresAt:   res.ClientSecretExpiresAt,
        }
}

// credentialsToResolution is the inverse of resolutionToCredentials. The
// RedirectURI is recovered from the persisted Key so consumers that read it
// off the resolution (e.g. pkg/authserver/runner.consumeResolution, which
// writes it back onto a run-config copy when the caller left it empty)
// see the canonical value.
//
// ClientIDIssuedAt is left zero because it is not persisted. Callers that
// care about it (none today) must read it directly from the live RFC 7591
// response, not from a cached resolution.
func credentialsToResolution(creds *storage.DCRCredentials) *Resolution {
        if creds == nil {
                return nil
        }
        return &Resolution{
                ClientID:                creds.ClientID,
                ClientSecret:            creds.ClientSecret,
                AuthorizationEndpoint:   creds.AuthorizationEndpoint,
                TokenEndpoint:           creds.TokenEndpoint,
                RegistrationAccessToken: creds.RegistrationAccessToken,
                RegistrationClientURI:   creds.RegistrationClientURI,
                TokenEndpointAuthMethod: creds.TokenEndpointAuthMethod,
                RedirectURI:             creds.Key.RedirectURI,
                ClientSecretExpiresAt:   creds.ClientSecretExpiresAt,
                CreatedAt:               creds.CreatedAt,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package dcr

import (
        "context"
        "crypto/rand"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "net/url"
        "regexp"
        "strings"
        "time"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth/oautherr"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// registrationProbeTimeout bounds the request that checks a cached
// registration, so an unresponsive upstream delays resolution by at most this
// long before the cached registration is used as-is.
const registrationProbeTimeout = 10 * time.Second

// authorizationProbeBodyLimit caps how much of an authorization endpoint's
// direct error response the probe reads to find an RFC 6749 error code. An
// error object is a few hundred bytes; anything larger is a page, not a
// verdict. The body is never logged.
const authorizationProbeBodyLimit = 16 << 10

// authorizationProbeRejectionCodes are the RFC 6749 error codes a direct
// (non-redirected) authorization response may carry when the server could
// not validate client_id or redirect_uri. invalid_request is included because
// §4.1.2.1 names no code for that case and servers report an unregistered
// redirect_uri that way (a 2026 vendor incident did exactly this), while a
// valid client with a malformed request would be answered by redirect.
var authorizationProbeRejectionCodes = map[string]struct{}{
        "invalid_client":      {},
        "unauthorized_client": {},
        "invalid_request":     {},
}

// authorizationProbeErrorCodePattern bounds what the probe is willing to
// repeat from a server's error code in its own error and logs.
var authorizationProbeErrorCodePattern = regexp.MustCompile(`^[a-z_]{1,40}$`)

// registrationStatus is the upstream's verdict on a cached registration.
type registrationStatus int

const (
        // registrationUnverified means the probe was skipped or inconclusive.
        registrationUnverified registrationStatus = iota
        // registrationLive means the upstream authenticated the cached client.
        registrationLive
        // registrationRejected means the upstream no longer recognises the
        // cached client.
        registrationRejected
)

// errProbeRedirectRefused stops the token-endpoint probe from following a
// redirect, which would re-send the client credentials to another origin.
var errProbeRedirectRefused = errors.New("dcr: token endpoint returned a redirect; refusing to follow")

// errProbeUnsupportedAuthMethod reports a cached registration whose token
// endpoint auth method the probe cannot authenticate with.
var errProbeUnsupportedAuthMethod = errors.New("dcr: token endpoint auth method cannot be probed")

// errProbeDisabled reports that the caller configured ProbeNone, so the cached
// registration is used without asking the upstream.
var errProbeDisabled = errors.New("dcr: registration probe disabled by configuration")

// verifyCachedRegistration asks the upstream whether it still recognises the
// cached client, using the strategy req.RegistrationProbe selects. An
// upstream can drop a registration that the resolver still caches: RFC 7591
// lets client_secret_expires_at be 0 while the server evicts idle clients
// anyway (a ToolHive auth server does exactly this), or it can lose its
// client storage. The browser is the only party that sees the resulting
// authorize error, so the resolver must ask.
//
// ProbeTokenEndpoint (the default) is verifyViaTokenEndpoint;
// ProbeAuthorizationEndpoint is verifyViaAuthorizationEndpoint. Each returns
// registrationUnverified for every outcome that is not a verdict (transport
// failures, 5xx, rate limiting, shapes the strategy cannot read), and callers
// then keep the cached registration so an unreachable upstream never blocks
// startup. ProbeNone returns registrationUnverified with errProbeDisabled
// without contacting the upstream.
//
// The returned error is non-nil only for an unverified status and says why;
// it never carries the upstream's response body (see probeErrorAttrs).
func verifyCachedRegistration(ctx context.Context, req *Request, cached *Resolution) (registrationStatus, error) {
        switch probe := req.RegistrationProbe.effective(); probe {
        case ProbeTokenEndpoint:
                return verifyViaTokenEndpoint(ctx, req, cached)
        case ProbeAuthorizationEndpoint:
                return verifyViaAuthorizationEndpoint(ctx, req, cached)
        case ProbeNone:
                return registrationUnverified, errProbeDisabled
        default:
                // validateResolveInputs rejects unknown values before resolution starts.
                return registrationUnverified, fmt.Errorf("dcr: unknown registration probe %q", string(probe))
        }
}

// verifyViaTokenEndpoint redeems a random authorization code with the cached
// client credentials. Servers that authenticate the client before validating
// the grant (as fosite does, for public and confidential clients alike)
// answer invalid_client when the registration is gone and invalid_grant when
// it is alive. No token can be issued for the random code. RFC 6749 §4.1.3
// only requires confidential clients to authenticate, so for a public client
// some servers validate the code first and answer invalid_grant for any
// client_id; against those servers this probe cannot detect a deleted
// registration (ProbeAuthorizationEndpoint can, when the server follows RFC
// 6749 §4.1.2.1). This probe skips private_key_jwt clients: it does not sign
// for them (the authorization-endpoint probe sends no credentials and does
// cover them).
//
// The configured token endpoint wins over the cached one, mirroring
// applyExplicitEndpointOverrides, so a tampered cached endpoint cannot
// receive the client secret when the operator configured the endpoint.
func verifyViaTokenEndpoint(ctx context.Context, req *Request, cached *Resolution) (registrationStatus, error) {
        authStyle, ok := probeAuthStyle(cached)
        if !ok {
                return registrationUnverified, errProbeUnsupportedAuthMethod
        }

        tokenEndpoint, operatorSupplied := req.TokenEndpoint, true
        if tokenEndpoint == "" {
                tokenEndpoint = cached.TokenEndpoint
                operatorSupplied = networking.AuthorityMatchesAny(tokenEndpoint, req.DiscoveryURL, req.RegistrationEndpoint)
        }
        client, err := newProbeHTTPClient(req, tokenEndpoint, "token_endpoint", !operatorSupplied)
        if err != nil {
                return registrationUnverified, fmt.Errorf("dcr: build token endpoint client: %w", err)
        }
        // The request carries client credentials; never re-send them to another origin.
        client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
                return errProbeRedirectRefused
        }

        clientSecret := cached.ClientSecret
        if cached.TokenEndpointAuthMethod == oauthproto.TokenEndpointAuthMethodNone {
                clientSecret = ""
        }
        cfg := &oauth2.Config{
                ClientID:     cached.ClientID,
                ClientSecret: clientSecret,
                Endpoint:     oauth2.Endpoint{TokenURL: tokenEndpoint, AuthStyle: authStyle},
                RedirectURL:  cached.RedirectURI,
        }
        probeCtx, cancel := context.WithTimeout(ctx, registrationProbeTimeout)
        defer cancel()
        probeCtx = context.WithValue(probeCtx, oauth2.HTTPClient, client)
        _, err = cfg.Exchange(probeCtx, rand.Text(), oauth2.VerifierOption(rand.Text()+rand.Text()))

        status := classifyProbe(err)
        if status != registrationUnverified {
                return status, nil
        }
        return status, err
}

// verifyViaAuthorizationEndpoint sends an authorization request for the cached
// client and reads the first response without following redirects. RFC 6749
// §4.1.2.1 only says that a server MUST NOT redirect when client_id or
// redirect_uri is invalid; the converse does not hold, since anything in
// front of the server (a gateway wanting a login, an IP policy, bot
// protection) can answer a direct 4xx for a live client. A 2xx (a login or
// consent page) or 3xx (to the IdP, or back to redirect_uri with an error
// such as invalid_scope) shows the client is known. A direct 4xx is a verdict
// only when the server itself answered: status 400, 401 or 403 with an RFC
// 6749 JSON error object whose error is one of
// authorizationProbeRejectionCodes. A server that had validated client_id
// and redirect_uri would have delivered that error by redirect, so the
// direct structured error is the server reporting it could not. The request
// carries no credentials: only the public client_id, the registered
// redirect_uri, the configured scopes and a throwaway state and PKCE
// challenge, so the server records at most an abandoned authorization
// request.
//
// Everything else is unverified and keeps the cached registration: other
// 4xx, a 4xx without a JSON error object (Keycloak and Auth0 render HTML),
// a JSON error outside the rejection set, 5xx and transport failures. A
// server that renders its unknown-client error with 200, or redirects to its
// own error page, reads as alive. Up to authorizationProbeBodyLimit of the
// body is read for the error code and is never logged; only a code matching
// authorizationProbeErrorCodePattern is repeated.
//
// A false rejection is costly: the resolver re-registers at every start (so
// users of an upstream that binds refresh tokens to the client sign in again
// each time) and, when registration also fails, startup fails. Hence the
// narrow verdict.
//
// The configured authorization endpoint wins over the cached one, mirroring
// verifyViaTokenEndpoint.
func verifyViaAuthorizationEndpoint(ctx context.Context, req *Request, cached *Resolution) (registrationStatus, error) {
        authzEndpoint, operatorSupplied := req.AuthorizationEndpoint, true
        if authzEndpoint == "" {
                authzEndpoint = cached.AuthorizationEndpoint
                operatorSupplied = networking.AuthorityMatchesAny(authzEndpoint, req.DiscoveryURL, req.RegistrationEndpoint)
        }
        client, err := newProbeHTTPClient(req, authzEndpoint, "authorization_endpoint", !operatorSupplied)
        if err != nil {
                return registrationUnverified, fmt.Errorf("dcr: build authorization endpoint client: %w", err)
        }
        // A redirect is the "known client" answer; read it, do not follow it.
        client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
                return http.ErrUseLastResponse
        }

        redirectURI := cached.RedirectURI
        if redirectURI == "" {
                if redirectURI, err = resolveUpstreamRedirectURI(req.RedirectURI, req.Issuer); err != nil {
                        return registrationUnverified, fmt.Errorf("dcr: authorization endpoint probe: %w", err)
                }
        }
        probeURL, err := authorizationProbeURL(authzEndpoint, cached.ClientID, redirectURI, req.Scopes)
        if err != nil {
                return registrationUnverified, err
        }

        probeCtx, cancel := context.WithTimeout(ctx, registrationProbeTimeout)
        defer cancel()
        httpReq, err := http.NewRequestWithContext(probeCtx, http.MethodGet, probeURL, nil)
        if err != nil {
                return registrationUnverified, fmt.Errorf("dcr: authorization endpoint probe: %w", err)
        }
        resp, err := client.Do(httpReq)
        if err != nil {
                return registrationUnverified, fmt.Errorf("dcr: authorization endpoint probe: %w", err)
        }
        defer func() {
                _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, authorizationProbeBodyLimit))
                _ = resp.Body.Close()
        }()
        body, err := io.ReadAll(io.LimitReader(resp.Body, authorizationProbeBodyLimit))
        if err != nil {
                return registrationUnverified, fmt.Errorf("dcr: authorization endpoint probe: read response: %w", err)
        }

        status, errorCode := classifyAuthorizationProbe(resp.StatusCode, body)
        //nolint:gosec // G706: client_id is public metadata per RFC 7591.
        slog.Debug("dcr: authorization endpoint probe answered",
                "client_id", cached.ClientID,
                "status", resp.StatusCode,
                "error_code", errorCode,
                "registration_live", status == registrationLive,
                "registration_rejected", status == registrationRejected,
        )
        if status != registrationUnverified {
                return status, nil
        }
        if errorCode != "" {
                return status, fmt.Errorf("dcr: authorization endpoint answered HTTP %d (error %q)", resp.StatusCode, errorCode)
        }
        return status, fmt.Errorf("dcr: authorization endpoint answered HTTP %d", resp.StatusCode)
}

// authorizationProbeURL builds the probe's authorization request URL,
// preserving any query the configured endpoint already carries (some
// providers key tenants that way).
func authorizationProbeURL(endpoint, clientID, redirectURI string, scopes []string) (string, error) {
        u, err := url.Parse(endpoint)
        if err != nil {
                return "", fmt.Errorf("dcr: parse authorization endpoint: %w", err)
        }
        q := u.Query()
        q.Set("response_type", "code")
        q.Set("client_id", clientID)
        q.Set("redirect_uri", redirectURI)
        q.Set("state", rand.Text())
        q.Set("code_challenge", oauth2.S256ChallengeFromVerifier(rand.Text()+rand.Text()))
        q.Set("code_challenge_method", "S256")
        if len(scopes) > 0 {
                q.Set("scope", strings.Join(scopes, " "))
        }
        u.RawQuery = q.Encode()
        return u.String(), nil
}

// classifyAuthorizationProbe maps the authorization endpoint's first response
// to a registrationStatus; see verifyViaAuthorizationEndpoint for the rules.
// The second result is the RFC 6749 error code found in a direct error
// response, filtered through authorizationProbeErrorCodePattern, or "".
func classifyAuthorizationProbe(statusCode int, body []byte) (registrationStatus, string) {
        if statusCode >= http.StatusOK && statusCode < http.StatusBadRequest {
                return registrationLive, ""
        }
        switch statusCode {
        case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden:
        default:
                return registrationUnverified, ""
        }
        var errResp struct {
                Error string `json:"error"`
        }
        if json.Unmarshal(body, &errResp) != nil || !authorizationProbeErrorCodePattern.MatchString(errResp.Error) {
                return registrationUnverified, ""
        }
        if _, rejected := authorizationProbeRejectionCodes[errResp.Error]; rejected {
                return registrationRejected, errResp.Error
        }
        return registrationUnverified, errResp.Error
}

// probeErrorAttrs renders why a probe was unverified as log attributes. An
// OAuth error response is reduced to its HTTP status and RFC 6749 error
// code: oauth2.RetrieveError's message includes the server's
// error_description, or its raw response body when there is no error code,
// and none of that belongs in our logs. Other errors (transport failures,
// refused redirects, an unusable endpoint, a non-verdict status from the
// authorization endpoint) originate on our side and go through
// SanitizeErrorForLog.
func probeErrorAttrs(err error) []any {
        if retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err); ok {
                attrs := []any{"error_code", retrieveErr.ErrorCode}
                if retrieveErr.Response != nil {
                        attrs = append(attrs, "status", retrieveErr.Response.StatusCode)
                }
                return attrs
        }
        return []any{"error", SanitizeErrorForLog(err)}
}

// probeAuthStyle maps the cached token endpoint auth method onto the oauth2
// client's auth style. It reports false for methods the probe cannot use
// (private_key_jwt needs a signed assertion) and for secret-based methods
// whose secret is missing.
func probeAuthStyle(cached *Resolution) (oauth2.AuthStyle, bool) {
        switch cached.TokenEndpointAuthMethod {
        case oauthproto.TokenEndpointAuthMethodClientSecretBasic:
                return oauth2.AuthStyleInHeader, cached.ClientSecret != ""
        case oauthproto.TokenEndpointAuthMethodClientSecretPost:
                return oauth2.AuthStyleInParams, cached.ClientSecret != ""
        case oauthproto.TokenEndpointAuthMethodNone:
                return oauth2.AuthStyleInParams, true
        default:
                return oauth2.AuthStyleAutoDetect, false
        }
}

// newProbeHTTPClient builds the guarded client for a probe against endpoint,
// validated under label. It reuses the resolver's SSRF policy
// (newGuardedDCRClient); callers set CheckRedirect for their strategy.
func newProbeHTTPClient(req *Request, endpoint, label string, serverSupplied bool) (*http.Client, error) {
        if err := validateUpstreamEndpointURL(endpoint, label); err != nil {
                return nil, err
        }
        host, err := hostFromURL(endpoint)
        if err != nil {
                return nil, err
        }
        return newGuardedDCRClient(host, req.AllowPrivateIPs, req.ServerSuppliedEndpoints || serverSupplied, req.CAFilePath)
}

// classifyProbe turns the token-endpoint probe's exchange error into a
// registrationStatus. Only a structured RFC 6749 error code is a verdict;
// transient shapes (5xx, 429, non-OAuth bodies) and transport errors are
// unverified.
func classifyProbe(err error) registrationStatus {
        if err == nil {
                // The upstream issued tokens for a random code; the client is
                // authenticated, whatever else is wrong with that server.
                return registrationLive
        }
        retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err)
        if !ok || !oautherr.IsPermanentCredentialError(err) {
                return registrationUnverified
        }
        switch retrieveErr.ErrorCode {
        case "invalid_client":
                return registrationRejected
        case "invalid_grant":
                return registrationLive
        default:
                return registrationUnverified
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package discovery provides authentication discovery utilities for detecting
// authentication requirements from remote servers.
//
// Supported Authentication Types:
// - OAuth 2.0 with PKCE (Proof Key for Code Exchange)
// - OIDC (OpenID Connect) discovery
// - Manual OAuth endpoint configuration
// - RFC 9728 Protected Resource Metadata
package discovery

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

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/dcr"
        "github.com/stacklok/toolhive/pkg/auth/oauth"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// Default timeout constants for authentication operations
const (
        DefaultOAuthTimeout      = 5 * time.Minute
        DefaultHTTPTimeout       = 30 * time.Second
        DefaultAuthDetectTimeout = 10 * time.Second
        MaxRetryAttempts         = 3
        RetryBaseDelay           = 2 * time.Second
        MaxResponseBodyDrain     = 1 * 1024 * 1024 // 1 MB - limit response body draining to prevent resource exhaustion
)

// AuthInfo contains authentication information extracted from WWW-Authenticate header
type AuthInfo struct {
        Realm            string
        Type             string
        ResourceMetadata string
        Error            string
        ErrorDescription string
}

// AuthServerInfo contains information about a validated authorization server
type AuthServerInfo struct {
        Issuer                            string
        AuthorizationURL                  string
        TokenURL                          string
        RegistrationEndpoint              string
        ClientIDMetadataDocumentSupported bool
}

// Config holds configuration for authentication discovery
type Config struct {
        Timeout               time.Duration
        TLSHandshakeTimeout   time.Duration
        ResponseHeaderTimeout time.Duration
        EnablePOSTDetection   bool // Whether to try POST requests for detection
}

// DefaultDiscoveryConfig returns a default discovery configuration
func DefaultDiscoveryConfig() *Config {
        return &Config{
                Timeout:               DefaultAuthDetectTimeout,
                TLSHandshakeTimeout:   5 * time.Second,
                ResponseHeaderTimeout: 5 * time.Second,
                EnablePOSTDetection:   true,
        }
}

// newDetectionClient builds the HTTP client used to probe a target server for
// WWW-Authenticate. The remote MCP server is untrusted, so it refuses
// cross-host / scheme-downgrade redirects to prevent the server driving the
// host into an SSRF (CWE-918). Extracted so the pool bounds
// networking.SetIdleConnBounds applies stay unit-testable.
func newDetectionClient(config *Config) *http.Client {
        transport := &http.Transport{
                TLSHandshakeTimeout:   config.TLSHandshakeTimeout,
                ResponseHeaderTimeout: config.ResponseHeaderTimeout,
        }
        networking.SetIdleConnBounds(transport)
        return &http.Client{
                Timeout:       config.Timeout,
                Transport:     transport,
                CheckRedirect: networking.SameHostRedirectPolicy(),
        }
}

// DetectAuthenticationFromServer attempts to detect authentication requirements from the target server
func DetectAuthenticationFromServer(ctx context.Context, targetURI string, config *Config) (*AuthInfo, error) {
        if config == nil {
                config = DefaultDiscoveryConfig()
        }

        // Create a context with timeout for auth detection
        detectCtx, cancel := context.WithTimeout(ctx, config.Timeout)
        defer cancel()

        // Make a test request to the target server to see if it returns WWW-Authenticate.
        client := newDetectionClient(config)

        // First try a GET request
        authInfo, err := detectAuthWithRequest(detectCtx, client, targetURI, http.MethodGet, nil)
        if err != nil {
                return nil, err
        }
        if authInfo != nil {
                return authInfo, nil
        }

        // If no auth detected with GET and POST detection is enabled, try a POST request with JSON-RPC initialize
        // Some servers only return WWW-Authenticate on specific requests
        if config.EnablePOSTDetection {
                postBody := strings.NewReader(`{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}`)
                authInfo, err = detectAuthWithRequest(detectCtx, client, targetURI, http.MethodPost, postBody)
                if err != nil {
                        return nil, err
                }
                if authInfo != nil {
                        return authInfo, nil
                }
        }

        // NEW: Well-known URI fallback per MCP specification
        // When no WWW-Authenticate header found, try well-known URIs
        slog.Debug("No WWW-Authenticate header found, attempting well-known URI discovery")

        wellKnownAuthInfo, err := tryWellKnownDiscovery(detectCtx, client, targetURI)
        if err != nil {
                slog.Debug("Well-known URI discovery failed", "error", err)
                return nil, nil // Not an error, just no auth detected
        }

        if wellKnownAuthInfo != nil {
                slog.Debug("Discovered authentication via well-known URI")
                return wellKnownAuthInfo, nil
        }

        return nil, nil // No authentication required
}

// detectAuthWithRequest makes a specific HTTP request and checks for authentication requirements
func detectAuthWithRequest(
        ctx context.Context,
        client *http.Client,
        targetURI string,
        method string,
        body *strings.Reader,
) (*AuthInfo, error) {
        var req *http.Request
        var err error

        if body != nil {
                req, err = http.NewRequestWithContext(ctx, method, targetURI, body)
                if err != nil {
                        return nil, fmt.Errorf("failed to create %s request: %w", method, err)
                }
                req.Header.Set("Content-Type", "application/json")
        } else {
                req, err = http.NewRequestWithContext(ctx, method, targetURI, nil)
                if err != nil {
                        return nil, fmt.Errorf("failed to create %s request: %w", method, err)
                }
        }

        // #nosec G704 -- targetURI is server-controlled; the client refuses cross-host and
        // HTTPS->HTTP redirects (networking.SameHostRedirectPolicy) to contain SSRF.
        resp, err := client.Do(req)
        if err != nil {
                return nil, fmt.Errorf("failed to make %s request: %w", method, err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("Failed to close response body", "error", err)
                }
        }()

        // Check if we got a 401 Unauthorized with WWW-Authenticate header
        if resp.StatusCode == http.StatusUnauthorized {
                wwwAuth := resp.Header.Get("WWW-Authenticate")
                if wwwAuth != "" {
                        return ParseWWWAuthenticate(wwwAuth)
                }
        }

        return nil, nil
}

// buildWellKnownURI constructs a well-known URI for OAuth Protected Resource metadata
// per RFC 9728 Section 3.1 and MCP specification
func buildWellKnownURI(parsedURL *url.URL, endpointSpecific bool) string {
        baseURL := url.URL{
                Scheme: parsedURL.Scheme,
                Host:   parsedURL.Host,
        }

        if endpointSpecific && parsedURL.Path != "" && parsedURL.Path != "/" {
                // Endpoint-specific: /.well-known/oauth-protected-resource/<original-path>
                // Remove leading slash from original path to avoid double slashes
                cleanPath := strings.TrimPrefix(parsedURL.Path, "/")
                baseURL.Path = path.Join(oauthproto.WellKnownOAuthResourcePath, cleanPath)
        } else {
                // Root-level: /.well-known/oauth-protected-resource
                baseURL.Path = oauthproto.WellKnownOAuthResourcePath
        }

        return baseURL.String()
}

// checkWellKnownURIExists returns true if a well-known URI is accessible and returns application/json
// Per RFC 9728, protected resource metadata MUST be queried using HTTP GET and MUST return application/json
func checkWellKnownURIExists(ctx context.Context, client *http.Client, uri string) bool {
        req, err := http.NewRequestWithContext(ctx, http.MethodGet, uri, nil)
        if err != nil {
                //nolint:gosec // G706: uri is from server endpoint discovery
                slog.Debug("Failed to create GET request", "uri", uri, "error", err)
                return false
        }

        req.Header.Set("Accept", "application/json")

        // #nosec G704 -- uri is server-controlled; the client refuses cross-host and
        // HTTPS->HTTP redirects (networking.SameHostRedirectPolicy) to contain SSRF.
        resp, err := client.Do(req)
        if err != nil {
                //nolint:gosec // G706: uri is from server endpoint discovery
                slog.Debug("Failed to check well-known URI", "uri", uri, "error", err)
                return false
        }
        defer func() {
                // Drain and close response body to enable connection reuse
                // Limit draining to MaxResponseBodyDrain to prevent resource exhaustion from large responses
                _, _ = io.CopyN(io.Discard, resp.Body, MaxResponseBodyDrain)
                _ = resp.Body.Close()
        }()

        // RFC 9728 requires 200 OK status code - metadata endpoints must be publicly accessible
        if resp.StatusCode != http.StatusOK {
                return false
        }

        // RFC 9728 requires Content-Type to be application/json
        contentType := strings.ToLower(resp.Header.Get("Content-Type"))
        if !strings.Contains(contentType, "application/json") {
                //nolint:gosec // G706: content type from server response is safe to log
                slog.Debug("Well-known URI returned unexpected content type",
                        "uri", uri, "content_type", contentType)
                return false
        }

        return true
}

// tryWellKnownDiscovery attempts to discover authentication requirements via well-known URIs
// per MCP specification Section: Protected Resource Metadata Discovery Requirements.
// Tries endpoint-specific path first, then root-level path.
func tryWellKnownDiscovery(ctx context.Context, client *http.Client, targetURI string) (*AuthInfo, error) {
        parsedURL, err := url.Parse(targetURI)
        if err != nil {
                return nil, fmt.Errorf("invalid target URI: %w", err)
        }

        // Build well-known URIs to try (in priority order per MCP spec)
        wellKnownURIs := []string{
                // 1. Endpoint-specific: /.well-known/oauth-protected-resource/<path>
                buildWellKnownURI(parsedURL, true),
                // 2. Root-level: /.well-known/oauth-protected-resource
                buildWellKnownURI(parsedURL, false),
        }

        // Try each well-known URI in order
        for _, wellKnownURI := range wellKnownURIs {
                //nolint:gosec // G706: well-known URIs are built from server endpoint
                slog.Debug("Trying well-known URI", "uri", wellKnownURI)

                // Check if the URI exists before attempting to fetch
                if !checkWellKnownURIExists(ctx, client, wellKnownURI) {
                        //nolint:gosec // G706: well-known URIs are built from server endpoint
                        slog.Debug("Well-known URI not found", "uri", wellKnownURI)
                        continue
                }

                // URI exists - return AuthInfo with ResourceMetadata set
                // Downstream handler will use FetchResourceMetadata to get the actual metadata
                //nolint:gosec // G706: well-known URIs are built from server endpoint
                slog.Debug("Found well-known URI", "uri", wellKnownURI)
                return &AuthInfo{
                        Type:             "OAuth",
                        ResourceMetadata: wellKnownURI,
                }, nil
        }

        return nil, nil // No well-known metadata found
}

// ParseWWWAuthenticate parses the WWW-Authenticate header to extract authentication information
// Supports multiple authentication schemes and complex header formats
func ParseWWWAuthenticate(header string) (*AuthInfo, error) {
        // Trim whitespace and handle empty headers
        header = strings.TrimSpace(header)
        if header == "" {
                return nil, fmt.Errorf("empty WWW-Authenticate header")
        }

        // Check for OAuth/Bearer authentication
        // Note: We don't split by comma because Bearer parameters can contain commas in quoted values
        if strings.HasPrefix(header, "Bearer") {
                authInfo := &AuthInfo{Type: "Bearer"}

                // Extract parameters after "Bearer"
                params := strings.TrimSpace(strings.TrimPrefix(header, "Bearer"))
                if params != "" {
                        // Parse parameters (realm, scope, resource_metadata, etc.)
                        realm := ExtractParameter(params, "realm")
                        if realm != "" {
                                authInfo.Realm = realm
                        }

                        // RFC 9728: Check for resource_metadata parameter
                        resourceMetadata := ExtractParameter(params, "resource_metadata")
                        if resourceMetadata != "" {
                                authInfo.ResourceMetadata = resourceMetadata
                        }

                        // Extract error information if present
                        errorParam := ExtractParameter(params, "error")
                        if errorParam != "" {
                                authInfo.Error = errorParam
                        }

                        errorDesc := ExtractParameter(params, "error_description")
                        if errorDesc != "" {
                                authInfo.ErrorDescription = errorDesc
                        }
                }

                return authInfo, nil
        }

        // Check for OAuth-specific schemes
        if strings.HasPrefix(header, "OAuth") {
                authInfo := &AuthInfo{Type: "OAuth"}

                // Extract parameters after "OAuth"
                params := strings.TrimSpace(strings.TrimPrefix(header, "OAuth"))
                if params != "" {
                        // Parse parameters (realm, scope, etc.)
                        realm := ExtractParameter(params, "realm")
                        if realm != "" {
                                authInfo.Realm = realm
                        }

                        // RFC 9728: Check for resource_metadata parameter
                        resourceMetadata := ExtractParameter(params, "resource_metadata")
                        if resourceMetadata != "" {
                                authInfo.ResourceMetadata = resourceMetadata
                        }
                }

                return authInfo, nil
        }

        // Currently only OAuth-based authentication is supported
        // Basic and Digest authentication are not implemented
        if strings.HasPrefix(header, "Basic") || strings.HasPrefix(header, "Digest") {
                //nolint:gosec // G706: auth scheme name (Basic/Digest) is safe to log
                slog.Debug("Unsupported authentication scheme", "header", header)
                return nil, fmt.Errorf("unsupported authentication scheme: %s", strings.Split(header, " ")[0])
        }

        return nil, fmt.Errorf("no supported authentication type found in header: %s", header)
}

// DeriveIssuerFromURL attempts to derive the OAuth issuer from the remote URL using general patterns
func DeriveIssuerFromURL(remoteURL string) string {
        // Parse the URL to extract the domain
        parsedURL, err := url.Parse(remoteURL)
        if err != nil {
                slog.Debug("Failed to parse remote URL", "error", err)
                return ""
        }

        host := parsedURL.Hostname()
        if host == "" {
                return ""
        }

        // Append port if explicitly present in the original URL
        port := parsedURL.Port()
        if port != "" {
                host = fmt.Sprintf("%s:%s", host, port)
        }

        // For localhost, preserve the original scheme (HTTP or HTTPS)
        // This supports local development and testing scenarios
        scheme := networking.HttpsScheme
        if networking.IsLocalhost(host) && parsedURL.Scheme != "" {
                scheme = parsedURL.Scheme
        }

        // General pattern: use the domain as the issuer
        // This works for most OAuth providers that use their domain as the issuer
        issuer := fmt.Sprintf("%s://%s", scheme, host)

        //nolint:gosec // G706: derived issuer URL is from server configuration
        slog.Debug("Derived issuer from URL", "remote_url", remoteURL, "issuer", issuer)
        return issuer
}

// ExtractParameter extracts a parameter value from an authentication header
// Handles both quoted and unquoted values according to RFC 2617 and RFC 6750
func ExtractParameter(params, paramName string) string {
        // Parameters can be separated by comma or space
        // Handle both paramName=value and paramName="value" formats

        // First try to find the parameter with equals sign
        searchStr := paramName + "="
        idx := strings.Index(params, searchStr)
        if idx == -1 {
                return ""
        }

        // Extract the value after the equals sign
        valueStart := idx + len(searchStr)
        if valueStart >= len(params) {
                return ""
        }

        remainder := params[valueStart:]

        // Check if the value is quoted
        if strings.HasPrefix(remainder, `"`) {
                // Find the closing quote
                endIdx := 1
                for endIdx < len(remainder) {
                        if remainder[endIdx] == '"' && (endIdx == 1 || remainder[endIdx-1] != '\\') {
                                // Found unescaped closing quote
                                value := remainder[1:endIdx]
                                // Unescape any escaped quotes
                                value = strings.ReplaceAll(value, `\"`, `"`)
                                return value
                        }
                        endIdx++
                }
                // No closing quote found, return empty
                return ""
        }

        // Unquoted value - find the end (comma, space, or end of string)
        endIdx := 0
        for endIdx < len(remainder) {
                if remainder[endIdx] == ',' || remainder[endIdx] == ' ' {
                        break
                }
                endIdx++
        }

        return strings.TrimSpace(remainder[:endIdx])
}

// DeriveIssuerFromRealm attempts to derive the OAuth issuer from the realm parameter
// According to RFC 8414, the issuer MUST be a URL using the "https" scheme with no query or fragment
func DeriveIssuerFromRealm(realm string) string {
        if realm == "" {
                return ""
        }

        // Check if realm is already a valid HTTPS URL
        parsedURL, err := url.Parse(realm)
        if err != nil {
                slog.Debug("Realm is not a valid URL", "error", err)
                return ""
        }

        // RFC 8414: The issuer identifier MUST be a URL using the "https" scheme
        // with no query or fragment components
        if parsedURL.Scheme != "https" && !networking.IsLocalhost(parsedURL.Host) {
                slog.Debug("Realm is not using HTTPS scheme", "realm", realm)
                return ""
        }

        // Normalize the path to prevent path traversal attacks
        if parsedURL.Path != "" {
                // Clean the path to resolve . and .. elements
                cleanPath := path.Clean(parsedURL.Path)
                // Ensure the path doesn't escape the root
                if !strings.HasPrefix(cleanPath, "/") {
                        cleanPath = "/" + cleanPath
                }
                parsedURL.Path = cleanPath
        }

        if parsedURL.RawQuery != "" || parsedURL.Fragment != "" {
                slog.Debug("Realm contains query or fragment components", "realm", realm)
                // Remove query and fragment to make it a valid issuer
                parsedURL.RawQuery = ""
                parsedURL.Fragment = ""
        }

        issuer := parsedURL.String()
        //nolint:gosec // G706: realm is from WWW-Authenticate header of configured remote
        slog.Debug("Derived issuer from realm", "realm", realm, "issuer", issuer)
        return issuer
}

// OAuthFlowConfig contains configuration for performing OAuth flows
type OAuthFlowConfig struct {
        ClientID             string
        ClientSecret         string //nolint:gosec // G117: field legitimately holds sensitive data
        AuthorizeURL         string // Manual OAuth endpoint (optional)
        TokenURL             string // Manual OAuth endpoint (optional)
        RegistrationEndpoint string // Manual registration endpoint (optional)
        Scopes               []string
        CallbackPort         int
        Timeout              time.Duration
        SkipBrowser          bool
        Resource             string // RFC 8707 resource indicator (optional)
        OAuthParams          map[string]string
        ScopeParamName       string // Override scope query parameter name (e.g., "user_scope" for Slack)

        // AllowPrivateIPs permits the Dynamic Client Registration calls (discovery
        // fetch and registration POST) to reach private/loopback/link-local
        // addresses. Callers set it from networking.TargetIsPrivate on the remote
        // target: when ToolHive is pointed at a private remote the upstream IdP may
        // legitimately be private too, so its DCR calls must be allowed; otherwise
        // they are guarded to contain SSRF (CWE-918). This mirrors the
        // blockPrivateIPs decision the flow already applies to its other discovery
        // fetches. Defaults to false (guarded).
        AllowPrivateIPs bool

        // TokenEndpointTrusted reports that the token endpoint's authority was
        // supplied by the operator — configured directly, or named by the
        // operator-configured issuer's own metadata document. It permits
        // private/loopback dialing of the token endpoint; the zero value is the
        // guarded posture. Never set it for an endpoint whose authority a remote
        // MCP server chose.
        TokenEndpointTrusted bool

        // IssuerTrusted reports that the issuer URL came from operator
        // configuration rather than from the remote MCP server's WWW-Authenticate
        // realm / resource metadata.
        IssuerTrusted bool

        // DCR renewal metadata — populated by handleDynamicRegistration and threaded
        // into OAuthFlowResult so callers can persist the data for RFC 7592 operations.
        SecretExpiry            time.Time // zero means the secret never expires
        RegistrationAccessToken string    //nolint:gosec // G117: field legitimately holds sensitive data
        RegistrationClientURI   string
        TokenEndpointAuthMethod string
        RegisteredCallbackPort  int
}

// OAuthFlowResult contains the result of an OAuth flow
type OAuthFlowResult struct {
        TokenSource oauth2.TokenSource
        Config      *oauth.Config

        // Token details for persistence across restarts
        AccessToken  string //nolint:gosec // G117: field legitimately holds sensitive data
        RefreshToken string //nolint:gosec // G117: field legitimately holds sensitive data
        Expiry       time.Time

        // DCR client credentials for persistence (obtained during Dynamic Client Registration)
        ClientID     string
        ClientSecret string //nolint:gosec // G117: field legitimately holds sensitive data

        // DCR renewal metadata (RFC 7591 §3.2.1 / RFC 7592).
        // SecretExpiry is zero when the provider did not issue an expiring secret.
        // RegistrationAccessToken and RegistrationClientURI are empty when the
        // provider does not support RFC 7592 management operations.
        SecretExpiry            time.Time
        RegistrationAccessToken string //nolint:gosec // G117: field legitimately holds sensitive data
        RegistrationClientURI   string
        TokenEndpointAuthMethod string
        RegisteredCallbackPort  int
}

func shouldDynamicallyRegisterClient(config *OAuthFlowConfig) bool {
        return config.ClientID == "" && config.ClientSecret == ""
}

// PerformOAuthFlow performs an OAuth authentication flow with the given configuration
func PerformOAuthFlow(ctx context.Context, issuer string, config *OAuthFlowConfig) (*OAuthFlowResult, error) {
        slog.Debug("Starting OAuth authentication flow", "issuer", issuer)

        if config == nil {
                return nil, fmt.Errorf("OAuth flow config cannot be nil")
        }

        // Resolve port availability before registration. DCR clients allow port fallback
        // because the actual port is registered after selection. Pre-registered and CIMD
        // clients require the configured port to be available as-is — it is already
        // published in their IdP application or metadata document redirect URI.
        if shouldDynamicallyRegisterClient(config) {
                // For dynamic registration, we can allow fallback to alternative ports
                // since we can register the client with the actual port we'll use
                port, err := networking.FindOrUsePort(config.CallbackPort)
                if err != nil {
                        return nil, fmt.Errorf("failed to find available port: %w", err)
                }

                if port != config.CallbackPort {
                        slog.Warn("Specified auth callback port is unavailable", "requested_port", config.CallbackPort, "actual_port", port)
                }
                config.CallbackPort = port
        } else {
                // For pre-registered clients and CIMD, use strict port checking.
                // The port is either configured in the IdP app or baked into the
                // redirect URI in the hosted metadata document.
                if !networking.IsAvailable(config.CallbackPort) {
                        return nil, fmt.Errorf(
                                "specified auth callback port %d is not available - please choose a different port or ensure it's not in use",
                                config.CallbackPort,
                        )
                }
        }

        // Handle dynamic client registration if needed
        if shouldDynamicallyRegisterClient(config) {
                if err := handleDynamicRegistration(ctx, issuer, config); err != nil {
                        return nil, err
                }
        }

        // Create OAuth configuration
        oauthConfig, err := createOAuthConfig(ctx, issuer, config)
        if err != nil {
                return nil, fmt.Errorf("failed to create OAuth config: %w", err)
        }

        // Create and execute OAuth flow
        return newOAuthFlow(ctx, oauthConfig, config)
}

// handleDynamicRegistration handles the dynamic client registration process.
//
// Persistence model (option (b) from #5219): the resolver runs against an
// in-memory dcr.CredentialStore for the duration of a single CLI
// invocation, so within one PerformOAuthFlow call concurrent goroutines
// share the singleflight (proving the inherited "singleflight
// deduplication" property). Cross-invocation persistence is handled
// outside the resolver by pkg/auth/remote/handler.go, which reads
// CachedClientID / CachedClientSecretRef before this code path runs and
// short-circuits to a refresh-token flow when a usable cached client
// exists.
//
// One consequence of option (b) is that the resolver's RFC 7591 §3.2.1
// expiry-driven refetch does NOT participate in the CLI's cross-invocation
// persistence loop: each PerformOAuthFlow call builds a fresh in-memory store,
// so a cached entry from a previous invocation never reaches the resolver.
// Cross-invocation client-secret expiry is handled instead by the remote
// handler, which consults CachedSecretExpiry and renews through RFC 7592 before
// cached credentials are used.
//
// Wrapping the remote handler's secretProvider into a dcr.CredentialStore
// adapter (option (a)) would close that loop and is the natural follow-up;
// it was rejected here as out-of-scope churn for sub-issue 4b.
func handleDynamicRegistration(ctx context.Context, issuer string, config *OAuthFlowConfig) error {
        discoveredDoc, err := getDiscoveryDocument(ctx, issuer, config)
        if err != nil {
                return fmt.Errorf("failed to discover registration endpoint: %w", err)
        }

        // Check if the provider supports Dynamic Client Registration before
        // invoking the resolver. The CLI-flag hint below is intentional: this
        // function is CLI-facing (pkg/auth/discovery is not a protocol-level
        // package) and the flags named here are the correct fallback for
        // operators who need to supply credentials manually. The
        // protocol-neutral version of this message lives in
        // pkg/oauthproto.handleHTTPResponse for the HTTP 404/405/501 paths.
        if discoveredDoc.RegistrationEndpoint == "" {
                return fmt.Errorf("this provider does not support Dynamic Client Registration (DCR). " +
                        "Please configure OAuth client credentials using --remote-auth-client-id and --remote-auth-client-secret flags, " +
                        "or register a client manually with the provider")
        }

        resolution, err := resolveDCRCredentials(ctx, issuer, config, discoveredDoc)
        if err != nil {
                return err
        }

        // Update config with registered client credentials. The remote handler
        // at pkg/auth/remote/handler.go reads ClientID / ClientSecret off
        // OAuthFlowResult and persists them into CachedClientID /
        // CachedClientSecretRef for the next invocation.
        config.ClientID = resolution.ClientID
        config.ClientSecret = resolution.ClientSecret

        // Surface the resolved authorization / token endpoints to the OAuth
        // flow. The resolver returns the endpoints it used for registration
        // (caller-supplied if specified, discovered otherwise), so
        // downstream OAuth-config construction can rely on these fields
        // being populated.
        if resolution.AuthorizationEndpoint != "" {
                config.AuthorizeURL = resolution.AuthorizationEndpoint
        }
        if resolution.TokenEndpoint != "" {
                // A registration response — or the metadata document behind it — can name
                // a token endpoint on any authority. Trust survives only while the new
                // authority is one the operator already named (CWE-918); this can clear
                // the flag but never set it.
                if !networking.AuthorityMatchesAny(resolution.TokenEndpoint, config.TokenURL, issuer) {
                        config.TokenEndpointTrusted = false
                }
                config.TokenURL = resolution.TokenEndpoint
        }

        // Store DCR renewal metadata for RFC 7592 operations.
        // A zero ClientSecretExpiresAt means the secret never expires (RFC 7591 §3.2.1).
        config.SecretExpiry = resolution.ClientSecretExpiresAt
        config.RegistrationAccessToken = resolution.RegistrationAccessToken
        config.RegistrationClientURI = resolution.RegistrationClientURI
        config.TokenEndpointAuthMethod = resolution.TokenEndpointAuthMethod
        config.RegisteredCallbackPort = config.CallbackPort

        if resolution.RegistrationAccessToken != "" {
                slog.Debug("DCR response includes registration access token for RFC 7592 operations")
        }

        return nil
}

// resolveDCRCredentials routes the CLI-flow DCR registration through the
// shared pkg/auth/dcr resolver, inheriting its singleflight deduplication,
// S256 PKCE gating, RFC 7591 §3.2.1 expiry-driven refetch, bearer-token
// transport with redirect refusal, and panic recovery.
//
// The redirect URI is a loopback per RFC 8252 §7.3 — this is the CLI's
// existing public-client model, and PublicClient=true tells the resolver
// to register with token_endpoint_auth_method=none and to refuse when
// S256 PKCE is not advertised (rather than silently downgrading).
//
// Discovery fetch — multi-URL fallback. oauthproto.FetchAuthorizationServerMetadata
// tries the three well-known URL forms defined by RFC 8414 §3.1 and OIDC
// Discovery in priority order: RFC 8414 path-insertion
// (scheme://host/.well-known/oauth-authorization-server/{path}), OIDC
// issuer-suffix ({issuer}/.well-known/openid-configuration), and bare
// RFC 8414 (scheme://host/.well-known/oauth-authorization-server). Using
// this function rather than constructing a single OIDC URL ensures that
// authorization servers with non-root issuer paths whose metadata lives
// at the path-insertion URL are correctly reached (fixes #5356).
// The fetched code_challenge_methods_supported is forwarded to the
// resolver via dcr.Request so the S256 PKCE gate fires without a
// second discovery round-trip inside the resolver.
//
// Threading code_challenge_methods_supported through AuthServerInfo (and
// updating every caller of ValidateAndDiscoverAuthServer) would eliminate
// this fetch on the pre-discovered path and is a natural follow-up.
func resolveDCRCredentials(
        ctx context.Context,
        issuer string,
        config *OAuthFlowConfig,
        discoveredDoc *oauthproto.OIDCDiscoveryDocument,
) (*dcr.Resolution, error) {
        redirectURI := fmt.Sprintf("http://localhost:%d/callback", config.CallbackPort)

        // dcr.Request.Issuer carries the caller's logical scope for cache
        // keying. The CLI flow has no separate logical issuer of its own, so
        // it deliberately reuses the upstream's issuer URL here. This is
        // safe because the resolver's cache key is (Issuer, UpstreamID,
        // RedirectURI, ScopesHash) and the CLI always supplies an explicit
        // loopback RedirectURI per RFC 8252 §7.3 — even if a future
        // embedded-authserver upstream happened to share Issuer with a CLI
        // invocation, the distinct RedirectURI keeps the cache keys apart.
        // UpstreamID (derived by the resolver from the RegistrationEndpoint set
        // below, or a DiscoveryURL) further distinguishes distinct upstreams.
        // See the Issuer field doc on dcr.Request for the wider semantics.
        req := &dcr.Request{
                Issuer:                issuer,
                RedirectURI:           redirectURI,
                Scopes:                config.Scopes,
                AuthorizationEndpoint: discoveredDoc.AuthorizationEndpoint,
                TokenEndpoint:         discoveredDoc.TokenEndpoint,
                PublicClient:          true,
                // Carry the flow's private-IP decision so DCR shares the same SSRF
                // posture as the flow's other discovery fetches; without this the DCR
                // calls would always be guarded and a legitimately private upstream
                // would be refused. See OAuthFlowConfig.AllowPrivateIPs.
                AllowPrivateIPs: config.AllowPrivateIPs,
                // An operator-configured issuer is not server-supplied input; only a
                // realm or resource-metadata URL named by the remote MCP server is.
                ServerSuppliedEndpoints: !config.IssuerTrusted,
        }

        // Fetch AS metadata using the multi-URL fallback so non-root issuers
        // are correctly reached (see function-level doc). discoveredDoc.Issuer
        // is non-empty for every reachable caller; the else branch is
        // defence-in-depth for a future refactor that produces an empty-issuer
        // doc and preserves the pre-existing RegistrationEndpoint-direct path.
        if discoveredDoc.Issuer != "" {
                // Guard this fetch the same way pkg/auth/dcr's own outbound calls are
                // guarded (CWE-918): discoveredDoc.Issuer can be untrusted server input
                // on some discovery branches (see discoverIssuerAndScopes in
                // pkg/auth/remote/handler.go), so a nil client here would reopen the
                // discovery-indirection and DNS-rebinding vectors this PR closes
                // elsewhere. Which builder applies depends on who named the issuer: an
                // operator-configured one gets the host-scoped policy, a server-supplied
                // one the strict policy where localhost and
                // INSECURE_DISABLE_URL_VALIDATION cannot widen the private-IP gate.
                metaHost, parseErr := url.Parse(discoveredDoc.Issuer)
                if parseErr != nil {
                        return nil, fmt.Errorf("dynamic client registration failed: parse issuer for http client: %w", parseErr)
                }
                metaBuilder := func() *networking.HttpClientBuilder {
                        if config.IssuerTrusted {
                                return networking.NewHostScopedClientBuilder(metaHost.Host, config.AllowPrivateIPs, false)
                        }
                        return networking.NewServerSuppliedHostClientBuilder(metaHost.Host, config.AllowPrivateIPs, false)
                }()
                metaClient, clientErr := metaBuilder.
                        WithDisableKeepAlives(true).
                        Build()
                if clientErr != nil {
                        return nil, fmt.Errorf("dynamic client registration failed: build metadata http client: %w", clientErr)
                }
                metaClient.CheckRedirect = networking.SameHostRedirectPolicy()

                fullMeta, metaErr := oauthproto.FetchAuthorizationServerMetadata(ctx, discoveredDoc.Issuer, metaClient)
                if metaErr != nil && !errors.Is(metaErr, oauthproto.ErrRegistrationEndpointMissing) {
                        return nil, fmt.Errorf("dynamic client registration failed: discover authorization server metadata: %w", metaErr)
                }
                if fullMeta == nil {
                        // Contract violation guard — FetchAuthorizationServerMetadata
                        // guarantees non-nil metadata alongside ErrRegistrationEndpointMissing.
                        return nil, fmt.Errorf("dynamic client registration failed: authorization server metadata unexpectedly nil")
                }
                regEndpoint := fullMeta.RegistrationEndpoint
                if regEndpoint == "" {
                        // Synthesise per nanobot/Hydra convention, mirroring the resolver's
                        // own synthesiseRegistrationEndpoint for the DiscoveryURL branch.
                        u, parseErr := url.Parse(discoveredDoc.Issuer)
                        if parseErr != nil {
                                return nil, fmt.Errorf("dynamic client registration failed: parse issuer for endpoint synthesis: %w", parseErr)
                        }
                        regEndpoint = (&url.URL{
                                Scheme: u.Scheme,
                                Host:   u.Host,
                                Path:   strings.TrimRight(u.Path, "/") + "/register",
                        }).String()
                }
                req.RegistrationEndpoint = regEndpoint
                req.CodeChallengeMethodsSupported = fullMeta.CodeChallengeMethodsSupported
        } else {
                req.RegistrationEndpoint = discoveredDoc.RegistrationEndpoint
        }

        store := dcr.NewInMemoryStore()
        defer func() {
                // Close returns nil today (inMemoryStore.Close is sync.Once-guarded
                // over an in-memory map), but a future change to the underlying
                // MemoryStorage.Close — e.g., timeout-aware cleanup goroutine
                // teardown — could surface a real error. Log it at debug rather
                // than dropping it so a regression is visible without elevating
                // every CLI-flow shutdown to WARN.
                if err := store.Close(); err != nil {
                        slog.Debug("dcr: in-memory store close failed", "error", err)
                }
        }()

        resolution, err := dcr.ResolveCredentials(ctx, req, store)
        if err != nil {
                // Surface the structured error to the boundary log so operators
                // see one slog.Error record with the step / issuer / redirect_uri
                // attributes, then wrap with the CLI-facing prefix.
                dcr.LogStepError(issuer, err)
                return nil, fmt.Errorf("dynamic client registration failed: %w", err)
        }
        return resolution, nil
}

// getDiscoveryDocument retrieves the OIDC discovery document.
//
// The "pre-discovered" short-circuit synthesises a document from
// OAuthFlowConfig fields populated by earlier discovery in the remote
// handler (pkg/auth/remote/handler.go via ValidateAndDiscoverAuthServer).
// On that path the synthesised document carries only the endpoint URLs,
// NOT the server-capability fields (code_challenge_methods_supported,
// token_endpoint_auth_methods_supported, scopes_supported) — those
// remain at their zero values. resolveDCRCredentials handles this by
// re-issuing a discovery fetch through the resolver (see its doc for
// the rationale and trade-off).
func getDiscoveryDocument(
        ctx context.Context,
        issuer string,
        config *OAuthFlowConfig,
) (*oauthproto.OIDCDiscoveryDocument, error) {
        // If we already have the registration endpoint from earlier discovery, use it
        if config.RegistrationEndpoint != "" && config.AuthorizeURL != "" && config.TokenURL != "" {
                slog.Debug("Using pre-discovered OAuth endpoints for dynamic registration")
                return &oauthproto.OIDCDiscoveryDocument{
                        AuthorizationServerMetadata: oauthproto.AuthorizationServerMetadata{
                                Issuer:                issuer,
                                AuthorizationEndpoint: config.AuthorizeURL,
                                TokenEndpoint:         config.TokenURL,
                                RegistrationEndpoint:  config.RegistrationEndpoint,
                        },
                }, nil
        }

        // Fall back to discovering endpoints. This fetch precedes the DCR
        // registration this function's callers are about to perform, so it shares
        // the same private-IP policy (CWE-918) rather than defaulting to unguarded.
        return oauth.DiscoverOIDCEndpoints(ctx, issuer, !config.AllowPrivateIPs)
}

// createOAuthConfig creates the OAuth configuration based on available endpoints
func createOAuthConfig(ctx context.Context, issuer string, config *OAuthFlowConfig) (*oauth.Config, error) {
        // Check if we have OAuth endpoints configured
        if config.AuthorizeURL != "" && config.TokenURL != "" {
                slog.Debug("Using OAuth endpoints",
                        "authorize_url", config.AuthorizeURL, "token_url", config.TokenURL)

                cfg, err := oauth.CreateOAuthConfigManual(
                        config.ClientID,
                        config.ClientSecret,
                        config.AuthorizeURL,
                        config.TokenURL,
                        config.Scopes,
                        true, // Enable PKCE by default for security
                        config.CallbackPort,
                        config.Resource,
                        config.OAuthParams,
                        config.ScopeParamName,
                )
                if err != nil {
                        return nil, err
                }
                // The token URL here is the operator's own, so the flag carries over as-is.
                cfg.TokenEndpointTrusted = config.TokenEndpointTrusted
                return cfg, nil
        }

        // Fall back to OIDC discovery
        slog.Debug("Using OIDC discovery")
        cfg, err := oauth.CreateOAuthConfigFromOIDC(
                ctx,
                issuer,
                config.ClientID,
                config.ClientSecret,
                config.Scopes,
                true, // Enable PKCE by default for security
                config.CallbackPort,
                config.Resource,
                !config.AllowPrivateIPs,
        )
        if err != nil {
                return nil, err
        }
        cfg.ScopeParamName = config.ScopeParamName
        // cfg.TokenURL came out of the issuer's discovery document. The issuer naming
        // its own authority stays within the operator's decision to configure it;
        // naming any other authority does not (CWE-918).
        cfg.TokenEndpointTrusted = config.TokenEndpointTrusted &&
                networking.AuthorityMatchesAny(cfg.TokenURL, config.TokenURL, issuer)
        return cfg, nil
}

func newOAuthFlow(ctx context.Context, oauthConfig *oauth.Config, config *OAuthFlowConfig) (*OAuthFlowResult, error) {
        flow, err := oauth.NewFlow(oauthConfig)
        if err != nil {
                return nil, fmt.Errorf("failed to create OAuth flow: %w", err)
        }

        // Create a context with timeout for the OAuth flow
        oauthTimeout := config.Timeout
        if oauthTimeout <= 0 {
                oauthTimeout = DefaultOAuthTimeout
        }

        oauthCtx, cancel := context.WithTimeout(ctx, oauthTimeout)
        defer cancel()

        // Start OAuth flow
        tokenResult, err := flow.Start(oauthCtx, config.SkipBrowser)
        if err != nil {
                if errors.Is(oauthCtx.Err(), context.DeadlineExceeded) {
                        return nil, fmt.Errorf("OAuth flow timed out after %v - user did not complete authentication", oauthTimeout)
                }
                return nil, fmt.Errorf("OAuth flow failed: %w", err)
        }

        slog.Debug("OAuth authentication successful")

        // Log token info (without exposing the actual token)
        if tokenResult.Claims != nil {
                if sub, ok := tokenResult.Claims["sub"].(string); ok {
                        slog.Debug("Authenticated as subject", "sub", sub)
                }
                if email, ok := tokenResult.Claims["email"].(string); ok {
                        slog.Debug("Authenticated with email", "email", email)
                }
        }

        source := flow.TokenSource()
        return buildOAuthFlowResult(source, oauthConfig, tokenResult, config), nil
}

func buildOAuthFlowResult(
        tokenSource oauth2.TokenSource,
        oauthConfig *oauth.Config,
        tokenResult *oauth.TokenResult,
        config *OAuthFlowConfig,
) *OAuthFlowResult {
        return &OAuthFlowResult{
                TokenSource:  tokenSource,
                Config:       oauthConfig,
                AccessToken:  tokenResult.AccessToken,
                RefreshToken: tokenResult.RefreshToken,
                Expiry:       tokenResult.Expiry,
                ClientID:     oauthConfig.ClientID,
                ClientSecret: oauthConfig.ClientSecret,
                // DCR renewal metadata — populated only when dynamic registration was performed.
                SecretExpiry:            config.SecretExpiry,
                RegistrationAccessToken: config.RegistrationAccessToken,
                RegistrationClientURI:   config.RegistrationClientURI,
                TokenEndpointAuthMethod: config.TokenEndpointAuthMethod,
                RegisteredCallbackPort:  config.RegisteredCallbackPort,
        }
}

// newResourceMetadataTransport builds the transport for RFC 9728 metadata
// fetches. The metadataURL is server-supplied, so the caller pairs it with a
// same-host redirect policy; when blockPrivateIPs is true it also refuses to
// dial private/loopback/link-local addresses on every hop (and disables
// keep-alive so a pooled connection cannot skip the per-dial check).
// Extracted so the pool bounds networking.SetIdleConnBounds applies stay
// unit-testable.
func newResourceMetadataTransport(blockPrivateIPs bool) *http.Transport {
        transport := &http.Transport{
                TLSHandshakeTimeout:   5 * time.Second,
                ResponseHeaderTimeout: 5 * time.Second,
        }
        networking.SetIdleConnBounds(transport)
        if blockPrivateIPs {
                transport.DialContext = networking.NewPrivateIPBlockingDialContext()
                transport.DisableKeepAlives = true
        }
        return transport
}

// FetchResourceMetadata fetches RFC 9728 protected-resource metadata from a
// server-supplied URL.
//
// The metadataURL originates from untrusted remote-server discovery (the
// WWW-Authenticate resource_metadata parameter), so the client refuses
// cross-host and HTTPS->HTTP redirects (CWE-918). When blockPrivateIPs is true
// it additionally refuses to dial private/loopback/link-local addresses on
// every hop. Callers set blockPrivateIPs when the operator-configured target is
// public; when the operator deliberately targets an internal network it is
// false so legitimately-internal auth metadata stays reachable.
func FetchResourceMetadata(ctx context.Context, metadataURL string, blockPrivateIPs bool) (*auth.RFC9728AuthInfo, error) {
        if metadataURL == "" {
                return nil, fmt.Errorf("metadata URL is empty")
        }

        // Validate URL
        parsedURL, err := url.Parse(metadataURL)
        if err != nil {
                return nil, fmt.Errorf("invalid metadata URL: %w", err)
        }

        // RFC 9728: Must use HTTPS (except for localhost in development)
        if parsedURL.Scheme != "https" && parsedURL.Hostname() != "localhost" && parsedURL.Hostname() != "127.0.0.1" {
                return nil, fmt.Errorf("metadata URL must use HTTPS: %s", metadataURL)
        }

        transport := newResourceMetadataTransport(blockPrivateIPs)
        client := &http.Client{
                Timeout:       DefaultHTTPTimeout,
                Transport:     transport,
                CheckRedirect: networking.SameHostRedirectPolicy(),
        }

        req, err := http.NewRequestWithContext(ctx, http.MethodGet, metadataURL, nil)
        if err != nil {
                return nil, fmt.Errorf("failed to create request: %w", err)
        }

        req.Header.Set("Accept", "application/json")

        // #nosec G704 -- metadataURL is server-controlled; the client refuses cross-host and
        // HTTPS->HTTP redirects (networking.SameHostRedirectPolicy) to contain SSRF.
        resp, err := client.Do(req)
        if err != nil {
                return nil, fmt.Errorf("failed to fetch metadata: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("Failed to close response body", "error", err)
                }
        }()

        if resp.StatusCode != http.StatusOK {
                return nil, fmt.Errorf("metadata request failed with status %d", resp.StatusCode)
        }

        // Check content type
        contentType := strings.ToLower(resp.Header.Get("Content-Type"))
        if !strings.Contains(contentType, "application/json") {
                return nil, fmt.Errorf("unexpected content type: %s", contentType)
        }

        // Parse the metadata
        const maxResponseSize = 1024 * 1024 // 1MB limit
        var metadata auth.RFC9728AuthInfo
        if err := json.NewDecoder(io.LimitReader(resp.Body, maxResponseSize)).Decode(&metadata); err != nil {
                return nil, fmt.Errorf("failed to parse metadata: %w", err)
        }

        // RFC 9728 Section 3.3: Validate that the resource value matches
        // For now we just check it's not empty
        if metadata.Resource == "" {
                return nil, fmt.Errorf("metadata missing required 'resource' field")
        }

        return &metadata, nil
}

// ValidateAndDiscoverAuthServer attempts to validate if a URL is an authorization server
// and discover its actual issuer by fetching its metadata.
// This handles the case where the URL used to fetch metadata differs from the actual issuer
// (e.g., Stripe's case where https://mcp.stripe.com hosts metadata for https://marketplace.stripe.com)
//
// When blockPrivateIPs is true the underlying discovery refuses to dial
// private/loopback/link-local addresses (SSRF guard for server-supplied issuer
// URLs); callers pass false when the issuer is operator-configured or the
// configured target is itself internal.
func ValidateAndDiscoverAuthServer(ctx context.Context, potentialIssuer string, blockPrivateIPs bool) (*AuthServerInfo, error) {
        // Use DiscoverActualIssuer which doesn't validate issuer match
        // This allows us to discover the real issuer even when it differs from the metadata URL
        doc, err := oauth.DiscoverActualIssuer(ctx, potentialIssuer, blockPrivateIPs)
        if err == nil && doc != nil && doc.Issuer != "" {
                // Found valid authorization server metadata, return the actual issuer and endpoints
                if doc.Issuer != potentialIssuer {
                        slog.Debug("Discovered actual issuer", "issuer", doc.Issuer, "metadata_url", potentialIssuer)
                } else {
                        slog.Debug("Validated authorization server", "issuer", potentialIssuer)
                }

                return &AuthServerInfo{
                        Issuer:                            doc.Issuer,
                        AuthorizationURL:                  doc.AuthorizationEndpoint,
                        TokenURL:                          doc.TokenEndpoint,
                        RegistrationEndpoint:              doc.RegistrationEndpoint,
                        ClientIDMetadataDocumentSupported: doc.ClientIDMetadataDocumentSupported,
                }, nil
        }

        // If that fails, the URL might not be a valid authorization server
        return nil, fmt.Errorf("could not validate %s as an authorization server: %w", potentialIssuer, err)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

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

        "github.com/golang-jwt/jwt/v5"
        "golang.org/x/time/rate"

        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// GitHubTokenCheckURL is the base URL pattern for GitHub OAuth token validation
//
//nolint:gosec // This is a URL pattern, not a credential
const GitHubTokenCheckURL = "api.github.com/applications"

// GitHubProvider implements token introspection for GitHub.com's OAuth token validation API
// GitHub uses a non-standard token validation endpoint that differs from RFC 7662
// Endpoint: POST /applications/{client_id}/token
// Auth: Basic (client_id:client_secret)
// Body: {"access_token": "gho_..."}
//
// Note: This provider is designed for GitHub.com only, not GitHub Enterprise Server
type GitHubProvider struct {
        client       *http.Client
        clientID     string
        clientSecret string
        baseURL      string
        rateLimiter  *rate.Limiter
}

// NewGitHubProvider creates a new GitHub token introspection provider
// Parameters:
//   - introspectURL: GitHub token validation endpoint (must be api.github.com with HTTPS)
//   - clientID: OAuth App client ID
//   - clientSecret: OAuth App client secret
//   - caCertPath: Path to CA certificate bundle (optional)
//   - allowPrivateIP: Allow private IP addresses (should be false for production)
func NewGitHubProvider(
        introspectURL, clientID, clientSecret, caCertPath string,
        allowPrivateIP bool,
) (*GitHubProvider, error) {
        return newGitHubProviderWithClient(introspectURL, clientID, clientSecret, caCertPath, allowPrivateIP, nil)
}

// newGitHubProviderWithClient creates a new GitHub provider with custom client (for testing)
func newGitHubProviderWithClient(
        introspectURL, clientID, clientSecret, caCertPath string,
        allowPrivateIP bool,
        customClient *http.Client,
) (*GitHubProvider, error) {
        var client *http.Client
        var err error

        if customClient != nil {
                // Use provided client (for testing)
                client = customClient
        } else {
                // Create secured HTTP client
                // Note: insecureAllowHTTP is always false for GitHub.com (requires HTTPS)
                client, err = networking.NewHttpClientBuilder().
                        WithCABundle(caCertPath).
                        WithPrivateIPs(allowPrivateIP).
                        Build()
                if err != nil {
                        return nil, fmt.Errorf("failed to create HTTP client: %w", err)
                }
        }

        // Create rate limiter: 100 requests per second with burst of 200
        // GitHub API allows 5,000 requests/hour, but we rate limit locally to prevent abuse
        limiter := rate.NewLimiter(100, 200)

        return &GitHubProvider{
                client:       client,
                clientID:     clientID,
                clientSecret: clientSecret,
                baseURL:      introspectURL,
                rateLimiter:  limiter,
        }, nil
}

// Name returns the provider name
func (*GitHubProvider) Name() string {
        return "github"
}

// CanHandle returns true if this provider can handle the given introspection URL
// This validates that the URL is a legitimate GitHub.com token validation endpoint
// Note: GitHub Enterprise Server is NOT supported - use corporate IdP instead
func (*GitHubProvider) CanHandle(introspectURL string) bool {
        // Parse URL to validate structure
        u, err := url.Parse(introspectURL)
        if err != nil {
                return false
        }

        // Validate scheme (must be HTTPS)
        if u.Scheme != "https" {
                return false
        }

        // Validate host - must be exactly api.github.com (GitHub.com only, no enterprise)
        if u.Host != "api.github.com" {
                return false
        }

        // Validate path structure: /applications/{client_id}/token
        path := u.Path
        return strings.Contains(path, "/applications/") && strings.HasSuffix(path, "/token")
}

// IntrospectToken introspects a GitHub OAuth token and returns JWT claims
// This calls GitHub's token validation API to verify the token and extract user information
func (g *GitHubProvider) IntrospectToken(ctx context.Context, token string) (jwt.MapClaims, error) {
        //nolint:gosec // G706 - baseURL is a configured GitHub API endpoint
        slog.Debug("using GitHub token validation provider", "url", g.baseURL)

        // Apply rate limiting to prevent DoS and respect GitHub API limits
        if err := g.rateLimiter.Wait(ctx); err != nil {
                return nil, fmt.Errorf("rate limit wait failed: %w", err)
        }

        // Create request body with the access token
        reqBody := map[string]string{"access_token": token}
        bodyBytes, err := json.Marshal(reqBody)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal request body: %w", err)
        }

        // Create POST request
        //nolint:gosec // G704 - URL is configured GitHub API endpoint
        req, err := http.NewRequestWithContext(ctx, "POST", g.baseURL, bytes.NewReader(bodyBytes))
        if err != nil {
                return nil, fmt.Errorf("failed to create GitHub validation request: %w", err)
        }

        // Set headers
        req.Header.Set("Content-Type", "application/json")
        req.Header.Set("Accept", "application/json")
        req.Header.Set("User-Agent", oauthproto.UserAgent)

        // GitHub requires Basic Auth with OAuth App credentials
        req.SetBasicAuth(g.clientID, g.clientSecret)

        // Make the request
        resp, err := g.client.Do(req) // #nosec G704 -- URL is the configured GitHub API endpoint
        if err != nil {
                return nil, fmt.Errorf("github validation request failed: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        // Read the response with a reasonable limit to prevent DoS attacks
        const maxResponseSize = 64 * 1024 // 64KB should be more than enough
        limitedReader := io.LimitReader(resp.Body, maxResponseSize)
        body, err := io.ReadAll(limitedReader)
        if err != nil {
                return nil, fmt.Errorf("failed to read GitHub validation response: %w", err)
        }

        // Check for HTTP errors
        if resp.StatusCode == http.StatusNotFound {
                // 404 means token is invalid or doesn't belong to this OAuth App
                return nil, ErrInvalidToken
        }
        if resp.StatusCode == http.StatusTooManyRequests {
                // 429 means we've hit GitHub's rate limit
                retryAfter := resp.Header.Get("Retry-After")
                remaining := resp.Header.Get("X-RateLimit-Remaining")
                reset := resp.Header.Get("X-RateLimit-Reset")
                //nolint:gosec // G706: rate limit headers are public HTTP metadata
                slog.Warn("github rate limit exceeded",
                        "retry_after", retryAfter, "remaining", remaining, "reset", reset)
                return nil, fmt.Errorf("github rate limit exceeded, retry after: %s", retryAfter)
        }
        if resp.StatusCode != http.StatusOK {
                return nil, fmt.Errorf("github validation failed with status %d: %s", resp.StatusCode, string(body))
        }

        // Parse the GitHub response and convert to JWT claims
        //nolint:gosec // G706: HTTP status code is not sensitive
        slog.Debug("successfully validated GitHub token", "status", resp.StatusCode)
        return g.parseGitHubResponse(body)
}

// parseGitHubResponse parses GitHub's token validation response and converts it to JWT claims
func (*GitHubProvider) parseGitHubResponse(body []byte) (jwt.MapClaims, error) {
        // Parse GitHub's response format
        // Reference: https://docs.github.com/en/rest/apps/oauth-applications#check-a-token
        var githubResp struct {
                ID    int64  `json:"id"`
                Token string `json:"token"`
                User  struct {
                        Login     string `json:"login"`
                        ID        int64  `json:"id"`
                        NodeID    string `json:"node_id"`
                        Email     string `json:"email"`
                        Name      string `json:"name"`
                        Type      string `json:"type"`
                        SiteAdmin bool   `json:"site_admin"`
                } `json:"user"`
                Scopes    []string `json:"scopes"`
                CreatedAt string   `json:"created_at"`
                UpdatedAt string   `json:"updated_at"`
                App       struct {
                        Name     string `json:"name"`
                        URL      string `json:"url"`
                        ClientID string `json:"client_id"`
                } `json:"app"`
        }

        if err := json.Unmarshal(body, &githubResp); err != nil {
                return nil, fmt.Errorf("failed to decode GitHub response: %w", err)
        }

        // Convert to JWT MapClaims format
        claims := jwt.MapClaims{
                "iss": "https://github.com", // Fixed issuer for GitHub
                "aud": "https://github.com", // Use issuer as audience
                // Mark token as active (consistent with RFC 7662 behavior)
                "active": true,
        }

        // Subject (sub) - use GitHub user ID as the unique identifier
        if githubResp.User.ID != 0 {
                claims["sub"] = fmt.Sprintf("%d", githubResp.User.ID)
        } else {
                return nil, fmt.Errorf("missing user ID in GitHub response")
        }

        // User information
        if githubResp.User.Login != "" {
                claims["preferred_username"] = githubResp.User.Login
                claims["login"] = githubResp.User.Login // GitHub-specific claim
        }
        if githubResp.User.Email != "" {
                claims["email"] = githubResp.User.Email
        }
        if githubResp.User.Name != "" {
                claims["name"] = githubResp.User.Name
        }

        // Parse created_at for iat (issued at) claim
        if githubResp.CreatedAt != "" {
                if t, err := time.Parse(time.RFC3339, githubResp.CreatedAt); err == nil {
                        claims["iat"] = float64(t.Unix())
                }
        }

        // Add scopes - GitHub returns them as an array
        if len(githubResp.Scopes) > 0 {
                claims["scopes"] = githubResp.Scopes
                // Also add as space-separated string for compatibility
                claims["scope"] = strings.Join(githubResp.Scopes, " ")
        }

        // GitHub-specific claims for advanced policies
        if githubResp.User.Type != "" {
                claims["user_type"] = githubResp.User.Type
        }
        if githubResp.User.SiteAdmin {
                claims["site_admin"] = true
        }
        if githubResp.App.Name != "" {
                claims["app_name"] = githubResp.App.Name
        }

        // Note: GitHub OAuth tokens don't have a standard expiration
        // They remain valid until revoked by the user or the app
        // We rely on the introspection call to validate token freshness

        return claims, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

import (
        "encoding/json"
        "fmt"

        "github.com/stacklok/toolhive-core/audit"
        upstreamtoken "github.com/stacklok/toolhive/pkg/auth/upstreamtoken"
)

// PrincipalInfo contains the non-sensitive identity fields safe for external consumption.
// This is the canonical projection of Identity for webhook payloads, audit logs, and
// any context where credentials must not appear — not even in redacted form.
//
// Identity embeds this type, so fields are accessible directly on Identity
// (e.g., identity.Subject, identity.Email) while keeping the credential-free
// subset available as a first-class type for external APIs.
type PrincipalInfo struct {
        // Subject is the unique identifier for the principal (from 'sub' claim).
        // This is always required per OIDC Core 1.0 spec § 5.1.
        Subject string `json:"sub,omitempty"`

        // PlatformUserID is the platform's canonical user identifier.
        //
        // It defaults to the sub claim (see claimsToIdentity). That default is
        // correct for standalone OSS use (where sub IS the canonical local User.ID)
        // and for AS-bearer paths under an enterprise directory binding (where the
        // AS-minted sub equals the directory user_id). It is INCORRECT for any
        // middleware that validates a JWT issued by a different IdP whose sub is not
        // the platform-canonical user identifier (e.g. a corporate IdP whose sub
        // rotates per-application). Such middleware MUST override PlatformUserID
        // (typically via a resolution call into a directory service) before calling
        // WithIdentity. On request-serving paths storage reads this value via
        // CanonicalUserFromContext (which falls back to this field when no dedicated
        // platform-user key is set); an unresolved PlatformUserID from a corporate-IdP
        // bearer will silently mis-key writes.
        //
        // Only claimsToIdentity populates this field today. Other Identity
        // constructors in this repo (local.go, anonymous.go) intentionally leave it
        // unset for now: standalone OSS keys upstream-token storage on the session,
        // not PlatformUserID, so those paths have no canonical-user reader to
        // satisfy. Populating them is deferred until a storage layer that reads
        // PlatformUserID on those paths exists.
        PlatformUserID string `json:"platform_user_id,omitempty"`

        // Name is the human-readable name (from 'name' claim).
        Name string `json:"name,omitempty"`

        // Email is the email address (from 'email' claim, if available).
        Email string `json:"email,omitempty"`

        // Groups are the groups this identity belongs to.
        //
        // NOTE: This field is intentionally NOT populated by authentication middleware.
        // Authorization logic MUST extract groups from the Claims map, as group claim
        // names vary by provider (e.g., "groups", "roles", "cognito:groups").
        Groups []string `json:"groups,omitempty"`

        // Claims contains additional claims from the auth token.
        // This preserves all JWT claims for authorization policies.
        Claims map[string]any `json:"claims,omitempty"`

        // DelegationChain is the parsed RFC 8693 "act" claim: the ordered list of
        // parties that acted on behalf of this identity's subject (outermost/most
        // recent first). Nil when the token carries no act claim. The "delegation"
        // tag matches the audit-event wire key so webhook payloads and audit logs
        // share one spelling of the same structure.
        DelegationChain *audit.DelegationChain `json:"delegation,omitempty"`
}

// Identity represents an authenticated user or service account.
// This is the primary type for representing authenticated principals throughout ToolHive.
//
// It embeds PrincipalInfo (the credential-free subset) and adds sensitive fields
// (Token, TokenType) and internal metadata that must never be externalized.
//
// # Field-completeness contract
//
// A non-nil *Identity is always COMPLETE: all credential fields that are relevant
// for its authentication path are populated by the constructor that created it.
// Anonymous and "no principal known" states are represented by nil, not by a
// struct with empty credential fields. Code that receives a non-nil *Identity is
// entitled to assume it is fully initialized for its session-binding semantics.
//
// In particular, do NOT construct &Identity{Subject: …} with Token and
// UpstreamTokens unset as a substitute for nil when no live bearer token is
// available (e.g. during session restore). Pass nil instead and let downstream
// consumers read the identity from the per-request context, where
// TokenValidator.Middleware places a fully-populated identity on every incoming
// request.
type Identity struct {
        PrincipalInfo

        // Token is the original authentication token (for pass-through scenarios).
        // This is redacted in String() and MarshalJSON() to prevent leakage.
        Token string

        // TokenType is the type of token (e.g., "Bearer", "JWT").
        TokenType string

        // Metadata stores additional identity information.
        Metadata map[string]string

        // SessionlessGrant is true only when TokenValidator validated this token for
        // the configured embedded authorization-server issuer, no upstream session
        // credentials were loaded, and the token affirmatively proves that it was
        // minted by a session-less grant. It is internal runtime provenance and is
        // intentionally not included in PrincipalInfo or MarshalJSON.
        SessionlessGrant bool

        // UpstreamTokens maps upstream provider names to their access tokens.
        // This is populated by the auth middleware when an embedded auth server
        // is active and the JWT contains a token session ID (tsid claim).
        // Redacted in MarshalJSON() to prevent token leakage.
        //
        // State semantics:
        //   - nil            — no upstream credentials were loaded for this identity.
        //                      Either the incoming JWT carried no tsid claim, or the
        //                      Identity came from a constructor that has no upstream
        //                      session at all: local.go, anonymous.go, or middleware
        //                      validating a JWT from an IdP other than the embedded
        //                      auth server (see the PlatformUserID note above).
        //   - empty map      — never produced by TokenValidator, which answers 401
        //                      when a tsid session holds no usable access token;
        //                      a consumer that meets one denies it like a session
        //                      lacking the pinned provider's credential.
        //   - populated map  — keys are upstream provider names; values are the
        //                      stored access token strings.
        //
        // A nil map therefore says only that no credential was loaded, never why.
        // Authorization code must use SessionlessGrant, which TokenValidator sets
        // only for a validated embedded-auth-server token that affirmatively proves
        // it was minted without an upstream session.
        //
        // MUST NOT be mutated after the Identity is placed in the publicly-reachable
        // request context. It MAY be mutated while the Identity is reachable only via
        // a load-scoped ctx, provided the loader does not share that ctx with
        // concurrent code. See TokenValidator.Middleware for the canonical pattern.
        UpstreamTokens map[string]string

        // UpstreamIDTokens maps upstream provider names to their ID tokens.
        // This is populated by the auth middleware when an embedded auth server
        // is active and the JWT contains a token session ID (tsid claim).
        // Each value is the rotated ID token when a refresh produced one
        // (OIDC Core 1.0 §12.2), otherwise the original JWT captured at the initial
        // OIDC login.
        //
        // State semantics:
        //   - nil            — no tsid claim was present on the incoming JWT
        //                      (middleware did not attempt to load credentials).
        //   - empty map      — tsid claim was valid but no providers had an
        //                      ID token stored for the session.
        //   - populated map  — keys are upstream provider names; values are the
        //                      stored ID token JWTs, which may be expired.
        //
        // Whether `exp` matters depends on what the caller does with the token. Nothing
        // in this package enforces it; the two current consumers use the token
        // differently and neither validates it locally:
        //
        //   - PRESENTING it as a credential — e.g. as the subject_token of an RFC 8693
        //     exchange (pkg/vmcp/auth/strategies) — is subject to expiry, because the
        //     party receiving it will reject an expired assertion. That consumer
        //     deliberately does not pre-check `exp` either: it lets the IdP reject the
        //     exchange and surfaces the resulting `invalid_grant`. So such a caller must
        //     either check `exp` first or be prepared to handle that failure.
        //   - READING identity claims out of it — as the Cedar authorizer does for the
        //     profile claims an upstream access token omits — is not subject to expiry at
        //     all. The token is evidence that the upstream asserted something at login,
        //     and expiry says nothing about whether that assertion happened. Rejecting it
        //     once expired would flip an authorization decision mid-session, since
        //     sessions outlive ID tokens by days; and the claims are no staler than the
        //     alternative, since the profile claims mirrored into the ToolHive-issued
        //     token are likewise captured once at login and never refreshed.
        //
        // Either way the token is NOT re-validated for signature or freshness here, so
        // treat its claims as facts the upstream asserted when it issued the token
        // rather than current state.
        //
        // When populated by upstreamtoken.InProcessService (see NewInProcessService),
        // an expired ID token triggers a refresh even while the access token is still
        // valid, so a provider that rotates the ID token on refresh yields a currently
        // valid one here. It can still be expired when the provider omits id_token on
        // refresh (OIDC Core 1.0 §12.2) and the old one is carried forward, when that
        // refresh failed and the stored token was returned, or when the row has no
        // refresh token; a consumer that needs a fresh token must check `exp` itself.
        //
        // Redacted in MarshalJSON() to prevent token leakage.
        // MUST NOT be mutated after the Identity is placed in the request context.
        UpstreamIDTokens map[string]string
}

// String returns a string representation of the Identity with sensitive fields redacted.
// This prevents accidental token leakage when the Identity is logged or printed.
func (i *Identity) String() string {
        if i == nil {
                return "<nil>"
        }

        return fmt.Sprintf("Identity{Subject:%q}", i.Subject)
}

// MarshalJSON implements json.Marshaler to redact sensitive fields during JSON serialization.
// This prevents accidental token leakage in structured logs, API responses, or audit logs.
func (i *Identity) MarshalJSON() ([]byte, error) {
        if i == nil {
                return []byte("null"), nil
        }

        // Create a safe representation with lowercase field names and redacted token
        type SafeIdentity struct {
                Subject          string                 `json:"subject"`
                PlatformUserID   string                 `json:"platformUserId,omitempty"`
                Name             string                 `json:"name"`
                Email            string                 `json:"email"`
                Groups           []string               `json:"groups"`
                Claims           map[string]any         `json:"claims"`
                DelegationChain  *audit.DelegationChain `json:"delegation,omitempty"`
                Token            string                 `json:"token"`
                TokenType        string                 `json:"tokenType"`
                Metadata         map[string]string      `json:"metadata"`
                UpstreamTokens   map[string]string      `json:"upstreamTokens,omitempty"`
                UpstreamIDTokens map[string]string      `json:"upstreamIDTokens,omitempty"`
        }

        const redacted = "REDACTED"

        token := i.Token
        if token != "" {
                token = redacted
        }

        claims := i.Claims
        if _, hasTsid := claims[upstreamtoken.TokenSessionIDClaimKey]; hasTsid {
                claims = make(map[string]any, len(i.Claims))
                for k, v := range i.Claims {
                        if k != upstreamtoken.TokenSessionIDClaimKey {
                                claims[k] = v
                        }
                }
        }

        // Redact upstream tokens: preserve keys, replace non-empty values
        var redactedUpstreamTokens map[string]string
        // Guard with len() > 0 (not != nil) so that both nil and empty maps
        // produce a nil redactedUpstreamTokens, which omitempty then omits.
        if len(i.UpstreamTokens) > 0 {
                redactedUpstreamTokens = make(map[string]string, len(i.UpstreamTokens))
                for k, v := range i.UpstreamTokens {
                        if v != "" {
                                redactedUpstreamTokens[k] = redacted
                        } else {
                                redactedUpstreamTokens[k] = ""
                        }
                }
        }

        // Redact upstream ID tokens with the same pattern as access tokens.
        var redactedUpstreamIDTokens map[string]string
        if len(i.UpstreamIDTokens) > 0 {
                redactedUpstreamIDTokens = make(map[string]string, len(i.UpstreamIDTokens))
                for k, v := range i.UpstreamIDTokens {
                        if v != "" {
                                redactedUpstreamIDTokens[k] = redacted
                        } else {
                                redactedUpstreamIDTokens[k] = ""
                        }
                }
        }

        return json.Marshal(&SafeIdentity{
                Subject:          i.Subject,
                PlatformUserID:   i.PlatformUserID,
                Name:             i.Name,
                Email:            i.Email,
                Groups:           i.Groups,
                Claims:           claims,
                DelegationChain:  i.DelegationChain,
                Token:            token,
                TokenType:        i.TokenType,
                Metadata:         i.Metadata,
                UpstreamTokens:   redactedUpstreamTokens,
                UpstreamIDTokens: redactedUpstreamIDTokens,
        })
}

// GetPrincipalInfo returns a copy of the credential-free PrincipalInfo suitable
// for external consumption (webhook payloads, audit logs, etc.).
// Token, TokenType, and Metadata are structurally excluded.
func (i *Identity) GetPrincipalInfo() *PrincipalInfo {
        if i == nil {
                return nil
        }

        pi := i.PrincipalInfo
        return &pi
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

import (
        "net/http"
        "time"

        "github.com/golang-jwt/jwt/v5"
)

// LocalUserMiddleware creates an HTTP middleware that sets up local user identity.
// This allows specifying a local username while still bypassing authentication.
//
// This middleware is useful for development and testing scenarios where you want
// to simulate a specific user without going through the full authentication flow.
// Like AnonymousMiddleware, this is heavily discouraged in production settings.
func LocalUserMiddleware(username string) func(http.Handler) http.Handler {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // Create local user claims with the specified username
                        claims := jwt.MapClaims{
                                "sub":   username,
                                "iss":   "toolhive-local",
                                "aud":   "toolhive",
                                "exp":   time.Now().Add(24 * time.Hour).Unix(), // Valid for 24 hours
                                "iat":   time.Now().Unix(),
                                "nbf":   time.Now().Unix(),
                                "email": username + "@localhost",
                                "name":  "Local User: " + username,
                        }

                        // Create Identity from claims
                        identity := &Identity{
                                PrincipalInfo: PrincipalInfo{
                                        Subject: username,
                                        Name:    "Local User: " + username,
                                        Email:   username + "@localhost",
                                        Claims:  claims,
                                },
                                Token:     "", // No token for local auth
                                TokenType: "Bearer",
                        }

                        // Add the Identity to the request context
                        ctx := WithIdentity(r.Context(), identity)
                        next.ServeHTTP(w, r.WithContext(ctx))
                })
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package auth

import (
        "context"
        "encoding/json"
        "fmt"
        "net/http"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// Middleware type constant
const (
        MiddlewareType = "auth"
)

// MiddlewareParams represents the parameters for authentication middleware
type MiddlewareParams struct {
        OIDCConfig *TokenValidatorConfig `json:"oidc_config,omitempty"`
        // EmbeddedAuthServerIssuer binds session-less-grant provenance to the local
        // embedded authorization server. It is derived from runner configuration,
        // not supplied by an incoming token.
        EmbeddedAuthServerIssuer string `json:"embedded_auth_server_issuer,omitempty"`
}

// Middleware wraps authentication middleware functionality
type Middleware struct {
        middleware      types.MiddlewareFunction
        authInfoHandler http.Handler
}

// Handler returns the middleware function used by the proxy.
func (m *Middleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (*Middleware) Close() error {
        // Auth middleware doesn't need cleanup
        return nil
}

// AuthInfoHandler returns the authentication info handler.
func (m *Middleware) AuthInfoHandler() http.Handler {
        return m.authInfoHandler
}

// CreateMiddleware factory function for authentication middleware
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {

        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal auth middleware parameters: %w", err)
        }

        var opts []TokenValidatorOption
        if reader := runner.GetUpstreamTokenReader(); reader != nil {
                opts = append(opts, WithUpstreamTokenReader(reader))
        }
        if provider := runner.GetKeyProvider(); provider != nil {
                opts = append(opts, WithKeyProvider(provider))
        }
        if params.EmbeddedAuthServerIssuer != "" {
                opts = append(opts, WithEmbeddedAuthServerIssuer(params.EmbeddedAuthServerIssuer))
        }

        middleware, authInfoHandler, err := GetAuthenticationMiddleware(context.Background(), params.OIDCConfig, opts...)
        if err != nil {
                return fmt.Errorf("failed to create authentication middleware: %w", err)
        }

        authMw := &Middleware{
                middleware:      middleware,
                authInfoHandler: authInfoHandler,
        }

        // Add middleware to runner
        runner.AddMiddleware(config.Type, authMw)

        // Set auth info handler if present
        if authInfoHandler != nil {
                runner.SetAuthInfoHandler(authInfoHandler)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package auth

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "net"
        "os"
        "strconv"
        "sync"
        "syscall"
        "time"

        "github.com/cenkalti/backoff/v7"
        "golang.org/x/oauth2"
        "golang.org/x/sync/singleflight"

        "github.com/stacklok/toolhive/pkg/auth/oautherr"
        "github.com/stacklok/toolhive/pkg/container/runtime"
)

const (
        // tokenRefreshInitialRetryInterval is the default starting interval for
        // exponential backoff when a token refresh fails during background monitoring.
        // Override with TOOLHIVE_TOKEN_REFRESH_INITIAL_RETRY_INTERVAL (e.g. "10s", "1m").
        tokenRefreshInitialRetryInterval = 10 * time.Second
        // tokenRefreshMaxRetryInterval is the default cap on the exponential growth
        // of the retry interval.
        // Override with TOOLHIVE_TOKEN_REFRESH_MAX_RETRY_INTERVAL (e.g. "2m", "10m").
        tokenRefreshMaxRetryInterval = 2 * time.Minute
        // tokenRefreshMaxTries is the default maximum number of retry attempts.
        // Override with TOOLHIVE_TOKEN_REFRESH_MAX_TRIES (e.g. "10").
        tokenRefreshMaxTries = 5
        // tokenRefreshMaxElapsedTime is the default maximum elapsed time for all retry attempts.
        // Override with TOOLHIVE_TOKEN_REFRESH_MAX_ELAPSED_TIME (e.g. "10m").
        tokenRefreshMaxElapsedTime = 5 * time.Minute
        // authRetryingTickInterval is the default cadence between background refresh
        // attempts once the short retry has exhausted on a transient error. The
        // short retry inside transientRefresher.retry already performs exponential
        // backoff; this cross-tick layer uses a fixed cadence to avoid layering
        // exponential on top of exponential and to give operators a predictable
        // retry signature when investigating long outages.
        // Override with TOOLHIVE_TOKEN_AUTH_RETRYING_TICK_INTERVAL (e.g. "5m", "30m").
        authRetryingTickInterval = 10 * time.Minute
        // authRetryingMaxElapsed is the ceiling on how long the monitor will stay
        // in the AuthRetrying transient-failure window before giving up and
        // marking the workload unauthenticated.
        // Override with TOOLHIVE_TOKEN_AUTH_RETRYING_MAX_ELAPSED (e.g. "12h", "48h").
        authRetryingMaxElapsed = 24 * time.Hour
)

const (
        // #nosec G101 — not credentials, just initial retry interval
        tokenRefreshInitialRetryIntervalEnv = "TOOLHIVE_TOKEN_REFRESH_INITIAL_RETRY_INTERVAL"
        // #nosec G101 — not credentials, just max retry interval
        tokenRefreshMaxRetryIntervalEnv = "TOOLHIVE_TOKEN_REFRESH_MAX_RETRY_INTERVAL"
        // #nosec G101 — not credentials, just max elapsed time
        tokenRefreshMaxElapsedTimeEnv = "TOOLHIVE_TOKEN_REFRESH_MAX_ELAPSED_TIME"
        // #nosec G101 — not credentials, just max tries
        tokenRefreshMaxTriesEnv = "TOOLHIVE_TOKEN_REFRESH_MAX_TRIES"
        // #nosec G101 — not credentials, just retry tick interval
        authRetryingTickIntervalEnv = "TOOLHIVE_TOKEN_AUTH_RETRYING_TICK_INTERVAL"
        // #nosec G101 — not credentials, just max elapsed time
        authRetryingMaxElapsedEnv = "TOOLHIVE_TOKEN_AUTH_RETRYING_MAX_ELAPSED"
)

// resolveTokenRefreshInitialRetryInterval returns the initial retry interval for
// token refresh backoff, reading from TOOLHIVE_TOKEN_REFRESH_INITIAL_RETRY_INTERVAL
// if set, otherwise returning the default.
func resolveTokenRefreshInitialRetryInterval() time.Duration {
        return resolveDurationEnv(
                tokenRefreshInitialRetryIntervalEnv,
                tokenRefreshInitialRetryInterval,
        )
}

// resolveTokenRefreshMaxRetryInterval returns the max retry interval for token
// refresh backoff, reading from TOOLHIVE_TOKEN_REFRESH_MAX_RETRY_INTERVAL if
// set, otherwise returning the default.
func resolveTokenRefreshMaxRetryInterval() time.Duration {
        return resolveDurationEnv(
                tokenRefreshMaxRetryIntervalEnv,
                tokenRefreshMaxRetryInterval,
        )
}

// resolveTokenRefreshMaxTries returns the maximum number of retry attempts for
// token refresh backoff, reading from TOOLHIVE_TOKEN_REFRESH_MAX_TRIES if
// set, otherwise returning the default.
func resolveTokenRefreshMaxTries() uint {
        v := os.Getenv(tokenRefreshMaxTriesEnv)
        if v == "" {
                return uint(tokenRefreshMaxTries)
        }
        n, err := strconv.ParseUint(v, 10, strconv.IntSize)
        if err != nil {
                return uint(tokenRefreshMaxTries)
        }
        return uint(n)
}

// resolveTokenRefreshMaxElapsedTime returns the maximum elapsed time for all retry attempts for
// token refresh backoff, reading from TOOLHIVE_TOKEN_REFRESH_MAX_ELAPSED_TIME if
// set, otherwise returning the default.
func resolveTokenRefreshMaxElapsedTime() time.Duration {
        return resolveDurationEnv(
                tokenRefreshMaxElapsedTimeEnv,
                tokenRefreshMaxElapsedTime,
        )
}

// resolveAuthRetryingTickInterval returns the cadence between monitor refresh
// attempts during the AuthRetrying transient-failure window, reading from
// TOOLHIVE_TOKEN_AUTH_RETRYING_TICK_INTERVAL if set, otherwise the default.
func resolveAuthRetryingTickInterval() time.Duration {
        return resolveDurationEnv(authRetryingTickIntervalEnv, authRetryingTickInterval)
}

// resolveAuthRetryingMaxElapsed returns the ceiling on how long the monitor
// will stay in AuthRetrying before giving up, reading from
// TOOLHIVE_TOKEN_AUTH_RETRYING_MAX_ELAPSED if set, otherwise the default.
func resolveAuthRetryingMaxElapsed() time.Duration {
        return resolveDurationEnv(authRetryingMaxElapsedEnv, authRetryingMaxElapsed)
}

// resolveDurationEnv reads a duration from the given environment variable.
// Returns defaultVal if the variable is unset or its value is not a valid
// positive duration.
func resolveDurationEnv(envVar string, defaultVal time.Duration) time.Duration {
        v := os.Getenv(envVar)
        if v == "" {
                return defaultVal
        }
        d, err := time.ParseDuration(v)
        if err != nil || d <= 0 {
                slog.Warn("invalid duration env var, using default",
                        "env_var", envVar, "value", v, "default", defaultVal)
                return defaultVal
        }
        slog.Debug("using custom token refresh interval", "env_var", envVar, "value", d)
        return d
}

// StatusUpdater is an interface for updating workload authentication status.
// This abstraction allows the monitored token source to work with any status management system
// without creating import cycles.
type StatusUpdater interface {
        SetWorkloadStatus(ctx context.Context, workloadName string, status runtime.WorkloadStatus, reason string) error
}

// transientRefresher deduplicates concurrent token fetches during transient
// network failures and retries with exponential backoff. It is owned by
// MonitoredTokenSource and can be tested in isolation.
type transientRefresher struct {
        group    singleflight.Group
        source   oauth2.TokenSource
        workload string

        // upstream identifies the upstream authorization server that issued the
        // token, and clientID is the OAuth 2.0 client_id used by this workload.
        // Both are optional and are only surfaced in structured logs (notably the
        // DCR/CIMD remediation warning emitted from MonitoredTokenSource on the
        // transition to Unauthenticated). Empty strings are acceptable and are
        // omitted from log output by the slog handler.
        upstream string
        clientID string

        // newBackOff constructs the exponential backoff applied during transient-error
        // retries. Always non-nil; set by the constructor.
        newBackOff func() backoff.BackOff
}

// Refresh deduplicates concurrent callers via singleflight and retries the
// underlying token source with exponential backoff until the context is
// cancelled or a non-transient error is returned.
func (r *transientRefresher) Refresh(ctx context.Context, origErr error) (*oauth2.Token, error) {
        v, err, _ := r.group.Do("token-refresh", func() (interface{}, error) {
                return r.retry(ctx, origErr)
        })
        if err != nil {
                return nil, err
        }
        return v.(*oauth2.Token), nil
}

func newTransientRefresher(
        source oauth2.TokenSource,
        workload, upstream, clientID string,
        newBackOff func() backoff.BackOff,
) *transientRefresher {
        return &transientRefresher{
                source:     source,
                workload:   workload,
                upstream:   upstream,
                clientID:   clientID,
                newBackOff: newBackOff,
        }
}

func defaultBackOff() backoff.BackOff {
        eb := backoff.NewExponentialBackOff()
        eb.InitialInterval = resolveTokenRefreshInitialRetryInterval()
        eb.MaxInterval = resolveTokenRefreshMaxRetryInterval()
        eb.Reset()
        return eb
}

func (r *transientRefresher) retry(ctx context.Context, origErr error) (*oauth2.Token, error) {
        slog.Warn("token refresh failed due to transient network error, retrying with backoff",
                "workload", r.workload,
                "error", origErr,
        )

        return backoff.Retry(ctx, func() (*oauth2.Token, error) {
                t, tokenErr := r.source.Token()
                if tokenErr == nil {
                        return t, nil
                }
                if !isTransientNetworkError(tokenErr) {
                        return nil, backoff.Permanent(tokenErr)
                }
                return nil, tokenErr
        },
                backoff.WithBackOff(r.newBackOff()),
                backoff.WithNotify(func(retryErr error, d time.Duration) {
                        slog.Warn("token refresh retry failed",
                                "workload", r.workload,
                                "retry_in", d,
                                "error", retryErr,
                        )
                }),
                backoff.WithMaxTries(resolveTokenRefreshMaxTries()),
                backoff.WithMaxElapsedTime(resolveTokenRefreshMaxElapsedTime()),
        )
}

// MonitoredTokenSource is a wrapper around an oauth2.TokenSource that monitors
// authentication failures and surfaces them through workload status transitions.
// It provides both per-request token retrieval and background monitoring.
//
// Failure handling has two layers:
//
//   - Short retry (inside transientRefresher.retry): exponential backoff up to
//     a ~5min total window. Retries transient failures (DNS errors, dropped
//     connections, OAuth server 5xx/429, WAF 4xx-with-HTML) without changing
//     workload status. Failures that resolve within the window never leave
//     Running; failures that outlast the window fall through to AuthRetrying.
//
//   - AuthRetrying ceiling (this struct): after the short retry exhausts on a
//     still-transient error, the workload transitions to AuthRetrying and the
//     monitor keeps running on a longer cadence (default 10min). On the next
//     successful refresh, the workload returns to Running. After a ceiling
//     (default 24h) the workload is finally marked Unauthenticated and the
//     monitor stops.
//
// Hot callers — request-path Token() calls, e.g. from the token injection
// middleware serving a live MCP request — fast-fail with the cached error
// during AuthRetrying so they return 503+Retry-After immediately rather
// than blocking on a singleflight retry that would exhaust again on every
// call.
//
// Hot callers and the background monitor can both drive state transitions
// independently (the monitor via its tick, hot callers via short-retry
// exhaustion in Token). When the upstream alternates between transient
// failures and brief recoveries, the workload status oscillates with it.
type MonitoredTokenSource struct {
        tokenSource    oauth2.TokenSource
        workloadName   string
        upstream       string
        clientID       string
        statusUpdater  StatusUpdater
        monitoringCtx  context.Context
        stopMonitoring chan struct{}
        stopOnce       sync.Once
        // dcrWarnOnce gates the DCR/CIMD remediation Warn so it fires at most
        // once per workload lifetime. Kept separate from stopOnce so the gate
        // is independent of which code path first closes stopMonitoring: a
        // zero-expiry exit in onTick may close the channel without ever
        // transitioning to Unauthenticated, and a later markAsUnauthenticated
        // must still be able to emit the Warn for a permanent 4xx.
        dcrWarnOnce sync.Once
        refresher   *transientRefresher

        // stopped is closed when monitorLoop exits, regardless of the reason.
        stopped chan struct{}

        timer *time.Timer

        // AuthRetrying state, guarded by mu. A non-zero transientStartedAt
        // means the workload is in the AuthRetrying window (short retry
        // exhausted, ceiling not yet reached). lastTransientErr is the error
        // returned to hot callers via Token() during the window so they fail
        // fast instead of joining a singleflight that would block for the
        // full short-retry duration on every tick. Written by the monitor
        // goroutine and by Token() callers who reach the short-retry
        // exhaustion branch; read by Token() callers during the fast-fail
        // check. mu is held only across pure field reads/writes — never
        // across a statusUpdater call, which does I/O.
        mu                 sync.Mutex
        transientStartedAt time.Time
        lastTransientErr   error
}

// NewMonitoredTokenSource creates a new MonitoredTokenSource that wraps the provided
// oauth2.TokenSource and monitors it for authentication failures.
//
// upstream and clientID annotate structured logs emitted by the token source,
// most importantly the DCR/CIMD remediation warning fired on the transition
// to Unauthenticated when the token endpoint returns a permanent 4xx (which
// frequently indicates stale cached credentials). Pass empty strings when
// the workload does not use DCR/CIMD or the upstream issuer is unknown; the
// remediation log will be suppressed when clientID is empty since its
// operator-correlation field would be blank.
//
// The fields are fixed at construction time rather than exposed via a setter
// so there is no data race between a late writer and the readers in Token()
// / transientRefresher.retry() — both of which may run on the background
// monitor goroutine started by StartBackgroundMonitoring.
func NewMonitoredTokenSource(
        ctx context.Context,
        tokenSource oauth2.TokenSource,
        workloadName string,
        upstream string,
        clientID string,
        statusUpdater StatusUpdater,
) *MonitoredTokenSource {
        return newMonitoredTokenSourceWithBackOff(ctx, tokenSource, workloadName, upstream, clientID, statusUpdater, defaultBackOff)
}

func newMonitoredTokenSourceWithBackOff(
        ctx context.Context,
        tokenSource oauth2.TokenSource,
        workloadName string,
        upstream string,
        clientID string,
        statusUpdater StatusUpdater,
        newBackOff func() backoff.BackOff,
) *MonitoredTokenSource {
        return &MonitoredTokenSource{
                tokenSource:    tokenSource,
                workloadName:   workloadName,
                upstream:       upstream,
                clientID:       clientID,
                statusUpdater:  statusUpdater,
                monitoringCtx:  ctx,
                stopMonitoring: make(chan struct{}),
                stopped:        make(chan struct{}),
                refresher:      newTransientRefresher(tokenSource, workloadName, upstream, clientID, newBackOff),
        }
}

// Stopped returns a channel that is closed when background monitoring has stopped,
// regardless of the reason (context cancellation, auth failure, or clean shutdown).
func (mts *MonitoredTokenSource) Stopped() <-chan struct{} {
        return mts.stopped
}

// Token retrieves a token. On a non-transient error, it marks the workload
// unauthenticated and returns immediately. On a transient error, behavior
// depends on monitor state — see the paragraphs below. Context cancellation
// (workload removal) stops any in-flight retry without marking the workload
// unauthenticated. See isTransientNetworkError for the classification rule.
//
// When the monitor is in the AuthRetrying transient-failure window
// (short retry exhausted, ceiling not yet reached), Token() fast-fails
// with the cached error rather than joining a singleflight retry that
// would hang for the full short-retry duration on every call. Hot
// callers see 503+Retry-After until the next monitor tick observes
// upstream recovery and clears the state — only the monitor's onTick
// (which calls the raw token source directly) can do that.
//
// If a hot caller's own short retry exhausts on a still-transient error,
// the workload transitions to AuthRetrying and the monitor stays alive —
// subsequent hot callers fast-fail until the next monitor tick clears the
// state or extends it. If the monitor has already stopped (a prior
// permanent error closed it), enterAuthRetrying is a no-op; the hot
// caller returns the error without changing workload status.
//
// During the short retry on a transient failure, concurrent callers
// joining at the same time are deduplicated via singleflight so that
// only one retry loop runs at a time. Callers that arrive AFTER the
// leader's singleflight call has returned start their own retry — they
// are not deduplicated against past calls. This is why hot-caller-
// driven AuthRetrying entry (after the retry exhausts) is load-bearing:
// without it, sequential hot callers would each pay the full short-
// retry duration against a broken endpoint.
func (mts *MonitoredTokenSource) Token() (*oauth2.Token, error) {
        if err := mts.fastFailIfAuthRetrying(); err != nil {
                return nil, err
        }

        tok, err := mts.tokenSource.Token()
        if err == nil {
                return tok, nil
        }

        if !isTransientNetworkError(err) {
                mts.markAsUnauthenticated(
                        fmt.Sprintf("Token retrieval failed: %v", err),
                        isPermanentTokenEndpointError(err),
                )
                return nil, err
        }

        // If the monitor has already stopped (prior permanent error or 24h
        // ceiling), do not enter the short retry. The workload is terminal;
        // making hot callers each pay the full short-retry duration against
        // a known-broken endpoint is exactly the pathology AuthRetrying was
        // introduced to avoid.
        select {
        case <-mts.stopMonitoring:
                return nil, err
        default:
        }

        // Transient network error — funnel all concurrent callers through a
        // single retry loop so we don't hammer the token endpoint.
        tok, err = mts.refresher.Refresh(mts.monitoringCtx, err)
        if err == nil {
                return tok, nil
        }
        if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
                return nil, err
        }
        if !isTransientNetworkError(err) {
                mts.markAsUnauthenticated(
                        fmt.Sprintf("Token refresh failed after retries: %v", err),
                        isPermanentTokenEndpointError(err),
                )
                return nil, err
        }
        // Short retry exhausted on a still-transient error → enter AuthRetrying
        // and keep the monitor alive. Subsequent hot callers fast-fail; the
        // next monitor tick will decide whether to recover or extend.
        mts.enterAuthRetrying(err)
        return nil, err
}

// StartBackgroundMonitoring starts the background monitoring goroutine that checks
// token validity at expiry time and marks the workload as unauthenticated on the failure.
func (mts *MonitoredTokenSource) StartBackgroundMonitoring() {
        if mts.timer == nil {
                mts.timer = time.NewTimer(time.Millisecond) // kick immediately
        }
        go mts.monitorLoop()
}

func (mts *MonitoredTokenSource) monitorLoop() {
        defer close(mts.stopped)
        for {
                select {
                case <-mts.monitoringCtx.Done():
                        mts.stopTimer()
                        return
                case <-mts.stopMonitoring:
                        mts.stopTimer()
                        return
                case <-mts.timer.C:
                        shouldStop, next := mts.onTick()
                        if shouldStop {
                                mts.stopTimer()
                                return
                        }
                        mts.resetTimer(next)
                }
        }
}

func (mts *MonitoredTokenSource) stopTimer() {
        if mts.timer != nil && !mts.timer.Stop() {
                select {
                case <-mts.timer.C:
                default:
                }
        }
}

func (mts *MonitoredTokenSource) resetTimer(d time.Duration) {
        mts.stopTimer()
        mts.timer.Reset(d)
}

// onTick performs a token refresh attempt on the background monitor goroutine
// and returns whether to stop monitoring and the next tick delay.
//
// State transitions handled here:
//
//   - Success: clear any AuthRetrying state (→ Running) and schedule the
//     next tick at the new token's expiry.
//   - Permanent error: mark Unauthenticated immediately and stop monitoring
//     (existing semantics).
//   - Transient error while in AuthRetrying: check the ceiling. If exceeded,
//     mark Unauthenticated. Otherwise refresh the cached error and reschedule
//     at the AuthRetrying cadence.
//   - Transient error while NOT in AuthRetrying: run the existing short retry
//     once (singleflight + exponential backoff). If it succeeds, return to
//     normal expiry-based scheduling. If it exhausts on a still-transient
//     error, enter AuthRetrying.
func (mts *MonitoredTokenSource) onTick() (bool, time.Duration) {
        // Call the raw token source directly — NOT Token() — so the monitor
        // bypasses the AuthRetrying fast-fail and the singleflight retry. The
        // monitor owns those state transitions itself.
        tok, err := mts.tokenSource.Token()

        if err == nil {
                mts.exitAuthRetrying()
                if tok == nil || tok.Expiry.IsZero() {
                        // Token has no expiry, so there's nothing for the monitor to
                        // schedule against — exit cleanly. Close stopMonitoring so the
                        // Token() hot path's stopMonitoring gate sees the monitor is
                        // gone: without this, a future transient error would burn the
                        // full short-retry window before enterAuthRetrying's gate
                        // (under mu) finally aborted, leaving the workload in
                        // AuthRetrying with no monitor to drive the ceiling. We do
                        // NOT emit Unauthenticated here — the current token is valid.
                        mts.stopMonitoringOnce()
                        return true, 0
                }
                return false, waitUntilExpiry(tok.Expiry)
        }

        if !isTransientNetworkError(err) {
                mts.markAsUnauthenticated(
                        fmt.Sprintf("Token retrieval failed: %v", err),
                        isPermanentTokenEndpointError(err),
                )
                return true, 0
        }

        if mts.inAuthRetrying() {
                if mts.ceilingExceeded() {
                        mts.markAsUnauthenticated(
                                fmt.Sprintf("Token refresh failed transiently for over %s: %v",
                                        resolveAuthRetryingMaxElapsed(), err),
                                false, // transient by definition → DCR Warn correctly silent
                        )
                        return true, 0
                }
                mts.enterAuthRetrying(err)
                return false, resolveAuthRetryingTickInterval()
        }

        // First transient failure on this tick: run the existing short retry once
        // so we benefit from singleflight + exponential backoff for brief blips.
        tok, err = mts.refresher.Refresh(mts.monitoringCtx, err)
        if err == nil {
                // A hot caller may have entered AuthRetrying between this tick's
                // inAuthRetrying() check and Refresh returning; clear that state
                // now so subsequent hot Token() calls don't keep fast-failing
                // against the cached transient error until the next monitor tick
                // reaches the inAuthRetrying branch. No-op when not in
                // AuthRetrying.
                mts.exitAuthRetrying()
                if tok == nil || tok.Expiry.IsZero() {
                        mts.stopMonitoringOnce()
                        return true, 0
                }
                return false, waitUntilExpiry(tok.Expiry)
        }
        if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
                return true, 0
        }
        if !isTransientNetworkError(err) {
                mts.markAsUnauthenticated(
                        fmt.Sprintf("Token refresh failed after retries: %v", err),
                        isPermanentTokenEndpointError(err),
                )
                return true, 0
        }

        // Short retry exhausted on a still-transient error → enter AuthRetrying.
        mts.enterAuthRetrying(err)
        return false, resolveAuthRetryingTickInterval()
}

// waitUntilExpiry returns the duration until expiry, clamped to a minimum
// of one second so the monitor doesn't spin on a stale or near-past expiry.
func waitUntilExpiry(expiry time.Time) time.Duration {
        wait := time.Until(expiry)
        if wait < time.Second {
                wait = time.Second
        }
        return wait
}

// isTransientNetworkError reports whether err represents a transient
// condition that is likely to resolve on its own. The categories are:
//
//   - Network-level failures: DNS lookup errors, TCP transport errors,
//     timeouts.
//   - OAuth token-endpoint responses classified as transient by
//     isTransientRetrieveError.
//   - Unparsable token responses on a 2xx status (typically an HTML page
//     from a load balancer or CDN).
//
// All other errors return false, causing the workload to be marked
// unauthenticated. TLS failures (certificate verification, handshake
// failure) are intentionally non-transient even though they surface
// through net.OpError like transport-level errors.
//
// The function is side-effect free; callers that want to emit a DCR
// remediation hint on a permanent 4xx must do so themselves at the
// state-transition boundary using isPermanentTokenEndpointError to
// classify, so a tight Token() loop does not spam the same record.
func isTransientNetworkError(err error) bool {
        if err == nil ||
                errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
                return false
        }

        if retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err); ok {
                return isTransientRetrieveError(retrieveErr)
        }

        // Non-JSON responses from the OAuth server (e.g. load balancer HTML pages).
        // The oauth2 library returns a plain error (not *RetrieveError) when the
        // HTTP status is 2xx but the body cannot be parsed as JSON.
        if isOAuthParseError(err) {
                return true
        }

        // DNS lookup failures — covers VPN-disconnect scenarios where the corporate DNS
        // resolver is unreachable.
        if _, ok := errors.AsType[*net.DNSError](err); ok {
                return true
        }

        // *net.OpError covers both transport-level errors (connection refused, network
        // unreachable) AND TLS errors (certificate invalid, handshake failure). Only the
        // former are transient; TLS errors do not wrap syscall errors, so we use that
        // to distinguish them.
        if opErr, ok := errors.AsType[*net.OpError](err); ok {
                _, isSyscall := errors.AsType[*os.SyscallError](opErr)
                _, isErrno := errors.AsType[syscall.Errno](opErr)
                return isSyscall || isErrno
        }

        // Generic net.Error timeout (catches any remaining net.Error implementations).
        if netErr, ok := errors.AsType[net.Error](err); ok && netErr.Timeout() {
                return true
        }

        return false
}

// isPermanentTokenEndpointError reports whether the OAuth server itself
// rendered a verdict on the cached credentials (invalid_grant, invalid_client,
// etc.). Used at state-transition boundaries to decide whether to emit a
// DCR/CIMD remediation hint alongside the unauthentication. The rules live in
// pkg/auth/oautherr; see IsPermanentCredentialError for why a non-spec-compliant
// 4xx is deliberately not "permanent".
func isPermanentTokenEndpointError(err error) bool {
        return oautherr.IsPermanentCredentialError(err)
}

// isTransientRetrieveError reports whether an *oauth2.RetrieveError should
// be treated as transient. The rules live in pkg/auth/oautherr so this
// package and the LLM-gateway token source cannot drift on them.
func isTransientRetrieveError(retrieveErr *oauth2.RetrieveError) bool {
        return oautherr.IsTransientRetrieveError(retrieveErr)
}

// isOAuthParseError detects errors from the oauth2 library that indicate the
// token endpoint returned an unparsable response body on a 2xx status. The
// detection lives in pkg/auth/oautherr.
func isOAuthParseError(err error) bool {
        return oautherr.IsParseError(err)
}

// stopMonitoringOnce idempotently closes the stopMonitoring channel. It
// does NOT emit any workload status transition; the caller is responsible
// for that (or for explicitly skipping it, as the zero-expiry exit path
// does). The DCR/CIMD remediation Warn is gated by a separate sync.Once
// (dcrWarnOnce in markAsUnauthenticated) so callers that close the
// channel without transitioning to Unauthenticated do not consume the
// Warn slot.
func (mts *MonitoredTokenSource) stopMonitoringOnce() {
        mts.stopOnce.Do(func() {
                close(mts.stopMonitoring)
        })
}

// markAsUnauthenticated marks the workload as unauthenticated and stops
// background monitoring. If permanent4xx is true and the workload was
// constructed with a non-empty client_id, a one-shot DCR/CIMD remediation
// hint is emitted alongside the stop transition. The Warn is gated on
// dcrWarnOnce so a later call cannot re-emit it after the workload has
// already transitioned to Unauthenticated, and so the Warn remains
// emittable on the first real Unauthenticated transition even if an
// earlier code path (such as the zero-expiry exit in onTick) consumed
// the stopOnce slot. dcrWarnOnce is kept separate from stopOnce for the
// reason captured in the dcrWarnOnce field comment.
//
// Ordering matters: stopMonitoring is closed first so any future
// enterAuthRetrying call sees the gate closed before it acquires the
// field mutex. That alone only narrows the in-memory field race — an
// enterAuthRetrying that already passed the gate and is waiting on mu
// can still set transientStartedAt/lastTransientErr after we clear
// them. The residual race is closed by the defensive re-clear under
// mu below (the load-bearing step), so a hot caller cannot leave
// AuthRetrying state populated after we've written Unauthenticated.
// The Unauthenticated status emit follows the re-clear.
func (mts *MonitoredTokenSource) markAsUnauthenticated(reason string, permanent4xx bool) {
        mts.stopMonitoringOnce()

        mts.mu.Lock()
        mts.transientStartedAt = time.Time{}
        mts.lastTransientErr = nil
        mts.mu.Unlock()

        _ = mts.statusUpdater.SetWorkloadStatus(
                context.Background(),
                mts.workloadName,
                runtime.WorkloadStatusUnauthenticated,
                reason,
        )

        // Emit the DCR/CIMD remediation hint at most once per workload lifetime.
        // A permanent 4xx from the token endpoint commonly indicates the
        // cached client (DCR or CIMD) is no longer recognised — but the same
        // branch fires for revoked consent, disabled accounts, and statically
        // configured clients, so the message has to be honest about the
        // variability. Gating on clientID != "" suppresses the log entirely
        // for workloads where no client_id context is available; the
        // operator-correlation it provides would be empty.
        //
        // Gate the dcrWarnOnce.Do call on permanent4xx && clientID != ""
        // OUTSIDE the lambda: otherwise a non-permanent invocation (e.g. a
        // ceiling-triggered markAsUnauthenticated, permanent4xx=false) would
        // enter Do, early-return, and consume the once-slot — silently
        // suppressing the Warn for a later legitimate permanent 4xx.
        if permanent4xx && mts.clientID != "" {
                mts.dcrWarnOnce.Do(func() {
                        //nolint:gosec // G706: client_id is public metadata per RFC 7591.
                        slog.Warn(
                                "token endpoint returned a permanent error; if this workload uses "+
                                        "cached DCR or CIMD credentials they may be stale — delete the "+
                                        "cached credentials and restart to re-register.",
                                "workload", mts.workloadName,
                                "upstream", mts.upstream,
                                "client_id", mts.clientID,
                        )
                })
        }
}

// inAuthRetrying reports whether the monitor is currently in the
// AuthRetrying transient-failure window.
func (mts *MonitoredTokenSource) inAuthRetrying() bool {
        mts.mu.Lock()
        defer mts.mu.Unlock()
        return !mts.transientStartedAt.IsZero()
}

// fastFailIfAuthRetrying returns the cached transient error when the
// monitor is in AuthRetrying so hot callers (e.g. the token injection
// middleware) return 503+Retry-After immediately rather than blocking
// on a singleflight retry that will exhaust again.
func (mts *MonitoredTokenSource) fastFailIfAuthRetrying() error {
        mts.mu.Lock()
        defer mts.mu.Unlock()
        if mts.transientStartedAt.IsZero() {
                return nil
        }
        // enterAuthRetrying sets transientStartedAt and lastTransientErr together
        // under mu, so under normal flow lastTransientErr is non-nil here. Guard
        // against a future refactor that breaks the invariant; fmt.Errorf("%w", nil)
        // returns a non-nil error whose formatted message is misleading
        // (%!w(<nil>)).
        if mts.lastTransientErr == nil {
                return fmt.Errorf("auth retrying since %s",
                        mts.transientStartedAt.Format(time.RFC3339))
        }
        return fmt.Errorf("auth retrying since %s: %w",
                mts.transientStartedAt.Format(time.RFC3339), mts.lastTransientErr)
}

// ceilingExceeded reports whether the workload has been in AuthRetrying
// for longer than the configured ceiling. Returns false if not in
// AuthRetrying.
func (mts *MonitoredTokenSource) ceilingExceeded() bool {
        mts.mu.Lock()
        defer mts.mu.Unlock()
        if mts.transientStartedAt.IsZero() {
                return false
        }
        return time.Since(mts.transientStartedAt) > resolveAuthRetryingMaxElapsed()
}

// enterAuthRetrying marks the workload as AuthRetrying on the first call
// and refreshes the cached error on subsequent calls. The status
// transition is emitted only on first entry to avoid spamming the
// status file every tick.
//
// The gate check is performed under mts.mu so that markAsUnauthenticated
// (which closes the channel before acquiring mu) cannot interleave to
// re-populate transientStartedAt/lastTransientErr after clearing them.
// A narrower disk-write inversion is still possible (auth_retrying
// written briefly after Unauthenticated); the in-memory state is
// correct, and the disk entry resolves on the next workload restart
// (runner.go resets non-Unauthenticated statuses to Running).
func (mts *MonitoredTokenSource) enterAuthRetrying(err error) {
        mts.mu.Lock()
        select {
        case <-mts.stopMonitoring:
                mts.mu.Unlock()
                return
        default:
        }
        firstEntry := mts.transientStartedAt.IsZero()
        if firstEntry {
                mts.transientStartedAt = time.Now()
        }
        elapsed := time.Since(mts.transientStartedAt)
        mts.lastTransientErr = err
        mts.mu.Unlock()

        if firstEntry {
                slog.Warn("token refresh entering AuthRetrying after short retry exhaustion; will retry on a longer cadence",
                        "workload", mts.workloadName,
                        "tick_interval", resolveAuthRetryingTickInterval(),
                        "max_elapsed", resolveAuthRetryingMaxElapsed(),
                        "error", err,
                )
                _ = mts.statusUpdater.SetWorkloadStatus(
                        context.Background(),
                        mts.workloadName,
                        runtime.WorkloadStatusAuthRetrying,
                        fmt.Sprintf("Token refresh failing transiently; retrying every %s: %v",
                                resolveAuthRetryingTickInterval(), err),
                )
                return
        }
        slog.Warn("token refresh still failing transiently",
                "workload", mts.workloadName,
                "elapsed", elapsed,
                "error", err,
        )
}

// exitAuthRetrying clears the AuthRetrying window and emits the Running
// status transition if the workload was previously in AuthRetrying. No-op
// if not in AuthRetrying.
func (mts *MonitoredTokenSource) exitAuthRetrying() {
        mts.mu.Lock()
        wasInAuthRetrying := !mts.transientStartedAt.IsZero()
        if wasInAuthRetrying {
                mts.transientStartedAt = time.Time{}
                mts.lastTransientErr = nil
        }
        mts.mu.Unlock()

        if !wasInAuthRetrying {
                return
        }
        slog.Info("token refresh recovered; exiting AuthRetrying state",
                "workload", mts.workloadName,
        )
        _ = mts.statusUpdater.SetWorkloadStatus(
                context.Background(),
                mts.workloadName,
                runtime.WorkloadStatusRunning,
                "",
        )
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package oauth provides OAuth 2.0 and OIDC authentication functionality.
package oauth

import (
        "context"
        "crypto/rand"
        "encoding/base64"
        "errors"
        "fmt"
        "html"
        "log/slog"
        "net/http"
        "net/url"
        "os"
        "os/signal"
        "strings"
        "syscall"
        "time"

        "github.com/golang-jwt/jwt/v5"
        "github.com/pkg/browser"
        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// Config contains configuration for OAuth authentication
type Config struct {
        // ClientID is the OAuth client ID
        ClientID string

        // ClientSecret is the OAuth client secret (optional for PKCE flow)
        ClientSecret string //nolint:gosec // G117: field legitimately holds sensitive data

        // RedirectURL is the redirect URL for the OAuth flow
        RedirectURL string

        // AuthURL is the authorization endpoint URL
        AuthURL string

        // TokenURL is the token endpoint URL
        TokenURL string

        // Scopes are the OAuth scopes to request
        Scopes []string

        // UsePKCE enables PKCE (Proof Key for Code Exchange) for enhanced security
        UsePKCE bool

        // CallbackPort is the port for the OAuth callback server (optional, 0 means auto-select)
        CallbackPort int

        // RequireExactCallbackPort prevents fallback when CallbackPort is unavailable.
        // Use it when the callback URI is pre-registered with the identity provider.
        RequireExactCallbackPort bool

        // IntrospectionEndpoint is the optional introspection endpoint for validating tokens
        IntrospectionEndpoint string

        // Resource is the OAuth 2.0 resource indicator (RFC 8707).
        Resource string

        // OAuthParams are additional parameters to pass to the authorization URL
        OAuthParams map[string]string

        // ScopeParamName overrides the query parameter name used to send scopes in the
        // authorization URL. When empty (default), the standard "scope" parameter is used.
        // Some providers use non-standard parameter names (e.g., Slack uses "user_scope"
        // for user-token scopes). When set, scopes are sent under this parameter name
        // instead of "scope", and the standard "scope" parameter is cleared.
        ScopeParamName string

        // TokenEndpointTrusted reports that the token endpoint's authority was
        // supplied by the operator — either configured directly, or named by the
        // operator-configured issuer's own metadata document. Because the operator
        // chose that authority, the endpoint may be a private, loopback, or
        // link-local address. The zero value is the guarded posture and is what an
        // endpoint named by a remote MCP server, or by a metadata document pointing
        // at some other authority, must keep.
        TokenEndpointTrusted bool
}

// Flow handles the OAuth authentication flow
type Flow struct {
        config       *Config
        oauth2Config *oauth2.Config
        server       *http.Server
        port         int

        // PKCE parameters
        codeVerifier  string
        codeChallenge string
        state         string

        tokenSource oauth2.TokenSource
        httpClient  *http.Client
}

// TokenResult contains the result of the OAuth flow
type TokenResult struct {
        AccessToken  string //nolint:gosec // G117: field legitimately holds sensitive data
        RefreshToken string //nolint:gosec // G117: field legitimately holds sensitive data
        TokenType    string
        Expiry       time.Time
        Claims       jwt.MapClaims
        IDToken      string // The OIDC ID token (JWT), if present
}

// NewFlow creates a new OAuth flow
func NewFlow(config *Config) (*Flow, error) {
        if config == nil {
                return nil, errors.New("OAuth config cannot be nil")
        }

        if config.ClientID == "" {
                return nil, errors.New("client ID is required")
        }

        if config.AuthURL == "" {
                return nil, errors.New("authorization URL is required")
        }

        if config.TokenURL == "" {
                return nil, errors.New("token URL is required")
        }

        requestedPort := config.CallbackPort
        port := requestedPort
        if config.RequireExactCallbackPort && requestedPort != 0 {
                if !networking.IsAvailable(requestedPort) {
                        return nil, &networking.CallbackPortInUseError{Port: requestedPort}
                }
        } else {
                var err error
                port, err = networking.FindOrUsePort(requestedPort)
                if err != nil {
                        return nil, fmt.Errorf("failed to find available port: %w", err)
                }
        }

        // Set default redirect URL if not provided
        redirectURL := config.RedirectURL
        if redirectURL == "" {
                redirectURL = fmt.Sprintf("http://localhost:%d/callback", port)
        }

        // Public clients (no secret) must use AuthStyleInParams: strict OAuth 2.1 servers
        // (e.g. Datadog) reject Basic Auth for token_endpoint_auth_method=none clients and
        // consume the single-use auth code in doing so, causing a retry to fail with
        // invalid_grant. Confidential clients use AutoDetect so servers that mandate
        // client_secret_basic are not broken.
        authStyle := oauth2.AuthStyleInParams
        if config.ClientSecret != "" {
                authStyle = oauth2.AuthStyleAutoDetect
        }

        // Create OAuth2 config
        oauth2Config := &oauth2.Config{
                ClientID:     config.ClientID,
                ClientSecret: config.ClientSecret,
                RedirectURL:  redirectURL,
                Scopes:       config.Scopes,
                Endpoint: oauth2.Endpoint{
                        AuthURL:   config.AuthURL,
                        TokenURL:  config.TokenURL,
                        AuthStyle: authStyle,
                },
        }

        httpClient, err := NewTokenHTTPClient(config.TokenURL, config.TokenEndpointTrusted)
        if err != nil {
                return nil, fmt.Errorf("build token endpoint HTTP client: %w", err)
        }

        flow := &Flow{
                config:       config,
                oauth2Config: oauth2Config,
                port:         port,
                httpClient:   httpClient,
        }

        // Generate PKCE parameters if enabled
        if config.UsePKCE {
                flow.generatePKCEParams()
        }

        // Generate state parameter
        if err := flow.generateState(); err != nil {
                return nil, fmt.Errorf("failed to generate state parameter: %w", err)
        }

        return flow, nil
}

// NewTokenHTTPClient creates the client used for OAuth token exchanges and
// refreshes. A trusted endpoint (see Config.TokenEndpointTrusted) may be dialed
// at a private, loopback, or link-local address, because the operator named its
// authority. An untrusted endpoint is refused those addresses (CWE-918) and
// gets a plain-HTTP exception only for localhost. Both clients disable
// connection reuse so the dial-time guard re-runs per request, and follow only
// same-host redirects.
func NewTokenHTTPClient(tokenURL string, trusted bool) (*http.Client, error) {
        endpoint, err := url.Parse(tokenURL)
        if err != nil {
                return nil, fmt.Errorf("parse token endpoint URL %q: %w", tokenURL, err)
        }
        if endpoint.Scheme == "" || endpoint.Host == "" {
                return nil, fmt.Errorf("token endpoint URL %q is missing a scheme or host", tokenURL)
        }

        builder := func() *networking.HttpClientBuilder {
                if trusted {
                        return networking.NewHostScopedClientBuilder(endpoint.Host, true, false)
                }
                return networking.NewHttpClientBuilder().
                        WithPrivateIPs(false).
                        WithInsecureAllowHTTP(networking.IsLocalhost(endpoint.Host))
        }()
        client, err := builder.WithDisableKeepAlives(true).Build()
        if err != nil {
                return nil, err
        }
        client.Transport = &oauthproto.UserAgentTransport{Base: client.Transport}
        client.CheckRedirect = networking.SameHostRedirectPolicy()
        return client, nil
}

// generatePKCEParams generates PKCE code verifier and challenge using
// the standard oauth2 library functions.
func (f *Flow) generatePKCEParams() {
        // Generate code verifier using oauth2 stdlib (43-128 characters, RFC 7636)
        f.codeVerifier = oauth2.GenerateVerifier()

        // Use S256 method for enhanced security (RFC 7636 recommendation)
        f.codeChallenge = oauth2.S256ChallengeFromVerifier(f.codeVerifier)
}

// generateState generates a random state parameter
func (f *Flow) generateState() error {
        stateBytes := make([]byte, 16)
        if _, err := rand.Read(stateBytes); err != nil {
                return fmt.Errorf("failed to generate state: %w", err)
        }
        f.state = base64.RawURLEncoding.EncodeToString(stateBytes)
        return nil
}

// listenAddr returns the address the callback server binds to. The redirect
// URL points at localhost, so the listener is bound to the IPv4 loopback
// instead of all interfaces: a wildcard bind would expose the callback
// endpoint to every host on the LAN for the lifetime of the login.
func (f *Flow) listenAddr() string {
        return fmt.Sprintf("127.0.0.1:%d", f.port)
}

// Start starts the OAuth authentication flow
func (f *Flow) Start(ctx context.Context, skipBrowser bool) (*TokenResult, error) {
        // Create channels for communication
        tokenChan := make(chan *oauth2.Token, 1)
        errorChan := make(chan error, 1)

        // Set up HTTP server for handling the callback
        mux := http.NewServeMux()
        mux.HandleFunc("/callback", f.handleCallback(tokenChan, errorChan))
        mux.HandleFunc("/", f.handleRoot())

        f.server = &http.Server{
                Addr:              f.listenAddr(),
                Handler:           mux,
                ReadHeaderTimeout: 10 * time.Second,
        }

        // Start the server in a goroutine
        go func() {
                slog.Debug("Starting OAuth callback server", "port", f.port)
                if err := f.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
                        errorChan <- fmt.Errorf("failed to start callback server: %w", err)
                }
        }()

        // Ensure server cleanup
        defer func() {
                // Use Background context for server shutdown. This cleanup operation runs after
                // the OAuth flow completes (or fails). The parent context may already be cancelled,
                // so we need a fresh context with its own timeout to ensure the server shuts down
                // gracefully regardless of the parent context state.
                shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
                defer cancel()
                if err := f.server.Shutdown(shutdownCtx); err != nil {
                        slog.Warn("Failed to shutdown OAuth callback server", "error", err)
                }
        }()

        // Build authorization URL
        authURL := f.buildAuthURL()

        // Open browser or display URL
        if !skipBrowser {
                fmt.Fprintf(os.Stderr, "Opening browser: %s\n", authURL)
                if err := browser.OpenURL(authURL); err != nil {
                        slog.Warn("Failed to open browser", "error", err)
                        fmt.Fprintf(os.Stderr, "Please manually open this URL in your browser: %s\n", authURL)
                }
        } else {
                fmt.Fprintf(os.Stderr, "Please open this URL in your browser: %s\n", authURL)
        }

        fmt.Fprintln(os.Stderr, "Waiting for OAuth callback")

        // Set up signal handling for graceful shutdown
        sigChan := make(chan os.Signal, 1)
        signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

        // Wait for token, error, or cancellation
        select {
        case token := <-tokenChan:
                slog.Debug("OAuth flow completed successfully")
                return f.processToken(ctx, token), nil
        case err := <-errorChan:
                return nil, fmt.Errorf("OAuth flow failed: %w", err)
        case <-ctx.Done():
                return nil, fmt.Errorf("OAuth flow cancelled: %w", ctx.Err())
        case sig := <-sigChan:
                return nil, fmt.Errorf("OAuth flow interrupted by signal: %v", sig)
        }
}

// buildAuthURL builds the authorization URL with appropriate parameters
func (f *Flow) buildAuthURL() string {
        opts := []oauth2.AuthCodeOption{
                oauth2.SetAuthURLParam("state", f.state),
        }

        if f.config.Resource != "" {
                opts = append(opts, oauth2.SetAuthURLParam("resource", f.config.Resource))
        }

        if f.config.OAuthParams != nil {
                for key, value := range f.config.OAuthParams {
                        opts = append(opts, oauth2.SetAuthURLParam(key, value))
                }
        }

        // When a custom scope parameter name is configured, move scopes from the
        // standard "scope" parameter to the custom one. This supports OAuth providers
        // that use non-standard parameter names (e.g., Slack's "user_scope").
        // We temporarily nil out oauth2Config.Scopes so the library omits the standard
        // "scope" parameter entirely (an empty scope= would violate RFC 6749 §3.3).
        // Scopes are restored via defer so token refresh requests still work correctly.
        if f.config.ScopeParamName != "" && len(f.oauth2Config.Scopes) > 0 {
                scopeValue := strings.Join(f.oauth2Config.Scopes, " ")
                savedScopes := f.oauth2Config.Scopes
                f.oauth2Config.Scopes = nil
                defer func() { f.oauth2Config.Scopes = savedScopes }()
                opts = append(opts,
                        oauth2.SetAuthURLParam(f.config.ScopeParamName, scopeValue),
                )
        }

        // Add PKCE parameters if enabled
        if f.config.UsePKCE {
                opts = append(opts,
                        oauth2.SetAuthURLParam("code_challenge", f.codeChallenge),
                        oauth2.SetAuthURLParam("code_challenge_method", oauthproto.PKCEMethodS256),
                )
        }

        return f.oauth2Config.AuthCodeURL(f.state, opts...)
}

// handleCallback handles the OAuth callback
func (f *Flow) handleCallback(tokenChan chan<- *oauth2.Token, errorChan chan<- error) http.HandlerFunc {
        return func(w http.ResponseWriter, r *http.Request) {
                // Parse query parameters
                query := r.URL.Query()

                // Check for error
                if errParam := query.Get("error"); errParam != "" {
                        errDesc := query.Get("error_description")
                        err := fmt.Errorf("OAuth error: %s - %s", errParam, errDesc)
                        f.writeErrorPage(w, err)
                        errorChan <- err
                        return
                }

                // Validate state parameter
                state := query.Get("state")
                if state != f.state {
                        err := errors.New("invalid state parameter")
                        f.writeErrorPage(w, err)
                        errorChan <- err
                        return
                }

                // Get authorization code
                code := query.Get("code")
                if code == "" {
                        err := errors.New("missing authorization code")
                        f.writeErrorPage(w, err)
                        errorChan <- err
                        return
                }

                // Exchange code for token using the request context to respect cancellation
                ctx := context.WithValue(r.Context(), oauth2.HTTPClient, f.httpClient)
                opts := []oauth2.AuthCodeOption{}

                // Add PKCE verifier if enabled
                if f.config.UsePKCE {
                        opts = append(opts, oauth2.SetAuthURLParam("code_verifier", f.codeVerifier))
                }

                if f.config.Resource != "" {
                        opts = append(opts, oauth2.SetAuthURLParam("resource", f.config.Resource))
                }

                token, err := f.oauth2Config.Exchange(ctx, code, opts...)
                if err != nil {
                        err = fmt.Errorf("failed to exchange code for token: %w", err)
                        f.writeErrorPage(w, err)
                        errorChan <- err
                        return
                }

                // Write success page
                f.writeSuccessPage(w)

                // Send token
                tokenChan <- token
        }
}

// setSecurityHeaders sets common security headers for all responses
func (*Flow) setSecurityHeaders(w http.ResponseWriter) {
        w.Header().Set("Content-Type", "text/html; charset=utf-8")
        w.Header().Set("X-Content-Type-Options", "nosniff")
        w.Header().Set("X-Frame-Options", "DENY")
        w.Header().Set("X-XSS-Protection", "1; mode=block")
        w.Header().Set("Referrer-Policy", "strict-origin-when-cross-origin")
        w.Header().Set("Content-Security-Policy", "default-src 'self'; style-src 'unsafe-inline'; script-src 'none'; object-src 'none';")
}

// handleRoot handles requests to the root path
func (f *Flow) handleRoot() http.HandlerFunc {
        return func(w http.ResponseWriter, r *http.Request) {
                // Only allow GET requests
                if r.Method != http.MethodGet {
                        http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
                        return
                }

                f.setSecurityHeaders(w)
                htmlContent := `
<!DOCTYPE html>
<html>
<head>
    <title>ToolHive OAuth</title>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <style>
        body { font-family: Arial, sans-serif; margin: 40px; text-align: center; }
        .container { max-width: 600px; margin: 0 auto; }
        .message { padding: 20px; border-radius: 5px; margin: 20px 0; }
        .info { background-color: #e7f3ff; border: 1px solid #b3d9ff; color: #0066cc; }
    </style>
</head>
<body>
    <div class="container">
        <h1>ToolHive OAuth Authentication</h1>
        <div class="message info">
            <p>OAuth callback server is running. Please complete the authentication flow in your browser.</p>
        </div>
    </div>
</body>
</html>`
                if _, err := w.Write([]byte(htmlContent)); err != nil {
                        slog.Warn("Failed to write HTML content", "error", err)
                }
        }
}

// writeSuccessPage writes a success page to the response
func (f *Flow) writeSuccessPage(w http.ResponseWriter) {
        f.setSecurityHeaders(w)
        htmlContent := `
<!DOCTYPE html>
<html>
<head>
    <title>Authentication Successful</title>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <style>
        body { font-family: Arial, sans-serif; margin: 40px; text-align: center; }
        .container { max-width: 600px; margin: 0 auto; }
        .message { padding: 20px; border-radius: 5px; margin: 20px 0; }
        .success { background-color: #e7f6e7; border: 1px solid #b3e6b3; color: #006600; }
    </style>
</head>
<body>
    <div class="container">
        <h1>Authentication Successful!</h1>
        <div class="message success">
            <p>You have successfully authenticated with ToolHive. You can now close this window and return to the terminal.</p>
        </div>
    </div>
</body>
</html>`
        if _, err := w.Write([]byte(htmlContent)); err != nil {
                slog.Warn("Failed to write HTML content", "error", err)
        }
}

// writeErrorPage writes an error page to the response
func (*Flow) writeErrorPage(w http.ResponseWriter, err error) {
        w.Header().Set("Content-Type", "text/html; charset=utf-8")
        w.Header().Set("X-Content-Type-Options", "nosniff")
        w.Header().Set("X-Frame-Options", "DENY")
        w.Header().Set("X-XSS-Protection", "1; mode=block")
        w.WriteHeader(http.StatusBadRequest)

        // HTML escape the error message to prevent XSS
        escapedError := html.EscapeString(err.Error())
        htmlContent := fmt.Sprintf(`
<!DOCTYPE html>
<html>
<head>
    <title>Authentication Failed</title>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <style>
        body { font-family: Arial, sans-serif; margin: 40px; text-align: center; }
        .container { max-width: 600px; margin: 0 auto; }
        .message { padding: 20px; border-radius: 5px; margin: 20px 0; }
        .error { background-color: #ffe7e7; border: 1px solid #ffb3b3; color: #cc0000; }
    </style>
</head>
<body>
    <div class="container">
        <h1>Authentication Failed</h1>
        <div class="message error">
            <p>%s</p>
            <p>Please try again or contact support if the problem persists.</p>
        </div>
    </div>
</body>
</html>`, escapedError)
        if _, err := w.Write([]byte(htmlContent)); err != nil {
                slog.Warn("Failed to write HTML content", "error", err)
        }
}

// processToken processes the received token and extracts claims
func (f *Flow) processToken(_ context.Context, token *oauth2.Token) *TokenResult {
        result := &TokenResult{
                AccessToken:  token.AccessToken,
                RefreshToken: token.RefreshToken,
                TokenType:    token.TokenType,
                Expiry:       token.Expiry,
        }

        // Create a base token source using the original token with a background context.
        // We use context.Background() instead of the passed ctx because the TokenSource
        // is long-lived and will be used for token refresh operations long after the
        // initial OAuth flow completes. Using the original ctx would cause "context canceled"
        // errors when attempting to refresh tokens, as that context gets cancelled when
        // the OAuth callback server shuts down.
        var base oauth2.TokenSource
        if f.config.Resource != "" {
                // Use resourceTokenSource wrapper to add resource parameter to refresh requests (RFC 8707)
                base = NewResourceTokenSource(f.oauth2Config, token, f.config.Resource, f.httpClient)
        } else {
                ctx := context.WithValue(context.Background(), oauth2.HTTPClient, f.httpClient)
                base = f.oauth2Config.TokenSource(ctx, token)
        }

        // ReuseTokenSource ensures that refresh happens only when needed
        f.tokenSource = oauth2.ReuseTokenSource(token, base)

        // Prefer extracting claims from the ID token if present (OIDC, e.g., Google)
        if idToken, ok := token.Extra("id_token").(string); ok && idToken != "" {
                result.IDToken = idToken
                if claims, err := f.extractJWTClaims(idToken); err == nil {
                        result.Claims = claims
                        slog.Debug("Successfully extracted JWT claims from ID token")
                } else {
                        slog.Debug("Could not extract JWT claims from ID token", "error", err)
                }
        } else {
                // Fallback: try to extract claims from the access token (e.g., Keycloak)
                if claims, err := f.extractJWTClaims(token.AccessToken); err == nil {
                        result.Claims = claims
                        slog.Debug("Successfully extracted JWT claims from access token")
                } else {
                        slog.Debug("Could not extract JWT claims from access token (may be opaque token)", "error", err)
                }
        }

        return result
}

// TokenSource returns the OAuth2 token source for refreshing tokens
func (f *Flow) TokenSource() oauth2.TokenSource {
        return f.tokenSource
}

// extractJWTClaims attempts to extract claims from a JWT token without validation
func (*Flow) extractJWTClaims(tokenString string) (jwt.MapClaims, error) {
        // Parse without verification to extract claims
        parser := jwt.NewParser(jwt.WithoutClaimsValidation())
        token, _, err := parser.ParseUnverified(tokenString, jwt.MapClaims{})
        if err != nil {
                return nil, err
        }

        claims, ok := token.Claims.(jwt.MapClaims)
        if !ok {
                return nil, errors.New("failed to extract claims")
        }

        return claims, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package oauth provides OAuth 2.0 and OIDC authentication functionality.
package oauth

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/networking"
)

// CreateOAuthConfigManual creates an OAuth config with manually provided endpoints
func CreateOAuthConfigManual(
        clientID, clientSecret string,
        authURL, tokenURL string,
        scopes []string,
        usePKCE bool,
        callbackPort int,
        resource string,
        oauthParams map[string]string,
        scopeParamName string,
) (*Config, error) {
        if clientID == "" {
                return nil, fmt.Errorf("client ID is required")
        }
        if authURL == "" {
                return nil, fmt.Errorf("authorization URL is required")
        }
        if tokenURL == "" {
                return nil, fmt.Errorf("token URL is required")
        }

        // Validate URLs
        if err := networking.ValidateEndpointURL(authURL); err != nil {
                return nil, fmt.Errorf("invalid authorization URL: %w", err)
        }
        if err := networking.ValidateEndpointURL(tokenURL); err != nil {
                return nil, fmt.Errorf("invalid token URL: %w", err)
        }

        // Default scopes for regular OAuth (don't assume OIDC scopes)
        if len(scopes) == 0 {
                scopes = []string{}
        }

        return &Config{
                ClientID:       clientID,
                ClientSecret:   clientSecret,
                AuthURL:        authURL,
                TokenURL:       tokenURL,
                Scopes:         scopes,
                UsePKCE:        usePKCE,
                CallbackPort:   callbackPort,
                Resource:       resource,
                OAuthParams:    oauthParams,
                ScopeParamName: scopeParamName,
        }, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauth

import (
        "context"
        "fmt"
        "log/slog"
        "net/http"
        "strings"
        "sync"
        "time"

        "golang.org/x/oauth2"
)

// NonCachingRefresher is an oauth2.TokenSource that always performs a network
// refresh when Token() is called — it holds no internal token cache.
//
// This is the correct innermost source for a preemptive-refresh chain:
// the outer oauth2.ReuseTokenSource provides caching; the inner source must
// always refresh when asked so that one network round-trip occurs per
// preemptive window instead of looping indefinitely.
//
// It handles both standard OAuth 2.0 refresh (resource == "") and RFC 8707
// resource-indicator refresh (resource != "") in a single type.
//
// Token() is safe for concurrent use. mu serializes access to refreshToken,
// which is updated in place when the IdP rotates it.
//
// When the IdP omits a new refresh token the previous token is preserved so
// the session survives providers that do not rotate on every refresh.
type NonCachingRefresher struct {
        mu           sync.Mutex
        cfg          *oauth2.Config
        resource     string // RFC 8707 resource indicator; empty for standard OAuth 2.0
        refreshToken string
        httpClient   *http.Client
}

// NewNonCachingRefresher creates a NonCachingRefresher that refreshes using
// cfg and the given refresh token. resource is the RFC 8707 resource indicator;
// pass "" for standard OAuth 2.0 refresh. httpClient carries the token
// endpoint's SSRF policy and must be built by NewTokenHTTPClient.
func NewNonCachingRefresher(
        cfg *oauth2.Config,
        refreshToken, resource string,
        httpClient *http.Client,
) *NonCachingRefresher {
        if cfg == nil {
                panic("oauth.NewNonCachingRefresher: config must not be nil")
        }
        if httpClient == nil {
                panic("oauth.NewNonCachingRefresher: HTTP client must not be nil")
        }
        return &NonCachingRefresher{
                cfg:          cfg,
                resource:     resource,
                refreshToken: refreshToken,
                httpClient:   httpClient,
        }
}

// Token always performs a token-endpoint refresh. It updates the stored refresh
// token when the IdP rotates it so callers (e.g. PersistingTokenSource) can
// detect the change and persist it.
func (n *NonCachingRefresher) Token() (*oauth2.Token, error) {
        n.mu.Lock()
        defer n.mu.Unlock()

        if n.refreshToken == "" {
                return nil, fmt.Errorf("no refresh token available")
        }

        ctx := context.WithValue(context.Background(), oauth2.HTTPClient, n.httpClient)

        var (
                tok *oauth2.Token
                err error
        )
        if n.resource != "" {
                tok, err = n.refreshWithResource(ctx)
        } else {
                tok, err = n.refreshStandard(ctx)
        }
        if err != nil {
                return nil, err
        }

        if tok.RefreshToken == "" {
                tok.RefreshToken = n.refreshToken
        } else {
                n.refreshToken = tok.RefreshToken
        }

        slog.Debug("Token refreshed", "resource", n.resource)
        return tok, nil
}

// refreshStandard uses the library's native tokenRefresher path, which sends
// grant_type=refresh_token without the empty code= parameter that Exchange
// always appends. A fresh TokenSource is constructed on each call so there is
// no internal cache. Scopes are not sent explicitly; servers must preserve them
// per RFC 6749 §6 (MUST).
func (n *NonCachingRefresher) refreshStandard(ctx context.Context) (*oauth2.Token, error) {
        // Passing an always-expired token forces the ReuseTokenSource returned by
        // cfg.TokenSource to call the inner tokenRefresher immediately.
        ts := n.cfg.TokenSource(ctx, &oauth2.Token{
                RefreshToken: n.refreshToken,
                Expiry:       time.Unix(1, 0),
        })
        tok, err := ts.Token()
        if err != nil {
                return nil, fmt.Errorf("token refresh failed: %w", err)
        }
        return tok, nil
}

// refreshWithResource uses cfg.Exchange with overridden grant_type and
// refresh_token parameters to send the RFC 8707 resource= indicator.
// golang.org/x/oauth2 has no native support for resource indicators, so the
// Exchange workaround is unavoidable here. The side-effect empty code=
// parameter from Exchange is acceptable: resource-indicator IdPs are required
// to dispatch on grant_type first and typically tolerate extra parameters.
// Scopes are sent explicitly because some resource-indicator IdPs do not
// preserve them when omitted despite the RFC 6749 §6 MUST requirement.
func (n *NonCachingRefresher) refreshWithResource(ctx context.Context) (*oauth2.Token, error) {
        opts := []oauth2.AuthCodeOption{
                oauth2.SetAuthURLParam("grant_type", "refresh_token"),
                oauth2.SetAuthURLParam("refresh_token", n.refreshToken),
                oauth2.SetAuthURLParam("resource", n.resource),
        }
        if len(n.cfg.Scopes) > 0 {
                opts = append(opts, oauth2.SetAuthURLParam("scope", strings.Join(n.cfg.Scopes, " ")))
        }
        tok, err := n.cfg.Exchange(ctx, "", opts...)
        if err != nil {
                return nil, fmt.Errorf("token refresh failed: %w", err)
        }
        return tok, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package oauth provides OAuth 2.0 and OIDC authentication functionality.
package oauth

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "net/url"
        "path"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// DiscoverOIDCEndpoints discovers OAuth endpoints from an OIDC issuer.
//
// issuer originates from untrusted remote-server discovery in some callers
// (e.g. the CLI DCR flow's well-known fallback) and from an operator-supplied
// flag in others; blockPrivateIPs lets each caller apply its own SSRF policy
// (CWE-918) rather than this function silently choosing one for everyone.
func DiscoverOIDCEndpoints(
        ctx context.Context,
        issuer string,
        blockPrivateIPs bool,
) (*oauthproto.OIDCDiscoveryDocument, error) {
        return discoverOIDCEndpointsWithClient(ctx, issuer, nil, false, blockPrivateIPs)
}

// DiscoverActualIssuer discovers the actual issuer from a URL that might be different from the issuer itself
// This is useful when the resource metadata points to a URL that hosts the authorization server metadata
// but the actual issuer identifier is different (e.g., Stripe's case)
//
// metadataURL originates from untrusted remote-server discovery. When
// blockPrivateIPs is true the default client refuses to dial private, loopback,
// or link-local addresses on every hop (SSRF guard); callers pass false when the
// operator-configured target is itself internal.
func DiscoverActualIssuer(
        ctx context.Context,
        metadataURL string,
        blockPrivateIPs bool,
) (*oauthproto.OIDCDiscoveryDocument, error) {
        return discoverOIDCEndpointsWithClientAndValidation(ctx, metadataURL, nil, false, false, blockPrivateIPs)
}

// discoverOIDCEndpointsWithClient discovers OAuth endpoints from an OIDC
// issuer with a custom HTTP client (private for testing). blockPrivateIPs is
// only consulted when client is nil — a caller-supplied client carries its
// own dial policy.
func discoverOIDCEndpointsWithClient(
        ctx context.Context,
        issuer string,
        client networking.HTTPClient,
        insecureAllowHTTP bool,
        blockPrivateIPs bool,
) (*oauthproto.OIDCDiscoveryDocument, error) {
        return discoverOIDCEndpointsWithClientAndValidation(ctx, issuer, client, true, insecureAllowHTTP, blockPrivateIPs)
}

// newOIDCDiscoveryTransport builds the transport used to fetch OIDC/OAuth
// discovery documents from an untrusted, remote-server-supplied issuer. The
// caller pairs it with a same-host redirect policy to prevent a 30x driving the
// host into an SSRF (CWE-918); when blockPrivateIPs is true it also refuses to
// dial private/loopback/link-local addresses on every hop (and disables
// keep-alive so a pooled connection cannot skip the per-dial check). Extracted
// so the pool bounds networking.SetIdleConnBounds applies stay unit-testable.
func newOIDCDiscoveryTransport(blockPrivateIPs bool) *http.Transport {
        transport := &http.Transport{
                TLSHandshakeTimeout:   10 * time.Second,
                ResponseHeaderTimeout: 10 * time.Second,
        }
        networking.SetIdleConnBounds(transport)
        if blockPrivateIPs {
                transport.DialContext = networking.NewPrivateIPBlockingDialContext()
                transport.DisableKeepAlives = true
        }
        return transport
}

// discoverOIDCEndpointsWithClientAndValidation discovers OAuth endpoints with optional issuer validation
//
//nolint:gocyclo // Function complexity justified by comprehensive OIDC discovery logic
func discoverOIDCEndpointsWithClientAndValidation(
        ctx context.Context,
        issuer string,
        client networking.HTTPClient,
        validateIssuer bool,
        insecureAllowHTTP bool,
        blockPrivateIPs bool,
) (*oauthproto.OIDCDiscoveryDocument, error) {

        oidcURL, oauthURL, err := buildWellKnownURLs(issuer, insecureAllowHTTP)
        if err != nil {
                return nil, fmt.Errorf("failed to build well-known URLs: %w", err)
        }

        if client == nil {
                client = &http.Client{
                        Timeout:       30 * time.Second,
                        Transport:     newOIDCDiscoveryTransport(blockPrivateIPs),
                        CheckRedirect: networking.SameHostRedirectPolicy(),
                }
        }

        try := func(urlStr string, oidc bool) (*oauthproto.OIDCDiscoveryDocument, error) {
                req, err := http.NewRequestWithContext(ctx, http.MethodGet, urlStr, nil)
                if err != nil {
                        return nil, fmt.Errorf("build request: %w", err)
                }
                req.Header.Set("User-Agent", oauthproto.UserAgent)
                req.Header.Set("Accept", "application/json")

                resp, err := client.Do(req)
                if err != nil {
                        return nil, fmt.Errorf("GET %s: %w", urlStr, err)
                }
                defer func() {
                        if err := resp.Body.Close(); err != nil {
                                slog.Debug("Failed to close response body", "error", err)
                        }
                }()

                if resp.StatusCode != http.StatusOK {
                        return nil, fmt.Errorf("%s: HTTP %d", urlStr, resp.StatusCode)
                }
                ct := strings.ToLower(resp.Header.Get("Content-Type"))
                if !strings.Contains(ct, "application/json") {
                        return nil, fmt.Errorf("%s: unexpected content-type %q", urlStr, ct)
                }

                // Limit response size to prevent DoS
                const maxResponseSize = 1024 * 1024 // 1MB
                var doc oauthproto.OIDCDiscoveryDocument
                if err := json.NewDecoder(io.LimitReader(resp.Body, maxResponseSize)).Decode(&doc); err != nil {
                        return nil, fmt.Errorf("%s: unexpected response: %w", urlStr, err)
                }
                expectedIssuer := issuer
                if !validateIssuer && doc.Issuer != "" {
                        // When not validating, use the discovered issuer as the expected one
                        expectedIssuer = doc.Issuer
                }
                if err := validateOIDCDocument(&doc, expectedIssuer, oidc, insecureAllowHTTP); err != nil {
                        return nil, fmt.Errorf("%s: invalid metadata: %w", urlStr, err)
                }
                return &doc, nil
        }

        doc, err := try(oidcURL, true)
        if err == nil {
                // OIDC discovery succeeded, but check if we need registration_endpoint
                // If it's missing, try OAuth authorization server well-known URL as fallback
                if doc.RegistrationEndpoint == "" {
                        slog.Debug("OIDC discovery succeeded but registration_endpoint is missing, " +
                                "trying OAuth authorization server well-known URL")
                        oauthDoc, oauthErr := try(oauthURL, false)
                        if oauthErr == nil && oauthDoc.RegistrationEndpoint != "" {
                                // Validate issuer matches before merging
                                if oauthDoc.Issuer == doc.Issuer {
                                        doc.RegistrationEndpoint = oauthDoc.RegistrationEndpoint
                                        slog.Debug("Found registration_endpoint in OAuth authorization server metadata", "endpoint", doc.RegistrationEndpoint)
                                        // Merge CIMD support flag — some servers (e.g. Granola) only advertise
                                        // client_id_metadata_document_supported in the OAuth AS metadata, not
                                        // in the OIDC discovery document.
                                        if oauthDoc.ClientIDMetadataDocumentSupported {
                                                doc.ClientIDMetadataDocumentSupported = true
                                        }
                                } else {
                                        slog.Warn("Issuer mismatch between OIDC and OAuth discovery documents, not merging registration_endpoint",
                                                "oidc_issuer", doc.Issuer, "oauth_issuer", oauthDoc.Issuer)
                                }
                        }
                }
                return doc, nil
        }
        oidcErr := err
        doc, err = try(oauthURL, false)
        if err == nil {
                return doc, nil
        }
        oauthErr := err

        return nil, fmt.Errorf("unable to discover OIDC endpoints at %q or %q: OIDC error: %v, OAuth error: %v",
                oidcURL, oauthURL, oidcErr, oauthErr)
}

// validateOIDCDocument validates the OIDC discovery document.
// It delegates basic field presence validation to the shared doc.Validate() method,
// then performs additional security checks (issuer mismatch, URL HTTPS validation).
func validateOIDCDocument(
        doc *oauthproto.OIDCDiscoveryDocument,
        expectedIssuer string,
        oidc bool,
        insecureAllowHTTP bool,
) error {
        // Delegate basic field presence validation to the shared method. We pass
        // oidc=false because this legacy path handles jwks_uri separately below and
        // does not require response_types_supported. It nevertheless always creates
        // an interactive authorization request, so authorization_endpoint is required
        // regardless of the remote server's advertised grant types.
        if err := doc.Validate(false); err != nil {
                return err
        }
        if doc.AuthorizationEndpoint == "" {
                return oauthproto.ErrMissingAuthorizationEndpoint
        }

        // Require jwks_uri for OIDC (with specific error message)
        if oidc && doc.JWKSURI == "" {
                return fmt.Errorf("missing jwks_uri (OIDC requires it)")
        }

        // Security check: issuer must match expected value
        if doc.Issuer != expectedIssuer {
                return fmt.Errorf("issuer mismatch: expected %s, got %s", expectedIssuer, doc.Issuer)
        }

        // Security check: validate endpoint URLs (HTTPS required unless insecureAllowHTTP)
        endpoints := map[string]string{
                "authorization_endpoint": doc.AuthorizationEndpoint,
                "token_endpoint":         doc.TokenEndpoint,
                "jwks_uri":               doc.JWKSURI,
                "introspection_endpoint": doc.IntrospectionEndpoint,
        }

        if doc.UserinfoEndpoint != "" {
                endpoints["userinfo_endpoint"] = doc.UserinfoEndpoint
        }
        for name, endpoint := range endpoints {
                if endpoint != "" {
                        if err := networking.ValidateEndpointURLWithInsecure(endpoint, insecureAllowHTTP); err != nil {
                                return fmt.Errorf("invalid %s: %w", name, err)
                        }
                }
        }
        return nil
}

// CreateOAuthConfigFromOIDC creates an OAuth config from OIDC discovery
func CreateOAuthConfigFromOIDC(
        ctx context.Context,
        issuer, clientID, clientSecret string,
        scopes []string,
        usePKCE bool,
        callbackPort int,
        resource string,
        blockPrivateIPs bool,
) (*Config, error) {
        return createOAuthConfigFromOIDCWithClient(
                ctx, issuer, clientID, clientSecret, scopes, usePKCE, callbackPort, resource, nil, blockPrivateIPs,
        )
}

// createOAuthConfigFromOIDCWithClient creates an OAuth config from OIDC discovery with a custom HTTP client (private for testing)
func createOAuthConfigFromOIDCWithClient(
        ctx context.Context,
        issuer, clientID, clientSecret string,
        scopes []string,
        usePKCE bool,
        callbackPort int,
        resource string,
        client networking.HTTPClient,
        blockPrivateIPs bool,
) (*Config, error) {
        // Discover OIDC endpoints (insecureAllowHTTP is false for OAuth config creation).
        doc, err := discoverOIDCEndpointsWithClient(ctx, issuer, client, false, blockPrivateIPs)
        if err != nil {
                return nil, fmt.Errorf("failed to discover OIDC endpoints: %w", err)
        }

        // Default scopes for OIDC if none provided
        if len(scopes) == 0 {
                scopes = []string{"openid", "profile", "email"}
        }

        // Enable PKCE by default if supported
        if !usePKCE && len(doc.CodeChallengeMethodsSupported) > 0 {
                for _, method := range doc.CodeChallengeMethodsSupported {
                        if method == oauthproto.PKCEMethodS256 {
                                usePKCE = true
                                break
                        }
                }
        }

        return &Config{
                ClientID:              clientID,
                ClientSecret:          clientSecret,
                AuthURL:               doc.AuthorizationEndpoint,
                IntrospectionEndpoint: doc.IntrospectionEndpoint,
                TokenURL:              doc.TokenEndpoint,
                Scopes:                scopes,
                UsePKCE:               usePKCE,
                CallbackPort:          callbackPort,
                Resource:              resource,
        }, nil
}

func buildWellKnownURLs(issuer string, insecureAllowHTTP bool) (oidcURL, oauthURL string, err error) {
        // Parse issuer
        issuerURL, err := url.Parse(issuer)
        if err != nil {
                return "", "", fmt.Errorf("invalid issuer URL: %w", err)
        }

        // Enforce HTTPS except localhost or explicit override
        if issuerURL.Scheme != "https" &&
                !networking.IsLocalhost(issuerURL.Host) &&
                !insecureAllowHTTP {
                return "", "", fmt.Errorf("issuer must use HTTPS: %s (use insecureAllowHTTP for testing only)", issuer)
        }

        // Extract tenant/realm path (may be nested)
        tenant := strings.Trim(issuerURL.EscapedPath(), "/")

        // Base = scheme + host
        base := &url.URL{
                Scheme: issuerURL.Scheme,
                Host:   issuerURL.Host,
        }

        //
        // Build OIDC Discovery URL
        //   /{tenant}/.well-known/openid-configuration
        //
        oidc := *base
        oidc.Path = path.Join("/", tenant, oauthproto.WellKnownOIDCPath)
        oidcURL = oidc.String()

        //
        // Build OAuth AS Metadata URL
        //   /.well-known/oauth-authorization-server/{tenant}
        //
        oauth := *base
        oauth.Path = path.Join(oauthproto.WellKnownOAuthServerPath, tenant)
        oauthURL = oauth.String()

        return oidcURL, oauthURL, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauth

import (
        "net/http"

        "golang.org/x/oauth2"
)

// resourceTokenSource wraps a NonCachingRefresher with an internal token cache
// and adds the resource parameter to token refresh requests per RFC 8707.
// Token() returns the cached token when still valid and delegates to the
// inner NonCachingRefresher only when a refresh is needed.
type resourceTokenSource struct {
        ncr   *NonCachingRefresher
        token *oauth2.Token
}

// NewResourceTokenSource creates a token source that includes the resource parameter
// in all token requests, including refresh requests.
// The resource parameter must be non-empty (caller should check before calling).
// httpClient carries the token endpoint's SSRF policy and must be built by
// NewTokenHTTPClient.
func NewResourceTokenSource(
        config *oauth2.Config,
        token *oauth2.Token,
        resource string,
        httpClient *http.Client,
) oauth2.TokenSource {
        if token == nil {
                panic("oauth.NewResourceTokenSource: token must not be nil")
        }
        return &resourceTokenSource{
                ncr:   NewNonCachingRefresher(config, token.RefreshToken, resource, httpClient),
                token: token,
        }
}

// Token returns a valid token, refreshing it if necessary.
// When refreshing, it delegates to NonCachingRefresher which adds the resource
// parameter per RFC 8707.
func (r *resourceTokenSource) Token() (*oauth2.Token, error) {
        if r.token.Valid() {
                return r.token, nil
        }

        newToken, err := r.ncr.Token()
        if err != nil {
                return nil, err
        }

        r.token = newToken
        return newToken, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package oautherr classifies OAuth token-endpoint failures as transient or
// permanent.
//
// The distinction matters wherever a cached credential is exchanged without a
// human present: a transient failure (the IdP is down, rate-limiting, or a WAF
// answered instead of it) is worth retrying, while a permanent one means the
// OAuth server itself rendered a verdict on the credential and no amount of
// retrying will help — the only way forward is a fresh interactive login.
//
// It is a leaf package so both the token-source construction path and the
// workload auth monitor can share one implementation instead of each carrying
// its own copy of the rules.
package oautherr

import (
        "errors"
        "net/http"
        "strings"

        "golang.org/x/oauth2"
)

// IsTransientRetrieveError reports whether an *oauth2.RetrieveError should be
// treated as transient. The classification rules are:
//
//   - nil Response: non-transient. There is no signal to act on, so callers
//     fall through to the unauthenticated path rather than retry blindly.
//   - 5xx status: transient (server-side issue, likely to resolve).
//   - 429 Too Many Requests: transient regardless of body (HTTP standard).
//   - 4xx with an empty ErrorCode: transient. The oauth2 library populates
//     ErrorCode from the RFC 6749 'error' field in a JSON response body. An
//     empty ErrorCode means the response was not a parseable OAuth error —
//     typically an HTML page from a WAF, CDN, or reverse proxy that
//     intercepted the request before it reached the OAuth server. These
//     infrastructure errors (Cloudflare blocks, residential-IP allowlist
//     misses, transient bad-config deploys) commonly resolve on their own.
//   - 4xx with a populated ErrorCode: permanent. The OAuth server returned
//     a structured error code (invalid_grant, invalid_client, etc.) telling
//     us specifically what's wrong; retrying won't help.
func IsTransientRetrieveError(retrieveErr *oauth2.RetrieveError) bool {
        if retrieveErr == nil || retrieveErr.Response == nil {
                return false
        }
        statusCode := retrieveErr.Response.StatusCode

        if statusCode >= 500 {
                return true
        }
        if statusCode == http.StatusTooManyRequests {
                return true
        }
        if retrieveErr.ErrorCode == "" {
                return true
        }
        return false
}

// IsPermanentCredentialError reports whether err is an *oauth2.RetrieveError
// whose response carries a structured RFC 6749 'error' code, implying the
// OAuth server itself rendered a verdict on the cached credentials
// (invalid_grant, invalid_client, etc.).
//
// This is the strict inverse of IsTransientRetrieveError on the
// *oauth2.RetrieveError branch: a response is "permanent" iff the classifier
// would NOT call it transient. Concretely, it is true only when ErrorCode is
// populated. 4xx responses without an OAuth error code (HTML pages from a WAF,
// CDN, or reverse proxy) — like 5xx, 429, 408, and nil-Response shapes — are
// treated as non-permanent because there is no OAuth-protocol verdict to act
// on. Telling a user their stored credential is dead based on a
// non-spec-compliant response would frequently mislead operators whose real
// problem is upstream of the OAuth server.
//
// "Permanent" here means only "retrying will not help". It does NOT mean the
// stored credential is dead — invalid_client and unauthorized_client indict the
// client registration, not the refresh token. Callers deciding whether to tell
// a user to log in again want IsRejectedRefreshGrant instead.
func IsPermanentCredentialError(err error) bool {
        retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err)
        if !ok || retrieveErr.Response == nil {
                return false
        }
        return !IsTransientRetrieveError(retrieveErr)
}

// rejectedGrantCode is the RFC 6749 'error' code a token endpoint returns when
// the grant it was handed is invalid, expired, revoked, or was issued to a
// different client. On a refresh exchange it is the one code that indicts the
// stored refresh token itself.
const rejectedGrantCode = "invalid_grant"

// IsRejectedRefreshGrant reports whether err is an *oauth2.RetrieveError in
// which the OAuth server rejected the refresh token itself: an RFC 6749
// 'invalid_grant' on a response IsTransientRetrieveError does not already
// excuse.
//
// This is deliberately narrower than IsPermanentCredentialError, and the two
// answer different questions. Both describe failures retrying cannot fix, but
// only invalid_grant means a fresh interactive login is the remedy.
// invalid_client, unauthorized_client, and invalid_scope are just as permanent
// and just as pointless to retry, yet they indict the client registration or
// the requested scopes — running the login flow again against the same broken
// configuration reproduces them exactly, so telling the user to log in again
// sends them in a circle. Callers rendering a re-login remediation must key off
// this predicate; callers that only need "stop retrying" want the broader one.
func IsRejectedRefreshGrant(err error) bool {
        retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err)
        if !ok || retrieveErr.ErrorCode != rejectedGrantCode {
                return false
        }
        return IsPermanentCredentialError(err)
}

// IsParseError detects errors from the oauth2 library that indicate the token
// endpoint returned an unparsable response body on a 2xx status. This
// typically happens when a load balancer, CDN, or reverse proxy intercepts the
// request and returns its own HTML page instead of the expected JSON token
// response. The oauth2 library uses fmt.Errorf with %v (not %w) for these
// errors, so string matching is the only reliable detection method.
func IsParseError(err error) bool {
        if err == nil {
                return false
        }
        msg := err.Error()
        return strings.Contains(msg, "oauth2: cannot parse json") ||
                strings.Contains(msg, "oauth2: cannot parse response")
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package obo

import "errors"

// ErrEnterpriseRequired is returned by every default OBO dispatch point — the
// controllerutil handler hook, the vMCP converter stub, and the middleware
// stub — when no out-of-tree handler/factory has been registered. Callers must
// use errors.Is to compare; the error wraps cleanly through
// fmt.Errorf("...: %w", ...).
//
// Lives in pkg/auth/obo (a leaf package) so that callers in
// cmd/thv-operator/... and pkg/vmcp/... can share the same sentinel without
// either layer importing the other. To register an out-of-tree handler, see
// controllerutil.RegisterOBOHandler (for the operator dispatch points) and
// obo.RegisterFactory (for the proxy middleware factory).
var ErrEnterpriseRequired = errors.New(
        "on-behalf-of (OBO) external auth type requires an enterprise build")

// ValidationError is the typed error an OBO handler returns when its input
// is genuinely malformed and the user must fix the spec for the failure to
// clear. It is the contract for the "permanent, user-fix" bucket in the
// OBOHandler error triage:
//
//   - errors.Is(err, ErrEnterpriseRequired) → not licensed; permanent until an
//     out-of-tree handler is registered.
//   - errors.As(err, &*ValidationError) → permanent until the user edits the
//     spec; the operator writes condition.Reason=InvalidConfig and does not
//     requeue.
//   - anything else → treated as transient by the reconciler, which returns
//     the error so controller-runtime requeues with backoff.
//
// Handler authors must use ValidationError for any condition the user can
// fix by editing the spec (missing field, malformed URL, schema violation)
// and must return a non-ValidationError for transient I/O failures (Secret
// not yet available, JWKS unreachable, webhook 5xx) so the reconciler
// retries instead of locking the resource into a permanent InvalidConfig
// state.
//
// The Message field is written verbatim into condition.Message — handler
// authors are responsible for ensuring it is safe to expose (no Secret
// names, no internal addressing, no credential fragments).
type ValidationError struct {
        Message string
}

// Error returns the user-facing message verbatim. The OBO branch in the
// MCPExternalAuthConfig reconciler writes this string into the Valid
// condition's Message field.
func (e *ValidationError) Error() string { return e.Message }

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package obo provides the proxy-runtime middleware factory hook for the
// on-behalf-of (OBO) external auth type. The default factory produces a
// stub middleware that responds 503 to every request. An out-of-tree build
// replaces the factory by calling RegisterFactory once during init().
package obo

import (
        "net/http"
        "sync"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// MiddlewareType is the type identifier used in MiddlewareConfig.Type for
// OBO middleware. Matches the ExternalAuthType constant value "obo".
const MiddlewareType = "obo"

// stubMessage is the body returned by the default 503 handler. It is shared
// between the default handler and tests in this package so the exact literal
// is only spelled once. Unexported because no caller outside this package has
// a use for it — tests reference it directly.
const stubMessage = "obo requires a build with a registered OBO middleware factory"

// stub is the default Middleware implementation. Its handler responds 503
// with a body explaining that the OBO middleware factory has not been
// registered. It is never instantiated except by the default CreateMiddleware
// factory, which an out-of-tree build replaces via RegisterFactory.
type stub struct{}

// Handler returns a MiddlewareFunction that wraps any downstream handler with
// a 503 response. The downstream handler is intentionally never called — the
// stub is the user-visible signal that the OBO middleware factory has not
// been registered.
func (*stub) Handler() types.MiddlewareFunction {
        return func(http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
                        http.Error(w, stubMessage, http.StatusServiceUnavailable)
                })
        }
}

// Close is a no-op for the stub middleware.
func (*stub) Close() error { return nil }

// factoryMu guards reads and writes of currentFactory. Each call to
// CreateMiddleware takes the read lock; RegisterFactory takes the write lock.
// The mutex makes hot-reload / re-registration safe in case a future caller
// (e.g. an admission webhook) wires it up, even though production today only
// ever writes during init().
var factoryMu sync.RWMutex

// currentFactory holds the actual middleware factory dispatched to by
// CreateMiddleware. The default returns a 503 stub; an out-of-tree build
// replaces it via RegisterFactory.
var currentFactory types.MiddlewareFactory = DefaultFactory

// DefaultFactory adds a stub middleware whose handler responds 503 to every
// request. Exposed primarily so external test code (e.g. pkg/runner) can pass
// it to RegisterFactory in a t.Cleanup to restore the package's default
// behavior after a test mutates currentFactory.
func DefaultFactory(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        runner.AddMiddleware(config.Type, &stub{})
        return nil
}

// CreateMiddleware is the package-level middleware factory. It is a stable
// indirection over currentFactory: each call dispatches to whatever factory
// is registered at call time, so out-of-tree builds replacing the factory
// via RegisterFactory take effect on subsequent calls even if a caller has
// already captured CreateMiddleware itself (e.g. pkg/runner builds its
// factory map once and reuses it across runner instances). The default
// produces a 503 stub.
//
// Declared as a function (matching sibling middleware packages such as
// awssts, upstreamswap, and oauthproto/tokenexchange) so RegisterFactory is
// the only mutation path — there is no second escape hatch via direct
// assignment to CreateMiddleware.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        factoryMu.RLock()
        f := currentFactory
        factoryMu.RUnlock()
        return f(config, runner)
}

// RegisterFactory replaces the underlying middleware factory. Calling it
// more than once is allowed and last-write-wins, matching the existing
// pkg/config.RegisterProviderFactory precedent. Panics if f is nil — a nil
// factory would dispatch into a nil function on the next CreateMiddleware
// call, far from the registration site; surface the problem at init() time
// instead.
func RegisterFactory(f types.MiddlewareFactory) {
        if f == nil {
                panic("obo.RegisterFactory: factory is nil")
        }
        factoryMu.Lock()
        defer factoryMu.Unlock()
        currentFactory = f
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package remote

import (
        "time"

        "golang.org/x/oauth2"
)

// BearerTokenSource implements oauth2.TokenSource for static bearer tokens.
// It returns a token with the bearer token value as the access token.
type BearerTokenSource struct {
        token string
}

// NewBearerTokenSource creates a new BearerTokenSource with the provided bearer token.
func NewBearerTokenSource(bearerToken string) *BearerTokenSource {
        return &BearerTokenSource{
                token: bearerToken,
        }
}

// Token returns an oauth2.Token with the bearer token as the access token.
// For static bearer tokens, this always returns the same token.
func (b *BearerTokenSource) Token() (*oauth2.Token, error) {
        return &oauth2.Token{
                AccessToken: b.token,
                TokenType:   "Bearer",
                Expiry:      time.Time{}, // No expiry for static bearer tokens
        }, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package remote

import (
        "encoding/json"
        "errors"
        "fmt"
        "strings"
)

const cachedRefreshTokenEnvelopeVersion = 1

// errLegacyCachedRefreshToken marks a decode failure caused by the cached
// value not being JSON at all — i.e. a bare refresh token cached before the
// envelope format existed, as opposed to a corrupted or tampered envelope.
var errLegacyCachedRefreshToken = errors.New("legacy cached refresh token: not in envelope format")

// cachedRefreshTokenEnvelope binds a refresh token to the authorization server
// and token endpoint that issued it. It is intentionally stored in the secret
// manager as one opaque value so existing public persistence APIs remain stable.
type cachedRefreshTokenEnvelope struct {
        Version  int    `json:"version"`
        Token    string `json:"token"`
        Issuer   string `json:"issuer"`
        TokenURL string `json:"token_url"`
}

func encodeCachedRefreshToken(token, issuer, tokenURL string) (string, error) {
        if token == "" || issuer == "" || tokenURL == "" {
                return "", fmt.Errorf("cached refresh token envelope requires token, issuer, and token endpoint")
        }

        envelope, err := json.Marshal(cachedRefreshTokenEnvelope{
                Version:  cachedRefreshTokenEnvelopeVersion,
                Token:    token,
                Issuer:   issuer,
                TokenURL: tokenURL,
        })
        if err != nil {
                return "", fmt.Errorf("marshal cached refresh token envelope: %w", err)
        }
        return string(envelope), nil
}

func decodeCachedRefreshToken(value string) (*cachedRefreshTokenEnvelope, error) {
        var envelope cachedRefreshTokenEnvelope
        if err := json.Unmarshal([]byte(value), &envelope); err != nil {
                // A value that isn't even shaped like a JSON object was never an
                // envelope — it's a pre-upgrade bare refresh token. A value that
                // looks like a JSON object but fails to parse is a corrupted or
                // tampered envelope and must not be treated as legacy.
                if !looksLikeJSONObject(value) {
                        return nil, fmt.Errorf("%w: %w", errLegacyCachedRefreshToken, err)
                }
                return nil, fmt.Errorf("invalid cached refresh token envelope: %w", err)
        }
        if envelope.Version != cachedRefreshTokenEnvelopeVersion {
                return nil, fmt.Errorf("unsupported cached refresh token envelope version")
        }
        if envelope.Token == "" || envelope.Issuer == "" || envelope.TokenURL == "" {
                return nil, fmt.Errorf("incomplete cached refresh token envelope")
        }
        return &envelope, nil
}

func looksLikeJSONObject(value string) bool {
        return strings.HasPrefix(strings.TrimSpace(value), "{")
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package remote

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "net/url"
        "strings"
        "time"

        registry "github.com/stacklok/toolhive-core/registry/types"
        httpval "github.com/stacklok/toolhive-core/validation/http"
)

// Config holds authentication configuration for remote MCP servers.
// Supports OAuth/OIDC-based authentication with automatic discovery.
type Config struct {
        ClientID         string        `json:"client_id,omitempty" yaml:"client_id,omitempty"`
        ClientSecret     string        `json:"client_secret,omitempty" yaml:"client_secret,omitempty"` //nolint:gosec // G117
        ClientSecretFile string        `json:"client_secret_file,omitempty" yaml:"client_secret_file,omitempty"`
        Scopes           []string      `json:"scopes,omitempty" yaml:"scopes,omitempty"`
        SkipBrowser      bool          `json:"skip_browser,omitempty" yaml:"skip_browser,omitempty"`
        Timeout          time.Duration `json:"timeout,omitempty" yaml:"timeout,omitempty" swaggertype:"string" example:"5m"`
        CallbackPort     int           `json:"callback_port,omitempty" yaml:"callback_port,omitempty"`
        UsePKCE          bool          `json:"use_pkce" yaml:"use_pkce"`

        // Resource is the OAuth 2.0 resource indicator (RFC 8707).
        Resource string `json:"resource,omitempty" yaml:"resource,omitempty"`

        // OAuth endpoint configuration (from registry)
        Issuer       string `json:"issuer,omitempty" yaml:"issuer,omitempty"`
        AuthorizeURL string `json:"authorize_url,omitempty" yaml:"authorize_url,omitempty"`
        TokenURL     string `json:"token_url,omitempty" yaml:"token_url,omitempty"`

        // Headers for HTTP requests
        Headers []*registry.Header `json:"headers,omitempty" yaml:"headers,omitempty" swaggerignore:"true"`

        // Environment variables for the client
        EnvVars []*registry.EnvVar `json:"env_vars,omitempty" yaml:"env_vars,omitempty" swaggerignore:"true"`

        // OAuth parameters for server-specific customization
        OAuthParams map[string]string `json:"oauth_params,omitempty" yaml:"oauth_params,omitempty"`

        // ScopeParamName overrides the query parameter name used to send scopes in the
        // authorization URL. When empty, the standard "scope" parameter is used.
        // Some providers require a non-standard name (e.g., Slack uses "user_scope").
        ScopeParamName string `json:"scope_param_name,omitempty" yaml:"scope_param_name,omitempty"`

        // Bearer token configuration (alternative to OAuth)
        BearerToken     string `json:"bearer_token,omitempty" yaml:"bearer_token,omitempty"` //nolint:gosec // G117
        BearerTokenFile string `json:"bearer_token_file,omitempty" yaml:"bearer_token_file,omitempty"`

        // Cached OAuth token reference for persistence across restarts.
        // The refresh token is stored securely in the secret manager, and this field
        // contains the reference to retrieve it (e.g., "OAUTH_REFRESH_TOKEN_workload").
        // This enables session restoration without requiring a new browser-based login.
        CachedRefreshTokenRef string    `json:"cached_refresh_token_ref,omitempty" yaml:"cached_refresh_token_ref,omitempty"`
        CachedTokenExpiry     time.Time `json:"cached_token_expiry,omitempty" yaml:"cached_token_expiry,omitempty"`

        // Cached DCR client credentials for persistence across restarts.
        // These are obtained during Dynamic Client Registration and needed to refresh tokens.
        // ClientID is stored as plain text since it's public information.
        CachedClientID        string `json:"cached_client_id,omitempty" yaml:"cached_client_id,omitempty"`
        CachedClientSecretRef string `json:"cached_client_secret_ref,omitempty" yaml:"cached_client_secret_ref,omitempty"`
        // ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server).
        // A zero value means the secret does not expire.
        CachedSecretExpiry time.Time `json:"cached_secret_expiry,omitempty" yaml:"cached_secret_expiry,omitempty"`
        // CachedRegTokenRef is a secret manager reference to the registration_access_token
        // returned in the DCR response. Used for RFC 7592 client update operations.
        // Stored as a secret reference since it's sensitive.
        CachedRegTokenRef string `json:"cached_reg_token_ref,omitempty" yaml:"cached_reg_token_ref,omitempty"`

        // CachedCIMDClientID stores the CIMD metadata URL used as client_id when CIMD
        // authentication was used. Kept separate from CachedClientID (which holds
        // DCR-issued IDs) so the two can have independent lifecycles — DCR credential
        // rotation clears CachedClientID without touching the stable CIMD URL.
        // Read by resolveClientCredentials to send the correct client_id on token refresh.
        CachedCIMDClientID string `json:"cached_cimd_client_id,omitempty" yaml:"cached_cimd_client_id,omitempty"`
        // CachedRegClientURI is the registration_client_uri from the DCR response.
        // This is the endpoint used for RFC 7592 client read/update/delete operations.
        // Stored as plain text since it is not sensitive.
        CachedRegClientURI string `json:"cached_reg_client_uri,omitempty" yaml:"cached_reg_client_uri,omitempty"`
        // CachedTokenEndpointAuthMethod is the auth method used for the token endpoint
        // (e.g., "client_secret_basic", "none"). Persisted for RFC 7592 updates.
        CachedTokenEndpointAuthMethod string `json:"cached_token_auth_method,omitempty" yaml:"cached_token_auth_method,omitempty"`
        // CachedDCRCallbackPort is the callback port that was actually registered
        // during DCR. It may differ from CallbackPort when the requested port was
        // unavailable and a fallback port was selected.
        CachedDCRCallbackPort int `json:"cached_dcr_callback_port,omitempty" yaml:"cached_dcr_callback_port,omitempty"`

        // Resolved credentials are runtime-only. The exported fields retain their
        // persistable values so copying Config cannot leak plaintext into state.
        resolvedClientSecret string
        clientSecretResolved bool
        resolvedBearerToken  string
        bearerTokenResolved  bool
}

// ClearResolvedSecrets removes runtime-only credential values.
func (c *Config) ClearResolvedSecrets() {
        c.resolvedClientSecret = ""
        c.clientSecretResolved = false
        c.resolvedBearerToken = ""
        c.bearerTokenResolved = false
}

// SetResolvedClientSecret records a client secret resolved at runtime.
func (c *Config) SetResolvedClientSecret(value string) {
        c.resolvedClientSecret = value
        c.clientSecretResolved = true
}

// SetResolvedBearerToken records a bearer token resolved at runtime.
func (c *Config) SetResolvedBearerToken(value string) {
        c.resolvedBearerToken = value
        c.bearerTokenResolved = true
}

func (c *Config) clientSecret() string {
        if c.clientSecretResolved {
                return c.resolvedClientSecret
        }
        return c.ClientSecret
}

func (c *Config) bearerToken() string {
        if c.bearerTokenResolved {
                return c.resolvedBearerToken
        }
        return c.BearerToken
}

// BearerTokenEnvVarName is the environment variable name used for bearer token authentication.
// The bearer token will be read from this environment variable if not provided via flag or file.
// #nosec G101 - this is an environment variable name, not a credential
const BearerTokenEnvVarName = "TOOLHIVE_REMOTE_AUTH_BEARER_TOKEN"

// UnmarshalJSON implements custom JSON unmarshaling for backward compatibility
// This handles both the old PascalCase format and the new snake_case format
func (r *Config) UnmarshalJSON(data []byte) error {
        // Parse the JSON to check which format is being used
        var raw map[string]interface{}
        if err := json.Unmarshal(data, &raw); err != nil {
                return err
        }

        // Check if this is the old PascalCase format by looking for old field name
        // if one old field is present, then it's the old format
        if _, isOld := raw["ClientID"]; isOld {
                // Unmarshal using old PascalCase format
                var oldFormat struct {
                        ClientID         string             `json:"ClientID,omitempty"`
                        ClientSecret     string             `json:"ClientSecret,omitempty"` //nolint:gosec // G117
                        ClientSecretFile string             `json:"ClientSecretFile,omitempty"`
                        Scopes           []string           `json:"Scopes,omitempty"`
                        SkipBrowser      bool               `json:"SkipBrowser,omitempty"`
                        Timeout          time.Duration      `json:"Timeout,omitempty"`
                        CallbackPort     int                `json:"CallbackPort,omitempty"`
                        UsePKCE          bool               `json:"UsePKCE,omitempty"`
                        Issuer           string             `json:"Issuer,omitempty"`
                        AuthorizeURL     string             `json:"AuthorizeURL,omitempty"`
                        TokenURL         string             `json:"TokenURL,omitempty"`
                        Headers          []*registry.Header `json:"Headers,omitempty"`
                        EnvVars          []*registry.EnvVar `json:"EnvVars,omitempty"`
                        OAuthParams      map[string]string  `json:"OAuthParams,omitempty"`
                        BearerToken      string             `json:"BearerToken,omitempty"` //nolint:gosec // G117
                        BearerTokenFile  string             `json:"BearerTokenFile,omitempty"`
                }

                if err := json.Unmarshal(data, &oldFormat); err != nil {
                        return fmt.Errorf("failed to unmarshal Config in old format: %w", err)
                }

                // Copy from old format to new format
                r.ClientID = oldFormat.ClientID
                r.ClientSecret = oldFormat.ClientSecret
                r.ClientSecretFile = oldFormat.ClientSecretFile
                r.Scopes = oldFormat.Scopes
                r.SkipBrowser = oldFormat.SkipBrowser
                r.Timeout = oldFormat.Timeout
                r.CallbackPort = oldFormat.CallbackPort
                r.UsePKCE = oldFormat.UsePKCE
                r.Issuer = oldFormat.Issuer
                r.AuthorizeURL = oldFormat.AuthorizeURL
                r.TokenURL = oldFormat.TokenURL
                r.Headers = oldFormat.Headers
                r.EnvVars = oldFormat.EnvVars
                r.OAuthParams = oldFormat.OAuthParams
                r.BearerToken = oldFormat.BearerToken
                r.BearerTokenFile = oldFormat.BearerTokenFile
                r.ClearResolvedSecrets()
                return nil
        }

        // Use the new snake_case format
        type Alias Config
        alias := (*Alias)(r)
        if err := json.Unmarshal(data, alias); err != nil {
                return err
        }
        r.ClearResolvedSecrets()
        return nil
}

// DefaultCallbackPort is the default port for the OAuth callback server
const DefaultCallbackPort = 8666

// HasValidCachedTokens returns true if the config has a cached token reference that can be used
// to create a TokenSource without requiring a new OAuth flow.
// Note: This only checks if a refresh token reference exists, not if the token is actually valid.
// The actual validity will be determined when the token is used.
func (c *Config) HasValidCachedTokens() bool {
        // We need at least a refresh token reference to restore the session
        return c.CachedRefreshTokenRef != ""
}

// ClearCachedTokens removes any cached OAuth token references from the config.
// Note: This does not delete the actual secret from the secret manager.
func (c *Config) ClearCachedTokens() {
        c.CachedRefreshTokenRef = ""
        c.CachedTokenExpiry = time.Time{}
}

// HasCachedClientCredentials returns true if the config has cached DCR client credentials.
func (c *Config) HasCachedClientCredentials() bool {
        return c.CachedClientID != ""
}

// HasCachedCIMDClientID returns true if a CIMD client_id was cached from a prior session.
func (c *Config) HasCachedCIMDClientID() bool {
        return c.CachedCIMDClientID != ""
}

// ClearCachedClientCredentials removes any cached DCR client credential references from the config.
// It does not clear CachedCIMDClientID — the CIMD URL is a stable constant that does not
// need to be rotated alongside DCR secrets.
func (c *Config) ClearCachedClientCredentials() {
        c.CachedClientID = ""
        c.CachedClientSecretRef = ""
        c.CachedSecretExpiry = time.Time{}
        c.CachedRegTokenRef = ""
        c.CachedRegClientURI = ""
        c.CachedTokenEndpointAuthMethod = ""
        c.CachedDCRCallbackPort = 0
}

// LogContext returns the upstream issuer and resolved client_id for use as
// log-correlation fields on the MonitoredTokenSource. Returns ("", "")
// when c is nil so callers do not need to guard the call. Mirrors the
// precedence applied at runtime when sending the client_id on token
// refresh: cached CIMD URL > cached DCR client_id > statically configured
// client_id. Lives next to resolveClientCredentials so the precedence has
// a single home — adding a fourth cached field updates both call sites
// in one place.
func (c *Config) LogContext() (upstream, clientID string) {
        if c == nil {
                return "", ""
        }
        clientID = func() string {
                if c.CachedCIMDClientID != "" {
                        return c.CachedCIMDClientID
                }
                if c.CachedClientID != "" {
                        return c.CachedClientID
                }
                return c.ClientID
        }()
        return c.Issuer, clientID
}

// DefaultResourceIndicator derives the resource indicator (RFC 8707) from the remote server URL.
// This function should only be called when the user has not explicitly provided a resource indicator.
// If the resource indicator cannot be derived, it returns an empty string.
func DefaultResourceIndicator(remoteServerURL string) string {
        // Normalize the remote server URL
        normalized, err := normalizeResourceURI(remoteServerURL)
        if err != nil {
                // Normalization failed - log warning and leave resource empty
                slog.Warn("Failed to normalize resource indicator from remote server URL", "url", remoteServerURL, "error", err)
                return ""
        }

        // Validate the normalized result
        if err := httpval.ValidateResourceURI(normalized); err != nil {
                // Validation failed - log warning and leave resource empty
                slog.Warn("Normalized resource indicator is invalid", "resource", normalized, "error", err)
                return ""
        }

        return normalized
}

// normalizeResourceURI normalizes a resource URI to conform to MCP specification requirements.
// This function performs the following normalizations:
// - Lowercase scheme and host
// - Strip fragments
func normalizeResourceURI(resourceURI string) (string, error) {
        if resourceURI == "" {
                return "", fmt.Errorf("resource URI cannot be empty")
        }

        // Parse the URI
        parsed, err := url.Parse(resourceURI)
        if err != nil {
                return "", fmt.Errorf("invalid resource URI: %w", err)
        }

        // Normalize: lowercase scheme and host
        parsed.Scheme = strings.ToLower(parsed.Scheme)
        parsed.Host = strings.ToLower(parsed.Host)

        // Strip fragment if present (fragments are not allowed in resource indicators)
        parsed.Fragment = ""

        return parsed.String(), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package remote

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "strings"
        "time"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth/discovery"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// Handler handles authentication for remote MCP servers.
// Supports OAuth/OIDC-based authentication with automatic discovery.
type Handler struct {
        config                     *Config
        tokenPersister             TokenPersister
        clientCredentialsPersister ClientCredentialsPersister
        secretProvider             secrets.Provider
        httpClient                 networking.HTTPClient
}

// NewHandler creates a new remote authentication handler
func NewHandler(config *Config) *Handler {
        return &Handler{
                config: config,
        }
}

// SetTokenPersister sets a callback function that will be called whenever
// OAuth tokens are refreshed. This enables token persistence across restarts.
func (h *Handler) SetTokenPersister(persister TokenPersister) {
        h.tokenPersister = persister
}

// SetSecretProvider sets the secret provider used to store and retrieve cached tokens.
func (h *Handler) SetSecretProvider(provider secrets.Provider) {
        h.secretProvider = provider
}

// SetClientCredentialsPersister sets a callback function that will be called
// when DCR client credentials are obtained and need to be persisted.
func (h *Handler) SetClientCredentialsPersister(persister ClientCredentialsPersister) {
        h.clientCredentialsPersister = persister
}

// SetHTTPClient sets the HTTP client used for RFC 7592 registration management requests.
func (h *Handler) SetHTTPClient(client networking.HTTPClient) {
        h.httpClient = client
}

// Authenticate is the main entry point for remote MCP server authentication
func (h *Handler) Authenticate(ctx context.Context, remoteURL string) (oauth2.TokenSource, error) {
        // Priority 1: Bearer token authentication (if configured)
        bearerToken := h.config.bearerToken()
        if bearerToken != "" {
                slog.Debug("Using bearer token authentication")
                return NewBearerTokenSource(bearerToken), nil
        }

        // Detect authentication requirements once (used by both cached token restore and fresh OAuth)
        authInfo, err := discovery.DetectAuthenticationFromServer(ctx, remoteURL, nil)
        if err != nil {
                slog.Debug("Could not detect authentication from server", "error", err)
                return nil, nil // Not an error, just no auth detected
        }

        if authInfo == nil {
                return nil, nil // No authentication required
        }

        slog.Debug("Detected authentication requirement from server",
                "type", authInfo.Type, "realm", authInfo.Realm, "resource_metadata", authInfo.ResourceMetadata)

        // Check if we need to handle Bearer token requirement
        if err := h.validateBearerRequirement(authInfo); err != nil {
                return nil, err
        }

        // Only proceed with OAuth if the auth type supports it
        if authInfo.Type != "OAuth" && authInfo.Type != "Bearer" {
                slog.Error("Unsupported authentication type", "type", authInfo.Type)
                return nil, nil
        }

        // Discover OAuth endpoints once (used by both cached token restore and fresh OAuth)
        issuer, scopes, authServerInfo, allowPrivateIPs, err := h.discoverIssuerAndScopes(ctx, authInfo, remoteURL)
        if err != nil {
                return nil, err
        }

        // Priority 2: Try to use cached OAuth tokens (if available)
        if h.config.HasValidCachedTokens() {
                tokenSource, err := h.tryRestoreFromCachedTokens(ctx, issuer, scopes, authServerInfo)
                if err != nil {
                        slog.Warn("Failed to restore from cached tokens, will perform fresh OAuth flow", "error", err)
                        h.clearCachedSessionMetadata()
                } else if tokenSource != nil {
                        slog.Debug("Successfully restored OAuth session from cached tokens")
                        return tokenSource, nil
                }
        }

        // Priority 3: Fresh OAuth authentication flow. Reuse the same target-private
        // decision the discovery fetches above already computed, so DCR (discovery +
        // registration) is allowed to reach a private upstream only when the
        // operator-configured remote is itself private, and guarded otherwise
        // (CWE-918) — without a second, later TargetIsPrivate DNS lookup that could
        // disagree with the first.
        return h.performOAuthFlow(ctx, issuer, scopes, authServerInfo, allowPrivateIPs)
}

// validateBearerRequirement checks if Bearer auth is required without OAuth fallback
func (*Handler) validateBearerRequirement(authInfo *discovery.AuthInfo) error {
        if authInfo.Type != "Bearer" {
                return nil
        }

        // For backward compatibility, fall back to OAuth flow if realm or resource_metadata is present
        // Many servers use Bearer header but support OAuth flow
        if authInfo.Realm != "" || authInfo.ResourceMetadata != "" {
                slog.Warn("Bearer header without token, attempting OAuth flow for backward compatibility",
                        "realm_present", authInfo.Realm != "", "resource_metadata_present", authInfo.ResourceMetadata != "")
                return nil
        }

        // No realm or resource_metadata - likely requires static bearer token
        return fmt.Errorf("server requires bearer token authentication but no bearer token is configured. "+
                "Please provide a bearer token using --remote-auth-bearer-token flag or %s environment variable", BearerTokenEnvVarName)
}

// performOAuthFlow executes the OAuth authentication flow
func (h *Handler) performOAuthFlow(
        ctx context.Context,
        issuer string,
        scopes []string,
        authServerInfo *discovery.AuthServerInfo,
        allowPrivateIPs bool,
) (oauth2.TokenSource, error) {
        slog.Debug("Starting OAuth authentication flow", "issuer", issuer)

        // Client registration priority (MCP spec: stored credentials → CIMD → DCR):
        // Priority 1: Pre-configured credentials — set by buildOAuthFlowConfig from h.config's effective values.
        // Priority 2: CIMD — AS advertises support and no credentials are set; use metadata URL as client_id.
        // Priority 3: DCR — PerformOAuthFlow handles this when ClientID is still empty after the above.
        flowConfig := h.buildOAuthFlowConfig(scopes, authServerInfo, allowPrivateIPs)
        if shouldUseCIMD(authServerInfo, flowConfig) {
                flowConfig.ClientID = oauthproto.ToolHiveClientMetadataDocumentURL
                slog.Debug("Using CIMD client_id", "url", oauthproto.ToolHiveClientMetadataDocumentURL)
        }

        result, err := discovery.PerformOAuthFlow(ctx, issuer, flowConfig)
        if err != nil {
                // If we used CIMD and it was rejected, we need to retry with DCR.
                if flowConfig.ClientID == oauthproto.ToolHiveClientMetadataDocumentURL && isCIMDRejectionError(err) {
                        slog.Warn("CIMD client_id rejected by AS, retrying with DCR", "issuer", issuer, "error", err)
                        flowConfig.ClientID = ""
                        result, err = discovery.PerformOAuthFlow(ctx, issuer, flowConfig)
                }
        }
        if err != nil {
                return nil, err
        }

        return h.wrapWithPersistence(result, h.effectiveIssuer(issuer, authServerInfo)), nil
}

// buildOAuthFlowConfig creates the OAuth flow configuration
func (h *Handler) buildOAuthFlowConfig(
        scopes []string,
        authServerInfo *discovery.AuthServerInfo,
        allowPrivateIPs bool,
) *discovery.OAuthFlowConfig {
        flowConfig := &discovery.OAuthFlowConfig{
                ClientID:        h.config.ClientID,
                ClientSecret:    h.config.clientSecret(),
                AuthorizeURL:    h.config.AuthorizeURL,
                TokenURL:        h.config.TokenURL,
                Scopes:          scopes,
                CallbackPort:    h.config.CallbackPort,
                Timeout:         h.config.Timeout,
                SkipBrowser:     h.config.SkipBrowser,
                Resource:        h.config.Resource,
                OAuthParams:     h.config.OAuthParams,
                ScopeParamName:  h.config.ScopeParamName,
                AllowPrivateIPs: allowPrivateIPs,
                IssuerTrusted:   h.config.Issuer != "",
        }

        // If we have discovered endpoints from the authorization server metadata,
        // use them instead of trying to discover them again
        if authServerInfo != nil && h.config.AuthorizeURL == "" && h.config.TokenURL == "" {
                flowConfig.AuthorizeURL = authServerInfo.AuthorizationURL
                flowConfig.TokenURL = authServerInfo.TokenURL
                flowConfig.RegistrationEndpoint = authServerInfo.RegistrationEndpoint
                slog.Debug("Using discovered OAuth endpoints",
                        "authorize", authServerInfo.AuthorizationURL,
                        "token", authServerInfo.TokenURL,
                        "registration", authServerInfo.RegistrationEndpoint)
        }

        // Derived after the override above so the flag describes the URL that will
        // actually be dialed, not the one the operator may have been overridden out of.
        flowConfig.TokenEndpointTrusted = h.tokenEndpointTrusted(flowConfig.TokenURL)

        return flowConfig
}

// tokenEndpointTrusted reports whether tokenURL's authority was supplied by the
// operator — either configured directly, or named by the operator-configured
// issuer's own metadata document.
//
// An empty tokenURL means no endpoint has been resolved yet: OIDC discovery will
// fill it in from the issuer, and createOAuthConfig re-checks the discovered
// authority against that issuer before the trust survives.
func (h *Handler) tokenEndpointTrusted(tokenURL string) bool {
        if tokenURL == "" {
                return h.config.Issuer != ""
        }
        return networking.AuthorityMatchesAny(tokenURL, h.config.TokenURL, h.config.Issuer)
}

// clearCachedSessionMetadata prevents a failed restore from reusing any cached
// session credential. It deliberately leaves the secret manager unchanged.
func (h *Handler) clearCachedSessionMetadata() {
        h.config.ClearCachedTokens()
        h.config.ClearCachedClientCredentials()
        h.config.CachedCIMDClientID = ""
}

// wrapWithPersistence wraps the OAuth result with token persistence.
func (h *Handler) wrapWithPersistence(result *discovery.OAuthFlowResult, issuer string) oauth2.TokenSource {
        endpoint := ""
        if result.Config != nil {
                endpoint = result.Config.TokenURL
        }
        persister := h.tokenPersisterFor(issuer, endpoint)

        // The client identity is written before the refresh token. A cached refresh
        // token is only usable with the client credentials it was issued to, so if
        // the process dies between the two writes we want the harmless order: client
        // credentials with no cached token (a fresh flow follows) rather than a
        // cached token with no credentials to redeem it.
        if err := h.persistClientIdentity(result); err != nil {
                slog.Warn("Failed to persist client credentials; refresh token will not be cached", "error", err)
                return result.TokenSource
        }

        // Persist the refresh token for future restarts.
        if persister != nil && result.RefreshToken != "" {
                if err := persister(result.RefreshToken, result.Expiry); err != nil {
                        slog.Warn("Failed to persist OAuth tokens", "error", err)
                } else {
                        slog.Debug("Successfully persisted OAuth tokens for future restarts")
                }
        }

        // Wrap the token source to persist refreshed tokens
        tokenSource := result.TokenSource
        if persister != nil {
                tokenSource = NewPersistingTokenSource(result.TokenSource, persister)
        }

        return tokenSource
}

// persistClientIdentity stores the client credentials the flow ended up using,
// so a later restore has an identity to redeem its cached refresh token with.
// It returns nil when there is nothing to store.
func (h *Handler) persistClientIdentity(result *discovery.OAuthFlowResult) error {
        // CIMD client IDs (HTTPS URLs) are stable constants, not DCR-issued
        // credentials, so they are recorded on the config instead.
        if oauthproto.IsClientIDMetadataDocumentURL(result.ClientID) {
                h.config.CachedCIMDClientID = result.ClientID
                slog.Debug("Persisted CIMD client_id for future restarts", "url", result.ClientID)
                return nil
        }

        // Only persist if client_id exists - client_secret may be empty for PKCE flows.
        if h.clientCredentialsPersister == nil || result.ClientID == "" {
                return nil
        }

        if err := h.clientCredentialsPersister(
                result.ClientID,
                result.ClientSecret,
                result.SecretExpiry,
                result.RegistrationAccessToken,
                result.RegistrationClientURI,
                result.TokenEndpointAuthMethod,
                result.RegisteredCallbackPort,
        ); err != nil {
                return err
        }

        slog.Debug("Successfully persisted DCR client credentials for future restarts")
        return nil
}

func (h *Handler) tokenPersisterFor(issuer, endpoint string) TokenPersister {
        if h.tokenPersister == nil {
                return nil
        }
        return func(token string, expiry time.Time) error {
                envelope, err := encodeCachedRefreshToken(token, issuer, endpoint)
                if err != nil {
                        return err
                }
                return h.tokenPersister(envelope, expiry)
        }
}

func (*Handler) effectiveIssuer(issuer string, authServerInfo *discovery.AuthServerInfo) string {
        if authServerInfo != nil && authServerInfo.Issuer != "" {
                return authServerInfo.Issuer
        }
        return issuer
}

// restoreCachedRefreshToken resolves the cached refresh token and the
// authorization server it should be bound to. The returned migrated flag is
// true when the cached value was a pre-upgrade bare string rather than a
// bound envelope; callers must persist a bound envelope once the token is
// verified so the migration only happens once per session.
func (h *Handler) restoreCachedRefreshToken(
        ctx context.Context,
        issuer string,
        authServerInfo *discovery.AuthServerInfo,
) (token, effectiveIssuer, tokenURL string, migrated bool, err error) {
        tokenURL = h.config.TokenURL
        if tokenURL == "" && authServerInfo != nil {
                tokenURL = authServerInfo.TokenURL
        }
        refreshTokenValue, err := h.secretProvider.GetSecret(ctx, h.config.CachedRefreshTokenRef)
        if err != nil {
                return "", "", "", false, fmt.Errorf("failed to retrieve cached refresh token: %w", err)
        }
        effectiveIssuer = h.effectiveIssuer(issuer, authServerInfo)
        envelope, err := decodeCachedRefreshToken(refreshTokenValue)
        if err != nil {
                if !errors.Is(err, errLegacyCachedRefreshToken) {
                        return "", "", "", false, fmt.Errorf("unsafe legacy cached refresh token: %w", err)
                }
                // Pre-upgrade bare-string token. An operator-pinned issuer/token
                // endpoint was never exposed to GHSA-pc64-52fx-5v22 — the attacker
                // cannot change what the operator configured — so binding to it is
                // not a trust extension. effectiveIssuer already resolves to the
                // operator-configured issuer when authServerInfo is unavailable, and
                // to the issuer validated against the operator-configured metadata
                // URL otherwise — never to a bare-token-time guess.
                if h.config.Issuer != "" || h.config.TokenURL != "" {
                        return refreshTokenValue, effectiveIssuer, tokenURL, true, nil
                }
                // Auto-discovery deployment: this is exactly the class the advisory
                // describes. A bare legacy token carries no independently trusted
                // issuer to bind to, so migrating it would mean trusting whatever
                // the remote currently advertises — replaying the original exploit
                // instead of closing it. Force a fresh OAuth flow.
                return "", "", "", false,
                        fmt.Errorf("legacy cached refresh token has no trusted authorization server binding, re-authentication required")
        }
        if envelope.Issuer != effectiveIssuer || envelope.TokenURL != tokenURL {
                return "", "", "", false, fmt.Errorf("cached refresh token authorization server does not match current discovery")
        }
        return envelope.Token, effectiveIssuer, tokenURL, false, nil
}

// resolveClientCredentials returns the client ID and secret to use, preferring
// cached DCR credentials over statically configured ones.
func (h *Handler) resolveClientCredentials(ctx context.Context) (clientID, clientSecret string) {
        // First try to use statically configured credentials
        clientID = h.config.ClientID
        clientSecret = h.config.clientSecret()

        // If CIMD was used in a prior session, use the cached metadata URL as client_id.
        // CIMD clients have no secret (token_endpoint_auth_method=none).
        // Checked before DCR so that DCR credential rotation does not change which
        // client_id is sent on token refresh.
        if h.config.HasCachedCIMDClientID() {
                slog.Debug("Using cached CIMD client_id", "url", h.config.CachedCIMDClientID)
                return h.config.CachedCIMDClientID, ""
        }

        // If we have cached DCR client credentials, use those instead
        if h.config.HasCachedClientCredentials() {
                // ClientID is stored as plain text (it's public information)
                clientID = h.config.CachedClientID
                slog.Debug("Using cached DCR client credentials", "client_id", clientID)

                // Client secret is stored securely and may be empty for PKCE flows
                if h.config.CachedClientSecretRef != "" && h.secretProvider != nil {
                        cachedClientSecret, err := h.secretProvider.GetSecret(ctx, h.config.CachedClientSecretRef)
                        if err != nil {
                                slog.Warn("Failed to retrieve cached client secret", "error", err)
                        } else {
                                clientSecret = cachedClientSecret
                        }
                }
        }

        return clientID, clientSecret
}

// tryRestoreFromCachedTokens attempts to create a TokenSource from cached tokens
func (h *Handler) tryRestoreFromCachedTokens(
        ctx context.Context,
        issuer string,
        scopes []string,
        authServerInfo *discovery.AuthServerInfo,
) (oauth2.TokenSource, error) {
        // Resolve the refresh token from the secret manager
        if h.secretProvider == nil {
                return nil, fmt.Errorf("secret provider not configured, cannot restore cached tokens")
        }

        // Decode and bind the refresh token before using any cached credential.
        // This must precede DCR renewal and client credential retrieval.
        refreshToken, effectiveIssuer, tokenURL, migrated, err := h.restoreCachedRefreshToken(ctx, issuer, authServerInfo)
        if err != nil {
                return nil, err
        }

        // Check if the cached client secret is expired before attempting token refresh.
        // If it has fully expired and renewal also fails we must force a fresh OAuth flow.
        if h.isSecretExpiredOrExpiringSoon() {
                slog.Debug("Cached client secret is expiring or expired; attempting renewal before token restore",
                        "expiry", h.config.CachedSecretExpiry)
                if renewErr := h.renewClientSecret(ctx, effectiveIssuer); renewErr != nil {
                        slog.Warn("Client secret renewal failed", "error", renewErr)
                        // Hard-fail only when the secret is already past its expiry.
                        // If we are still in the buffer window the existing secret may work.
                        if !h.config.CachedSecretExpiry.IsZero() && time.Now().After(h.config.CachedSecretExpiry) {
                                return nil, fmt.Errorf(
                                        "client secret expired at %v and renewal failed: %w",
                                        h.config.CachedSecretExpiry, renewErr)
                        }
                        // Still within buffer — log and continue with the existing (still-valid) secret
                        slog.Warn("Proceeding with expiring client secret after failed renewal attempt")
                } else {
                        slog.Debug("Successfully renewed client secret before token restore")
                }
        }

        // Resolve client credentials - prefer cached DCR credentials over config.
        clientID, clientSecret := h.resolveClientCredentials(ctx)

        // Public clients (no secret) must use AuthStyleInParams: strict OAuth 2.1 servers
        // (e.g. Datadog) reject Basic Auth for token_endpoint_auth_method=none clients and
        // consume the single-use auth code in doing so. Confidential clients (DCR or
        // statically configured) use AutoDetect so servers that mandate client_secret_basic
        // are not broken.
        authStyle := oauth2.AuthStyleInParams
        if clientSecret != "" {
                authStyle = oauth2.AuthStyleAutoDetect
        }

        // Build OAuth2 config for token refresh.
        oauth2Config := &oauth2.Config{
                ClientID:     clientID,
                ClientSecret: clientSecret,
                Scopes:       scopes,
                Endpoint: oauth2.Endpoint{
                        AuthURL:   h.config.AuthorizeURL,
                        TokenURL:  tokenURL,
                        AuthStyle: authStyle,
                },
        }

        // Use the discovered authorization endpoint if available.
        if authServerInfo != nil && h.config.AuthorizeURL == "" {
                oauth2Config.Endpoint.AuthURL = authServerInfo.AuthorizationURL
        }

        // Create token source from cached refresh token.
        // Passes resource for RFC 8707 compliance when configured.
        baseSource, err := CreateTokenSourceFromCached(
                oauth2Config,
                refreshToken,
                h.config.CachedTokenExpiry,
                h.config.Resource,
                h.tokenEndpointTrusted(oauth2Config.Endpoint.TokenURL),
        )
        if err != nil {
                return nil, err
        }

        // Try to get a token to verify the cached tokens are valid
        // This will trigger a refresh since we don't have an access token
        tok, err := baseSource.Token()
        if err != nil {
                return nil, fmt.Errorf("cached tokens are invalid or expired: %w", err)
        }

        slog.Debug("Restored OAuth session from cached tokens", "issuer", issuer)

        if migrated {
                h.persistMigratedRefreshToken(tok, refreshToken, effectiveIssuer, tokenURL)
        }

        // Wrap with a persisting token source to save refreshed tokens bound to the
        // same authorization server and token endpoint.
        if persister := h.tokenPersisterFor(effectiveIssuer, tokenURL); persister != nil {
                return NewPersistingTokenSource(baseSource, persister), nil
        }

        return baseSource, nil
}

// persistMigratedRefreshToken persists a bound envelope after a legacy bare-string
// token has been verified against issuer/tokenURL, rather than waiting for a
// future refresh. It uses the refresh token/expiry the AS just returned, falling
// back to the pre-verification values if it didn't rotate them.
func (h *Handler) persistMigratedRefreshToken(tok *oauth2.Token, refreshToken, issuer, tokenURL string) {
        newRefreshToken := tok.RefreshToken
        if newRefreshToken == "" {
                newRefreshToken = refreshToken
        }
        expiry := tok.Expiry
        if expiry.IsZero() {
                expiry = h.config.CachedTokenExpiry
        }
        if persister := h.tokenPersisterFor(issuer, tokenURL); persister != nil {
                if err := persister(newRefreshToken, expiry); err != nil {
                        slog.Warn("Failed to persist migrated cached refresh token", "error", err)
                }
        }
}

// discoverIssuerAndScopes attempts to discover the OAuth issuer and scopes from various sources
// following RFC 8414 and RFC 9728 standards
// If the issuer is not derived from Realm and Resource Metadata, it derives from the remote URL
func (h *Handler) discoverIssuerAndScopes(
        ctx context.Context,
        authInfo *discovery.AuthInfo,
        remoteURL string,
) (string, []string, *discovery.AuthServerInfo, bool, error) {
        // Decide once whether discovery fetches derived from untrusted server input
        // (realm, resource_metadata, authorization_servers) may reach private
        // addresses. If the operator-configured target is itself internal they may;
        // otherwise block them to contain SSRF (CWE-918). The configured-issuer path
        // below is exempt because that URL comes from the operator, not the server.
        // Callers also thread this decision into the DCR flow (see performOAuthFlow),
        // so it's computed exactly once per Authenticate call rather than being
        // re-derived from a second, later TargetIsPrivate DNS lookup that could
        // disagree with this one.
        allowPrivateIPs := networking.TargetIsPrivate(ctx, remoteURL)
        blockPrivateIPs := !allowPrivateIPs

        // Priority 1: Use configured issuer if available. Fetch discovery to populate
        // AuthServerInfo (including ClientIDMetadataDocumentSupported) even when the
        // issuer is pre-configured, so CIMD detection works on this path.
        if h.config.Issuer != "" {
                slog.Debug("Using configured issuer", "issuer", h.config.Issuer)
                authServerInfo, _ := discovery.ValidateAndDiscoverAuthServer(ctx, h.config.Issuer, false)
                return h.config.Issuer, h.config.Scopes, authServerInfo, allowPrivateIPs, nil
        }

        // Priority 2: Try to derive from realm (RFC 8414). Fetch discovery for the
        // same reason as Priority 1 — the realm path skips resource metadata discovery.
        if authInfo.Realm != "" {
                derivedIssuer := discovery.DeriveIssuerFromRealm(authInfo.Realm)
                if derivedIssuer != "" {
                        slog.Debug("Derived issuer from realm", "issuer", derivedIssuer)
                        authServerInfo, _ := discovery.ValidateAndDiscoverAuthServer(ctx, derivedIssuer, blockPrivateIPs)
                        return derivedIssuer, h.config.Scopes, authServerInfo, allowPrivateIPs, nil
                }
        }

        // Priority 3: Fetch from resource metadata (RFC 9728)
        if authInfo.ResourceMetadata != "" {
                issuer, scopes, authServerInfo, err := h.tryDiscoverFromResourceMetadata(ctx, authInfo.ResourceMetadata, blockPrivateIPs)
                if err == nil {
                        return issuer, scopes, authServerInfo, allowPrivateIPs, nil
                }
                slog.Debug("Resource metadata discovery failed, falling through to well-known discovery", "error", err)
        }

        // Priority 4: Try to discover actual issuer from the server's well-known endpoint
        // This handles cases where the issuer differs from the server URL (e.g., Atlassian)
        issuer, scopes, authServerInfo, err := h.tryDiscoverFromWellKnown(ctx, remoteURL, blockPrivateIPs)
        if err == nil {
                return issuer, scopes, authServerInfo, allowPrivateIPs, nil
        }
        slog.Debug("Could not discover from well-known endpoint", "error", err)

        // Priority 5: Last resort - derive issuer from URL without discovery
        derivedIssuer := discovery.DeriveIssuerFromURL(remoteURL)
        if derivedIssuer != "" {
                slog.Debug("Using derived issuer from URL", "issuer", derivedIssuer)
                return derivedIssuer, h.config.Scopes, nil, allowPrivateIPs, nil
        }

        // No issuer could be determined
        return "", nil, nil, false, fmt.Errorf("could not determine OAuth issuer. Please provide issuer in configuration, " +
                "or ensure the server provides a valid realm parameter or resource_metadata URL in the WWW-Authenticate header")
}

// tryDiscoverFromResourceMetadata attempts to discover issuer and scopes from resource metadata.
// blockPrivateIPs is propagated to the SSRF guard on the server-supplied metadata
// and authorization-server fetches.
func (h *Handler) tryDiscoverFromResourceMetadata(
        ctx context.Context,
        resourceMetadataURL string,
        blockPrivateIPs bool,
) (string, []string, *discovery.AuthServerInfo, error) {
        slog.Debug("Fetching resource metadata", "url", resourceMetadataURL)

        metadata, err := discovery.FetchResourceMetadata(ctx, resourceMetadataURL, blockPrivateIPs)
        if err != nil {
                slog.Debug("Failed to fetch resource metadata", "error", err)
                return "", nil, nil, fmt.Errorf("could not determine OAuth issuer")
        }

        if metadata == nil {
                return "", nil, nil, fmt.Errorf("could not determine OAuth issuer")
        }

        // Try to find a valid authorization server from the list
        authServerInfo, issuer := h.findValidAuthServer(ctx, metadata.AuthorizationServers, blockPrivateIPs)
        if authServerInfo == nil {
                if len(metadata.AuthorizationServers) > 0 {
                        slog.Warn("Resource metadata contained authorization_servers, " +
                                "but none could be validated as actual OAuth authorization servers")
                }
                return "", nil, nil, fmt.Errorf("could not determine OAuth issuer")
        }

        // Determine scopes - use configured or fall back to metadata
        scopes := h.config.Scopes
        if len(scopes) == 0 && len(metadata.ScopesSupported) > 0 {
                scopes = metadata.ScopesSupported
                slog.Debug("Using scopes from resource metadata", "scopes", scopes)
        }

        return issuer, scopes, authServerInfo, nil
}

// findValidAuthServer validates authorization servers and returns the first valid one
func (*Handler) findValidAuthServer(
        ctx context.Context,
        authServers []string,
        blockPrivateIPs bool,
) (*discovery.AuthServerInfo, string) {
        for _, authServer := range authServers {
                slog.Debug("Validating authorization server", "server", authServer)

                authServerInfo, err := discovery.ValidateAndDiscoverAuthServer(ctx, authServer, blockPrivateIPs)
                if err != nil {
                        slog.Debug("Authorization server validation failed", "server", authServer, "error", err)
                        continue
                }

                // Found a valid authorization server
                slog.Debug("Using validated authorization server",
                        "server", authServer, "issuer", authServerInfo.Issuer)
                return authServerInfo, authServerInfo.Issuer
        }

        return nil, ""
}

// tryDiscoverFromWellKnown attempts to discover the actual OAuth issuer
// by probing the server's well-known endpoints without validating issuer match
// This is useful when the issuer differs from the server URL (e.g., Atlassian case)
func (h *Handler) tryDiscoverFromWellKnown(
        ctx context.Context,
        remoteURL string,
        blockPrivateIPs bool,
) (string, []string, *discovery.AuthServerInfo, error) {
        // First try to derive a base URL from the remote URL
        derivedURL := discovery.DeriveIssuerFromURL(remoteURL)
        if derivedURL == "" {
                return "", nil, nil, fmt.Errorf("could not derive base URL from %s", remoteURL)
        }

        // Try to discover the actual issuer without validation
        // This uses DiscoverActualIssuer which doesn't validate issuer match
        authServerInfo, err := discovery.ValidateAndDiscoverAuthServer(ctx, derivedURL, blockPrivateIPs)
        if err != nil {
                return "", nil, nil, fmt.Errorf("well-known discovery failed: %w", err)
        }

        // Successfully discovered the actual issuer
        if authServerInfo.Issuer != derivedURL {
                slog.Debug("Discovered actual issuer",
                        "issuer", authServerInfo.Issuer, "server_url", derivedURL)
        }

        // Determine scopes - use configured or fall back to defaults
        scopes := h.config.Scopes
        if len(scopes) == 0 {
                // Use some reasonable defaults if no scopes configured
                scopes = []string{"openid", "profile"}
                slog.Debug("No scopes configured, using defaults", "scopes", scopes)
        }

        return authServerInfo.Issuer, scopes, authServerInfo, nil
}

// shouldUseCIMD reports whether the CIMD client_id should be presented to the AS.
// The AS must advertise CIMD support and no pre-configured credentials may be set.
// Mirrors shouldDynamicallyRegisterClient in pkg/auth/discovery for consistency.
func shouldUseCIMD(authServerInfo *discovery.AuthServerInfo, flowConfig *discovery.OAuthFlowConfig) bool {
        if authServerInfo == nil || !authServerInfo.ClientIDMetadataDocumentSupported {
                return false
        }
        return flowConfig.ClientID == "" && flowConfig.ClientSecret == ""
}

// isCIMDRejectionError returns true if err indicates the AS rejected the CIMD
// client_id. Only the RFC 6749 error codes invalid_client and unauthorized_client
// trigger a DCR retry; all other errors — including invalid_request and
// token-exchange failures — surface as-is.
//
// CIMD rejection can surface from two stages:
//   - Authorization endpoint: AS redirects to callback with error=invalid_client;
//     flow.go formats this as "OAuth error: <code> - <description>" (a plain error).
//   - Token endpoint: oauth2.RetrieveError with ErrorCode set.
func isCIMDRejectionError(err error) bool {
        if err == nil {
                return false
        }
        // Token endpoint rejection — structured error from golang.org/x/oauth2.
        var rerr *oauth2.RetrieveError
        if errors.As(err, &rerr) {
                switch rerr.ErrorCode {
                case "invalid_client", "unauthorized_client":
                        return true
                }
                return false
        }
        // Authorization endpoint rejection — flow.go formats callback errors as
        // "OAuth error: <code> - <description>". Check for the code after the prefix.
        msg := err.Error()
        return strings.HasPrefix(msg, "OAuth error: invalid_client") ||
                strings.HasPrefix(msg, "OAuth error: unauthorized_client")
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package remote

import (
        "context"
        "fmt"
        "log/slog"
        "sync"
        "time"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth/oauth"
)

// TokenPersister is a callback function that persists a cached OAuth session.
// It is called whenever tokens are refreshed. Only the cached session is
// persisted since the access token can be regenerated from it.
//
// The value passed is opaque and must be stored and returned unchanged:
// Handler may pass either a bare refresh token (legacy) or a versioned
// envelope binding the token to the authorization server that issued it
// (see cachedRefreshTokenEnvelope). Persisters must not assume the value is
// a directly-usable refresh token.
type TokenPersister func(cachedSession string, expiry time.Time) error

// ClientCredentialsPersister is called when DCR client credentials need to be persisted.
// This is used to store client_id, client_secret, and renewal metadata obtained during
// Dynamic Client Registration (RFC 7591) and needed for secret renewal (RFC 7592).
//
// Parameters:
//   - clientID: the registered client ID (public, stored as plain text)
//   - clientSecret: the registered client secret (sensitive, stored via secret manager)
//   - secretExpiry: when the client secret expires; zero value means it never expires
//   - registrationAccessToken: bearer token for RFC 7592 management operations (sensitive)
//   - registrationClientURI: endpoint for RFC 7592 client update/read operations (plain text)
//   - tokenEndpointAuthMethod: the auth method used for the token endpoint (e.g., "client_secret_basic", "none")
//   - registeredCallbackPort: the callback port used in the original DCR redirect URI
type ClientCredentialsPersister func(
        clientID string,
        clientSecret string,
        secretExpiry time.Time,
        registrationAccessToken string,
        registrationClientURI string,
        tokenEndpointAuthMethod string,
        registeredCallbackPort int,
) error

// PersistingTokenSource wraps an oauth2.TokenSource and persists tokens
// whenever they are refreshed. This enables session restoration across
// workload restarts without requiring a new browser-based OAuth flow.
type PersistingTokenSource struct {
        source    oauth2.TokenSource
        persister TokenPersister

        mu        sync.Mutex
        lastToken *oauth2.Token
}

// NewPersistingTokenSource creates a new PersistingTokenSource that wraps
// the given token source and calls the persister function whenever tokens
// are refreshed.
func NewPersistingTokenSource(source oauth2.TokenSource, persister TokenPersister) *PersistingTokenSource {
        return &PersistingTokenSource{
                source:    source,
                persister: persister,
        }
}

// Token returns a valid token, refreshing it if necessary.
// If the token was refreshed, it will be persisted using the configured persister.
func (p *PersistingTokenSource) Token() (*oauth2.Token, error) {
        token, err := p.source.Token()
        if err != nil {
                return nil, err
        }

        // Check if the refresh token changed - only persist when it actually differs
        // Refresh tokens are long-lived and usually don't change on every access token refresh
        p.mu.Lock()
        defer p.mu.Unlock()

        if token.RefreshToken != "" && p.persister != nil &&
                (p.lastToken == nil || token.RefreshToken != p.lastToken.RefreshToken) {
                // Refresh token changed, persist it
                if err := p.persister(token.RefreshToken, token.Expiry); err != nil {
                        // Log the error but don't fail the token retrieval
                        slog.Warn("Failed to persist refreshed OAuth token", "error", err)
                } else {
                        slog.Debug("Successfully persisted refreshed OAuth token")
                }
                p.lastToken = token
        }

        return token, nil
}

// CreateTokenSourceFromCached creates an oauth2.TokenSource from a cached refresh token.
// The returned token source will immediately refresh to get a new access token,
// then automatically refresh when it expires.
// If resource is non-empty, it is included in all refresh requests per RFC 8707.
// trusted reports that the token endpoint's authority came from operator
// configuration; see oauth.Config.TokenEndpointTrusted.
func CreateTokenSourceFromCached(
        config *oauth2.Config,
        refreshToken string,
        expiry time.Time,
        resource string,
        trusted bool,
) (oauth2.TokenSource, error) {
        if config == nil {
                return nil, fmt.Errorf("OAuth config must not be nil")
        }
        httpClient, err := oauth.NewTokenHTTPClient(config.Endpoint.TokenURL, trusted)
        if err != nil {
                return nil, fmt.Errorf("build token endpoint HTTP client: %w", err)
        }
        // Create a token with only the refresh token.
        // The access token is intentionally empty - ReuseTokenSource will detect
        // that the token is expired (since Expiry is in the past or AccessToken is empty)
        // and trigger a refresh using the refresh token.
        token := &oauth2.Token{
                AccessToken:  "", // Empty - will trigger immediate refresh
                RefreshToken: refreshToken,
                Expiry:       expiry,
                TokenType:    "Bearer",
        }

        // Use resource-aware token source if configured (RFC 8707)
        var base oauth2.TokenSource
        if resource != "" {
                base = oauth.NewResourceTokenSource(config, token, resource, httpClient)
        } else {
                ctx := context.WithValue(context.Background(), oauth2.HTTPClient, httpClient)
                base = config.TokenSource(ctx, token)
        }

        return oauth2.ReuseTokenSource(token, base), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package remote

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

        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// secretExpiryBuffer is the lead time before expiry at which we proactively
// renew the client secret (RFC 7592). Renewal is attempted when the secret
// expires within this window, not only after expiry.
const secretExpiryBuffer = 24 * time.Hour

// clientUpdateRequest is the body sent in a RFC 7592 §2.2 PUT request.
// Per the spec, all client metadata fields that were provided during
// registration must be included in the update request body.
type clientUpdateRequest struct {
        ClientID                string   `json:"client_id"`
        ClientName              string   `json:"client_name,omitempty"`
        RedirectURIs            []string `json:"redirect_uris,omitempty"`
        GrantTypes              []string `json:"grant_types,omitempty"`
        ResponseTypes           []string `json:"response_types,omitempty"`
        TokenEndpointAuthMethod string   `json:"token_endpoint_auth_method,omitempty"`
}

// clientUpdateResponse is the body returned by a RFC 7592 §2.1 response.
// The provider may rotate the registration_access_token; if present we must
// replace the stored one.
type clientUpdateResponse struct {
        // Required fields mirrored from registration response
        ClientID     string `json:"client_id"`
        ClientSecret string `json:"client_secret,omitempty"` //nolint:gosec // G117: field holds sensitive data

        // Expiry fields
        ClientSecretExpiresAt int64 `json:"client_secret_expires_at,omitempty"`

        // Management fields — registration_access_token may be rotated
        RegistrationAccessToken string `json:"registration_access_token,omitempty"` //nolint:gosec
        RegistrationClientURI   string `json:"registration_client_uri,omitempty"`
}

// isSecretExpiredOrExpiringSoon returns true when the cached client secret is
// already expired or will expire within secretExpiryBuffer.
// A zero CachedSecretExpiry means the secret never expires, so this returns false.
func (h *Handler) isSecretExpiredOrExpiringSoon() bool {
        expiry := h.config.CachedSecretExpiry
        if expiry.IsZero() {
                return false // Non-expiring secret
        }
        return time.Now().After(expiry.Add(-secretExpiryBuffer))
}

// renewClientSecret attempts to renew the client secret using RFC 7592 §2.2.
// It retrieves the stored registration_access_token and sends a PUT request
// to the registration_client_uri with the current client metadata.
//
// On success the handler's config is updated with the new secret, expiry, and
// (if rotated) the new registration_access_token.
//
// Callers should log a warning and continue if renewal fails — the existing
// secret may still be valid for some time, or the provider may not support renewal.
func (h *Handler) renewClientSecret(ctx context.Context, issuer string) error {
        if err := h.validateRenewalPrerequisites(); err != nil {
                return err
        }

        // Retrieve the registration access token from the secret manager
        regAccessToken, err := h.secretProvider.GetSecret(ctx, h.config.CachedRegTokenRef)
        if err != nil {
                return fmt.Errorf("failed to retrieve registration access token: %w", err)
        }

        slog.Debug("Attempting RFC 7592 client secret renewal",
                "registration_client_uri", h.config.CachedRegClientURI)

        // Validate the registration_client_uri before using it
        if err := validateRegistrationClientURI(h.config.CachedRegClientURI); err != nil {
                return fmt.Errorf("invalid registration_client_uri: %w", err)
        }

        // Build the update request body with the current client metadata.
        // Per RFC 7592 §2.2, the request MUST include all client metadata fields
        // that were provided during the initial registration.
        // ToolHive's CLI DCR flow registers this fixed metadata set; if future DCR
        // callers support custom metadata, persist and replay those original values.
        updateReq := clientUpdateRequest{
                ClientID:                h.config.CachedClientID,
                ClientName:              oauthproto.ToolHiveMCPClientName,
                RedirectURIs:            []string{fmt.Sprintf("http://localhost:%d/callback", h.registeredDCRCallbackPort())},
                GrantTypes:              []string{"authorization_code", "refresh_token"},
                ResponseTypes:           []string{"code"},
                TokenEndpointAuthMethod: h.config.CachedTokenEndpointAuthMethod,
        }

        reqBody, err := json.Marshal(updateReq)
        if err != nil {
                return fmt.Errorf("failed to marshal client update request: %w", err)
        }

        // Create PUT request per RFC 7592 §2.2
        req, err := http.NewRequestWithContext(
                ctx,
                http.MethodPut,
                h.config.CachedRegClientURI,
                bytes.NewReader(reqBody),
        )
        if err != nil {
                return fmt.Errorf("failed to create client update request: %w", err)
        }

        req.Header.Set("Content-Type", "application/json")
        req.Header.Set("Accept", "application/json")
        req.Header.Set("Authorization", "Bearer "+regAccessToken)

        httpClient := h.httpClient
        if httpClient == nil {
                httpClient = newRenewalHTTPClient(ctx, issuer)
        }

        resp, err := httpClient.Do(req)
        if err != nil {
                return fmt.Errorf("client update request failed: %w", err)
        }
        defer func() {
                if closeErr := resp.Body.Close(); closeErr != nil {
                        slog.Debug("Failed to close renewal response body", "error", closeErr)
                }
        }()

        if resp.StatusCode != http.StatusOK {
                errorBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
                _, _ = io.Copy(io.Discard, resp.Body)
                return fmt.Errorf("client update request returned HTTP %d: %s", resp.StatusCode, string(errorBody))
        }

        // Parse the renewal response
        const maxResponseSize = 1024 * 1024 // 1 MB
        var updateResp clientUpdateResponse
        if err := json.NewDecoder(io.LimitReader(resp.Body, maxResponseSize)).Decode(&updateResp); err != nil {
                return fmt.Errorf("failed to decode client update response: %w", err)
        }

        if updateResp.ClientID == "" {
                return fmt.Errorf("client update response missing client_id")
        }
        if updateResp.ClientSecret == "" {
                return fmt.Errorf("client update response missing client_secret")
        }

        return h.persistRenewedSecret(updateResp)
}

func (h *Handler) validateRenewalPrerequisites() error {
        if h.config.CachedRegClientURI == "" {
                return fmt.Errorf("registration_client_uri missing; cannot renew secret (RFC 7592 unsupported)")
        }
        if h.config.CachedRegTokenRef == "" {
                return fmt.Errorf("registration_access_token missing; cannot renew secret (RFC 7592 unsupported)")
        }
        if h.secretProvider == nil {
                return fmt.Errorf("secret provider not configured; cannot retrieve registration access token")
        }
        return nil
}

func (h *Handler) registeredDCRCallbackPort() int {
        if h.config.CachedDCRCallbackPort != 0 {
                return h.config.CachedDCRCallbackPort
        }
        return h.config.CallbackPort
}

func (h *Handler) persistRenewedSecret(updateResp clientUpdateResponse) error {
        if h.clientCredentialsPersister == nil {
                return fmt.Errorf("client credentials persister not configured; cannot save renewed secret")
        }

        var newExpiry time.Time
        if updateResp.ClientSecretExpiresAt > 0 {
                newExpiry = time.Unix(updateResp.ClientSecretExpiresAt, 0)
        }

        // Use the rotated registration_access_token if provided; fall back to existing.
        newRegToken := updateResp.RegistrationAccessToken
        newRegURI := updateResp.RegistrationClientURI
        if newRegURI != "" {
                if err := validateRegistrationClientURIRotation(h.config.CachedRegClientURI, newRegURI); err != nil {
                        return fmt.Errorf("invalid rotated registration_client_uri: %w", err)
                }
        }
        if newRegURI == "" {
                newRegURI = h.config.CachedRegClientURI
        }

        if err := h.clientCredentialsPersister(
                updateResp.ClientID,
                updateResp.ClientSecret,
                newExpiry,
                newRegToken,
                newRegURI,
                h.config.CachedTokenEndpointAuthMethod,
                h.registeredDCRCallbackPort(),
        ); err != nil {
                return fmt.Errorf("failed to persist renewed client secret: %w", err)
        }

        h.config.CachedClientID = updateResp.ClientID
        h.config.CachedSecretExpiry = newExpiry
        h.config.CachedRegClientURI = newRegURI
        if h.config.CachedDCRCallbackPort == 0 {
                h.config.CachedDCRCallbackPort = h.config.CallbackPort
        }

        slog.Debug("Successfully renewed client secret via RFC 7592",
                "client_id", updateResp.ClientID,
                "new_expiry_zero", newExpiry.IsZero(),
                "reg_token_rotated", newRegToken != "")

        return nil
}

// validateRegistrationClientURI validates that the registration_client_uri is
// a valid HTTPS URL (or localhost for development).
func validateRegistrationClientURI(registrationClientURI string) error {
        if registrationClientURI == "" {
                return fmt.Errorf("registration_client_uri is empty")
        }

        parsedURL, err := url.Parse(registrationClientURI)
        if err != nil {
                return fmt.Errorf("invalid registration_client_uri URL: %w", err)
        }
        if parsedURL.Scheme == "" || parsedURL.Host == "" {
                return fmt.Errorf("registration_client_uri must include a scheme and host: %s", registrationClientURI)
        }

        if parsedURL.Scheme != "https" && !networking.IsLocalhost(parsedURL.Host) {
                return fmt.Errorf("registration_client_uri must use HTTPS: %s", registrationClientURI)
        }
        cleanPath := path.Clean(parsedURL.Path)
        if parsedURL.Path == "" || cleanPath == "." || cleanPath == "/" {
                return fmt.Errorf("registration_client_uri must include a non-root path: %s", registrationClientURI)
        }

        return nil
}

// validateRegistrationClientURIRotation prevents a renewal response from
// moving the bearer-authenticated management endpoint to another origin.
func validateRegistrationClientURIRotation(currentURI, rotatedURI string) error {
        if err := validateRegistrationClientURI(rotatedURI); err != nil {
                return err
        }

        current, err := url.Parse(currentURI)
        if err != nil {
                return fmt.Errorf("invalid current registration_client_uri: %w", err)
        }
        rotated, err := url.Parse(rotatedURI)
        if err != nil {
                return fmt.Errorf("invalid rotated registration_client_uri: %w", err)
        }
        if !sameOrigin(current, rotated) {
                return fmt.Errorf("rotated registration_client_uri must remain on origin %s://%s", current.Scheme, current.Host)
        }

        return nil
}

func sameOrigin(first, second *url.URL) bool {
        return strings.EqualFold(first.Scheme, second.Scheme) &&
                strings.EqualFold(first.Hostname(), second.Hostname()) &&
                effectivePort(first) == effectivePort(second)
}

func effectivePort(parsedURL *url.URL) string {
        if port := parsedURL.Port(); port != "" {
                return port
        }
        if strings.EqualFold(parsedURL.Scheme, "https") {
                return "443"
        }
        if strings.EqualFold(parsedURL.Scheme, "http") {
                return "80"
        }
        return ""
}

// newRenewalHTTPClient derives renewal transport behavior from the shared
// OAuth client while adding protections required for a bearer-authenticated
// request to an endpoint supplied by an authorization server.
func newRenewalHTTPClient(ctx context.Context, issuer string) *http.Client {
        baseClient := oauthproto.DefaultHTTPClient()
        client := *baseClient
        client.CheckRedirect = networking.SameHostRedirectPolicy()

        // Internal IdPs are an intentional deployment mode. Base the exception on
        // the configured issuer rather than the server-returned management URI.
        if networking.TargetIsPrivate(ctx, issuer) {
                return &client
        }

        if transport, ok := baseClient.Transport.(*http.Transport); ok {
                protectedTransport := transport.Clone()
                // A proxy resolves and connects to the destination on ToolHive's behalf,
                // bypassing the destination-IP check below. Public management endpoints
                // therefore use direct dials; trusted private issuers retain base-client
                // proxy behavior through the branch above.
                protectedTransport.Proxy = nil
                protectedTransport.DialContext = networking.NewPrivateIPBlockingDialContext()
                protectedTransport.DisableKeepAlives = true
                client.Transport = protectedTransport
        }

        return &client
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package secrets provides generic secret management utilities for authentication.
// This package contains functions that can be used by any authentication method
// (OAuth, bearer tokens, etc.) to process secrets and store them in a secrets manager.
package secrets

import (
        "context"
        "crypto/rand"
        "errors"
        "fmt"
        "time"

        "github.com/stacklok/toolhive-core/env"
        "github.com/stacklok/toolhive/pkg/config"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// TokenType represents the type of authentication token/secret being processed
type TokenType string

const (
        // TokenTypeOAuthClientSecret represents an OAuth client secret
        // #nosec G101 - this is a type identifier, not a credential
        TokenTypeOAuthClientSecret TokenType = "oauth_client_secret"
        // TokenTypeBearerToken represents a bearer token
        TokenTypeBearerToken TokenType = "bearer_token"
        // TokenTypeOAuthRefreshToken represents a cached OAuth refresh token
        // #nosec G101 - this is a type identifier, not a credential
        TokenTypeOAuthRefreshToken TokenType = "oauth_refresh_token"
)

// tokenTypeConfig holds configuration for each token type
type tokenTypeConfig struct {
        prefix       string
        target       string
        errorContext string
}

var tokenTypeConfigs = map[TokenType]tokenTypeConfig{
        TokenTypeOAuthClientSecret: {
                prefix:       "OAUTH_CLIENT_SECRET_",
                target:       "oauth_secret",
                errorContext: "OAuth client secret",
        },
        TokenTypeBearerToken: {
                prefix:       "BEARER_TOKEN_",
                target:       "bearer_token",
                errorContext: "bearer token",
        },
        TokenTypeOAuthRefreshToken: {
                prefix:       "OAUTH_REFRESH_TOKEN_",
                target:       "oauth_refresh",
                errorContext: "OAuth refresh token",
        },
}

// ProcessSecret processes a secret, converting plain text to CLI format if needed.
// This is a generic function that can be used for any secret type.
// Parameters:
//   - workloadName: Name of the workload (used for secret naming)
//   - secretValue: The secret value (plain text or already in CLI format)
//   - tokenType: The type of token/secret (determines prefix, target, and error context)
//
// Returns the secret in CLI format: "secret-name,target=target_value"
func ProcessSecret(workloadName, secretValue string, tokenType TokenType) (string, error) {
        // Early return if no secret is provided - no need to access secrets manager
        if secretValue == "" {
                return "", nil
        }

        // Get the configuration for this token type
        tokenConfig, ok := tokenTypeConfigs[tokenType]
        if !ok {
                return "", fmt.Errorf("unknown token type: %s", tokenType)
        }

        // Return existing CLI secret references without requiring a configured secrets manager.
        if _, err := secrets.ParseSecretParameter(secretValue); err == nil {
                return secretValue, nil
        }

        // Get the secrets manager
        secretManager, err := GetSecretsManager()
        if err != nil {
                return "", fmt.Errorf("failed to get secrets manager: %w", err)
        }

        // Process the secret using the provider
        return processSecretWithProvider(workloadName, secretValue, secretManager, tokenConfig)
}

// ProcessSecretWithProvider processes a secret using the provided secret manager
// This version is testable with dependency injection and is used for testing
func ProcessSecretWithProvider(
        workloadName,
        secretValue string,
        secretManager secrets.Provider,
        tokenType TokenType,
) (string, error) {
        // Get the configuration for this token type
        tokenConfig, ok := tokenTypeConfigs[tokenType]
        if !ok {
                return "", fmt.Errorf("unknown token type: %s", tokenType)
        }

        return processSecretWithProvider(workloadName, secretValue, secretManager, tokenConfig)
}

// processSecretWithProvider processes a secret using the provided secret manager
// This is the internal implementation
func processSecretWithProvider(
        workloadName,
        secretValue string,
        secretManager secrets.Provider,
        tokenConfig tokenTypeConfig,
) (string, error) {
        if secretValue == "" {
                return "", nil
        }

        // Check if it's already in CLI format
        if _, err := secrets.ParseSecretParameter(secretValue); err == nil {
                return secretValue, nil // Already in CLI format
        }

        // It's plain text - convert to CLI format
        secretName, err := GenerateUniqueSecretNameWithPrefix(workloadName, tokenConfig.prefix, secretManager)
        if err != nil {
                return "", fmt.Errorf("failed to generate unique secret name: %w", err)
        }

        // Store the secret in the secret manager
        if err := StoreSecretInManagerWithProvider(context.Background(), secretName, secretValue, secretManager); err != nil {
                return "", fmt.Errorf("failed to store %s in manager: %w", tokenConfig.errorContext, err)
        }

        // Return CLI format reference
        return secrets.SecretParameter{Name: secretName, Target: tokenConfig.target}.ToCLIString(), nil
}

// GenerateUniqueSecretNameWithPrefix generates a unique secret name with a custom prefix
// This is a generic function that can be used for any secret type
func GenerateUniqueSecretNameWithPrefix(workloadName, prefix string, secretManager secrets.Provider) (string, error) {
        // Generate base name
        baseName := fmt.Sprintf("%s%s", prefix, workloadName)

        // Check if the base name is available
        _, err := secretManager.GetSecret(context.Background(), baseName)
        if err != nil {
                // Secret doesn't exist, we can use the base name
                return baseName, nil
        }

        // Secret exists, generate a unique name with timestamp
        // Use nanosecond precision + random suffix for collision resistance
        timestamp := time.Now().UnixNano()
        randomSuffix := make([]byte, 4)
        if _, err := rand.Read(randomSuffix); err != nil {
                return "", fmt.Errorf("failed to generate random suffix: %w", err)
        }
        uniqueName := fmt.Sprintf("%s_%d_%x", baseName, timestamp, randomSuffix)
        return uniqueName, nil
}

// StoreSecretInManagerWithProvider stores a secret using the provided secret manager
// This version is testable with dependency injection
func StoreSecretInManagerWithProvider(ctx context.Context, secretName, secretValue string, secretManager secrets.Provider) error {
        // Check if the provider supports writing secrets
        if caps := secretManager.Capabilities(); !caps.CanWrite {
                return fmt.Errorf("secrets provider (%s) does not support writing secrets", caps)
        }

        // Store the secret
        if err := secretManager.SetSecret(ctx, secretName, secretValue); err != nil {
                return fmt.Errorf("failed to store secret %s: %w", secretName, err)
        }

        return nil
}

// GetUserSecretsProvider returns a secrets provider suitable for user-facing
// callers (CLI, API, MCP tool server). It filters out system-reserved keys so
// user commands cannot accidentally read or overwrite internal secrets.
func GetUserSecretsProvider() (secrets.Provider, error) {
        configProvider := config.NewDefaultProvider()
        cfg := configProvider.GetConfig()

        if !cfg.Secrets.SetupCompleted {
                return nil, secrets.ErrSecretsNotSetup
        }

        providerType, err := cfg.Secrets.GetProviderType()
        if err != nil {
                return nil, fmt.Errorf("failed to get secrets provider type: %w", err)
        }

        provider, err := secrets.CreateProvider(providerType, secrets.WithUserFacing())
        if err != nil {
                return nil, fmt.Errorf("failed to create secrets provider: %w", err)
        }

        return provider, nil
}

// GetSystemSecretsProvider returns a raw (unscoped) secrets provider for
// advanced/emergency operations on system-managed keys. Unlike
// GetUserSecretsProvider it does not apply UserProvider filtering, so callers
// can read and delete __thv_* prefixed keys directly.
func GetSystemSecretsProvider() (secrets.Provider, error) {
        return getSystemSecretsProviderFromConfig(config.NewDefaultProvider(), &env.OSReader{})
}

// getSystemSecretsProviderFromConfig is the testable core of GetSystemSecretsProvider.
// It accepts an explicit config.Provider and env.Reader so tests can inject
// both without touching global environment state.
func getSystemSecretsProviderFromConfig(configProvider config.Provider, envReader env.Reader) (secrets.Provider, error) {
        cfg := configProvider.GetConfig()

        // GetProviderTypeWithEnv handles the TOOLHIVE_SECRETS_PROVIDER env var and
        // allows the "environment" provider even when SetupCompleted is false, so
        // Kubernetes and test deployments can skip interactive setup. When no env
        // var override is present and SetupCompleted is false, it returns
        // ErrSecretsNotSetup, which is propagated directly to the caller.
        providerType, err := cfg.Secrets.GetProviderTypeWithEnv(envReader)
        if err != nil {
                if errors.Is(err, secrets.ErrSecretsNotSetup) {
                        return nil, secrets.ErrSecretsNotSetup
                }
                return nil, fmt.Errorf("failed to get secrets provider type: %w", err)
        }

        provider, err := secrets.CreateProvider(providerType)
        if err != nil {
                return nil, fmt.Errorf("failed to create secrets provider: %w", err)
        }

        return provider, nil
}

// GetSecretsManager returns the secrets manager instance
// This is exported so it can be reused by other packages
func GetSecretsManager() (secrets.Provider, error) {
        configProvider := config.NewDefaultProvider()
        cfg := configProvider.GetConfig()

        // Check if secrets setup has been completed
        if !cfg.Secrets.SetupCompleted {
                return nil, secrets.ErrSecretsNotSetup
        }

        providerType, err := cfg.Secrets.GetProviderType()
        if err != nil {
                return nil, fmt.Errorf("failed to get secrets provider type: %w", err)
        }

        provider, err := secrets.CreateProvider(providerType, secrets.WithScope(secrets.ScopeWorkloads))
        if err != nil {
                return nil, fmt.Errorf("failed to create secrets provider: %w", err)
        }

        return provider, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package sessionbinding binds sessions to the validated issuer and subject, never
// to a bearer token. Authentication must run before ownership enforcement.
package sessionbinding

import (
        "context"
        "crypto/subtle"
        "errors"
        "net/http"
        "net/url"
        "strings"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/transport/session"
)

// MetadataKey is shared with persisted vMCP sessions for compatibility.
const MetadataKey = session.MetadataKeyIdentityBinding

// UnauthenticatedSentinel is reserved for an actually absent identity.
const UnauthenticatedSentinel = "unauthenticated"

var (
        // ErrInvalidBinding denotes malformed issuer/subject claims.
        ErrInvalidBinding = errors.New("invalid identity binding")
        // ErrNotFound deliberately conceals missing, unowned and foreign sessions.
        ErrNotFound = session.ErrSessionNotFound
)

// Format encodes a nonempty issuer and subject separated by exactly one NUL.
func Format(iss, sub string) (string, error) {
        if iss == "" || sub == "" || strings.ContainsRune(iss, '\x00') || strings.ContainsRune(sub, '\x00') {
                return "", ErrInvalidBinding
        }
        return iss + "\x00" + sub, nil
}

// Parse rejects malformed bindings and the unauthenticated sentinel.
func Parse(value string) (iss, sub string, ok bool) {
        iss, sub, ok = strings.Cut(value, "\x00")
        if !ok || iss == "" || sub == "" || strings.ContainsRune(sub, '\x00') {
                return "", "", false
        }
        return iss, sub, true
}

// FromIdentity accepts nil only for auth-disabled deployments. A synthetic
// anonymous/local identity with valid claims binds normally; malformed non-nil
// identities must never downgrade to the unauthenticated sentinel.
func FromIdentity(identity *auth.Identity) (string, error) {
        if identity == nil {
                return UnauthenticatedSentinel, nil
        }
        iss, _ := identity.Claims["iss"].(string)
        sub, _ := identity.Claims["sub"].(string)
        return Format(iss, sub)
}

// Validate fails closed for legacy records without an owner. Refreshed tokens
// for the same issuer and subject are accepted. It never mutates session state.
func Validate(stored string, identity *auth.Identity) error {
        if stored != UnauthenticatedSentinel {
                if _, _, ok := Parse(stored); !ok {
                        return ErrNotFound
                }
        }
        current, err := FromIdentity(identity)
        if err != nil || subtle.ConstantTimeCompare([]byte(stored), []byte(current)) != 1 {
                return ErrNotFound
        }
        return nil
}

// RequestID reads protocol session carriers, rejecting duplicate or conflicting
// values before a backend can interpret them differently. queryKey may be empty.
func RequestID(r *http.Request, queryKey string) (string, error) {
        headers := r.Header.Values("Mcp-Session-Id")
        if len(headers) > 1 {
                return "", ErrNotFound
        }
        id := r.Header.Get("Mcp-Session-Id")
        if queryKey == "" {
                return id, nil
        }
        query, err := url.ParseQuery(r.URL.RawQuery)
        if err != nil {
                return "", ErrNotFound
        }
        values := query[queryKey]
        if len(values) > 1 {
                return "", ErrNotFound
        }
        if len(values) == 1 {
                if values[0] == "" || (id != "" && id != values[0]) {
                        return "", ErrNotFound
                }
                id = values[0]
        }
        return id, nil
}

// Lookup reads authoritative owner metadata only, without restoring SDK sessions
// or connecting to backends. Return ErrNotFound for absent records, and propagate
// storage errors. Shared storage does not provide live-stream affinity.
type Lookup func(context.Context, string) (string, error)

// Check validates a session-bearing request using its authenticated context.
// Empty IDs are sessionless; callers enforce any protocol-required ID separately.
func Check(r *http.Request, id string, lookup Lookup) error {
        if id == "" {
                return nil
        }
        owner, err := lookup(r.Context(), id)
        if err != nil {
                return err
        }
        identity, _ := auth.IdentityFromContext(r.Context())
        return Validate(owner, identity)
}

// NewMiddleware constructs ownership enforcement. extract must classify requests
// without consuming their bodies and return an empty ID only for truly sessionless
// requests. reject must conceal ErrNotFound (404) and map storage failures to 503.
// Install inside authentication and outside stateful handlers/restoration.
func NewMiddleware(extract func(*http.Request) (string, error), lookup Lookup,
        reject func(http.ResponseWriter, *http.Request, error)) (func(http.Handler) http.Handler, error) {
        if extract == nil || lookup == nil || reject == nil {
                return nil, errors.New("session ownership middleware requires extractor, lookup and rejection handler")
        }
        return func(next http.Handler) http.Handler {
                if next == nil {
                        panic("session ownership middleware requires a handler")
                }
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        id, err := extract(r)
                        if err == nil {
                                err = Check(r, id, lookup)
                        }
                        if err != nil {
                                reject(w, r, err)
                                return
                        }
                        next.ServeHTTP(w, r)
                })
        }, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package sessionbinding

import (
        "context"
        "errors"
        "net/http"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/transport/session"
)

// BindOwner sets ownership before insertion, never on a published session.
func BindOwner(ctx context.Context, sess session.Session) error {
        identity, _ := auth.IdentityFromContext(ctx)
        owner, err := FromIdentity(identity)
        if err != nil {
                return err
        }
        sess.SetMetadata(MetadataKey, owner)
        return nil
}

// Store is the ordinary session manager surface required for conditional publication.
type Store interface {
        AddSession(session.Session) error
        LookupOwner(context.Context, string) (string, error)
}

// AddOwnedSession binds an ordinary proxy session before atomic publication.
// sess must be newly constructed and unpublished; its metadata is modified.
// A collision may reuse only the same owner's record, never replace metadata.
// Shared storage does not provide cross-replica live-stream delivery.
func AddOwnedSession(ctx context.Context, manager Store, sess session.Session) error {
        if err := BindOwner(ctx, sess); err != nil {
                return ErrNotFound
        }
        if err := manager.AddSession(sess); err != nil {
                if !errors.Is(err, session.ErrSessionAlreadyExists) {
                        return err
                }
                owner, err := manager.LookupOwner(ctx, sess.ID())
                if err != nil {
                        return err
                }
                identity, _ := auth.IdentityFromContext(ctx)
                return Validate(owner, identity)
        }
        return nil
}

// WriteOwnershipError uses the ordinary proxies' non-disclosing JSON-RPC 404.
func WriteOwnershipError(w http.ResponseWriter, requestID any, err error) {
        if errors.Is(err, ErrNotFound) {
                session.WriteNotFound(w, requestID)
                return
        }
        http.Error(w, "Session store unavailable", http.StatusServiceUnavailable)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

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

        "github.com/cenkalti/backoff/v7"
        "github.com/golang-jwt/jwt/v5"
        "github.com/jwx-go/jwkfetch/v4"
        "github.com/lestrrat-go/httprc/v3"
        "github.com/lestrrat-go/jwx/v4/jwk"

        "github.com/stacklok/toolhive-core/env"
        "github.com/stacklok/toolhive/pkg/auth/upstreamtoken"
        "github.com/stacklok/toolhive/pkg/authserver/server/keys"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// TokenIntrospector defines the interface for token introspection providers
type TokenIntrospector interface {
        // Name returns the provider name
        Name() string

        // CanHandle returns true if this provider can handle the given introspection URL
        CanHandle(introspectURL string) bool

        // IntrospectToken introspects an opaque token and returns JWT claims
        IntrospectToken(ctx context.Context, token string) (jwt.MapClaims, error)
}

// Registry maintains a list of available token introspection providers
type Registry struct {
        providers []TokenIntrospector
}

// NewRegistry creates a new provider registry
func NewRegistry() *Registry {
        return &Registry{
                providers: []TokenIntrospector{},
        }
}

// GetIntrospector returns the appropriate provider for the given introspection URL
func (r *Registry) GetIntrospector(introspectURL string) TokenIntrospector {
        for _, provider := range r.providers {
                if provider.CanHandle(introspectURL) {
                        //nolint:gosec // G706: provider name and URL are from server configuration
                        slog.Debug("selected provider for introspection", "provider", provider.Name(), "url", introspectURL)
                        return provider
                }
        }
        // Create a new fallback provider instance with the specific URL
        slog.Debug("using RFC7662 fallback provider for introspection", "url", introspectURL)
        return NewRFC7662Provider(introspectURL)
}

// AddProvider adds a new provider to the registry
func (r *Registry) AddProvider(provider TokenIntrospector) {
        r.providers = append(r.providers, provider)
}

// GoogleTokeninfoURL is the Google OAuth2 tokeninfo endpoint URL
const GoogleTokeninfoURL = "https://oauth2.googleapis.com/tokeninfo" //nolint:gosec

// GoogleProvider implements token introspection for Google's tokeninfo API
type GoogleProvider struct {
        client *http.Client
        url    string
}

// NewGoogleProvider creates a new Google token introspection provider
func NewGoogleProvider(introspectURL string) *GoogleProvider {
        return &GoogleProvider{
                client: http.DefaultClient,
                url:    introspectURL,
        }
}

// Name returns the provider name
func (*GoogleProvider) Name() string {
        return "google"
}

// CanHandle returns true if this provider can handle the given introspection URL
func (g *GoogleProvider) CanHandle(introspectURL string) bool {
        return introspectURL == g.url
}

// IntrospectToken introspects a Google opaque token and returns JWT claims
func (g *GoogleProvider) IntrospectToken(ctx context.Context, token string) (jwt.MapClaims, error) {
        slog.Debug("using Google tokeninfo provider for token introspection", "url", g.url)

        // Parse the URL and add query parameters
        u, err := url.Parse(g.url)
        if err != nil {
                return nil, fmt.Errorf("failed to parse introspection URL: %w", err)
        }

        // Add the access token as a query parameter
        query := u.Query()
        query.Set("access_token", token)
        u.RawQuery = query.Encode()

        // Create the GET request
        //nolint:gosec // G704 - URL from trusted OIDC discovery config
        req, err := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
        if err != nil {
                return nil, fmt.Errorf("failed to create Google tokeninfo request: %w", err)
        }

        req.Header.Set("Accept", "application/json")
        req.Header.Set("User-Agent", oauthproto.UserAgent)

        // Make the request
        resp, err := g.client.Do(req) // #nosec G704 -- URL is the configured OIDC introspection endpoint
        if err != nil {
                return nil, fmt.Errorf("google tokeninfo request failed: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        // Read the response with a reasonable limit to prevent DoS attacks
        const maxResponseSize = 64 * 1024 // 64KB should be more than enough for tokeninfo response
        limitedReader := io.LimitReader(resp.Body, maxResponseSize)
        body, err := io.ReadAll(limitedReader)
        if err != nil {
                return nil, fmt.Errorf("failed to read Google tokeninfo response: %w", err)
        }

        // Check for HTTP errors
        if resp.StatusCode != http.StatusOK {
                return nil, fmt.Errorf("google tokeninfo request failed with status %d: %s", resp.StatusCode, string(body))
        }

        // Parse the Google response and convert to JWT claims
        //nolint:gosec // G706: HTTP status code is not sensitive
        slog.Debug("successfully received Google tokeninfo response", "status", resp.StatusCode)
        return g.parseGoogleResponse(body)
}

// parseGoogleResponse parses Google's tokeninfo response and converts it to JWT claims
func (*GoogleProvider) parseGoogleResponse(body []byte) (jwt.MapClaims, error) {
        // Parse Google's response format
        var googleResp struct {
                // Standard OAuth fields
                Aud   string `json:"aud,omitempty"`
                Sub   string `json:"sub,omitempty"`
                Scope string `json:"scope,omitempty"`

                // Google returns Unix timestamp as string (RFC 7662 uses numeric)
                Exp string `json:"exp,omitempty"`

                // Google-specific fields
                Azp           string `json:"azp,omitempty"`
                ExpiresIn     string `json:"expires_in,omitempty"`
                Email         string `json:"email,omitempty"`
                EmailVerified string `json:"email_verified,omitempty"`
        }

        if err := json.Unmarshal(body, &googleResp); err != nil {
                return nil, fmt.Errorf("failed to decode Google tokeninfo JSON: %w", err)
        }

        // Convert to JWT MapClaims format
        claims := jwt.MapClaims{
                "iss": "https://accounts.google.com", // Default Google issuer
        }

        // Copy standard fields
        if googleResp.Sub != "" {
                claims["sub"] = googleResp.Sub
        }
        if googleResp.Aud != "" {
                claims["aud"] = googleResp.Aud
        }
        if googleResp.Scope != "" {
                claims["scope"] = googleResp.Scope
        }

        // Handle expiration - convert string timestamp to float64
        if googleResp.Exp != "" {
                expInt, err := strconv.ParseInt(googleResp.Exp, 10, 64)
                if err != nil {
                        return nil, fmt.Errorf("invalid exp format: %w", err)
                }
                claims["exp"] = float64(expInt) // JWT expects float64

                // Check if token is expired and return error if so (consistent with RFC 7662 behavior)
                isActive := time.Now().Unix() < expInt
                claims["active"] = isActive
                if !isActive {
                        return nil, ErrInvalidToken
                }
        } else {
                return nil, fmt.Errorf("missing exp field in Google response")
        }

        // Copy Google-specific fields
        if googleResp.Azp != "" {
                claims["azp"] = googleResp.Azp
        }
        if googleResp.ExpiresIn != "" {
                claims["expires_in"] = googleResp.ExpiresIn
        }
        if googleResp.Email != "" {
                claims["email"] = googleResp.Email
        }
        if googleResp.EmailVerified != "" {
                claims["email_verified"] = googleResp.EmailVerified
        }

        return claims, nil
}

// RFC7662Provider implements standard RFC 7662 OAuth 2.0 Token Introspection
type RFC7662Provider struct {
        client       *http.Client
        clientID     string
        clientSecret string
        url          string
}

// NewRFC7662Provider creates a new RFC 7662 token introspection provider
func NewRFC7662Provider(introspectURL string) *RFC7662Provider {
        return &RFC7662Provider{
                client: http.DefaultClient,
                url:    introspectURL,
        }
}

// NewRFC7662ProviderWithAuth creates a new RFC 7662 provider with client credentials
func NewRFC7662ProviderWithAuth(
        introspectURL, clientID, clientSecret, caCertPath, authTokenFile string, allowPrivateIP bool,
) (*RFC7662Provider, error) {
        // Create HTTP client with CA bundle and auth token support
        client, err := networking.NewHttpClientBuilder().
                WithCABundle(caCertPath).
                WithTokenFromFile(authTokenFile).
                WithPrivateIPs(allowPrivateIP).
                Build()
        if err != nil {
                return nil, fmt.Errorf("failed to create HTTP client: %w", err)
        }

        return &RFC7662Provider{
                client:       client,
                clientID:     clientID,
                clientSecret: clientSecret,
                url:          introspectURL,
        }, nil
}

// Name returns the provider name
func (*RFC7662Provider) Name() string {
        return "rfc7662"
}

// CanHandle returns true if this provider can handle the given introspection URL
// Returns true for any URL when no specific URL was configured (fallback behavior)
// or when the URL matches the configured URL
func (r *RFC7662Provider) CanHandle(introspectURL string) bool {
        // If no URL was configured, this is a fallback provider that handles everything
        if r.url == "" {
                return true
        }
        // Otherwise, only handle the specific configured URL
        return r.url == introspectURL
}

// IntrospectToken introspects a token using RFC 7662 standard
func (r *RFC7662Provider) IntrospectToken(ctx context.Context, token string) (jwt.MapClaims, error) {
        // Prepare form data for POST request
        formData := url.Values{}
        formData.Set("token", token)
        formData.Set("token_type_hint", "access_token")

        // Create POST request with form data
        //nolint:gosec // G704 - URL is configured introspection endpoint
        req, err := http.NewRequestWithContext(ctx, "POST", r.url, strings.NewReader(formData.Encode()))
        if err != nil {
                return nil, fmt.Errorf("failed to create introspection request: %w", err)
        }

        // Set headers
        req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
        req.Header.Set("User-Agent", oauthproto.UserAgent)
        req.Header.Set("Accept", "application/json")

        // Add client authentication if configured
        if r.clientID != "" && r.clientSecret != "" {
                req.SetBasicAuth(r.clientID, r.clientSecret)
        }

        // Make the request
        resp, err := r.client.Do(req) // #nosec G704 -- URL is the configured OIDC introspection endpoint
        if err != nil {
                return nil, fmt.Errorf("introspection request failed: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        // Read response body with a reasonable limit to prevent DoS attacks
        const maxResponseSize = 64 * 1024 // 64KB should be more than enough for introspection response
        limitedReader := io.LimitReader(resp.Body, maxResponseSize)
        body, err := io.ReadAll(limitedReader)
        if err != nil {
                return nil, fmt.Errorf("failed to read introspection response: %w", err)
        }

        // Check for HTTP errors
        if resp.StatusCode != http.StatusOK {
                if resp.StatusCode == http.StatusUnauthorized {
                        return nil, fmt.Errorf("introspection unauthorized")
                }
                return nil, fmt.Errorf("introspection failed with status %d: %s", resp.StatusCode, string(body))
        }

        // Parse RFC 7662 response - use the existing parseIntrospectionClaims function
        return parseIntrospectionClaims(strings.NewReader(string(body)))
}

// Common errors
var (
        ErrNoToken                 = errors.New("no token provided")
        ErrInvalidToken            = errors.New("invalid token")
        ErrTokenExpired            = errors.New("token expired")
        ErrInvalidIssuer           = errors.New("invalid issuer")
        ErrInvalidAudience         = errors.New("invalid audience")
        ErrMissingJWKSURL          = errors.New("missing JWKS URL")
        ErrFailedToFetchJWKS       = errors.New("failed to fetch JWKS")
        ErrFailedToDiscoverOIDC    = errors.New("failed to discover OIDC configuration")
        ErrMissingIssuerAndJWKSURL = errors.New("either issuer or JWKS URL must be provided")
)

// TokenValidator validates JWT or opaque tokens using OIDC configuration.
type TokenValidator struct {
        // OIDC configuration
        issuer            string
        audience          string
        jwksURL           string
        clientID          string
        clientSecret      string // Optional client secret for introspection
        jwksClient        *jwkfetch.Cache
        introspectURL     string       // Optional introspection endpoint
        client            *http.Client // HTTP client for making requests
        resourceURL       string       // (RFC 9728)
        scopes            []string     // OAuth scopes to advertise in WWW-Authenticate (RFC 6750 §3)
        registry          *Registry    // Token introspection providers
        insecureAllowHTTP bool         // Allow HTTP (non-HTTPS) OIDC issuers for development/testing

        // upstreamTokenReader loads upstream provider tokens for identity enrichment.
        // nil means no enrichment (no embedded auth server).
        upstreamTokenReader upstreamtoken.TokenReader

        // keyProvider provides in-process JWKS key lookups from the embedded auth
        // server's key provider. When set, getKeyFromJWKS resolves keys locally
        // before falling back to HTTP. Eliminates self-referential HTTP calls.
        // nil when no embedded auth server is configured.
        keyProvider keys.PublicKeyProvider

        // embeddedAuthServerIssuer binds session-less-grant provenance to tokens
        // validated for this embedded authorization server only.
        embeddedAuthServerIssuer string

        // Lazy JWKS registration
        jwksRegistered      bool
        jwksRegistrationMu  sync.Mutex
        jwksRegistrationErr error

        // Lazy OIDC discovery - allows deferring discovery until first validation request.
        // This is needed when the OIDC provider (auth server) is the same pod and starts after
        // the token validator is created.
        // Uses mutex+flag instead of sync.Once so that failed discovery can be retried
        // on subsequent ValidateToken calls (transient failures should not be permanent).
        oidcDiscoveryMu  sync.Mutex
        oidcDiscovered   bool
        oidcDiscoveryErr error
}

// TokenValidatorConfig contains configuration for the token validator.
type TokenValidatorConfig struct {
        // Issuer is the OIDC issuer URL (e.g., https://accounts.google.com)
        Issuer string

        // Audience is the expected audience for the token
        Audience string

        // JWKSURL is the URL to fetch the JWKS from
        JWKSURL string

        // ClientID is the OIDC client ID
        ClientID string

        // ClientSecret is the optional OIDC client secret for introspection
        ClientSecret string // #nosec G117 -- not a hardcoded credential, populated at runtime from config

        // CACertPath is the path to the CA certificate bundle for HTTPS requests
        CACertPath string

        // AuthTokenFile is the path to file containing bearer token for authentication
        AuthTokenFile string

        // AllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses
        AllowPrivateIP bool

        // InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing
        // WARNING: This is insecure and should NEVER be used in production
        InsecureAllowHTTP bool

        // IntrospectionURL is the optional introspection endpoint for validating tokens
        IntrospectionURL string

        // Store http client with the right config
        httpClient *http.Client

        // ResourceURL is the explicit resource URL for OAuth discovery (RFC 9728)
        ResourceURL string

        // Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728)
        // If empty, defaults to ["openid"]
        Scopes []string
}

// discoverOIDCConfiguration discovers OIDC configuration from the issuer's well-known endpoint
func discoverOIDCConfiguration(
        ctx context.Context,
        issuer string,
        client *http.Client,
        insecureAllowHTTP bool,
) (*oauthproto.OIDCDiscoveryDocument, error) {
        // Validate issuer URL scheme
        if err := networking.ValidateEndpointURLWithInsecure(issuer, insecureAllowHTTP); err != nil {
                return nil, fmt.Errorf("invalid issuer URL: %w", err)
        }

        // Construct the well-known endpoint URL
        wellKnownURL := strings.TrimSuffix(issuer, "/") + oauthproto.WellKnownOIDCPath

        // Create HTTP request
        req, err := http.NewRequestWithContext(ctx, http.MethodGet, wellKnownURL, nil)
        if err != nil {
                return nil, fmt.Errorf("failed to create request: %w", err)
        }

        // Set User-Agent header
        req.Header.Set("User-Agent", oauthproto.UserAgent)
        req.Header.Set("Accept", "application/json")

        resp, err := client.Do(req) // #nosec G704 -- URL is the configured OIDC issuer discovery endpoint
        if err != nil {
                return nil, fmt.Errorf("failed to fetch OIDC configuration: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        // Check response status
        if resp.StatusCode != http.StatusOK {
                return nil, fmt.Errorf("OIDC discovery endpoint returned status %d", resp.StatusCode)
        }

        // Parse the response
        var doc oauthproto.OIDCDiscoveryDocument
        if err := json.NewDecoder(resp.Body).Decode(&doc); err != nil {
                return nil, fmt.Errorf("failed to decode OIDC configuration: %w", err)
        }

        // Validate that we got the required fields
        if doc.JWKSURI == "" {
                return nil, fmt.Errorf("OIDC configuration missing jwks_uri")
        }

        return &doc, nil
}

// registerIntrospectionProviders creates and configures the provider registry
// for token introspection based on the configuration.
func registerIntrospectionProviders(config TokenValidatorConfig, clientSecret string) (*Registry, error) {
        registry := NewRegistry()

        // Add Google provider if the introspection URL matches
        if config.IntrospectionURL == GoogleTokeninfoURL {
                slog.Debug("registering Google tokeninfo provider", "url", config.IntrospectionURL)
                registry.AddProvider(NewGoogleProvider(config.IntrospectionURL))
        }

        // Add GitHub provider if the introspection URL matches GitHub's API pattern
        if strings.Contains(config.IntrospectionURL, GitHubTokenCheckURL) {
                slog.Debug("registering GitHub token validation provider", "url", config.IntrospectionURL)
                githubProvider, err := NewGitHubProvider(
                        config.IntrospectionURL,
                        config.ClientID,
                        clientSecret,
                        config.CACertPath,
                        config.AllowPrivateIP,
                )
                if err != nil {
                        return nil, fmt.Errorf("failed to create GitHub provider: %w", err)
                }
                registry.AddProvider(githubProvider)
        }

        // Add RFC7662 provider with auth if configured
        if config.ClientID != "" || clientSecret != "" {
                rfc7662Provider, err := NewRFC7662ProviderWithAuth(
                        config.IntrospectionURL, config.ClientID, clientSecret,
                        config.CACertPath, config.AuthTokenFile, config.AllowPrivateIP,
                )
                if err != nil {
                        return nil, fmt.Errorf("failed to create RFC7662 provider: %w", err)
                }
                registry.AddProvider(rfc7662Provider)
        }

        return registry, nil
}

// tokenValidatorOptions holds optional dependencies for NewTokenValidator.
type tokenValidatorOptions struct {
        envReader                env.Reader
        upstreamTokenReader      upstreamtoken.TokenReader
        keyProvider              keys.PublicKeyProvider
        embeddedAuthServerIssuer string
}

// TokenValidatorOption is a functional option for NewTokenValidator.
type TokenValidatorOption func(*tokenValidatorOptions)

// WithEnvReader overrides the default environment variable reader.
// This is primarily used in tests to avoid process-wide env var manipulation.
func WithEnvReader(reader env.Reader) TokenValidatorOption {
        return func(o *tokenValidatorOptions) {
                o.envReader = reader
        }
}

// WithUpstreamTokenReader configures the token validator to enrich Identity
// with upstream provider tokens. When set, the Middleware extracts the token
// session ID (tsid) from JWT claims and loads all upstream tokens into
// Identity.UpstreamTokens (access tokens) and Identity.UpstreamIDTokens
// (ID tokens) before placing the Identity in the request context. A token
// session whose upstream access token cannot be refreshed, or that holds no
// usable upstream access token at all, is answered with 401 and a
// WWW-Authenticate challenge instead, so the client signs in again.
func WithUpstreamTokenReader(reader upstreamtoken.TokenReader) TokenValidatorOption {
        return func(o *tokenValidatorOptions) {
                o.upstreamTokenReader = reader
        }
}

// WithEmbeddedAuthServerIssuer binds session-less-grant provenance to tokens
// validated for this embedded authorization server. This runtime-only value is
// deliberately separate from TokenValidatorConfig, which is user configuration.
func WithEmbeddedAuthServerIssuer(issuer string) TokenValidatorOption {
        return func(o *tokenValidatorOptions) {
                o.embeddedAuthServerIssuer = issuer
        }
}

// WithKeyProvider configures the token validator to use an in-process key
// provider for JWKS lookups instead of fetching keys over HTTP. This is used
// when the embedded auth server's key provider is available in the same process,
// eliminating self-referential HTTP calls and the need for insecureAllowHTTP
// and jwksAllowPrivateIP flags.
//
// Only PublicKeyProvider is required — the validator never signs tokens.
func WithKeyProvider(provider keys.PublicKeyProvider) TokenValidatorOption {
        return func(o *tokenValidatorOptions) {
                o.keyProvider = provider
        }
}

// resolveClientSecret returns the client secret from the config, falling back
// to the TOOLHIVE_OIDC_CLIENT_SECRET environment variable if not set.
func resolveClientSecret(configSecret string, envReader env.Reader) string {
        if configSecret != "" {
                return configSecret
        }
        if envSecret := envReader.Getenv("TOOLHIVE_OIDC_CLIENT_SECRET"); envSecret != "" {
                return envSecret
        }
        return ""
}

// resolveJWKSDiscovery determines the JWKS URL to use, and whether lazy OIDC
// discovery is needed. Discovery is deferred when JWKS URL is not provided
// but issuer is, allowing the validator to start before the OIDC provider is
// ready; ensureOIDCDiscovered populates jwksURL on first use in that case.
func resolveJWKSDiscovery(config TokenValidatorConfig, o *tokenValidatorOptions) (string, error) {
        jwksURL := config.JWKSURL

        // Skip discovery if TOOLHIVE_SKIP_OIDC_DISCOVERY is set (for testing only)
        skipDiscovery := o.envReader.Getenv("TOOLHIVE_SKIP_OIDC_DISCOVERY") == "true"
        if skipDiscovery && config.Issuer != "" {
                if jwksURL == "" {
                        return "", fmt.Errorf(
                                "TOOLHIVE_SKIP_OIDC_DISCOVERY=true requires explicit JWKSURL in config. " +
                                        "This env var is for testing only and cannot guess provider-specific JWKS URLs",
                        )
                }
                slog.Warn(
                        "OIDC discovery skipped (TOOLHIVE_SKIP_OIDC_DISCOVERY=true)",
                        "issuer", config.Issuer,
                )
        } else if jwksURL == "" && config.Issuer != "" {
                if o.keyProvider != nil {
                        slog.Debug("OIDC discovery deferred - failure non-fatal with local key provider",
                                "issuer", config.Issuer)
                } else {
                        slog.Debug("OIDC discovery deferred - will discover on first validation request",
                                "issuer", config.Issuer)
                }
        }

        // Ensure we have either an explicit JWKS URL, an issuer to discover from,
        // or a local key provider (embedded auth server).
        if jwksURL == "" && config.Issuer == "" && o.keyProvider == nil {
                return "", ErrMissingIssuerAndJWKSURL
        }

        return jwksURL, nil
}

// NewTokenValidator creates a new token validator.
func NewTokenValidator(ctx context.Context, config TokenValidatorConfig, opts ...TokenValidatorOption) (*TokenValidator, error) {
        // Apply functional options
        o := &tokenValidatorOptions{envReader: &env.OSReader{}}
        for _, opt := range opts {
                opt(o)
        }

        // Log warning if insecure HTTP is enabled
        if config.InsecureAllowHTTP {
                slog.Warn(
                        "insecure HTTP is enabled - "+
                                "HTTP OIDC URLs are allowed - this is INSECURE and should NEVER be used in production",
                        "issuer", config.Issuer,
                )
        }

        jwksURL, err := resolveJWKSDiscovery(config, o)
        if err != nil {
                return nil, err
        }

        if err := validateGoogleTokeninfoAudience(config); err != nil {
                return nil, err
        }

        // Create HTTP client with CA bundle and auth token support for JWKS
        httpClient, err := networking.NewHttpClientBuilder().
                WithCABundle(config.CACertPath).
                WithPrivateIPs(config.AllowPrivateIP).
                WithTokenFromFile(config.AuthTokenFile).
                WithInsecureAllowHTTP(config.InsecureAllowHTTP).
                Build()
        if err != nil {
                return nil, fmt.Errorf("failed to create HTTP client: %w", err)
        }
        config.httpClient = httpClient

        // Create a new JWKS client with auto-refresh.
        httprcClient := httprc.NewClient()
        cache, err := jwkfetch.NewCache(ctx, httprcClient, jwkfetch.WithHTTPClient(httpClient))
        if err != nil {
                return nil, fmt.Errorf("failed to create JWKS cache: %w", err)
        }

        // Skip synchronous JWKS registration - will be done lazily on first use

        // Resolve client secret from config or environment variable
        clientSecret := resolveClientSecret(config.ClientSecret, o.envReader)

        // Register introspection providers
        registry, err := registerIntrospectionProviders(config, clientSecret)
        if err != nil {
                return nil, err
        }

        validator := &TokenValidator{
                issuer:                   config.Issuer,
                audience:                 config.Audience,
                jwksURL:                  jwksURL,
                introspectURL:            config.IntrospectionURL,
                clientID:                 config.ClientID,
                clientSecret:             clientSecret,
                jwksClient:               cache,
                client:                   config.httpClient,
                resourceURL:              config.ResourceURL,
                scopes:                   config.Scopes,
                registry:                 registry,
                insecureAllowHTTP:        config.InsecureAllowHTTP,
                upstreamTokenReader:      o.upstreamTokenReader,
                keyProvider:              o.keyProvider,
                embeddedAuthServerIssuer: o.embeddedAuthServerIssuer,
        }

        return validator, nil
}

// validateGoogleTokeninfoAudience rejects a Google tokeninfo introspection
// config with no audience. Google's tokeninfo endpoint returns no iss claim -
// parseGoogleResponse synthesises iss locally, so a configured-issuer check
// against it is self-satisfying and proves nothing about which OAuth client
// the token was minted for. The only real binding to this deployment is the
// audience check, which is skipped when audience is empty. Without it, ANY
// valid Google access token (e.g. one minted for an attacker's own unrelated
// OAuth client) passes validation. Refuse the combination at startup instead
// of silently accepting cross-client tokens at runtime.
func validateGoogleTokeninfoAudience(config TokenValidatorConfig) error {
        if config.IntrospectionURL == GoogleTokeninfoURL && strings.TrimSpace(config.Audience) == "" {
                return fmt.Errorf("audience is required when using Google's tokeninfo endpoint for introspection: " +
                        "without it any valid Google access token is accepted regardless of the OAuth client it was minted for")
        }
        return nil
}

// ensureJWKSRegistered ensures that the JWKS URL is registered with the cache.
// This is called lazily on first use to avoid blocking startup.
// On failure, registration is retried on subsequent calls (transient failures
// should not permanently disable the validator).
func (v *TokenValidator) ensureJWKSRegistered(ctx context.Context) error {
        v.jwksRegistrationMu.Lock()
        defer v.jwksRegistrationMu.Unlock()

        // Already registered successfully - nothing to do
        if v.jwksRegistered {
                return nil
        }

        // Create context with 5-second timeout for JWKS registration
        registrationCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
        defer cancel()

        // The CA-aware client is configured on the cache so every fetch,
        // including background refreshes, uses the same transport policy.
        err := v.jwksClient.Register(registrationCtx, v.jwksURL)
        switch {
        case err == nil:
                // Registered and the first fetch succeeded.
        case errors.Is(err, httprc.ErrNotReady()):
                // The resource is registered and httprc will keep fetching it in the
                // background; only the first fetch has not completed within the
                // registration budget. Treat it as registered: Lookup returns a
                // not-ready error until a background fetch succeeds. Retrying Register
                // instead would fail with ErrResourceAlreadyExists forever.
                //nolint:gosec // G706: JWKS URL is from server configuration or OIDC discovery
                slog.Debug(
                        "JWKS URL registered but first fetch not ready; fetching continues in background",
                        "jwks_url", v.jwksURL, "error", err,
                )
        case errors.Is(err, httprc.ErrResourceAlreadyExists()):
                // The URL is already in httprc's resource map, e.g. a previous attempt
                // returned ErrNotReady before this state was tracked, or OIDC
                // re-discovery produced the same URL after resetting the flag.
        default:
                v.jwksRegistrationErr = fmt.Errorf("failed to register JWKS URL: %w", err)
                // Do NOT set jwksRegistered = true -- allow retry on next call
                return v.jwksRegistrationErr
        }

        v.jwksRegistered = true
        v.jwksRegistrationErr = nil
        return nil
}

// OIDC discovery retry configuration constants.
const (
        // oidcDiscoveryMaxAttempts is the maximum number of OIDC discovery attempts
        // (including the initial attempt). Set to 3 to allow sufficient time for
        // auth server startup (~3.5s with backoff).
        oidcDiscoveryMaxAttempts = 3

        // oidcDiscoveryInitialInterval is the starting delay for exponential backoff.
        oidcDiscoveryInitialInterval = 500 * time.Millisecond

        // oidcDiscoveryMaxInterval caps the maximum delay between retry attempts.
        oidcDiscoveryMaxInterval = 2 * time.Second

        // oidcDiscoveryAttemptTimeout is the timeout for each discovery attempt.
        oidcDiscoveryAttemptTimeout = 5 * time.Second
)

// ensureOIDCDiscovered performs lazy OIDC discovery with retry logic.
// This method is called on first token validation to discover the OIDC configuration
// (including JWKS URL) from the issuer. It uses exponential backoff to handle cases
// where the OIDC provider is starting up (e.g., same pod, load balancer warmup).
//
// Discovery is retried on each ValidateToken call until it succeeds. Once successful,
// subsequent calls return immediately. This allows recovery from transient failures
// (e.g., auth server not yet ready, temporary DNS issues, context cancellation).
func (v *TokenValidator) ensureOIDCDiscovered(ctx context.Context) error {
        // If there is no issuer to discover from, nothing to do.
        // v.issuer is immutable after construction, so this check is safe without a lock.
        if v.issuer == "" {
                return nil
        }

        v.oidcDiscoveryMu.Lock()
        defer v.oidcDiscoveryMu.Unlock()

        // Already discovered successfully (or JWKS URL was provided at construction) - return immediately.
        if v.oidcDiscovered || v.jwksURL != "" {
                return nil
        }

        // Configure exponential backoff with jitter (provided by the library)
        expBackoff := backoff.NewExponentialBackOff()
        expBackoff.InitialInterval = oidcDiscoveryInitialInterval
        expBackoff.MaxInterval = oidcDiscoveryMaxInterval
        expBackoff.Reset()

        // Define the discovery operation with per-attempt timeout
        operation := func() (*oauthproto.OIDCDiscoveryDocument, error) {
                attemptCtx, cancel := context.WithTimeout(ctx, oidcDiscoveryAttemptTimeout)
                defer cancel()

                return discoverOIDCConfiguration(
                        attemptCtx,
                        v.issuer,
                        v.client,
                        v.insecureAllowHTTP,
                )
        }

        // Execute with retry using cenkalti/backoff (consistent with ToolHive patterns)
        doc, err := backoff.Retry(ctx, operation,
                backoff.WithBackOff(expBackoff),
                backoff.WithMaxTries(oidcDiscoveryMaxAttempts),
                backoff.WithNotify(func(err error, duration time.Duration) {
                        slog.Debug(
                                "oidc discovery failed, retrying",
                                "issuer", v.issuer, "retry_in", duration, "error", err,
                        )
                }),
        )

        if err != nil {
                v.oidcDiscoveryErr = fmt.Errorf("%w: %w", ErrFailedToDiscoverOIDC, err)
                //nolint:gosec // G706: issuer is from server configuration
                slog.Error(
                        "oidc discovery failed after retries",
                        "issuer", v.issuer, "attempts", oidcDiscoveryMaxAttempts, "error", err,
                )
                // Do NOT set oidcDiscovered = true -- allow retry on next ValidateToken call
                return v.oidcDiscoveryErr
        }

        // Discovery succeeded - update the JWKS URL and mark as done
        v.jwksURL = doc.JWKSURI
        v.oidcDiscovered = true
        v.oidcDiscoveryErr = nil
        // Reset JWKS registration so it re-registers with the newly discovered URL.
        // Acquire jwksRegistrationMu to safely reset the flag, since ensureJWKSRegistered
        // reads it under that mutex. Lock ordering: oidcDiscoveryMu -> jwksRegistrationMu.
        v.jwksRegistrationMu.Lock()
        v.jwksRegistered = false
        v.jwksRegistrationMu.Unlock()
        //nolint:gosec // G706: issuer and JWKS URL are from OIDC discovery
        slog.Debug(
                "oidc discovery succeeded",
                "issuer", v.issuer, "jwks_url", doc.JWKSURI,
        )

        return nil
}

// getKeyFromLocalProvider attempts to find a verification key from the local
// key provider (embedded auth server). Returns (key, nil) on success,
// (nil, nil) to signal fallback to HTTP, or (nil, error) for hard failures.
// validateTokenHeader checks the signing method is supported (RSA or ECDSA) and
// extracts the key ID from the token header. Returns an error for unsupported
// methods or a missing kid claim.
func validateTokenHeader(token *jwt.Token) (string, error) {
        switch token.Method.(type) {
        case *jwt.SigningMethodRSA, *jwt.SigningMethodECDSA:
                // Supported signing methods
        default:
                return "", fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
        }

        kid, ok := token.Header["kid"].(string)
        if !ok {
                return "", fmt.Errorf("token header missing kid")
        }
        return kid, nil
}

func (v *TokenValidator) getKeyFromLocalProvider(ctx context.Context, token *jwt.Token) (interface{}, error) {
        if v.keyProvider == nil {
                return nil, nil
        }

        kid, err := validateTokenHeader(token)
        if err != nil {
                return nil, err
        }

        pubKeys, err := v.keyProvider.PublicKeys(ctx)
        if err != nil {
                slog.Debug("local JWKS provider failed, falling back to HTTP", "error", err)
                return nil, nil
        }

        for _, k := range pubKeys {
                if k.KeyID == kid {
                        slog.Debug("resolved JWKS key from embedded auth server", "kid", kid)
                        return k.PublicKey, nil
                }
        }

        // Key not found locally — fall back to HTTP JWKS
        slog.Debug("key not found in local JWKS provider, falling back to HTTP", "kid", kid)
        return nil, nil
}

// getKeyFromJWKS gets the key from the JWKS.
func (v *TokenValidator) getKeyFromJWKS(ctx context.Context, token *jwt.Token) (interface{}, error) {
        // Try local key provider first (embedded auth server in-process keys).
        // This avoids self-referential HTTP calls when the auth server and
        // token validator run in the same process.
        if key, err := v.getKeyFromLocalProvider(ctx, token); err != nil {
                return nil, err
        } else if key != nil {
                return key, nil
        }

        // Fall through to HTTP-based JWKS lookup.
        // When a local key provider is present, OIDC discovery may have been
        // skipped entirely, so jwksURL can legitimately be empty.
        if v.jwksURL == "" {
                if v.keyProvider != nil {
                        return nil, fmt.Errorf(
                                "local key provider could not resolve key and no JWKS URL is available: %w",
                                ErrMissingJWKSURL,
                        )
                }
                return nil, ErrMissingJWKSURL
        }

        // Ensure JWKS is registered before attempting to use it
        if err := v.ensureJWKSRegistered(ctx); err != nil {
                return nil, fmt.Errorf("JWKS registration failed: %w", err)
        }

        kid, err := validateTokenHeader(token)
        if err != nil {
                return nil, err
        }

        // Get the key set from the JWKS.
        keySet, err := v.jwksClient.Lookup(ctx, v.jwksURL)
        if err != nil {
                return nil, fmt.Errorf("failed to lookup JWKS: %w", err)
        }

        // Get the key with the matching key ID
        key, found := keySet.LookupKeyID(kid)
        if !found {
                return nil, fmt.Errorf("key ID %s not found in JWKS", kid)
        }

        // Export the raw key.
        rawKey, err := jwk.Export[any](key)
        if err != nil {
                return nil, fmt.Errorf("failed to export raw key: %w", err)
        }

        return rawKey, nil
}

// validateClaims validates the claims in the token.
func (v *TokenValidator) validateClaims(claims jwt.MapClaims) error {
        // Validate the issuer if provided
        if v.issuer != "" {
                issuerClaim, err := claims.GetIssuer()
                if err != nil {
                        return fmt.Errorf("failed to get issuer from claims: %w", err)
                }
                if strings.TrimSpace(issuerClaim) != strings.TrimSpace(v.issuer) {
                        return ErrInvalidIssuer
                }
        }
        // Validate the audience if provided
        if v.audience != "" {
                audiences, err := claims.GetAudience()
                if err != nil {
                        return ErrInvalidAudience
                }

                found := false
                for _, aud := range audiences {
                        if aud == v.audience {
                                found = true
                                break
                        }
                }

                if !found {
                        return ErrInvalidAudience
                }
        }

        // Validate the expiration time
        expirationTime, err := claims.GetExpirationTime()
        if err != nil || expirationTime == nil || expirationTime.Before(time.Now()) {
                return ErrTokenExpired
        }

        return nil
}

func parseIntrospectionClaims(r io.Reader) (jwt.MapClaims, error) {
        var j struct {
                Active bool        `json:"active"`
                Exp    *float64    `json:"exp,omitempty"`
                Sub    string      `json:"sub,omitempty"`
                Aud    interface{} `json:"aud,omitempty"`
                Scope  string      `json:"scope,omitempty"`
                Iss    string      `json:"iss,omitempty"`
        }
        var body json.RawMessage

        if err := json.NewDecoder(r).Decode(&body); err != nil {
                return nil, fmt.Errorf("failed to decode introspection JSON: %w", err)
        }
        if err := json.Unmarshal(body, &j); err != nil {
                return nil, fmt.Errorf("failed to decode introspection JSON: %w", err)
        }
        if !j.Active {
                return nil, ErrInvalidToken
        }

        claims := jwt.MapClaims{}
        var extra map[string]interface{}
        if err := json.Unmarshal(body, &extra); err != nil {
                return nil, fmt.Errorf("failed to decode introspection JSON: %w", err)
        }
        for k, v := range extra {
                switch k {
                case "active", "exp", "sub", "aud", "scope", "iss":
                        continue
                default:
                        claims[k] = v
                }
        }
        if j.Exp != nil {
                claims["exp"] = *j.Exp
        }
        if j.Sub != "" {
                claims["sub"] = strings.TrimSpace(j.Sub)
        }
        if j.Aud != nil {
                claims["aud"] = j.Aud
        }
        if j.Scope != "" {
                claims["scope"] = strings.TrimSpace(j.Scope)
        }
        if j.Iss != "" {
                claims["iss"] = strings.TrimSpace(j.Iss)
        }

        return claims, nil
}

// introspectOpaqueToken uses the provider pattern to introspect opaque tokens
func (v *TokenValidator) introspectOpaqueToken(ctx context.Context, tokenStr string) (jwt.MapClaims, error) {
        if v.introspectURL == "" {
                return nil, fmt.Errorf("no introspection endpoint available")
        }

        // Find appropriate provider for the introspection URL
        provider := v.registry.GetIntrospector(v.introspectURL)
        if provider == nil {
                return nil, fmt.Errorf("no provider available for introspection URL: %s", v.introspectURL)
        }

        // Use provider to introspect the token
        claims, err := provider.IntrospectToken(ctx, tokenStr)
        if err != nil {
                return nil, fmt.Errorf("%s introspection failed: %w", provider.Name(), err)
        }

        // Validate required claims (exp, iss, aud if configured)
        if err := v.validateClaims(claims); err != nil {
                return nil, err
        }

        return claims, nil
}

// ValidateToken validates a token.
func (v *TokenValidator) ValidateToken(ctx context.Context, tokenString string) (jwt.MapClaims, error) {
        // Ensure OIDC discovery is complete (lazy discovery with retry).
        // When a local key provider is configured (embedded auth server),
        // discovery failure is non-fatal — signing keys can be resolved
        // in-process and the issuer URL may not be reachable from within
        // the cluster. Mark discovery as done to avoid per-request retries.
        if err := v.ensureOIDCDiscovered(ctx); err != nil {
                if v.keyProvider == nil {
                        return nil, fmt.Errorf("OIDC discovery failed: %w", err)
                }
                slog.Warn("OIDC discovery failed, proceeding with local key provider",
                        "issuer", v.issuer, "error", err)
                v.oidcDiscoveryMu.Lock()
                v.oidcDiscovered = true
                v.oidcDiscoveryMu.Unlock()
        }

        // Parse the token
        token, err := jwt.Parse(tokenString, func(token *jwt.Token) (any, error) {
                return v.getKeyFromJWKS(ctx, token)
        })

        if err != nil {
                if errors.Is(err, jwt.ErrTokenMalformed) {
                        // check against introspection endpoint if available
                        claims, err := v.introspectOpaqueToken(ctx, tokenString)
                        if err != nil {
                                return nil, fmt.Errorf("failed to introspect opaque token: %w", err)
                        }
                        return claims, nil
                }
                return nil, fmt.Errorf("failed to parse token: %w", err)
        }

        // it is a jwt token
        // Check if the token is valid
        if !token.Valid {
                return nil, ErrInvalidToken
        }

        // Get the claims
        claims, ok := token.Claims.(jwt.MapClaims)
        if !ok {
                return nil, fmt.Errorf("failed to get claims from token")
        }

        // Validate the claims
        if err := v.validateClaims(claims); err != nil {
                return nil, err
        }

        return claims, nil
}

// buildWWWAuthenticate builds a RFC 6750 / RFC 9728 compliant value for the
// WWW-Authenticate header. It always includes realm and, if set, resource_metadata.
// When errorCode is non-empty ("invalid_request", "invalid_token", or "insufficient_scope"),
// it appends the error and optional error_description.
func (v *TokenValidator) buildWWWAuthenticate(errorCode string, errDescription string) string {
        var parts []string

        // realm (RFC 6750)
        if v.issuer != "" {
                parts = append(parts, fmt.Sprintf(`realm="%s"`, EscapeQuotes(v.issuer)))
        }

        // resource_metadata (RFC 9728)
        // Per RFC 9728 Section 3.1, the well-known URI is inserted between the host and path components
        // Example: https://resource.example.com/resource1 -> https://resource.example.com/.well-known/oauth-protected-resource/resource1
        if v.resourceURL != "" {
                parsedURL, err := url.Parse(v.resourceURL)
                if err == nil {
                        // Per RFC 9728 Section 3.1, remove any terminating slash from path
                        path := parsedURL.Path
                        if path == "/" {
                                path = ""
                        }

                        // Construct the metadata URL by inserting the well-known path between host and path
                        metadataURL := fmt.Sprintf("%s://%s/.well-known/oauth-protected-resource%s",
                                parsedURL.Scheme,
                                parsedURL.Host,
                                path)
                        parts = append(parts, fmt.Sprintf(`resource_metadata="%s"`, EscapeQuotes(metadataURL)))
                }
        }

        // scope (RFC 6750 §3) - only when explicitly configured
        if len(v.scopes) > 0 {
                parts = append(parts, fmt.Sprintf(`scope="%s"`, EscapeQuotes(strings.Join(v.scopes, " "))))
        }

        // error fields (RFC 6750 §3)
        if errorCode != "" {
                parts = append(parts, fmt.Sprintf(`error="%s"`, EscapeQuotes(errorCode)))
                if errDescription != "" {
                        parts = append(parts, fmt.Sprintf(`error_description="%s"`, EscapeQuotes(errDescription)))
                }
        }
        return "Bearer " + strings.Join(parts, ", ")
}

// RFC 6750 error code constants for Bearer token authentication.
const (
        OAuthErrInvalidRequest    = "invalid_request"
        OAuthErrInvalidToken      = "invalid_token"
        OAuthErrInsufficientScope = "insufficient_scope"
)

// upstreamReauthRequired is the error_description of the 401 the validator
// answers when a token session's upstream credential is unusable: the client's
// only cure is a new upstream login.
const upstreamReauthRequired = "upstream token is no longer valid; re-authentication required"

// RFC6750Error represents an RFC 6750 compliant OAuth error response body.
type RFC6750Error struct {
        Error            string `json:"error"`
        ErrorDescription string `json:"error_description"`
}

// writeOAuthError writes an RFC 6750 compliant JSON error response.
func writeOAuthError(w http.ResponseWriter, errorCode, description string) {
        body, err := json.Marshal(RFC6750Error{
                Error:            errorCode,
                ErrorDescription: description,
        })
        if err != nil {
                http.Error(w, "internal server error", http.StatusInternalServerError)
                return
        }
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusUnauthorized)
        _, _ = w.Write(body)
}

// writeUpstreamReauthRequired answers 401 with the RFC 6750 invalid_token
// challenge for a token session whose upstream credential is unusable, so the
// client re-authenticates and the login writes fresh upstream rows.
func (v *TokenValidator) writeUpstreamReauthRequired(w http.ResponseWriter) {
        w.Header().Set("WWW-Authenticate", v.buildWWWAuthenticate(OAuthErrInvalidToken, upstreamReauthRequired))
        writeOAuthError(w, OAuthErrInvalidToken, upstreamReauthRequired)
}

// loadUpstreamTokens extracts the token session ID from claims and loads
// all upstream provider credentials (access + ID tokens) for that session.
// Returns (nil, nil, nil) if no tsid claim exists. Returns a non-nil error
// when a tsid claim is present but token loading fails due to an infrastructure
// error (storage unavailable). A non-empty failed map means one or more
// providers' tokens could not be refreshed; the caller should return HTTP 401
// to trigger re-authentication.
func (v *TokenValidator) loadUpstreamTokens(
        ctx context.Context, claims jwt.MapClaims,
) (map[string]upstreamtoken.UpstreamCredential, map[string]error, error) {
        tsid, ok := claims[upstreamtoken.TokenSessionIDClaimKey].(string)
        if !ok || tsid == "" {
                return nil, nil, nil
        }

        creds, failed, err := v.upstreamTokenReader.GetAllUpstreamCredentials(ctx, tsid)
        if err != nil {
                // Log tsid at DEBUG only — it is credential-adjacent and must not
                // leak into WARN-level logs or returned error strings.
                slog.DebugContext(ctx, "load upstream credentials failed", "tsid", tsid, "error", err)
                return nil, nil, fmt.Errorf("load upstream credentials: %w", err)
        }
        return creds, failed, nil
}

// isSessionlessGrant reports whether a validated token from the configured
// embedded authorization server affirmatively proves it was minted without an
// upstream session. A nil upstream-token map alone is never sufficient.
//
// The gate requires the validator's own issuer to be the embedded authorization
// server's, rather than trusting the token's "iss" claim on its own. validateClaims
// only verifies "iss" when v.issuer is set, so on a JWKS-only configuration the
// claim is unverified beyond the signature and would be a self-assertion. Binding
// to v.issuer makes "validated for the configured embedded authorization-server
// issuer" true by construction instead of by deployment convention, and leaves the
// claim comparison below as a redundant restatement of what validateClaims proved.
func (v *TokenValidator) isSessionlessGrant(identity *Identity, claims jwt.MapClaims) bool {
        // Trim on both sides, as validateClaims does: an issuer configured with stray
        // whitespace must not silently turn this gate off and deny every session-less
        // grant with no diagnostic.
        embeddedIssuer := strings.TrimSpace(v.embeddedAuthServerIssuer)
        if embeddedIssuer == "" || strings.TrimSpace(v.issuer) != embeddedIssuer {
                return false
        }
        if identity.UpstreamTokens != nil {
                return false
        }
        issuer, ok := claims["iss"].(string)
        if !ok || strings.TrimSpace(issuer) != embeddedIssuer {
                return false
        }
        if chain := identity.DelegationChain; chain != nil && !chain.Malformed && len(chain.Chain) > 0 {
                return true
        }
        marker, ok := claims[upstreamtoken.NoUpstreamSessionClaimKey].(bool)
        return ok && marker
}

// Middleware creates an HTTP middleware that validates JWT tokens and creates Identity.
func (v *TokenValidator) Middleware(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                // Extract the bearer token from the Authorization header
                tokenString, err := ExtractBearerToken(r)
                if err != nil {
                        w.Header().Set("WWW-Authenticate", v.buildWWWAuthenticate(OAuthErrInvalidRequest, err.Error()))
                        writeOAuthError(w, OAuthErrInvalidRequest, err.Error())
                        return
                }

                // Validate the token
                claims, err := v.ValidateToken(r.Context(), tokenString)
                if err != nil {
                        w.Header().Set("WWW-Authenticate", v.buildWWWAuthenticate(OAuthErrInvalidToken, err.Error()))
                        writeOAuthError(w, OAuthErrInvalidToken, fmt.Sprintf("Invalid token: %v", err))
                        return
                }

                // Convert claims to Identity
                identity, err := claimsToIdentity(claims, tokenString)
                if err != nil {
                        slog.Error("failed to convert claims to identity", "error", err)
                        w.Header().Set("WWW-Authenticate", v.buildWWWAuthenticate(OAuthErrInvalidToken, err.Error()))
                        writeOAuthError(w, OAuthErrInvalidToken, "Invalid authentication claims")
                        return
                }

                // Enrich Identity with upstream provider tokens when an embedded auth
                // server is active (reader configured via WithUpstreamTokenReader). The
                // load runs against a temporary load-scoped context carrying the Identity,
                // so storage layers invoked during the load can resolve the canonical user
                // from context (via CanonicalUserFromContext, which falls back to the
                // Identity's PlatformUserID on this request-serving path). The in-place
                // UpstreamTokens mutation completes here, before the publicly-reachable
                // context is built below — so the Identity placed in the served context is
                // never mutated afterwards (see the UpstreamTokens doc comment in identity.go).
                if v.upstreamTokenReader != nil {
                        loadCtx := WithIdentity(r.Context(), identity)
                        creds, failed, loadErr := v.loadUpstreamTokens(loadCtx, claims)
                        if loadErr != nil {
                                slog.WarnContext(loadCtx, "upstream token storage unavailable",
                                        "error", loadErr,
                                )
                                http.Error(w, "authentication service temporarily unavailable", http.StatusServiceUnavailable)
                                return
                        }
                        // Conservative policy: any provider refresh failure rejects the entire
                        // request. In vMCP multi-provider sessions this may block requests that
                        // could partially succeed if only some backends need the failing provider.
                        // A per-backend provider check would require routing context not available
                        // at this layer; this is the correct tradeoff for now.
                        if len(failed) > 0 {
                                slog.WarnContext(loadCtx, "upstream token refresh failed; returning 401",
                                        "providers", slices.Sorted(maps.Keys(failed)))
                                v.writeUpstreamReauthRequired(w)
                                return
                        }
                        // Project the per-provider credential bundle onto the two
                        // consumer-facing Identity fields. creds is non-nil exactly when
                        // the token carried a tsid claim; without one both fields stay nil.
                        if creds != nil {
                                accessTokens := make(map[string]string, len(creds))
                                idTokens := make(map[string]string, len(creds))
                                for provider, cred := range creds {
                                        // Omit providers with no usable access token (mirrors the ID-token
                                        // filter below); a provider may carry an ID token but no access token.
                                        if cred.AccessToken != "" {
                                                accessTokens[provider] = cred.AccessToken
                                        }
                                        // Omit providers whose upstream login did not yield an
                                        // ID token so consumers see the same map shape as before
                                        // (keyed only on providers with a usable ID token).
                                        if cred.IDToken != "" {
                                                idTokens[provider] = cred.IDToken
                                        }
                                }
                                // A token session that holds no usable access token for any
                                // provider (every stored row has expired or been deleted, or
                                // none carried an access token) is in the same state as one
                                // whose refresh was refused: consumers of UpstreamTokens can
                                // only deny it, and a denial does not make the client sign in
                                // again. Answer the same 401 challenge here, before any
                                // consumer runs.
                                //
                                // Debug, unlike the refresh-failure WARN above: an expired row
                                // is ordinary lifecycle rather than an upstream fault, a client
                                // retrying on a stale token sets the rate, and the audit
                                // middleware records the 401. The line names the session (tsid
                                // stays at Debug, as in loadUpstreamTokens) and the providers
                                // whose rows were still stored without an access token.
                                if len(accessTokens) == 0 {
                                        tsid, _ := claims[upstreamtoken.TokenSessionIDClaimKey].(string)
                                        slog.DebugContext(loadCtx, "token session holds no upstream access token; returning 401",
                                                "tsid", tsid, "stored_providers", slices.Sorted(maps.Keys(creds)))
                                        v.writeUpstreamReauthRequired(w)
                                        return
                                }
                                identity.UpstreamTokens = accessTokens
                                identity.UpstreamIDTokens = idTokens
                        }
                }

                identity.SessionlessGrant = v.isSessionlessGrant(identity, claims)

                // Add the fully-populated Identity to the request context and serve.
                ctx := WithIdentity(r.Context(), identity)
                next.ServeHTTP(w, r.WithContext(ctx))
        })
}

// RFC9728AuthInfo represents the OAuth Protected Resource metadata as defined in RFC 9728
type RFC9728AuthInfo struct {
        Resource               string   `json:"resource"`
        AuthorizationServers   []string `json:"authorization_servers"`
        BearerMethodsSupported []string `json:"bearer_methods_supported"`
        JWKSURI                string   `json:"jwks_uri,omitempty"`
        ScopesSupported        []string `json:"scopes_supported"`
}

// NewAuthInfoHandler creates an HTTP handler that returns RFC-9728 compliant OAuth Protected Resource metadata
func NewAuthInfoHandler(issuer, resourceURL string, scopes []string) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                // Set CORS headers for all requests
                origin := r.Header.Get("Origin")
                if origin == "" {
                        // Allow all origins if none specified. This should be fine because this is a discovery endpoint.
                        origin = "*"
                }
                w.Header().Set("Access-Control-Allow-Origin", origin)
                w.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS")
                // At least mcp-inspector requires these headers to be set for CORS. It seems to be a little
                // off since this is a discovery endpoint, but let's make the inspector happy.
                w.Header().Set("Access-Control-Allow-Headers", "mcp-protocol-version, Content-Type, Authorization")
                w.Header().Set("Access-Control-Max-Age", "86400") // 24 hours

                if r.Method == http.MethodOptions {
                        w.WriteHeader(http.StatusNoContent)
                        return
                }

                // if resourceURL is not set, return 404 as we shouldn't presume a resource URL
                if resourceURL == "" {
                        w.WriteHeader(http.StatusNotFound)
                        return
                }

                // Use provided scopes or fall back to a minimal default.
                // When the embedded auth server is used, the caller provides
                // the AS's ScopesSupported explicitly (see config_builder.go).
                supportedScopes := scopes
                if len(supportedScopes) == 0 {
                        supportedScopes = []string{"openid"}
                }

                authInfo := RFC9728AuthInfo{
                        Resource:               resourceURL,
                        AuthorizationServers:   []string{issuer},
                        BearerMethodsSupported: []string{"header"},
                        ScopesSupported:        supportedScopes,
                }

                // Set content type
                w.Header().Set("Content-Type", "application/json")

                // Encode and send the response
                if err := json.NewEncoder(w).Encode(authInfo); err != nil {
                        slog.Error("failed to encode OAuth discovery response", "error", err)
                        http.Error(w, "Internal server error", http.StatusInternalServerError)
                        return
                }
        })
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package tokensource provides a shared OIDC-backed OAuth token source used by
// both the LLM gateway client and the registry authentication client.
//
// The shared OAuthTokenSource implements a four-tier token strategy:
//
//  1. In-memory cached oauth2.TokenSource (auto-refreshes transparently)
//  2. Secrets-provider cached access token (cross-process reuse to avoid racing)
//  3. Refresh token stored in the secrets provider (restores across CLI invocations)
//  4. Browser-based OIDC+PKCE flow (only when interactive is true)
//
// Callers parameterise the source via Options: an OIDC config struct, a key
// provider function (which determines the secrets-provider key and its prefix),
// and a ConfigPersister callback that persists the token reference into the
// application config.
package tokensource

import (
        "context"
        "crypto/sha256"
        "encoding/hex"
        "errors"
        "fmt"
        "log/slog"
        "slices"
        "strings"
        "sync"
        "time"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth/oauth"
        "github.com/stacklok/toolhive/pkg/auth/oautherr"
        "github.com/stacklok/toolhive/pkg/auth/remote"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// errCacheMiss is a sentinel used internally to distinguish "no token in cache"
// from real backend errors. It is never exposed to callers.
var errCacheMiss = errors.New("no cached refresh token")

// preemptiveRefreshWindow is the default distance before actual expiry at which
// a token is treated as expired, triggering a proactive refresh before the
// upstream rejects it. Callers that re-invoke the token source on a fixed
// cadence (e.g. the LLM gateway's apiKeyHelper) widen this via
// Options.PreemptiveRefreshWindow so a refresh always lands before expiry.
const preemptiveRefreshWindow = 30 * time.Second

// OIDCParams holds the OIDC connection parameters for a token source.
type OIDCParams struct {
        Issuer                   string
        ClientID                 string
        Scopes                   []string // "offline_access" is appended automatically if absent
        Audience                 string
        CallbackPort             int
        RequireExactCallbackPort bool
}

// ConfigPersister is called when the refresh token key or expiry changes —
// after a successful browser flow (initial login) or when the OIDC provider
// rotates the refresh token during a refresh. Callers wire this to their config
// persistence layer. It is NOT called on routine access-token refreshes where
// the refresh token is unchanged.
type ConfigPersister func(refreshTokenKey string, expiry time.Time)

// Options configures an OAuthTokenSource.
type Options struct {
        OIDC OIDCParams
        // SecretsProvider is used to persist and restore refresh and access tokens.
        // May be nil; in that case tier 2/3 return an actionable error rather than
        // the FallbackErr, so the caller sees the real cause.
        SecretsProvider secrets.Provider
        // Interactive controls whether the browser OIDC+PKCE flow is allowed.
        Interactive bool
        // SkipBrowser, when true, makes the interactive flow print the authorization
        // URL to stderr and wait for the callback instead of opening a browser. It
        // only has effect when Interactive is true. Useful for headless/SSH/CI
        // environments where no system browser is available.
        SkipBrowser bool
        // KeyProvider returns the secrets-provider key for the refresh token.
        // May be called multiple times per Token() invocation; must be deterministic
        // and free of side effects. Typically returns a cached config ref if set,
        // otherwise derives a key from url+issuer.
        KeyProvider func() string
        // ConfigPersister is called when a new refresh token is persisted (login or
        // rotation). May be nil to skip config persistence.
        ConfigPersister ConfigPersister
        // FallbackErr is returned in non-interactive mode when no cached credentials
        // exist and no actionable lastErr is available. Defaults to a generic error.
        FallbackErr error
        // PreemptiveRefreshWindow overrides how far before expiry a token is
        // proactively refreshed. Zero uses the default (preemptiveRefreshWindow,
        // 30 s); a negative value panics in New. Widen this when the caller polls
        // the token source on a fixed cadence (e.g. the LLM gateway apiKeyHelper) so
        // the window comfortably exceeds that cadence and a refresh always lands
        // before expiry.
        PreemptiveRefreshWindow time.Duration
}

// OAuthTokenSource provides OIDC-backed access tokens via a four-tier strategy.
// It is safe for concurrent use.
type OAuthTokenSource struct {
        opts        Options
        mu          sync.Mutex
        tokenSource oauth2.TokenSource
}

// New creates an OAuthTokenSource from the given Options.
// It panics if opts.KeyProvider is nil, as this is a required field, or if
// opts.PreemptiveRefreshWindow is negative.
// If opts.OIDC.CallbackPort is zero, it defaults to remote.DefaultCallbackPort.
// If opts.PreemptiveRefreshWindow is zero, it defaults to preemptiveRefreshWindow.
func New(opts Options) *OAuthTokenSource {
        if opts.KeyProvider == nil {
                panic("tokensource.New: Options.KeyProvider must not be nil")
        }
        if opts.PreemptiveRefreshWindow < 0 {
                panic("tokensource.New: Options.PreemptiveRefreshWindow must not be negative")
        }
        if opts.FallbackErr == nil {
                opts.FallbackErr = errors.New("authentication required: no cached credentials found; complete an interactive login first")
        }
        if opts.OIDC.CallbackPort == 0 {
                opts.OIDC.CallbackPort = remote.DefaultCallbackPort
        }
        return &OAuthTokenSource{opts: opts}
}

// Token returns a valid access token string.
// It is safe for concurrent use.
func (t *OAuthTokenSource) Token(ctx context.Context) (string, error) {
        t.mu.Lock()
        defer t.mu.Unlock()

        // lastErr tracks the most recent actionable error from tiers 1/2. In
        // non-interactive mode it is returned instead of FallbackErr so the caller
        // sees the real cause (e.g. invalid_grant, keyring locked).
        var lastErr error

        if tok, err := t.tryInMemoryToken(); err == nil {
                return tok, nil
        } else if !errors.Is(err, errCacheMiss) {
                lastErr = err
        }

        // Tier 1.5: secrets-cached access token — avoids IdP exchange when the
        // access token is still fresh. Concurrent short-lived processes share this
        // cached value instead of all racing to exchange the same refresh token,
        // preventing invalid_grant from OIDC providers that rotate refresh tokens.
        if tok, found := t.tryAccessTokenCache(ctx); found {
                return tok, nil
        }

        if tok, err := t.tryCachedToken(ctx); err == nil {
                return tok, nil
        } else if !errors.Is(err, errCacheMiss) {
                lastErr = err
        }

        // Tier 3: browser OIDC+PKCE flow — only in interactive mode.
        if !t.opts.Interactive {
                if lastErr != nil {
                        return "", t.classifyTerminalError(lastErr)
                }
                return "", t.opts.FallbackErr
        }
        if err := t.performBrowserFlow(ctx); err != nil {
                return "", fmt.Errorf("OIDC browser flow failed: %w", err)
        }
        tok, err := t.tokenSource.Token()
        if err != nil {
                return "", fmt.Errorf("failed to get token after browser flow: %w", err)
        }
        t.cacheAccessToken(ctx, tok.AccessToken, tok.Expiry)
        return tok.AccessToken, nil
}

// classifyTerminalError decides which error a non-interactive caller sees when
// every cache tier has failed.
//
// An RFC 6749 invalid_grant means the OAuth server rejected the stored refresh
// token itself: it expired, was revoked, or was rotated out from under us.
// Retrying cannot fix that, and only a fresh interactive login can. Returning
// the raw OAuth error there is actively misleading, because callers that
// distinguish "you must log in again" from "the provider is having a bad day"
// all key off FallbackErr, and a dead refresh token is precisely the case that
// needs the re-login remediation.
//
// Every other failure surfaces verbatim, including the permanent ones. A 5xx,
// a 429, a WAF page, or a locked keyring is worth retrying or fixing in place.
// invalid_client, unauthorized_client, and invalid_scope are not worth
// retrying, but they are verdicts on the client registration or the requested
// scopes rather than on the credential, so logging in again reproduces them
// unchanged — reporting them as FallbackErr would send the user in a circle
// while hiding the code that names the real problem.
//
// The stored credential is deliberately NOT deleted here. With an IdP that
// rotates refresh tokens, a sibling process may have just written a newer token
// under the same key, and the next Token() call re-reads the secrets provider;
// deleting on a rejection would destroy that cross-process recovery.
func (t *OAuthTokenSource) classifyTerminalError(lastErr error) error {
        if t.opts.FallbackErr == nil || !oautherr.IsRejectedRefreshGrant(lastErr) {
                return lastErr
        }
        return &credentialRejectedError{fallback: t.opts.FallbackErr, cause: lastErr}
}

// credentialRejectedError marks a non-interactive failure where the identity
// provider rejected the stored refresh token. It wraps BOTH the caller's
// FallbackErr (so errors.Is finds the re-login sentinel) and the underlying
// cause (so errors.As can still reach the *oauth2.RetrieveError).
//
// Error() interpolates nothing from the token endpoint, on purpose. The whole
// response is untrusted input: an *oauth2.RetrieveError's own Error() embeds the
// raw body, which can echo back bearer material, and even the parsed RFC 6749
// 'error' field is an arbitrary server-chosen string that may carry secrets or
// control characters into a log line. Since this type is only constructed for
// invalid_grant, the code carries no information the fixed sentence does not
// already convey. Diagnostics that need the exact verdict reach the cause
// through errors.As.
type credentialRejectedError struct {
        fallback error
        cause    error
}

func (e *credentialRejectedError) Error() string {
        return fmt.Sprintf("%s (the identity provider rejected the stored credential)", e.fallback)
}

// Unwrap exposes both the sentinel and the cause to errors.Is / errors.As.
func (e *credentialRejectedError) Unwrap() []error { return []error{e.fallback, e.cause} }

// tryInMemoryToken tries the in-memory token source (tier 1).
// Returns (token, nil) on success, ("", errCacheMiss) when no source is set,
// or ("", err) for a real token-endpoint error.
func (t *OAuthTokenSource) tryInMemoryToken() (string, error) {
        if t.tokenSource == nil {
                return "", errCacheMiss
        }
        tok, err := t.tokenSource.Token()
        if err == nil && tok.Valid() {
                return tok.AccessToken, nil
        }
        t.tokenSource = nil
        if err != nil {
                return "", err
        }
        return "", errCacheMiss
}

// tryCachedToken restores a token source from the secrets provider (tier 2/3)
// and tries to obtain a valid token from it.
// Returns (token, nil) on success, ("", errCacheMiss) on a cache miss,
// or ("", err) for a real error.
func (t *OAuthTokenSource) tryCachedToken(ctx context.Context) (string, error) {
        if err := t.tryRestoreFromCache(ctx); err != nil {
                return "", err // errCacheMiss propagates transparently
        }
        tok, err := t.tokenSource.Token()
        if err == nil && tok.Valid() {
                t.cacheAccessToken(ctx, tok.AccessToken, tok.Expiry)
                return tok.AccessToken, nil
        }
        t.tokenSource = nil
        if err != nil {
                return "", err
        }
        return "", errCacheMiss
}

// tryRestoreFromCache attempts to build a token source from the cached refresh
// token stored in the secrets provider.
func (t *OAuthTokenSource) tryRestoreFromCache(ctx context.Context) error {
        if t.opts.SecretsProvider == nil {
                return fmt.Errorf("no secrets provider available")
        }
        key := t.opts.KeyProvider()
        refreshToken, err := t.opts.SecretsProvider.GetSecret(ctx, key)
        if err != nil {
                if secrets.IsNotFoundError(err) {
                        return errCacheMiss
                }
                return fmt.Errorf("reading cached refresh token: %w", err)
        }
        if refreshToken == "" {
                return errCacheMiss
        }

        oauth2Cfg, err := t.buildOAuth2Config(ctx)
        if err != nil {
                return fmt.Errorf("building oauth2 config for cache restore: %w", err)
        }
        httpClient, err := oauth.NewTokenHTTPClient(oauth2Cfg.Endpoint.TokenURL, true)
        if err != nil {
                return fmt.Errorf("building token endpoint HTTP client for cache restore: %w", err)
        }

        // Use a non-caching refresher as the innermost source.
        //
        // oauth2Cfg.TokenSource caches internally and only refreshes when the real
        // expiry passes. When the outer ReuseTokenSource enters the preemptive window
        // (refreshWindow before real expiry) it calls preemptiveTokenSource, which calls
        // the inner source; the inner source sees the real token as still valid and returns
        // it unchanged; preemptiveTokenSource shifts the expiry back by the window, producing
        // an already-expired token; the outer ReuseTokenSource then re-enters the chain
        // on every subsequent call — an infinite non-refreshing loop.
        //
        // A non-caching refresher always performs a network round-trip when called, so
        // the first call inside the preemptive window obtains a fresh token with a
        // real-future expiry. The outer ReuseTokenSource caches the shifted result and
        // serves it until the next window — exactly one refresh per window.
        //
        // Target stacking: ReuseTokenSource(nil, preemptive{PersistingTokenSource{nonCachingRefresher}})
        rawRefresher := oauth.NewNonCachingRefresher(oauth2Cfg, refreshToken, t.opts.OIDC.Audience, httpClient)

        // Persist rotated refresh tokens back to the secrets provider so future
        // invocations can still restore the session if the provider invalidates the
        // old token on refresh (common with OIDC providers that rotate refresh tokens).
        base := remote.NewPersistingTokenSource(rawRefresher, t.makeTokenPersister(key))

        // Wrap with preemptive refresh so tokens are renewed before real expiry.
        t.tokenSource = withPreemptiveRefresh(base, t.refreshWindow())
        return nil
}

// startFlow runs the interactive OAuth flow to completion. It is a package var
// so tests can stub the browser/callback round-trip and assert how skipBrowser
// is forwarded without standing up a real browser or callback listener.
var startFlow = func(ctx context.Context, flow *oauth.Flow, skipBrowser bool) (*oauth.TokenResult, error) {
        return flow.Start(ctx, skipBrowser)
}

// performBrowserFlow runs the interactive OIDC+PKCE browser flow and persists
// the resulting refresh token for future non-interactive use.
func (t *OAuthTokenSource) performBrowserFlow(ctx context.Context) error {
        flowCfg, err := t.buildFlowConfig(ctx)
        if err != nil {
                return err
        }

        flow, err := oauth.NewFlow(flowCfg)
        if err != nil {
                return fmt.Errorf("creating OAuth flow: %w", err)
        }

        tokenResult, err := startFlow(ctx, flow, t.opts.SkipBrowser)
        if err != nil {
                return fmt.Errorf("OAuth flow start failed: %w", err)
        }

        // Build a non-caching refresher as the innermost source (see tryRestoreFromCache
        // for a detailed explanation of why caching inner sources cause an infinite loop
        // inside the preemptive window). Reuse the already-discovered flowCfg to avoid
        // a second OIDC round-trip.
        oauth2Cfg := oauth2ConfigFrom(flowCfg)
        initialToken := &oauth2.Token{
                AccessToken:  tokenResult.AccessToken,
                RefreshToken: tokenResult.RefreshToken,
                Expiry:       tokenResult.Expiry,
                TokenType:    tokenResult.TokenType,
        }

        httpClient, err := oauth.NewTokenHTTPClient(flowCfg.TokenURL, flowCfg.TokenEndpointTrusted)
        if err != nil {
                return fmt.Errorf("building token endpoint HTTP client: %w", err)
        }
        var base oauth2.TokenSource = oauth.NewNonCachingRefresher(
                oauth2Cfg, initialToken.RefreshToken, flowCfg.Resource, httpClient,
        )

        if t.opts.SecretsProvider != nil {
                key := t.opts.KeyProvider()
                base = remote.NewPersistingTokenSource(base, t.makeTokenPersister(key))
                if tokenResult.RefreshToken != "" {
                        if err := t.opts.SecretsProvider.SetSecret(ctx, key, tokenResult.RefreshToken); err != nil {
                                slog.Warn("failed to persist initial refresh token", "error", err)
                        } else {
                                t.persistConfig(key, tokenResult.Expiry)
                        }
                } else {
                        slog.Debug("OIDC provider did not return a refresh token; token will not be persisted")
                }
        }

        // Pre-seed the outer ReuseTokenSource with the shifted initial token so the
        // just-obtained access token is served without an immediate network round-trip.
        t.tokenSource = withPreemptiveRefreshFrom(initialToken, base, t.refreshWindow())
        return nil
}

// buildFlowConfig creates an oauth.Config for the interactive browser flow.
// PKCE (S256) is always enabled per OAuth 2.1 requirements for public clients.
func (t *OAuthTokenSource) buildFlowConfig(ctx context.Context) (*oauth.Config, error) {
        config, err := oauth.CreateOAuthConfigFromOIDC(
                ctx,
                t.opts.OIDC.Issuer,
                t.opts.OIDC.ClientID,
                "", // public client — no client secret
                EnsureOfflineAccess(t.opts.OIDC.Scopes),
                true, // always use PKCE
                t.opts.OIDC.CallbackPort,
                t.opts.OIDC.Audience,
                false, // issuer is operator-configured
        )
        if err != nil {
                return nil, err
        }
        config.TokenEndpointTrusted = true
        config.RequireExactCallbackPort = t.opts.OIDC.RequireExactCallbackPort
        return config, nil
}

// buildOAuth2Config creates a minimal oauth2.Config suitable for token refresh
// via the cached refresh token path (no browser required).
func (t *OAuthTokenSource) buildOAuth2Config(ctx context.Context) (*oauth2.Config, error) {
        flowCfg, err := t.buildFlowConfig(ctx)
        if err != nil {
                return nil, err
        }
        return oauth2ConfigFrom(flowCfg), nil
}

// oauth2ConfigFrom converts an oauth.Config (from OIDC discovery) into the
// minimal oauth2.Config used for headless token refresh.
func oauth2ConfigFrom(flowCfg *oauth.Config) *oauth2.Config {
        return &oauth2.Config{
                ClientID: flowCfg.ClientID,
                Scopes:   flowCfg.Scopes,
                Endpoint: oauth2.Endpoint{
                        AuthURL:   flowCfg.AuthURL,
                        TokenURL:  flowCfg.TokenURL,
                        AuthStyle: oauth2.AuthStyleInParams,
                },
        }
}

// makeTokenPersister returns a remote.TokenPersister that stores the refresh
// token in the secrets provider and calls ConfigPersister. It also invalidates
// the access-token cache so the next Token() call does not serve a token minted
// against the now-rotated refresh token.
func (t *OAuthTokenSource) makeTokenPersister(key string) remote.TokenPersister {
        return func(refreshToken string, expiry time.Time) error {
                ctx := context.Background()
                if err := t.opts.SecretsProvider.SetSecret(ctx, key, refreshToken); err != nil {
                        return fmt.Errorf("persisting refresh token: %w", err)
                }
                t.persistConfig(key, expiry)
                // Invalidate the access-token cache so the next call does not serve a
                // token minted against the now-rotated refresh token. Order is intentional:
                // the refresh token (durable) is written first, then the AT cache is cleared.
                // Use key+"_AT" directly rather than accessTokenCacheKey() to guarantee we
                // clear the same key that was just persisted — ConfigPersister may have
                // updated CachedRefreshTokenRef, which would change KeyProvider's return
                // value and cause accessTokenCacheKey() to resolve a different key.
                // We write "" rather than calling DeleteSecret for provider compatibility.
                if err := t.opts.SecretsProvider.SetSecret(ctx, key+"_AT", ""); err != nil {
                        slog.Warn("failed to invalidate access-token cache after refresh token rotation",
                                "error", err)
                }
                return nil
        }
}

// persistConfig calls the ConfigPersister (if set) to record the token ref and expiry.
func (t *OAuthTokenSource) persistConfig(key string, expiry time.Time) {
        if t.opts.ConfigPersister != nil {
                t.opts.ConfigPersister(key, expiry)
        }
}

// accessTokenCacheKey returns the secrets-provider key for the cached access token.
func (t *OAuthTokenSource) accessTokenCacheKey() string {
        return t.opts.KeyProvider() + "_AT"
}

// tryAccessTokenCache reads a previously cached access token from the secrets
// provider and returns it if still valid.
func (t *OAuthTokenSource) tryAccessTokenCache(ctx context.Context) (string, bool) {
        if t.opts.SecretsProvider == nil {
                return "", false
        }
        raw, err := t.opts.SecretsProvider.GetSecret(ctx, t.accessTokenCacheKey())
        if err != nil || raw == "" {
                return "", false
        }
        idx := strings.LastIndex(raw, "|")
        if idx <= 0 {
                return "", false
        }
        expiry, err := time.Parse(time.RFC3339, raw[idx+1:])
        if err != nil {
                return "", false
        }
        if time.Now().Before(expiry) {
                return raw[:idx], true
        }
        return "", false
}

// cacheAccessToken writes the access token and its expiry to the secrets
// provider so concurrent short-lived processes can reuse it. Errors are logged
// at debug level and suppressed — a write failure degrades gracefully to a full
// refresh on the next call.
func (t *OAuthTokenSource) cacheAccessToken(ctx context.Context, token string, expiry time.Time) {
        if t.opts.SecretsProvider == nil || token == "" || expiry.IsZero() {
                return
        }
        val := token + "|" + expiry.UTC().Format(time.RFC3339)
        if err := t.opts.SecretsProvider.SetSecret(ctx, t.accessTokenCacheKey(), val); err != nil {
                slog.Debug("failed to cache access token", "error", err)
        }
}

// DeriveSecretKey computes a secrets-provider key for an OAuth refresh token.
// The formula is: <prefix><8 hex chars> where the hex is derived from
// sha256(resourceURL + "\x00" + issuer)[:4].
func DeriveSecretKey(prefix, resourceURL, issuer string) string {
        h := sha256.Sum256([]byte(resourceURL + "\x00" + issuer))
        return prefix + hex.EncodeToString(h[:4])
}

// LLMAccessTokenEnvVar returns the environment-variable name under which the
// environment secrets provider expects a cached LLM access token for the given
// gateway and issuer URLs. The format is:
//
//        TOOLHIVE_SECRET___thv_llm_<DeriveSecretKey("LLM_OAUTH_", gateway, issuer)>_AT
//
// This is the canonical source of truth for the key construction used by both
// the proxy/token commands and the e2e tests that inject fake tokens.
func LLMAccessTokenEnvVar(gatewayURL, issuerURL string) string {
        key := DeriveSecretKey("LLM_OAUTH_", gatewayURL, issuerURL)
        scopedKey := secrets.SystemKeyPrefix + string(secrets.ScopeLLM) + "_" + key + "_AT"
        return secrets.EnvVarPrefix + scopedKey
}

// EnsureOfflineAccess returns scopes with "offline_access" appended if absent.
// This scope is required for the provider to return a refresh token.
func EnsureOfflineAccess(scopes []string) []string {
        if slices.Contains(scopes, "offline_access") {
                return scopes
        }
        return append(scopes[:len(scopes):len(scopes)], "offline_access")
}

// refreshWindow returns the configured preemptive refresh window, falling back
// to the package default when the caller left Options.PreemptiveRefreshWindow
// unset. New rejects negative values, so this is always > 0.
func (t *OAuthTokenSource) refreshWindow() time.Duration {
        if t.opts.PreemptiveRefreshWindow > 0 {
                return t.opts.PreemptiveRefreshWindow
        }
        return preemptiveRefreshWindow
}

// preemptiveTokenSource wraps an oauth2.TokenSource and shifts each returned
// token's expiry back by window.
type preemptiveTokenSource struct {
        inner  oauth2.TokenSource
        window time.Duration
}

func (p *preemptiveTokenSource) Token() (*oauth2.Token, error) {
        tok, err := p.inner.Token()
        if err != nil {
                return nil, err
        }
        if !tok.Expiry.IsZero() {
                shifted := *tok
                shifted.Expiry = tok.Expiry.Add(-p.window)
                return &shifted, nil
        }
        return tok, nil
}

// withPreemptiveRefresh wraps src so tokens appear expired window before they
// actually expire, then re-wraps with ReuseTokenSource so the refresh is only
// triggered once per window.
func withPreemptiveRefresh(src oauth2.TokenSource, window time.Duration) oauth2.TokenSource {
        return withPreemptiveRefreshFrom(nil, src, window)
}

// withPreemptiveRefreshFrom is like withPreemptiveRefresh but pre-seeds the outer
// ReuseTokenSource with a shifted copy of initial (if non-nil and valid).
func withPreemptiveRefreshFrom(initial *oauth2.Token, src oauth2.TokenSource, window time.Duration) oauth2.TokenSource {
        var seeded *oauth2.Token
        if initial != nil && initial.Valid() && !initial.Expiry.IsZero() {
                shifted := *initial
                shifted.Expiry = initial.Expiry.Add(-window)
                if shifted.Valid() {
                        seeded = &shifted
                }
        }
        return oauth2.ReuseTokenSource(seeded, &preemptiveTokenSource{inner: src, window: window})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package upstreamswap provides middleware for exchanging embedded auth server
// access tokens with upstream IdP tokens.
package upstreamswap

import (
        "encoding/json"
        "fmt"
        "net/http"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// MiddlewareType is the type identifier for upstream swap middleware.
const MiddlewareType = "upstreamswap"

// Header injection strategy constants
const (
        // HeaderStrategyReplace replaces the Authorization header with the upstream token.
        HeaderStrategyReplace = "replace"
        // HeaderStrategyCustom adds the upstream token to a custom header.
        HeaderStrategyCustom = "custom"
)

// Config holds configuration for upstream swap middleware.
type Config struct {
        // HeaderStrategy determines how to inject the token: "replace" (default) or "custom".
        HeaderStrategy string `json:"header_strategy,omitempty" yaml:"header_strategy,omitempty"`

        // CustomHeaderName is the header name when HeaderStrategy is "custom".
        CustomHeaderName string `json:"custom_header_name,omitempty" yaml:"custom_header_name,omitempty"`

        // ProviderName identifies which upstream provider's tokens to retrieve for injection.
        // This is required and must match a configured upstream provider name.
        ProviderName string `json:"provider_name" yaml:"provider_name"`
}

// MiddlewareParams represents the JSON parameters for the middleware factory.
type MiddlewareParams struct {
        Config *Config `json:"config,omitempty"`
}

// Middleware wraps the upstream swap middleware functionality.
type Middleware struct {
        middleware types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *Middleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (*Middleware) Close() error {
        return nil
}

// CreateMiddleware is the factory function for upstream swap middleware.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal upstream swap middleware parameters: %w", err)
        }

        // Config is optional; use defaults if not provided
        cfg := params.Config
        if cfg == nil {
                cfg = &Config{}
        }

        // Validate configuration
        if err := validateConfig(cfg); err != nil {
                return fmt.Errorf("invalid upstream swap configuration: %w", err)
        }

        middleware := createMiddlewareFunc(cfg)

        upstreamSwapMw := &Middleware{
                middleware: middleware,
        }

        runner.AddMiddleware(config.Type, upstreamSwapMw)
        return nil
}

// validateConfig validates the upstream swap configuration.
func validateConfig(cfg *Config) error {
        // ProviderName is required to identify which upstream provider's tokens to retrieve
        if cfg.ProviderName == "" {
                return fmt.Errorf("provider_name is required")
        }

        // Validate header strategy
        if cfg.HeaderStrategy != "" &&
                cfg.HeaderStrategy != HeaderStrategyReplace &&
                cfg.HeaderStrategy != HeaderStrategyCustom {
                return fmt.Errorf("invalid header_strategy: %s (valid values: '%s', '%s')",
                        cfg.HeaderStrategy, HeaderStrategyReplace, HeaderStrategyCustom)
        }

        // Custom header name is required when using custom strategy
        if cfg.HeaderStrategy == HeaderStrategyCustom && cfg.CustomHeaderName == "" {
                return fmt.Errorf("custom_header_name must be specified when header_strategy is '%s'", HeaderStrategyCustom)
        }

        return nil
}

// writeUpstreamAuthRequired writes a 401 response with a WWW-Authenticate Bearer
// challenge per RFC 6750 Section 3.1, signalling that the caller must re-authenticate
// with the upstream IdP.
func writeUpstreamAuthRequired(w http.ResponseWriter) {
        w.Header().Set("WWW-Authenticate",
                `Bearer error="invalid_token", error_description="upstream token is no longer valid; re-authentication required"`)
        http.Error(w, "upstream authentication required", http.StatusUnauthorized)
}

// injectionFunc is a function that injects a token into an HTTP request.
type injectionFunc func(*http.Request, string)

// createReplaceInjector creates an injection function that replaces the Authorization header.
func createReplaceInjector() injectionFunc {
        return func(r *http.Request, token string) {
                r.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
        }
}

// createCustomInjector creates an injection function that adds the token to a custom header.
//
// It also strips the client's original Authorization header (the ToolHive-issued
// JWT) so it is never forwarded to the backend: that token was minted for the
// proxy, not the upstream, and leaking it is credential passthrough (#5504). The
// "replace" strategy avoids this implicitly because it overwrites Authorization;
// the custom strategy must strip it explicitly. When the custom header name is
// itself Authorization (case-insensitive), the Set above already replaced the JWT
// with the upstream token, so it must be left in place.
func createCustomInjector(headerName string) injectionFunc {
        stripAuthorization := http.CanonicalHeaderKey(headerName) != "Authorization"
        return func(r *http.Request, token string) {
                r.Header.Set(headerName, fmt.Sprintf("Bearer %s", token))
                if stripAuthorization {
                        r.Header.Del("Authorization")
                }
        }
}

// createMiddlewareFunc creates the actual middleware function.
// It reads upstream tokens from Identity.UpstreamTokens, which are populated
// during JWT validation by the auth middleware (Step 3).
func createMiddlewareFunc(cfg *Config) types.MiddlewareFunction {
        // Determine injection strategy at startup time
        strategy := cfg.HeaderStrategy
        if strategy == "" {
                strategy = HeaderStrategyReplace
        }

        var injectToken injectionFunc
        switch strategy {
        case HeaderStrategyReplace:
                injectToken = createReplaceInjector()
        case HeaderStrategyCustom:
                injectToken = createCustomInjector(cfg.CustomHeaderName)
        default:
                // This shouldn't happen due to validation, but default to replace
                injectToken = createReplaceInjector()
        }

        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        identity, ok := auth.IdentityFromContext(r.Context())
                        if !ok {
                                next.ServeHTTP(w, r)
                                return
                        }

                        token, exists := identity.UpstreamTokens[cfg.ProviderName]
                        if !exists || token == "" {
                                writeUpstreamAuthRequired(w)
                                return
                        }

                        injectToken(r, token)
                        next.ServeHTTP(w, r)
                })
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package upstreamtoken

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "sync"
        "time"

        "github.com/golang-jwt/jwt/v5"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

const (
        // idTokenExpirySkew is how far ahead of the ID token's `exp` the
        // ID-token trigger fires, so a token is not handed out moments before it
        // expires.
        idTokenExpirySkew = 30 * time.Second

        // defaultIDTokenSuppression bounds how long the ID-token trigger stays
        // suppressed for a row whose access token asserts no expiry and whose
        // session expiry is unknown, so the suppression map cannot grow without
        // bound.
        defaultIDTokenSuppression = time.Hour

        // idTokenRefreshFailureBackoff is how long the ID-token trigger stays
        // suppressed after a failed ID-token-triggered refresh. It is short so a
        // transient upstream error does not reopen the stale-ID-token window for
        // the rest of the access token's lifetime.
        idTokenRefreshFailureBackoff = time.Minute

        // idTokenSuppressionSweepInterval rate-limits the full sweep of expired
        // suppression entries, which runs under mu. Entries are also dropped
        // lazily when read after expiry.
        idTokenSuppressionSweepInterval = time.Minute
)

// InProcessService implements the Service interface for in-process use.
// It composes storage (read) and refresher (refresh + persist) to provide a
// single GetValidTokens call. Concurrent-refresh deduplication is delegated to
// the refresher's own singleflight.Group keyed by an opaque storage-row identity,
// so the same Group covers both the handler's chain-walk path and the runtime
// token-swap path within this process.
//
// Refresh is triggered by an expired access token or by an expired ID token
// (its `exp` claim, read without re-verifying the signature it was validated
// with when stored). The ID-token trigger serves providers whose ID-token
// lifetime is shorter than their access-token lifetime (e.g. Microsoft Entra
// ID), which would otherwise hand consumers an expired ID token for part of
// every access-token lifetime. It never fires when the row has no refresh
// token, or for an empty, unparsable or `exp`-less ID token.
//
// It cannot help when the provider omits id_token on refresh (OIDC Core 1.0
// §12.2): the expired ID token is then carried forward, and the trigger is
// suppressed for that row until the access token next expires so reads do not
// loop on refresh. After a failed ID-token-triggered refresh the trigger backs
// off briefly (idTokenRefreshFailureBackoff). Consumers that need a currently
// valid ID token must therefore still check `exp` themselves.
//
// Suppression state is held in process memory: replicas sharing storage each
// keep their own, so each may perform one extra refresh per access-token
// lifetime for a provider that omits id_token. Likewise, refresh
// deduplication is per process; replicas that refresh the same row at the same
// moment can each redeem its refresh token (see upstreamTokenRefresher), which
// a provider enforcing single-use refresh-token rotation may treat as reuse.
// The ID-token trigger adds refresh events, but no new race shape.
type InProcessService struct {
        storage   storage.UpstreamTokenStorage
        refresher storage.UpstreamTokenRefresher

        // mu guards idTokenSuppressions and lastSuppressionSweep.
        mu sync.Mutex
        // idTokenSuppressions records rows for which an ID-token-triggered refresh
        // should not be attempted again until the access token next expires,
        // because a previous refresh left the ID token expired (the provider
        // omitted id_token, or the refresh failed). Keyed by suppressionKey.
        idTokenSuppressions  map[string]idTokenSuppression
        lastSuppressionSweep time.Time
}

// idTokenSuppression pins a suppression to the exact stale ID token it was
// recorded for, so a rotated ID token is never suppressed.
type idTokenSuppression struct {
        idToken string
        until   time.Time
}

// Compile-time checks.
var (
        _ Service     = (*InProcessService)(nil)
        _ TokenReader = (*InProcessService)(nil)
)

// NewInProcessService creates a new InProcessService.
// The refresher may be nil if upstream token refresh is not configured;
// expired tokens will return ErrNoRefreshToken in that case.
func NewInProcessService(
        stor storage.UpstreamTokenStorage,
        refresher storage.UpstreamTokenRefresher,
) *InProcessService {
        return &InProcessService{
                storage:             stor,
                refresher:           refresher,
                idTokenSuppressions: make(map[string]idTokenSuppression),
        }
}

// GetValidTokens returns a valid upstream credential for a session and provider.
// It transparently refreshes expired access tokens using the refresh token,
// and refreshes expired ID tokens where possible (see InProcessService).
func (s *InProcessService) GetValidTokens(ctx context.Context, sessionID, providerName string) (*UpstreamCredential, error) {
        tokens, err := s.storage.GetUpstreamTokens(ctx, sessionID, providerName)
        if err != nil {
                // ErrExpired returns tokens (including refresh token) alongside the error.
                // Attempt a refresh before giving up.
                if errors.Is(err, storage.ErrExpired) {
                        if tokens != nil {
                                return s.refreshOrFail(ctx, sessionID, providerName, tokens)
                        }
                        // Expired but storage returned nil tokens — can't refresh.
                        return nil, ErrNoRefreshToken
                }
                if errors.Is(err, storage.ErrNotFound) {
                        return nil, ErrSessionNotFound
                }
                if errors.Is(err, storage.ErrInvalidBinding) {
                        return nil, ErrInvalidBinding
                }
                return nil, fmt.Errorf("failed to get upstream tokens: %w", err)
        }

        // Defense in depth: some storage implementations may return tokens
        // without checking expiry (the interface does not require it).
        if !tokens.ExpiresAt.IsZero() && tokens.IsExpired(time.Now()) {
                return s.refreshOrFail(ctx, sessionID, providerName, tokens)
        }

        if s.shouldRefreshForIDToken(sessionID, providerName, tokens) {
                return s.refreshForIDToken(ctx, sessionID, providerName, tokens)
        }

        return &UpstreamCredential{AccessToken: tokens.AccessToken, IDToken: tokens.IDToken}, nil
}

// GetAllUpstreamCredentials returns access tokens and ID tokens for all upstream
// providers in a session in a single storage round-trip. Expired access tokens
// are refreshed transparently; providers whose access token cannot be refreshed
// are included in the returned failed map, keyed by provider name with the
// refresh error as the value, so downstream middleware can return a clean 401
// for the access-token use case. The IDToken on a returned entry is
// the rotated ID token when a refresh produced one (OIDC Core 1.0 §12.2),
// otherwise the original JWT captured at the initial OIDC login; it is not
// independently validated for freshness and may be empty if the upstream login
// never yielded one. An expired ID token also triggers a refresh (see
// InProcessService); a failed ID-token-triggered refresh does not put the
// provider in the failed map, since its access token is still valid. Callers
// that PRESENT it as a credential must expect expiry to be rejected; callers
// that only READ ITS CLAIMS are unaffected by expiry. See
// Identity.UpstreamIDTokens for both consumers.
//
// Returns an empty map and nil failed map (not error) for unknown sessions.
func (s *InProcessService) GetAllUpstreamCredentials(
        ctx context.Context, sessionID string,
) (map[string]UpstreamCredential, map[string]error, error) {
        allTokens, err := s.storage.GetAllUpstreamTokens(ctx, sessionID)
        if err != nil {
                return nil, nil, fmt.Errorf("bulk read upstream tokens: %w", err)
        }

        if len(allTokens) == 0 {
                return map[string]UpstreamCredential{}, nil, nil
        }

        result := make(map[string]UpstreamCredential, len(allTokens))
        var failed map[string]error
        fail := func(providerName string, err error) {
                if failed == nil {
                        failed = make(map[string]error)
                }
                failed[providerName] = err
        }
        // TODO(auth): Refresh providers in parallel using errgroup to avoid
        // worst-case latency of N refreshes when multiple providers need refresh.
        for providerName, tokens := range allTokens {
                if tokens == nil {
                        continue
                }

                // If token is not expired, use it directly.
                if tokens.ExpiresAt.IsZero() || !tokens.IsExpired(time.Now()) {
                        if s.shouldRefreshForIDToken(sessionID, providerName, tokens) {
                                cred, idErr := s.refreshForIDToken(ctx, sessionID, providerName, tokens)
                                if idErr != nil {
                                        slog.WarnContext(ctx, "upstream token row removed by failed refresh; provider will require re-authentication",
                                                "session_id", sessionID,
                                                "provider", providerName,
                                                "error", idErr,
                                        )
                                        fail(providerName, idErr)
                                        continue
                                }
                                result[providerName] = *cred
                                continue
                        }
                        result[providerName] = UpstreamCredential{
                                AccessToken: tokens.AccessToken,
                                IDToken:     tokens.IDToken,
                        }
                        continue
                }

                // Token is expired — attempt refresh.
                refreshed, refreshErr := s.refreshOrFail(ctx, sessionID, providerName, tokens)
                if refreshErr != nil {
                        slog.WarnContext(ctx, "upstream token refresh failed; provider will require re-authentication",
                                "session_id", sessionID,
                                "provider", providerName,
                                "error", refreshErr,
                        )
                        fail(providerName, refreshErr)
                        continue
                }
                // refreshOrFail carries through the rotated ID token when the provider
                // issued one, otherwise the original login ID token (see its doc).
                result[providerName] = *refreshed
        }

        // The returned ID tokens are deliberately not checked for expiry here. Neither
        // of today's two consumers wants that pre-check: one presents the token as an
        // RFC 8693 subject_token and relies on the IdP to reject an expired assertion
        // (pkg/vmcp/auth/strategies surfaces the resulting invalid_grant), and the other
        // only reads its claims, for which expiry is irrelevant. Enforcing it here would
        // break the second while duplicating what the IdP already does for the first.
        // A future consumer that presents the token to a sink which does NOT validate
        // `exp` would change that calculus — such a caller must check it itself.
        // An expired ID token is refreshed when the provider can rotate it, but this
        // still does not filter what it returns: an ID token the
        // provider declined to rotate is passed through expired.
        // See Identity.UpstreamIDTokens for both.
        return result, failed, nil
}

// refreshOrFail attempts a refresh via the shared refresher and maps errors to
// the service's sentinel errors. Deduplication of concurrent refreshes for the
// same logical upstream-token row is handled inside the refresher using an
// opaque storage-row identity.
func (s *InProcessService) refreshOrFail(
        ctx context.Context,
        sessionID string,
        providerName string,
        expired *storage.UpstreamTokens,
) (*UpstreamCredential, error) {
        if expired.RefreshToken == "" {
                return nil, ErrNoRefreshToken
        }

        if s.refresher == nil {
                slog.Debug("token refresher not configured, cannot refresh upstream tokens",
                        "session_id", sessionID,
                        "provider", providerName,
                )
                return nil, ErrNoRefreshToken
        }

        refreshed, err := s.refresher.RefreshAndStore(ctx, sessionID, expired)
        if err != nil {
                slog.Warn("upstream token refresh failed",
                        "session_id", sessionID,
                        "provider", providerName,
                        "error", err,
                )
                return nil, fmt.Errorf("%w: %w", ErrRefreshFailed, err)
        }

        if refreshed == nil {
                return nil, ErrRefreshFailed
        }

        // Prefer the ID token from the refresh response when the provider rotated
        // it (OIDC Core 1.0 §12.2 permits — but does not require — a new id_token on
        // refresh). Fall back to the original login ID token when the refresh response
        // omitted one. The primary carry-forward into storage is in
        // upstreamTokenRefresher.refreshAndStore; this fallback is defense-in-depth
        // so the caller never sees an empty subject token.
        idToken := refreshed.IDToken
        if idToken == "" {
                idToken = expired.IDToken
        }

        // A refresh that still leaves the ID token expired (the provider omitted id_token) must not be retried on every read.
        // Suppress the trigger for this exact ID token until the new access token
        // expires, when the access-token trigger refreshes anyway.
        if idTokenExpired(idToken, time.Now()) {
                s.suppressIDTokenRefresh(sessionID, providerName, idToken, refreshed.ExpiresAt, refreshed.SessionExpiresAt)
        }

        return &UpstreamCredential{AccessToken: refreshed.AccessToken, IDToken: idToken}, nil
}

// shouldRefreshForIDToken reports whether the ID-token trigger applies to a
// row whose access token is still valid: a refresh token exists, the stored ID token is expired, and the trigger is not
// suppressed for that ID token.
func (s *InProcessService) shouldRefreshForIDToken(
        sessionID, providerName string, tokens *storage.UpstreamTokens,
) bool {
        if tokens.RefreshToken == "" {
                return false
        }
        now := time.Now()
        if !idTokenExpired(tokens.IDToken, now) {
                return false
        }
        return !s.idTokenRefreshSuppressed(sessionID, providerName, tokens.IDToken, now)
}

// refreshForIDToken refreshes a row whose access token is still valid but whose
// ID token has expired. A failed refresh is deliberately not fatal: the stored
// credential is returned unchanged, since its access token is still usable and
// a transient upstream error should not force re-authentication. The trigger
// then backs off for idTokenRefreshFailureBackoff (or until the access token
// expires, if sooner) so a failing upstream is not hammered on every read; no
// backoff is recorded when the failure is the caller's own context ending.
//
// The one fatal case is a row the failed refresh removed: the refresher deletes
// it when the provider rotated the refresh token but the rotation could not be
// persisted, leaving the old one dead. Serving the stored copy would then hand
// out a credential that no longer exists, so the row is re-read and
// ErrRefreshFailed returned if it is gone.
func (s *InProcessService) refreshForIDToken(
        ctx context.Context, sessionID, providerName string, tokens *storage.UpstreamTokens,
) (*UpstreamCredential, error) {
        refreshed, err := s.refreshOrFail(ctx, sessionID, providerName, tokens)
        if err == nil {
                return refreshed, nil
        }
        // refreshOrFail has already logged the failure at WARN.
        if ctx.Err() != nil {
                return &UpstreamCredential{AccessToken: tokens.AccessToken, IDToken: tokens.IDToken}, nil
        }
        if _, readErr := s.storage.GetUpstreamTokens(ctx, sessionID, providerName); errors.Is(readErr, storage.ErrNotFound) {
                return nil, err
        }
        slog.DebugContext(ctx, "returning stored upstream tokens after failed refresh for expired ID token",
                "session_id", sessionID,
                "provider", providerName,
        )
        until := time.Now().Add(idTokenRefreshFailureBackoff)
        if !tokens.ExpiresAt.IsZero() && tokens.ExpiresAt.Before(until) {
                until = tokens.ExpiresAt
        }
        s.suppressIDTokenRefresh(sessionID, providerName, tokens.IDToken, until, time.Time{})
        return &UpstreamCredential{AccessToken: tokens.AccessToken, IDToken: tokens.IDToken}, nil
}

// suppressIDTokenRefresh records that the ID-token trigger must not fire for
// idToken on this row until the access token expires. When the access token
// asserts no expiry, the session expiry (or defaultIDTokenSuppression) bounds
// the entry instead. Expired entries are swept at most once per
// idTokenSuppressionSweepInterval, which keeps the map bounded by the rows
// suppressed within roughly the last interval plus those still live.
func (s *InProcessService) suppressIDTokenRefresh(
        sessionID, providerName, idToken string, accessExpiresAt, sessionExpiresAt time.Time,
) {
        now := time.Now()
        until := accessExpiresAt
        if until.IsZero() {
                until = sessionExpiresAt
        }
        if until.IsZero() {
                until = now.Add(defaultIDTokenSuppression)
        }

        s.mu.Lock()
        defer s.mu.Unlock()
        if now.Sub(s.lastSuppressionSweep) >= idTokenSuppressionSweepInterval {
                for k, v := range s.idTokenSuppressions {
                        if !now.Before(v.until) {
                                delete(s.idTokenSuppressions, k)
                        }
                }
                s.lastSuppressionSweep = now
        }
        s.idTokenSuppressions[suppressionKey(sessionID, providerName)] = idTokenSuppression{idToken: idToken, until: until}
}

// idTokenRefreshSuppressed reports whether the ID-token trigger is suppressed
// for idToken on this row at now. An entry found expired is dropped.
func (s *InProcessService) idTokenRefreshSuppressed(sessionID, providerName, idToken string, now time.Time) bool {
        key := suppressionKey(sessionID, providerName)
        s.mu.Lock()
        defer s.mu.Unlock()
        sup, ok := s.idTokenSuppressions[key]
        if !ok {
                return false
        }
        if !now.Before(sup.until) {
                delete(s.idTokenSuppressions, key)
                return false
        }
        return sup.idToken == idToken
}

func suppressionKey(sessionID, providerName string) string {
        return sessionID + "\x00" + providerName
}

// idTokenExpired reports whether idToken carries an `exp` claim that has
// passed, allowing idTokenExpirySkew. The token's signature is not verified:
// it was validated when stored, and only its expiry is read here. An empty or
// unparsable token, or one without `exp`, is not treated as expired, so the
// trigger never fires for a row it cannot reason about.
func idTokenExpired(idToken string, now time.Time) bool {
        if idToken == "" {
                return false
        }
        parsed, _, err := jwt.NewParser().ParseUnverified(idToken, jwt.MapClaims{})
        if err != nil {
                return false
        }
        exp, err := parsed.Claims.GetExpirationTime()
        if err != nil || exp == nil {
                return false
        }
        return !now.Add(idTokenExpirySkew).Before(exp.Time)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

import (
        "context"
        "errors"
        "log/slog"
        "net/http"
        "os/user"
        "strings"
)

// bearerTokenType defines the expected token type for Bearer authentication.
const bearerTokenType = "Bearer"

// Common Bearer token extraction errors
var (
        ErrAuthHeaderMissing       = errors.New("authorization header required")
        ErrInvalidAuthHeaderFormat = errors.New("invalid authorization header format, expected 'Bearer <token>'")
        ErrEmptyBearerToken        = errors.New("empty token in authorization header")
)

// ExtractBearerToken extracts and validates a Bearer token from the Authorization header.
// It performs the following validations:
//  1. Verifies the Authorization header is present
//  2. Checks for the "Bearer " prefix (case-sensitive per RFC 6750)
//  3. Ensures the token is not empty after removing the prefix
//
// The function returns the token string (without "Bearer " prefix) and any validation error.
// Callers are responsible for further token validation (JWT parsing, introspection, etc.)
// and for converting errors to appropriate HTTP responses.
//
// This function implements RFC 6750 Section 2.1 (Bearer Token Authorization Header).
// See: https://datatracker.ietf.org/doc/html/rfc6750#section-2.1
func ExtractBearerToken(r *http.Request) (string, error) {
        // Get the Authorization header
        authHeader := r.Header.Get("Authorization")
        if authHeader == "" {
                return "", ErrAuthHeaderMissing
        }

        // Check for the Bearer prefix (case-sensitive per RFC 6750)
        bearerPrefix := bearerTokenType + " "
        if !strings.HasPrefix(authHeader, bearerPrefix) {
                return "", ErrInvalidAuthHeaderFormat
        }

        // Extract the token
        tokenString := strings.TrimPrefix(authHeader, bearerPrefix)

        // Check for empty token (handles "Bearer " with no token or only whitespace)
        if strings.TrimSpace(tokenString) == "" {
                return "", ErrEmptyBearerToken
        }

        return tokenString, nil
}

// GetAuthenticationMiddleware returns the appropriate authentication middleware based on the configuration.
// If OIDC config is provided, it returns JWT validation middleware. Otherwise it returns local-user
// middleware and logs a WARNING: with no OIDC config, the proxy accepts every request and forwards it
// under a synthetic local-user identity with no token or credential check. This unauthenticated fallback
// is intended for local/development use; configure an OIDC provider to enforce authentication per request.
func GetAuthenticationMiddleware(ctx context.Context, oidcConfig *TokenValidatorConfig, opts ...TokenValidatorOption,
) (func(http.Handler) http.Handler, http.Handler, error) {
        if oidcConfig != nil {
                slog.Debug("oidc validation enabled")

                // Create JWT validator
                jwtValidator, err := NewTokenValidator(ctx, *oidcConfig, opts...)
                if err != nil {
                        return nil, nil, err
                }

                authInfoHandler := NewAuthInfoHandler(oidcConfig.Issuer, oidcConfig.ResourceURL, oidcConfig.Scopes)
                return jwtValidator.Middleware, authInfoHandler, nil
        }

        slog.Warn("no authentication configured: the proxy will accept every request unauthenticated and " +
                "assign each a synthetic local-user identity with no token or credential check; " +
                "configure an OIDC provider to enforce authentication per request")

        // Get current OS user
        currentUser, err := user.Current()
        if err != nil {
                slog.Warn("failed to get current user, using 'local' as default", "error", err)
                return LocalUserMiddleware("local"), nil, nil
        }

        slog.Debug("using local user authentication", "user", currentUser.Username)
        return LocalUserMiddleware(currentUser.Username), nil, nil
}

// EscapeQuotes escapes quotes in a string for use in a quoted-string context.
func EscapeQuotes(s string) string {
        // Simple escape of backslashes and quotes is sufficient for quoted-string.
        s = strings.ReplaceAll(s, `\`, `\\`)
        return strings.ReplaceAll(s, `"`, `\"`)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package auth provides authentication and authorization utilities.
package auth

import (
        "net/http"
        "strings"

        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// NewWellKnownHandler creates an HTTP handler that routes requests under the
// /.well-known/ path space to the appropriate handler.
//
// Per RFC 9728, the /.well-known/oauth-protected-resource endpoint and any subpaths
// under it must be accessible without authentication. This handler ensures proper
// routing of discovery requests while returning 404 for unknown paths.
//
// If authInfoHandler is nil, the returned handler responds with HTTP 404 and a
// JSON body for all /.well-known/ paths. This ensures OAuth discovery clients
// (e.g., Claude Code) receive a clean, parseable "not found" instead of falling
// through to the MCP handler, which would reject the GET with an HTTP 406
// JSON-RPC error that breaks OAuth error parsing.
//
// Usage:
//
//        authInfoHandler := auth.NewAuthInfoHandler(issuer, resourceURL, scopes)
//        wellKnownHandler := auth.NewWellKnownHandler(authInfoHandler)
//        mux.Handle("/.well-known/", wellKnownHandler)
//
// This handler matches:
//   - /.well-known/oauth-protected-resource (exact)
//   - /.well-known/oauth-protected-resource/* (subpaths)
//
// Returns 404 for other /.well-known/* paths or when auth is not configured.
func NewWellKnownHandler(authInfoHandler http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                // When auth is configured, route discovery requests to the auth handler.
                if authInfoHandler != nil &&
                        strings.HasPrefix(r.URL.Path, oauthproto.WellKnownOAuthResourcePath) {
                        authInfoHandler.ServeHTTP(w, r)
                        return
                }

                // No auth configured, or unknown .well-known path — return JSON 404
                // so OAuth discovery clients can parse the response cleanly.
                w.Header().Set("Content-Type", "application/json")
                w.WriteHeader(http.StatusNotFound)
                _, _ = w.Write([]byte(`{"error":"not_found"}`))
        })
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package authorizationdetails validates the RFC 9396 authorization_details
// parameter of a token request and asks per-type hooks what to grant.
//
// The package only parses and delegates. It does not persist the granted
// details, does not put them in tokens, and does not convey them to resource
// servers. Register only types whose enforcement the deployment performs
// itself.
//
// Usage has two phases. Registry.Parse runs early and cheaply, before any
// expensive checks, and rejects malformed input. Pending.Grant runs late, after
// the request is otherwise accepted, and calls the hooks.
package authorizationdetails

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "slices"
        "strings"
        "time"
        "unicode/utf8"
)

const (
        maxSize  = 64 * 1024
        maxDepth = 64
)

// Type registers one authorization-details type.
type Type struct {
        // GrantTypes lists the grant types this type is accepted on. It is required
        // and must be non-empty. An entry of this type on any other grant type is
        // rejected as invalid.
        GrantTypes []string

        // Grant decides what to grant for all of this type's entries in a request.
        // It is called at most once per request, after the other checks passed,
        // with every entry of the type in request order.
        //
        // Grant is decide-only: it may read state but must not commit it, because
        // token issuance can still fail after it returns. Returning Error rejects
        // the request with that code and reason; any other error is an internal
        // failure.
        Grant func(context.Context, Request) (Result, error)
}

// Request is the context a Grant hook receives. The caller supplies every
// field except Entries and GrantType, which the package fills.
//
// All hooks of one request share the same Request values: the slices, Subject
// and Subject.Claims are not copied per hook, so a hook must not modify them.
// Only Entries is specific to the hook's type.
type Request struct {
        // Entries holds this type's entries in request order.
        Entries []json.RawMessage
        // GrantType is the grant type of the token request.
        GrantType string
        // ClientID is the authenticated client.
        ClientID string
        // Actor is the resolved actor; it equals ClientID without an actor_token.
        Actor string
        // ActorTokenPresent reports whether the request carried an actor_token.
        ActorTokenPresent bool
        // RequestedTokenType is the requested_token_type of the request.
        RequestedTokenType string
        // RequestedScopes are the scopes the client asked for.
        RequestedScopes []string
        // GrantedAudience is the audience the server has already decided to grant.
        GrantedAudience []string
        // Subject describes the validated subject token.
        Subject Subject
}

// Subject describes a validated subject token.
type Subject struct {
        Issuer   string
        Subject  string
        ClientID string
        JWTID    string
        Audience []string
        IssuedAt time.Time
        Expiry   time.Time
        External bool // validated through a trusted issuer
        // Claims holds all verified claims, passed to hooks as given.
        Claims map[string]any
}

// Result is what a Grant hook decided.
type Result struct {
        // Entries are the granted entries. The list must be non-empty and every
        // entry must be a JSON object whose "type" is the hook's own type.
        Entries []json.RawMessage
}

// Registry holds the registered types. It is immutable and safe for concurrent
// use.
type Registry struct {
        types map[string]Type
}

// New copies and validates the registered types, keyed by exact,
// case-sensitive type name. It rejects empty or whitespace names, a nil Grant,
// and an empty or blank GrantTypes list. A registry with no types is valid; it
// accepts only the empty array.
func New(types map[string]Type) (*Registry, error) {
        copied := make(map[string]Type, len(types))
        for name, t := range types {
                if strings.TrimSpace(name) == "" {
                        return nil, fmt.Errorf("authorization details type name must not be empty or whitespace")
                }
                if t.Grant == nil {
                        return nil, fmt.Errorf("authorization details type %q must have a Grant function", name)
                }
                if len(t.GrantTypes) == 0 {
                        return nil, fmt.Errorf("authorization details type %q must list at least one grant type", name)
                }
                for _, grantType := range t.GrantTypes {
                        if strings.TrimSpace(grantType) == "" {
                                return nil, fmt.Errorf("authorization details type %q lists an empty grant type", name)
                        }
                }
                t.GrantTypes = slices.Clone(t.GrantTypes)
                copied[name] = t
        }
        return &Registry{types: copied}, nil
}

// Parse checks the raw authorization_details value for the given grant type.
// The value must be at most 64 KiB of valid UTF-8 holding one JSON array of
// objects, nested at most 64 containers deep, in which no object repeats a
// member name (also not differing only in letter case; that rule is ToolHive
// policy, not RFC 9396). Every entry needs a non-empty string "type" that is
// registered for grantType. The empty array "[]" is valid.
//
// Any failure returns ErrInvalid.
func (r *Registry) Parse(raw, grantType string) (*Pending, error) {
        data := []byte(raw)
        if len(data) > maxSize || !utf8.Valid(data) || !validArray(data) {
                return nil, ErrInvalid
        }
        // Unmarshalling copies the bytes, so each entry owns its backing array.
        var entries []json.RawMessage
        if err := json.Unmarshal(data, &entries); err != nil {
                return nil, ErrInvalid
        }

        pending := &Pending{
                grantType: grantType,
                types:     r.types,
                entries:   make(map[string][]json.RawMessage),
        }
        for _, entry := range entries {
                name, ok := entryType(entry)
                if !ok {
                        return nil, ErrInvalid
                }
                t, registered := r.types[name]
                if !registered || !slices.Contains(t.GrantTypes, grantType) {
                        return nil, ErrInvalid
                }
                if _, seen := pending.entries[name]; !seen {
                        pending.order = append(pending.order, name)
                }
                pending.entries[name] = append(pending.entries[name], entry)
        }
        return pending, nil
}

// Pending is a parsed, not yet granted authorization_details value. It is
// single-use: call Grant at most once.
type Pending struct {
        grantType string
        types     map[string]Type
        order     []string
        entries   map[string][]json.RawMessage
}

// Grant calls each type's hook once, in the order the types first appear in
// the request, and returns the granted entries as a JSON array: types in that
// order, each type's entries in the order its hook returned them ("[]" when
// the request held no entries). The package fills req.Entries and
// req.GrantType.
//
// A hook must return at least one entry, and every entry is checked as
// strictly as the request: valid UTF-8, within the size limit, no duplicate
// member names, and an object of the hook's own type. The combined array must
// also fit the size limit. The output is built with encoding/json, so entries
// are compacted and HTML-escaped (<, > and & become \u003c, \u003e and
// \u0026): equal to the hook's output as JSON, not byte-identical.
//
// Errors: Error is a hook rejection to pass to the client. A plain error means
// the package rejected hook output, and any other error is an untyped hook
// failure; both are internal failures whose text must not reach the client.
func (p *Pending) Grant(ctx context.Context, req Request) (json.RawMessage, error) {
        req.GrantType = p.grantType

        granted := []json.RawMessage{}
        for _, name := range p.order {
                typeReq := req
                typeReq.Entries = p.entries[name]
                result, err := p.types[name].Grant(ctx, typeReq)
                if err != nil {
                        return nil, hookError(name, err)
                }
                if len(result.Entries) == 0 {
                        return nil, fmt.Errorf("type %q returned no entries: %w", name, errInvalidHookResult)
                }
                for _, entry := range result.Entries {
                        if err := checkGranted(name, entry); err != nil {
                                return nil, err
                        }
                        granted = append(granted, entry)
                }
        }

        out, err := json.Marshal(granted)
        if err != nil || len(out) > maxSize {
                return nil, fmt.Errorf("combined hook output is not valid JSON or exceeds %d bytes: %w", maxSize, errInvalidHookResult)
        }
        return out, nil
}

func hookError(typeName string, err error) error {
        var typed Error
        if asError(err, &typed) {
                if !typed.valid() {
                        return fmt.Errorf("type %q returned an invalid error: %w", typeName, errInvalidHookResult)
                }
                return typed
        }
        return fmt.Errorf("type %q: %w", typeName, err)
}

// checkGranted applies the request-side checks to one entry a hook returned.
func checkGranted(typeName string, entry json.RawMessage) error {
        if len(entry) > maxSize || !utf8.Valid(entry) {
                return fmt.Errorf("type %q returned an entry that is oversized or not valid UTF-8: %w", typeName, errInvalidHookResult)
        }
        if got, ok := entryType(entry); !ok || got != typeName {
                return fmt.Errorf("type %q returned an entry that is not an object of that type: %w", typeName, errInvalidHookResult)
        }
        wrapped := make([]byte, 0, len(entry)+2)
        wrapped = append(append(append(wrapped, '['), entry...), ']')
        if !validArray(wrapped) {
                return fmt.Errorf("type %q returned an entry with duplicate member names or too deep nesting: %w",
                        typeName, errInvalidHookResult)
        }
        return nil
}

// asError extracts an Error returned by value or by pointer, possibly wrapped.
func asError(err error, out *Error) bool {
        var value Error
        if errors.As(err, &value) {
                *out = value
                return true
        }
        var pointer *Error
        if errors.As(err, &pointer) && pointer != nil {
                *out = *pointer
                return true
        }
        return false
}

// entryType returns the string "type" member of a JSON object.
func entryType(entry json.RawMessage) (string, bool) {
        var object map[string]json.RawMessage
        if err := json.Unmarshal(entry, &object); err != nil || object == nil {
                return "", false
        }
        var name string
        if err := json.Unmarshal(object["type"], &name); err != nil || name == "" {
                return "", false
        }
        return name, true
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authorizationdetails

import (
        "errors"
        "regexp"
)

var (
        // ErrInvalid reports that the authorization_details parameter is malformed:
        // bad JSON, over a size or depth limit, duplicate member names, an entry
        // that is not an object, or a missing, unknown or wrong-grant "type".
        // Callers should answer with a generic invalid_authorization_details error
        // and must not echo request data.
        ErrInvalid = errors.New("invalid authorization_details")

        // errInvalidHookResult reports that a type's Grant hook returned output the
        // package rejects: an empty result, an entry that is not an object of that
        // type, a limit violation, or a typed Error with an unknown Code or an
        // invalid Reason. It is a defect in the registered hook, not in the
        // request, so callers should treat it like any other hook failure.
        errInvalidHookResult = errors.New("invalid authorization details hook result")
)

// ErrorCode is the OAuth error code a hook may return to the client.
type ErrorCode string

const (
        // CodeInvalidAuthorizationDetails is the RFC 9396 invalid_authorization_details error.
        CodeInvalidAuthorizationDetails ErrorCode = "invalid_authorization_details"
        // CodeInvalidRequest is the RFC 6749 invalid_request error. It is for
        // problems outside the authorization details themselves, such as a missing
        // or unacceptable related parameter. Problems in the details' content
        // (unknown fields, wrong types, invalid values, missing required fields)
        // must use CodeInvalidAuthorizationDetails, per RFC 9396 section 5.
        CodeInvalidRequest ErrorCode = "invalid_request"
)

var reasonPattern = regexp.MustCompile(`^[a-z0-9_]{1,64}$`)

// Error is a client-visible rejection a Grant hook may return. Reason is shown
// to the client, so it must match [a-z0-9_]{1,64} and must not carry request
// data. Any other error a hook returns is treated as an internal failure.
type Error struct {
        Code   ErrorCode
        Reason string
}

// Error implements error.
func (e Error) Error() string {
        return string(e.Code) + ": " + e.Reason
}

func (e Error) valid() bool {
        switch e.Code {
        case CodeInvalidAuthorizationDetails, CodeInvalidRequest:
        default:
                return false
        }
        return reasonPattern.MatchString(e.Reason)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authorizationdetails

import (
        "bytes"
        "encoding/json"
        "io"
        "strings"
        "unicode"
)

type scanFrame struct {
        object    bool
        expectKey bool
        folded    map[string]struct{}
}

// arrayScanner tracks the container stack while validArray walks the tokens.
type arrayScanner struct {
        frames  []scanFrame
        started bool
        done    bool
}

// validArray reports whether raw is exactly one JSON array, nested at most
// maxDepth containers deep, in which no object repeats a member name.
//
// Member names are compared after JSON unescaping, so "type" and "\u0074ype"
// collide. Names that differ only in letter case also collide, using the same
// simple case folding encoding/json applies when it matches struct fields
// ("s" and "ſ" collide, as do "σ" and "ς"). RFC 9396 does not require that
// second rule: it is ToolHive policy, because a consumer that matches names
// case-insensitively would otherwise read a different value than one that
// matches exactly.
func validArray(raw []byte) bool {
        decoder := json.NewDecoder(bytes.NewReader(raw))
        decoder.UseNumber()
        s := &arrayScanner{frames: make([]scanFrame, 0, 8)}

        for {
                token, err := decoder.Token()
                if err == io.EOF {
                        return s.done
                }
                if err != nil || s.done || !s.accept(token) {
                        return false
                }
        }
}

// accept consumes one token and reports whether the array is still valid.
func (s *arrayScanner) accept(token json.Token) bool {
        delim, isDelim := token.(json.Delim)
        if !s.started {
                if !isDelim || delim != '[' {
                        return false
                }
                s.started = true
        }

        // Token enforces the JSON syntax: bracket pairing, string keys in key
        // position and separators. Only duplicates, depth and the single root
        // array remain to be checked here.
        if n := len(s.frames); n > 0 && s.frames[n-1].object && s.frames[n-1].expectKey && !isDelim {
                return s.frames[n-1].addKey(token)
        }
        return s.acceptValue(delim, isDelim)
}

// acceptValue handles a delimiter or scalar outside key position.
func (s *arrayScanner) acceptValue(delim json.Delim, isDelim bool) bool {
        if isDelim && (delim == '{' || delim == '[') {
                if len(s.frames) == maxDepth {
                        return false
                }
                s.frames = append(s.frames, newScanFrame(delim == '{'))
                return true
        }
        if isDelim {
                s.frames = s.frames[:len(s.frames)-1]
        }

        // A scalar or a whole container just ended.
        if n := len(s.frames); n == 0 {
                s.done = true
        } else if s.frames[n-1].object {
                s.frames[n-1].expectKey = true
        }
        return true
}

// addKey records an object member name and reports false on a duplicate.
func (f *scanFrame) addKey(token json.Token) bool {
        // Decoder.Token only returns strings in key position, so the assertion holds.
        key, _ := token.(string)
        // The folded set alone catches exact and case-variant duplicates.
        folded := foldKey(key)
        if _, dup := f.folded[folded]; dup {
                return false
        }
        f.folded[folded] = struct{}{}
        f.expectKey = false
        return true
}

func newScanFrame(object bool) scanFrame {
        frame := scanFrame{object: object}
        if object {
                frame.expectKey = true
                frame.folded = make(map[string]struct{})
        }
        return frame
}

// foldKey maps each rune to the smallest rune in its unicode.SimpleFold orbit.
// SimpleFold returns the next larger rune in the orbit and wraps to the
// smallest one, so the loop stops at the orbit minimum.
func foldKey(s string) string {
        var b strings.Builder
        for _, r := range s {
                for {
                        next := unicode.SimpleFold(r)
                        if next <= r {
                                r = next
                                break
                        }
                        r = next
                }
                b.WriteRune(r)
        }
        return b.String()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package authserver provides configuration and validation for the OAuth authorization server.
package authserver

import (
        "crypto/rand"
        "errors"
        "fmt"
        "log/slog"
        "math"
        "net/url"
        "regexp"
        "slices"
        "strings"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/auth/dcr"
        "github.com/stacklok/toolhive/pkg/authserver/authorizationdetails"
        "github.com/stacklok/toolhive/pkg/authserver/internal/consentpolicy"
        oauthserver "github.com/stacklok/toolhive/pkg/authserver/server"
        servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
        "github.com/stacklok/toolhive/pkg/authserver/server/handlers"
        "github.com/stacklok/toolhive/pkg/authserver/server/keys"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
        "github.com/stacklok/toolhive/pkg/authserver/upstreamfilter"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// CurrentSchemaVersion is the current version of the authserver RunConfig schema.
const CurrentSchemaVersion = "v0.1.0"

// TLSListenerRunConfig configures the auth server TLS listener on port 8443.
type TLSListenerRunConfig struct {
        CertFile string `json:"cert_file" yaml:"cert_file"`
        KeyFile  string `json:"key_file" yaml:"key_file"`
}

// RunConfig is the serializable configuration for the embedded auth server.
// It contains no secrets - only file paths and environment variable names
// that will be resolved at runtime.
//
// This follows the same pattern as pkg/runner.RunConfig - it's serializable,
// versioned, and portable. Secrets are referenced by file path or environment
// variable name, never embedded directly.
type RunConfig struct {
        // SchemaVersion is the version of the RunConfig schema.
        SchemaVersion string `json:"schema_version" yaml:"schema_version"`

        // Issuer is the issuer identifier for this authorization server.
        // This will be included in the "iss" claim of issued tokens.
        // Must be a valid HTTPS URL (or HTTP for localhost) without query, fragment, or trailing slash.
        Issuer string `json:"issuer" yaml:"issuer"`

        // AuthorizationEndpointBaseURL overrides the base URL used for the authorization_endpoint
        // in the OAuth discovery document. When set, the discovery document will advertise
        // `{authorization_endpoint_base_url}/oauth/authorize` instead of `{issuer}/oauth/authorize`.
        // All other endpoints remain derived from the issuer.
        //nolint:lll // field tags require full JSON+YAML names
        AuthorizationEndpointBaseURL string `json:"authorization_endpoint_base_url,omitempty" yaml:"authorization_endpoint_base_url,omitempty"`

        // SigningKeyConfig configures the signing key provider for JWT operations.
        // If nil or empty, an ephemeral signing key will be auto-generated (development only).
        SigningKeyConfig *SigningKeyRunConfig `json:"signing_key_config,omitempty" yaml:"signing_key_config,omitempty"`

        // HMACSecretFiles contains file paths to HMAC secrets for signing authorization codes
        // and refresh tokens (opaque tokens).
        // First file is the current secret (must be at least 32 bytes), subsequent files
        // are for rotation/verification of existing tokens.
        // If empty, an ephemeral secret will be auto-generated (development only).
        HMACSecretFiles []string `json:"hmac_secret_files,omitempty" yaml:"hmac_secret_files,omitempty"`

        // TokenLifespans configures the duration that various tokens are valid.
        // If nil, defaults are applied (access: 1h, refresh: 7d, authCode: 10m).
        TokenLifespans *TokenLifespanRunConfig `json:"token_lifespans,omitempty" yaml:"token_lifespans,omitempty"`

        // DelegationTokenLifespan is the maximum lifetime for delegated tokens issued
        // via RFC 8693 token exchange. Specified as a Go duration string (e.g., "15m").
        // If empty, defaults to 15 minutes.
        DelegationTokenLifespan string `json:"delegation_token_lifespan,omitempty" yaml:"delegation_token_lifespan,omitempty"`

        // Upstreams configures connections to upstream Identity Providers for
        // interactive authorization. It may be empty only when DelegateClients or a
        // TrustedIssuer with JWTBearerGrant enables token-only operation.
        // Multiple upstreams are supported for sequential authorization chains.
        Upstreams []UpstreamRunConfig `json:"upstreams" yaml:"upstreams"`

        // UpstreamFilter configures an optional group-based filter that narrows the
        // multi-upstream authorization chain per-request based on the first
        // upstream's resolved identity. When nil, every request walks the full
        // upstream chain in configured order (preserving behavior from before the
        // filter hook existed).
        //
        // The filter is applied exactly once per authorization, on the first leg's
        // callback. The mandatory first upstream (Upstreams[0]) is never filtered
        // out. See pkg/authserver/upstreamfilter for the configuration schema and
        // pkg/authserver/server/handlers for the runtime contract.
        UpstreamFilter *upstreamfilter.GroupBasedFilterConfig `json:"upstream_filter,omitempty" yaml:"upstream_filter,omitempty"`

        // ScopesSupported lists the OAuth 2.0 scope values advertised in discovery documents.
        // If empty, defaults to registration.DefaultScopes (["openid", "profile", "email", "offline_access"]).
        ScopesSupported []string `json:"scopes_supported,omitempty" yaml:"scopes_supported,omitempty"`

        // BaselineClientScopes is a baseline set of OAuth 2.0 scopes unioned into every
        // DCR registration. All values must appear in ScopesSupported; the auth server
        // rejects this RunConfig at startup otherwise. Empty means current behavior is
        // preserved (registered scope = client-requested, or the intersection of
        // DefaultScopes with ScopesSupported if the client requested none).
        // When ScopesSupported is empty, the subset check uses registration.DefaultScopes
        // (the same set applyDefaults would substitute at startup) — so
        // BaselineClientScopes containing standard OIDC scopes works without enumerating
        // ScopesSupported explicitly.
        //nolint:lll // field tags require full JSON+YAML names
        BaselineClientScopes []string `json:"baseline_client_scopes,omitempty" yaml:"baseline_client_scopes,omitempty"`

        // AllowedAudiences is the list of valid resource URIs that tokens can be issued for.
        // Per RFC 8707, the "resource" parameter in authorization and token requests is
        // validated against this list. Required for MCP compliance.
        AllowedAudiences []string `json:"allowed_audiences" yaml:"allowed_audiences"`

        // Storage configures the storage backend for the auth server.
        // If nil, defaults to in-memory storage.
        Storage *storage.RunConfig `json:"storage,omitempty" yaml:"storage,omitempty"`

        // DisableUpstreamTokenInjection prevents the upstream swap middleware from being added.
        // When true, the embedded auth server handles OAuth flows for clients, but instead of
        // injecting upstream IdP tokens the proxy strips the client's credential headers
        // (Authorization, Cookie, Proxy-Authorization) after the JWT is validated — the
        // backend receives an unauthenticated request. Incompatible with token exchange
        // and AWS STS, which would re-add credentials after the strip.
        //nolint:lll // field tags require full JSON+YAML names
        DisableUpstreamTokenInjection bool `json:"disable_upstream_token_injection,omitempty" yaml:"disable_upstream_token_injection,omitempty"`

        // CIMD controls client_id metadata document support. When enabled, the
        // embedded authorization server accepts HTTPS URLs as client_id values
        // and resolves them via the CIMD protocol instead of requiring DCR.
        CIMD *CIMDRunConfig `json:"cimd,omitempty" yaml:"cimd,omitempty"`

        // InsecureAllowHTTP permits an http:// issuer URL for non-localhost hosts.
        // Only set this for in-cluster Kubernetes deployments on a trusted network.
        // Production deployments reachable outside the cluster MUST use https://.
        //nolint:lll // field tags require full JSON+YAML names
        InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"`

        // TrustedIssuers lists external OIDC trust declarations.
        //
        // This legacy field is deprecated; RFC 8693 and JWT-bearer policies embedded in these entries
        // remain supported for compatibility. New configurations should put policy
        // under InboundGrants and reference a named trusted issuer.
        //
        // See tokenexchange.TrustedIssuer for the per-issuer field reference, and
        // docs/arch/17-token-exchange-delegation.md for the trust model, consent
        // signals, and operator-facing constraints (audience/scope bounding,
        // subject namespace qualification, required client binding) that aren't
        // visible from the config shape alone.
        //nolint:lll // field tags require full JSON+YAML names
        TrustedIssuers []tokenexchange.TrustedIssuer `json:"trusted_issuers,omitempty" yaml:"trusted_issuers,omitempty"`

        // AllowConfidentialClientRegistration permits Dynamic Client Registration
        // of confidential clients: when true, /oauth/register accepts
        // token_endpoint_auth_method values client_secret_basic and
        // client_secret_post in addition to "none" (still the default on
        // omission) and mints a client_secret returned exactly once. Confidential
        // clients are restricted to https non-loopback redirect URIs, and
        // registrations expire and must re-register when never used within
        // DefaultUnusedDCRClientTTL (24 hours) or idle for more than
        // DefaultDCRClientTTL (30 days). This gates registration only: disabling
        // it does not revoke or reject already-minted secrets at the token
        // endpoint.
        //
        // Security: /oauth/register is unauthenticated, so this issues client
        // secrets to any caller. Combining it with InsecureAllowHTTP is rejected
        // by Validate.
        //nolint:lll // field tags require full JSON+YAML names
        AllowConfidentialClientRegistration bool `json:"allow_confidential_client_registration,omitempty" yaml:"allow_confidential_client_registration,omitempty"`

        // AllowPrivateKeyJWTRegistration permits Dynamic Client Registration of
        // clients using private_key_jwt authentication. This is independent of
        // AllowConfidentialClientRegistration and defaults to false. Registration
        // behavior is controlled independently by the DCR handler and discovery
        // metadata.
        //
        // Security: /oauth/register is unauthenticated. Unlike
        // AllowConfidentialClientRegistration, this is NOT rejected when combined
        // with InsecureAllowHTTP: registration never returns a secret for a
        // private_key_jwt client, so there is nothing for cleartext HTTP to
        // expose.
        //nolint:lll // field tags require full JSON+YAML names
        AllowPrivateKeyJWTRegistration bool `json:"allow_private_key_jwt_registration,omitempty" yaml:"allow_private_key_jwt_registration,omitempty"`

        // DeviceFlowEnabled enables the RFC 8628 OAuth 2.0 Device Authorization
        // Grant: POST /oauth/device_authorization is mounted and
        // urn:ietf:params:oauth:grant-type:device_code is registered at the
        // token endpoint and advertised in discovery. The minimum polling
        // interval (RFC 8628 Section 3.5) is fixed at
        // oauthserver.DefaultDeviceCodeInterval; this is a deliberate
        // simplification to keep this config surface minimal — a future
        // increment may add an override.
        DeviceFlowEnabled bool `json:"device_flow_enabled,omitempty" yaml:"device_flow_enabled,omitempty"`

        // ForceConfidentialRedirectURIs lists redirect URIs that must be registered
        // as confidential clients regardless of the token_endpoint_auth_method the
        // DCR request declares. A registration whose redirect_uris contains an
        // EXACT match for one of these entries is issued a real client_secret and
        // reported back as token_endpoint_auth_method "client_secret_post", even
        // if the request said "none" or omitted the field.
        //
        // This exists for MCP clients (Perplexity is the known case) that declare
        // themselves public (token_endpoint_auth_method: "none") per RFC 7591 but
        // then refuse to proceed because the response carries no client_secret —
        // a self-contradictory request no conformant server can satisfy as
        // written. RFC 7591 §3.2.1 permits the server to substitute metadata, so
        // this takes such a client at its word that it wants a secret.
        //
        // Exact matching is deliberate: it is not a way to obtain a usable
        // credential for another client. An attacker who registers with someone
        // else's callback URI is issued a secret for a client whose authorization
        // codes are delivered to that someone else's redirect endpoint, not to
        // the attacker — the secret is useless without also controlling the
        // callback.
        //
        // Requires AllowConfidentialClientRegistration; every entry must be a
        // valid https non-loopback URI (Validate rejects loopback entries — the
        // same restriction AllowConfidentialClientRegistration itself enforces
        // exists so secrets do not land in distributed native apps, and this
        // override must not bypass it). Remove an entry once the client is fixed
        // to handle "none" registrations correctly.
        //nolint:lll // field tags require full JSON+YAML names
        ForceConfidentialRedirectURIs []string `json:"force_confidential_redirect_uris,omitempty" yaml:"force_confidential_redirect_uris,omitempty"`

        // InsecureAllowConfidentialOverLoopbackHTTP opts in to confidential clients
        // when Issuer is a plain-HTTP loopback URL. Without this flag, that
        // combination is rejected: a loopback http:// issuer is normally fine for
        // local development (the traffic never leaves the machine), but client
        // secrets would otherwise travel over cleartext. Defaults to false. Has no
        // effect when there are no confidential clients or Issuer is https.
        //
        // Applies identically to delegate clients and DCR-registered clients. The
        // Kubernetes CRD requires the explicit opt-in for a delegate client with an
        // HTTP issuer; the shared transport validator enforces that its host is
        // loopback — see EmbeddedAuthServerConfig's doc comment.
        //
        // private_key_jwt registration has no equivalent flag or transport
        // restriction: unlike confidential registration, it never returns a
        // client_secret (or any other secret) in the DCR response, so there is
        // nothing here for cleartext HTTP to expose.
        //nolint:lll // field tags require full JSON+YAML names
        InsecureAllowConfidentialOverLoopbackHTTP bool `json:"insecure_allow_confidential_over_loopback_http,omitempty" yaml:"insecure_allow_confidential_over_loopback_http,omitempty"`

        // DelegateClients declares confidential OAuth clients to register at
        // authorization-server startup, including clients intended for RFC 8693
        // token exchange.
        //
        // This legacy field is deprecated; use InboundGrants.TokenExchange.DelegateClients.
        //
        // Independent of AllowConfidentialClientRegistration: declaring a client
        // here does not require or enable self-service confidential DCR, and
        // setting that flag does not declare or enable any client here. They
        // govern different endpoints — this field is static configuration the
        // operator controls directly, while the flag is admission policy for the
        // unauthenticated /oauth/register endpoint.
        //
        // See DelegateClientRunConfig for the per-client field reference.
        DelegateClients []DelegateClientRunConfig `json:"delegate_clients,omitempty" yaml:"delegate_clients,omitempty"`

        // SPIFFETrustDomains declares SPIFFE trust roots. Each declaration must be
        // referenced by an InboundGrants.SPIFFEClientAuth entry.
        SPIFFETrustDomains []SPIFFETrustDomainRunConfig `json:"spiffe_trust_domains,omitempty" yaml:"spiffe_trust_domains,omitempty"`

        // InboundGrants declares canonical inbound grant configuration, including
        // SPIFFE client authentication, delegate clients, and issuer policy. A
        // non-nil value explicitly controls grant-family enablement.
        InboundGrants *InboundGrantsRunConfig `json:"inbound_grants,omitempty" yaml:"inbound_grants,omitempty"`

        // TLSListener serves the auth server routes over TLS on port 8443, in
        // addition to the MCP port. Required for spiffe_x509 client authentication.
        TLSListener *TLSListenerRunConfig `json:"tls_listener,omitempty" yaml:"tls_listener,omitempty"`

        // RegistrationRateLimit bounds the unauthenticated /oauth/register
        // endpoint. If omitted, defaults to 1 registration per second with a
        // burst of 5. The limiter is per process, so the effective rate is the
        // configured rate multiplied by the replica count. The cap on how many
        // registered clients may exist at once is Storage.MaxDCRClients; unused
        // clients live 24 hours, so a caller at the full rate keeps up to
        // rate x replicas x 86400 of them alive, which at the defaults exceeds the
        // default cap within about 2.8 hours. The operator sets it from the
        // embedded auth server's registrationRateLimit.
        //nolint:lll // field tags require full JSON+YAML names
        RegistrationRateLimit *RegistrationRateLimitRunConfig `json:"registration_rate_limit,omitempty" yaml:"registration_rate_limit,omitempty"`
}

// RegistrationRateLimitRunConfig is the token-bucket limit on /oauth/register.
type RegistrationRateLimitRunConfig struct {
        // RequestsPerSecond is the sustained registration rate. Must be positive.
        RequestsPerSecond float64 `json:"requests_per_second" yaml:"requests_per_second"`

        // Burst is the number of registrations admitted at once before the rate
        // applies. Must be positive.
        Burst int `json:"burst" yaml:"burst"`
}

// Validate rejects a non-positive rate or burst. A nil receiver is valid and
// means the default limit.
func (c *RegistrationRateLimitRunConfig) Validate() error {
        if c == nil {
                return nil
        }
        if c.RequestsPerSecond <= 0 || math.IsInf(c.RequestsPerSecond, 0) || math.IsNaN(c.RequestsPerSecond) {
                return fmt.Errorf("registration_rate_limit: requests_per_second must be a positive number, got %v",
                        c.RequestsPerSecond)
        }
        if c.Burst <= 0 {
                return fmt.Errorf("registration_rate_limit: burst must be positive, got %d", c.Burst)
        }
        return nil
}

// DelegateClientRunConfig declares a pre-provisioned confidential OAuth
// client for authorization-server startup, so it can act as the client in an
// RFC 8693 token-exchange request. The secret is always a reference (file or
// environment variable), never an inline literal. The grant type is fixed to
// RFC 8693 token exchange internally and is not configurable.
type DelegateClientRunConfig struct {
        // ClientID is the OAuth client_id this client presents at the token endpoint.
        ClientID string `json:"client_id" yaml:"client_id"`

        // ClientSecretFile is the path to a file containing the client secret.
        // If both this and ClientSecretEnvVar are set, the file takes precedence.
        ClientSecretFile string `json:"client_secret_file,omitempty" yaml:"client_secret_file,omitempty"`

        // ClientSecretEnvVar is the name of an environment variable containing
        // the client secret. One of ClientSecretFile or ClientSecretEnvVar is
        // required.
        //nolint:lll // field tags require full JSON+YAML names
        ClientSecretEnvVar string `json:"client_secret_env_var,omitempty" yaml:"client_secret_env_var,omitempty"`

        // Scopes are the OAuth scopes this client may request. Required, and
        // must be a subset of RunConfig.ScopesSupported: a declared client must
        // not receive every supported scope just because this was left empty.
        Scopes []string `json:"scopes" yaml:"scopes"`

        // Audiences are the RFC 8707 resource values this client may request a
        // token for. Required, and must be a subset of RunConfig.AllowedAudiences:
        // a declared client must not receive every allowed audience just because
        // this was left empty.
        Audiences []string `json:"audiences" yaml:"audiences"`
}

// Validate checks that the on-disk RunConfig is internally consistent. Called
// by the runner before resolving secrets and building the runtime Config; it
// catches operator-supplied misconfiguration early so server startup fails
// loudly instead of degrading silently at runtime.
func (c *RunConfig) Validate() error {
        if err := c.TLSListener.validate(); err != nil {
                return err
        }
        if err := c.validateClientAdmission(); err != nil {
                return err
        }
        // Also checked by Config.Validate() (same reasoning as
        // validateBaselineClientScopes below): buildUpstreamConfigs performs live
        // RFC 7591 registration against upstream IdPs before authserver.New ever
        // reaches Config.Validate(), so a malformed delegate-client config must
        // fail here, before that side-effecting work runs — not on a crash loop
        // after it, which would also orphan an upstream registration on every
        // restart with the default in-memory DCR store.
        if err := validateAllowedAudiences(c.AllowedAudiences); err != nil {
                return err
        }
        normalized, err := NormalizeInboundGrants(c)
        if err != nil {
                return err
        }
        if err := validateDelegateClients(normalized.DelegateClients, c.ScopesSupported, c.AllowedAudiences); err != nil {
                return err
        }
        if err := validateTrustedIssuers(normalized.TrustedIssuers, c.Issuer, c.AllowedAudiences); err != nil {
                return err
        }
        if err := ValidateConfidentialClientTransport(
                c.AllowConfidentialClientRegistration || len(normalized.DelegateClients) > 0, c.InsecureAllowHTTP,
                c.Issuer, c.InsecureAllowConfidentialOverLoopbackHTTP); err != nil {
                return err
        }
        if err := ValidateForceConfidentialRedirectURIs(
                c.ForceConfidentialRedirectURIs, c.AllowConfidentialClientRegistration); err != nil {
                return err
        }
        if err := ValidateSPIFFETrust(
                c.SPIFFETrustDomains, c.InboundGrants, c.ScopesSupported, c.AllowedAudiences,
        ); err != nil {
                return err
        }
        if err := validateSPIFFENotYetEnforced(c.SPIFFETrustDomains); err != nil {
                return err
        }
        if err := c.validateUpstreamFilter(); err != nil {
                return err
        }
        return c.validateBaselineClientScopes()
}

// validateClientAdmission checks the bounds on how unauthenticated callers
// can create clients: CIMD resolution, the /oauth/register rate, and the DCR
// client cap.
func (c *RunConfig) validateClientAdmission() error {
        if c.CIMD != nil {
                if err := c.CIMD.Validate(); err != nil {
                        return fmt.Errorf("cimd: %w", err)
                }
        }
        if err := c.RegistrationRateLimit.Validate(); err != nil {
                return err
        }
        if _, err := c.Storage.ResolvedMaxDCRClients(); err != nil {
                return fmt.Errorf("storage: %w", err)
        }
        return nil
}

// validate requires both certificate paths when a TLS listener is configured.
// A nil receiver means no listener and is valid.
func (l *TLSListenerRunConfig) validate() error {
        if l == nil {
                return nil
        }
        missing := make([]string, 0, 2)
        if l.CertFile == "" {
                missing = append(missing, "tls_listener.cert_file")
        }
        if l.KeyFile == "" {
                missing = append(missing, "tls_listener.key_file")
        }
        if len(missing) > 0 {
                return fmt.Errorf("%s must be set", strings.Join(missing, " and "))
        }
        return nil
}

// validateUpstreamFilter rejects an UpstreamFilter configuration whose rules or
// defaults reference upstream names that are not configured, that reference the
// mandatory first upstream, or that otherwise violate the filter's static
// invariants. It constructs the filter against the configured upstream names
// and discards the result — successful construction is the validation.
func (c *RunConfig) validateUpstreamFilter() error {
        if c.UpstreamFilter == nil {
                return nil
        }
        configured := make([]string, len(c.Upstreams))
        for i := range c.Upstreams {
                configured[i] = c.Upstreams[i].Name
        }
        if _, err := upstreamfilter.NewGroupBasedFilter(*c.UpstreamFilter, configured); err != nil {
                return fmt.Errorf("upstream_filter: %w", err)
        }
        return nil
}

// HasSPIFFEX509ClientAuth reports whether any SPIFFE client association
// enables spiffe_x509. A nil *RunConfig returns false.
func (c *RunConfig) HasSPIFFEX509ClientAuth() bool {
        if c == nil || c.InboundGrants == nil {
                return false
        }
        for _, association := range c.InboundGrants.SPIFFEClientAuth {
                for _, method := range association.Methods {
                        if method == SPIFFEAuthenticationMethodX509 {
                                return true
                        }
                }
        }
        return false
}

// validateSPIFFENotYetEnforced hard-rejects a non-empty SPIFFE trust
// configuration. ValidateSPIFFETrust above confirms the configuration is
// well-formed, but well-formed is not the same as enforced: nothing in this
// build ever verifies an X.509-SVID or JWT-SVID against the configured trust
// bundle, so a valid, non-empty SPIFFE trust configuration currently has no
// runtime authentication effect. Accepting it silently would let an operator
// believe SPIFFE client authentication is active when no credential is ever
// checked. This rejection must be removed by the future PR that adds real
// SVID verification against the configured trust bundle.
func validateSPIFFENotYetEnforced(trustDomains []SPIFFETrustDomainRunConfig) error {
        if len(trustDomains) == 0 {
                return nil
        }
        return fmt.Errorf(
                "spiffe_trust_domains: SPIFFE client authentication is not yet enforced by this build " +
                        "(no X.509-SVID or JWT-SVID is verified against the configured trust bundle); " +
                        "remove this configuration until a verification consumer lands")
}

// validateConfigSPIFFENotYetEnforced is validateSPIFFENotYetEnforced's
// Config-level counterpart: a caller that constructs Config directly (e.g.
// authserver.New) bypasses RunConfig.Validate() entirely, so the same
// fail-loud rejection must also apply to Config.SPIFFETrust -- otherwise a
// non-empty, well-formed SPIFFE trust policy could start a server through
// this path with no authentication consumer ever wired to it. Associations
// is nil-safe and empty for both a nil SPIFFETrust and one with no
// configured associations.
func validateConfigSPIFFENotYetEnforced(trust *SPIFFETrustConfig) error {
        if len(trust.Associations()) == 0 {
                return nil
        }
        return fmt.Errorf(
                "spiffe_trust: SPIFFE client authentication is not yet enforced by this build " +
                        "(no X.509-SVID or JWT-SVID is verified against the configured trust bundle); " +
                        "remove this configuration until a verification consumer lands")
}

// validateBaselineClientScopes ensures every entry in BaselineClientScopes is
// also present in ScopesSupported. If a baseline scope is not advertised by
// ScopesSupported, the embedded DCR handler would later try to register a
// client with a scope the server does not support, which fosite rejects at
// /oauth/authorize with invalid_scope.
//
// When ScopesSupported is empty, the check uses registration.DefaultScopes as
// the superset (matching what applyDefaults would substitute at startup), so
// operators can omit ScopesSupported and still configure standard OIDC scopes
// as baseline without error.
func (c *RunConfig) validateBaselineClientScopes() error {
        effective := c.ScopesSupported
        if len(effective) == 0 {
                effective = registration.DefaultScopes
        }
        return registration.ValidateScopeSubset(c.BaselineClientScopes, effective, "baseline_client_scopes")
}

func validateAllowedAudiences(audiences []string) error {
        for _, audience := range audiences {
                if audience == "" {
                        return fmt.Errorf("allowed_audiences must not contain an empty audience")
                }
                if err := oauthserver.ValidateAudienceURI(audience); err != nil {
                        return fmt.Errorf("allowed_audiences contains invalid audience %q: %w", audience, err)
                }
        }
        return nil
}

// validateDelegateClients validates the structural and narrowing invariants for
// RunConfig.DelegateClients. Scopes and audiences are required and must be
// subsets of the server-supported values so a declared client never defaults
// to every supported scope or audience.
func validateDelegateClients(
        clients []DelegateClientRunConfig, scopesSupported, allowedAudiences []string,
) error {
        effectiveScopes := scopesSupported
        if len(effectiveScopes) == 0 {
                effectiveScopes = registration.DefaultScopes
        }

        seen := make(map[string]struct{}, len(clients))
        for _, client := range clients {
                if client.ClientID == "" {
                        return fmt.Errorf("delegate_clients: client_id is required")
                }
                if _, ok := seen[client.ClientID]; ok {
                        return fmt.Errorf("delegate_clients: duplicate client_id %q", client.ClientID)
                }
                seen[client.ClientID] = struct{}{}
                // A URL-shaped client_id collides with CIMD (see
                // oauthproto.IsClientIDMetadataDocumentURL): when CIMD is enabled, the
                // CIMDStorageDecorator resolves any client_id matching this shape from
                // a live-fetched document instead of this pre-provisioned row, and its
                // write-through persistence (CIMDStorageDecorator.fetch) would then
                // overwrite this confidential delegate client in place — downgrading
                // it to a public client and discarding its hashed secret. Reject the
                // shape outright regardless of whether CIMD is enabled in this
                // particular deployment, since the config is portable across
                // deployments where it may be.
                if oauthproto.IsClientIDMetadataDocumentURL(client.ClientID) {
                        return fmt.Errorf(
                                "delegate_clients: client_id %q looks like a CIMD client metadata URL, "+
                                        "which collides with CIMD client resolution — pre-provisioned delegate "+
                                        "clients must use an opaque, non-URL client_id", client.ClientID)
                }

                if client.ClientSecretFile == "" && client.ClientSecretEnvVar == "" {
                        return fmt.Errorf(
                                "delegate_clients: client_id %q: client_secret_file or client_secret_env_var is required", client.ClientID)
                }
                if len(client.Scopes) == 0 {
                        return fmt.Errorf("delegate_clients: client_id %q: scopes is required", client.ClientID)
                }
                if err := registration.ValidateScopeSubset(
                        client.Scopes, effectiveScopes, fmt.Sprintf("delegate_clients[%s].scopes", client.ClientID)); err != nil {
                        return err
                }
                if len(client.Audiences) == 0 {
                        return fmt.Errorf("delegate_clients: client_id %q: audiences is required", client.ClientID)
                }
                for _, audience := range client.Audiences {
                        if !slices.Contains(allowedAudiences, audience) {
                                return fmt.Errorf(
                                        "delegate_clients: client_id %q: audience %q is not in allowed_audiences", client.ClientID, audience)
                        }
                }
        }
        return nil
}

// minDelegateClientSecretLength is the minimum accepted length for a resolved
// delegate-client secret. Delegate-client secrets come from the operator (a
// file or environment variable) rather than being minted by this server's own
// DCR path, so — unlike SHA256Hasher's DCR-issued-secret assumption (see that
// type's doc comment) — nothing else guarantees they carry meaningful
// entropy. 32 matches the byte length GenerateClientSecret uses for DCR
// secrets (32 bytes of crypto/rand, base64url-encoded), so this floor accepts
// anything at that scale without requiring a specific encoding.
const minDelegateClientSecretLength = 32

// validateResolvedDelegateClients validates resolved runtime delegate clients.
// The resolved secret must be nonempty and at least minDelegateClientSecretLength
// characters, because Config callers bypass secret resolution and
// RunConfig.Validate, and unlike DCR-minted secrets, a delegate-client secret's
// entropy is never otherwise checked.
func validateResolvedDelegateClients(
        clients []DelegateClient, scopesSupported, allowedAudiences []string,
) error {
        runClients := make([]DelegateClientRunConfig, len(clients))
        for i, client := range clients {
                if client.ClientSecret == "" {
                        return fmt.Errorf("delegate_clients: client_id %q: resolved client secret is required", client.ClientID)
                }
                if len(client.ClientSecret) < minDelegateClientSecretLength {
                        return fmt.Errorf(
                                "delegate_clients: client_id %q: resolved client secret must be at least %d characters",
                                client.ClientID, minDelegateClientSecretLength)
                }
                runClients[i] = DelegateClientRunConfig{
                        ClientID:           client.ClientID,
                        ClientSecretEnvVar: "resolved",
                        Scopes:             client.Scopes,
                        Audiences:          client.Audiences,
                }
        }
        return validateDelegateClients(runClients, scopesSupported, allowedAudiences)
}

// CIMDRunConfig controls client_id metadata document (CIMD) support.
type CIMDRunConfig struct {
        // Enabled activates CIMD client lookup when true.
        Enabled bool `json:"enabled" yaml:"enabled"`

        // CacheMaxSize is the maximum number of CIMD documents held in the LRU cache.
        // Defaults to 256 when Enabled is true and this field is zero.
        CacheMaxSize int `json:"cache_max_size,omitempty" yaml:"cache_max_size,omitempty"`

        // CacheFallbackTTL is the fixed TTL applied to every cached CIMD document.
        // Cache-Control header parsing is not yet implemented; all entries use this value.
        // Format: Go duration string (e.g. "5m", "10m", "1h").
        // Defaults to 5 minutes when Enabled is true and this field is omitted.
        CacheFallbackTTL string `json:"cache_fallback_ttl,omitempty" yaml:"cache_fallback_ttl,omitempty" example:"5m"`
}

// Validate checks that the CIMDRunConfig fields are internally consistent.
func (c *CIMDRunConfig) Validate() error {
        if !c.Enabled {
                return nil
        }
        if c.CacheMaxSize < 0 {
                return fmt.Errorf("cache_max_size must be non-negative when CIMD is enabled, got %d", c.CacheMaxSize)
        }
        if c.CacheFallbackTTL != "" {
                d, err := time.ParseDuration(c.CacheFallbackTTL)
                if err != nil {
                        return fmt.Errorf("cache_fallback_ttl: %w", err)
                }
                if d <= 0 {
                        return fmt.Errorf("cache_fallback_ttl must be positive when CIMD is enabled, got %s", c.CacheFallbackTTL)
                }
        }
        return nil
}

// SigningKeyRunConfig configures where to load signing keys from.
// Keys are loaded from PEM-encoded files on disk (typically mounted from secrets).
type SigningKeyRunConfig struct {
        // KeyDir is the directory containing PEM-encoded private key files.
        // All key filenames are relative to this directory.
        // In Kubernetes, this is typically a mounted Secret volume.
        KeyDir string `json:"key_dir,omitempty" yaml:"key_dir,omitempty"`

        // SigningKeyFile is the filename of the primary signing key (relative to KeyDir).
        // This key is used for signing new tokens.
        SigningKeyFile string `json:"signing_key_file,omitempty" yaml:"signing_key_file,omitempty"`

        // FallbackKeyFiles are filenames of additional keys for verification (relative to KeyDir).
        // These keys are included in the JWKS endpoint for token verification but are NOT
        // used for signing new tokens. Useful for key rotation.
        FallbackKeyFiles []string `json:"fallback_key_files,omitempty" yaml:"fallback_key_files,omitempty"`
}

// TokenLifespanRunConfig holds token lifetime configuration.
// All durations are specified as Go duration strings (e.g., "1h", "30m", "168h").
type TokenLifespanRunConfig struct {
        // AccessTokenLifespan is the duration that access tokens are valid.
        // If empty, defaults to 1 hour.
        AccessTokenLifespan string `json:"access_token_lifespan,omitempty" yaml:"access_token_lifespan,omitempty"`

        // RefreshTokenLifespan is the duration that refresh tokens are valid.
        // If empty, defaults to 7 days (168h).
        RefreshTokenLifespan string `json:"refresh_token_lifespan,omitempty" yaml:"refresh_token_lifespan,omitempty"`

        // AuthCodeLifespan is the duration that authorization codes are valid.
        // If empty, defaults to 10 minutes.
        AuthCodeLifespan string `json:"auth_code_lifespan,omitempty" yaml:"auth_code_lifespan,omitempty"`
}

// UpstreamProviderType identifies the type of upstream Identity Provider.
type UpstreamProviderType string

const (
        // UpstreamProviderTypeOIDC is for OIDC providers with discovery support.
        UpstreamProviderTypeOIDC UpstreamProviderType = "oidc"

        // UpstreamProviderTypeOAuth2 is for pure OAuth 2.0 providers with explicit endpoints.
        UpstreamProviderTypeOAuth2 UpstreamProviderType = "oauth2"
)

// DefaultUpstreamName is the name assigned to a single unnamed upstream.
const DefaultUpstreamName = "default"

// ResolveUpstreamName returns the canonical name for an upstream.
// An empty name is resolved to DefaultUpstreamName ("default").
func ResolveUpstreamName(name string) string {
        if name == "" {
                return DefaultUpstreamName
        }
        return name
}

// ResolveFirstUpstreamName returns the resolved name of the first element of
// names, or DefaultUpstreamName when names is empty. It is the single
// implementation of the "first upstream or default" pattern used wherever a
// subject-provider name must be derived from a list of configured upstreams.
func ResolveFirstUpstreamName(names []string) string {
        if len(names) > 0 {
                return ResolveUpstreamName(names[0])
        }
        return DefaultUpstreamName
}

// upstreamNameRegex validates upstream provider names.
// Names must be DNS-label-like to prevent delimiter injection in storage keys.
var upstreamNameRegex = regexp.MustCompile(`^[a-z0-9]([a-z0-9-]*[a-z0-9])?$`)

// UpstreamRunConfig configures an upstream identity provider.
type UpstreamRunConfig struct {
        // Name uniquely identifies this upstream.
        // Used for routing decisions and session binding in multi-upstream scenarios.
        // If empty when only one upstream is configured, defaults to "default".
        Name string `json:"name,omitempty" yaml:"name,omitempty"`

        // DisplayName is an optional human-readable label for this upstream, shown to
        // end users on the consent page in place of Name. It is display-only: never
        // compared, stored, or used to route. It is sanitized at render time and never
        // rejected by validation. It is shown as trusted text, so it must come from
        // operator-controlled configuration only.
        DisplayName string `json:"display_name,omitempty" yaml:"display_name,omitempty"`

        // Type specifies the provider type: "oidc" or "oauth2".
        Type UpstreamProviderType `json:"type" yaml:"type"`

        // OIDCConfig contains OIDC-specific configuration.
        // Required when Type is "oidc", must be nil when Type is "oauth2".
        OIDCConfig *OIDCUpstreamRunConfig `json:"oidc_config,omitempty" yaml:"oidc_config,omitempty"`

        // OAuth2Config contains OAuth 2.0-specific configuration.
        // Required when Type is "oauth2", must be nil when Type is "oidc".
        OAuth2Config *OAuth2UpstreamRunConfig `json:"oauth2_config,omitempty" yaml:"oauth2_config,omitempty"`
}

// OIDCUpstreamRunConfig contains OIDC provider configuration.
// OIDC providers support automatic endpoint discovery via the issuer URL.
type OIDCUpstreamRunConfig struct {
        // IssuerURL is the OIDC issuer URL for automatic endpoint discovery.
        // Must be a valid HTTPS URL.
        IssuerURL string `json:"issuer_url" yaml:"issuer_url"`

        // ClientID is the OAuth 2.0 client identifier registered with the upstream IDP.
        ClientID string `json:"client_id" yaml:"client_id"`

        // ClientSecretFile is the path to a file containing the OAuth 2.0 client secret.
        // Mutually exclusive with ClientSecretEnvVar. Optional for public clients using PKCE.
        ClientSecretFile string `json:"client_secret_file,omitempty" yaml:"client_secret_file,omitempty"`

        // ClientSecretEnvVar is the name of an environment variable containing the client secret.
        // Mutually exclusive with ClientSecretFile. Optional for public clients using PKCE.
        ClientSecretEnvVar string `json:"client_secret_env_var,omitempty" yaml:"client_secret_env_var,omitempty"`

        // RedirectURI is the callback URL where the upstream IDP will redirect after authentication.
        // For a DCR upstream (dcr_config set) an empty value is resolved to
        // `{issuer}/oauth/callback`; a pre-provisioned client (client_id set) must
        // specify it. Its hostname must match the browser-facing authorize URL
        // (issuer, or authorization_endpoint_base_url when set): /oauth/callback is
        // bound to the browser that started the login by a host-only cookie, so a
        // different host rejects every browser login (Config.Validate warns).
        RedirectURI string `json:"redirect_uri,omitempty" yaml:"redirect_uri,omitempty"`

        // Scopes are the OAuth scopes to request from the upstream IDP.
        // If not specified, defaults to ["openid", "offline_access"].
        // When using AdditionalAuthorizationParams with provider-specific refresh
        // token mechanisms (e.g., Google's access_type=offline), set explicit scopes
        // to avoid sending both offline_access and the provider-specific parameter.
        Scopes []string `json:"scopes,omitempty" yaml:"scopes,omitempty"`

        // UserInfoOverride allows customizing UserInfo fetching behavior for OIDC providers.
        // By default, the UserInfo endpoint is discovered automatically via OIDC discovery.
        UserInfoOverride *UserInfoRunConfig `json:"userinfo_override,omitempty" yaml:"userinfo_override,omitempty"`

        // AdditionalAuthorizationParams are extra query parameters to include in
        // authorization requests. Useful for provider-specific parameters like
        // Google's access_type=offline.
        //nolint:lll // field tags require full JSON+YAML names
        AdditionalAuthorizationParams map[string]string `json:"additional_authorization_params,omitempty" yaml:"additional_authorization_params,omitempty"`

        // AdditionalTokenParams are extra form-body parameters to include in
        // token requests (authorization code exchange and refresh). Useful for
        // providers that enforce RFC 8707 resource indicators on token requests.
        //nolint:lll // field tags require full JSON+YAML names
        AdditionalTokenParams map[string]string `json:"additional_token_params,omitempty" yaml:"additional_token_params,omitempty"`

        // OmitScopeOnRefresh stops the `scope` parameter from being sent on
        // refresh_token grant requests. Scopes are still sent on the authorization
        // request. Useful for providers that reject `scope` on refresh, such as
        // Salesforce. See upstream.CommonOAuthConfig.
        OmitScopeOnRefresh bool `json:"omit_scope_on_refresh,omitempty" yaml:"omit_scope_on_refresh,omitempty"`

        // SubjectClaim names the validated ID-token claim to use as the upstream
        // subject. Defaults to "sub" when empty. Set for IdPs where "sub" isn't
        // stable per user (e.g. Entra/Azure AD's "oid"). See upstream.OIDCConfig.
        SubjectClaim string `json:"subject_claim,omitempty" yaml:"subject_claim,omitempty"`

        // DCRConfig enables RFC 7591 Dynamic Client Registration against the
        // upstream authorization server. When set, the client credentials are
        // obtained at runtime rather than being pre-provisioned via ClientID /
        // ClientSecretFile / ClientSecretEnvVar, and ClientID must be left empty.
        // Mutually exclusive with ClientID.
        DCRConfig *DCRUpstreamConfig `json:"dcr_config,omitempty" yaml:"dcr_config,omitempty"`

        // CAFilePath is the path to a PEM CA bundle added to the system roots.
        CAFilePath string `json:"ca_file_path,omitempty" yaml:"ca_file_path,omitempty"`

        // AllowPrivateIPs permits the OIDC discovery and token HTTP clients to
        // connect to private IP ranges (RFC-1918, link-local). Use only when the
        // upstream is hosted inside the same cluster and has no public endpoint.
        // HTTP-scheme restrictions are unchanged — HTTPS is still required for
        // non-localhost hosts. Defaults to false.
        AllowPrivateIPs bool `json:"allow_private_ips,omitempty" yaml:"allow_private_ips,omitempty"`

        // InsecureAllowHTTP permits a plain-HTTP issuer URL and HTTP discovery
        // endpoints for this upstream. Only for in-cluster development environments
        // (e.g. Dex served over HTTP in a kind cluster) where TLS is not available.
        // Never set this in production.
        //nolint:lll // field tags require full JSON+YAML names
        InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"`
}

// OAuth2UpstreamRunConfig contains configuration for pure OAuth 2.0 providers.
// OAuth 2.0 providers require explicit endpoint configuration.
type OAuth2UpstreamRunConfig struct {
        // AuthorizationEndpoint is the URL for the OAuth authorization endpoint.
        AuthorizationEndpoint string `json:"authorization_endpoint" yaml:"authorization_endpoint"`

        // TokenEndpoint is the URL for the OAuth token endpoint.
        TokenEndpoint string `json:"token_endpoint" yaml:"token_endpoint"`

        // ClientID is the OAuth 2.0 client identifier registered with the upstream IDP.
        // Mutually exclusive with DCRConfig: when DCRConfig is set, ClientID is obtained
        // at runtime via RFC 7591 Dynamic Client Registration and must be left empty.
        ClientID string `json:"client_id" yaml:"client_id"`

        // ClientSecretFile is the path to a file containing the OAuth 2.0 client secret.
        // Mutually exclusive with ClientSecretEnvVar. Optional for public clients using PKCE.
        ClientSecretFile string `json:"client_secret_file,omitempty" yaml:"client_secret_file,omitempty"`

        // ClientSecretEnvVar is the name of an environment variable containing the client secret.
        // Mutually exclusive with ClientSecretFile. Optional for public clients using PKCE.
        ClientSecretEnvVar string `json:"client_secret_env_var,omitempty" yaml:"client_secret_env_var,omitempty"`

        // TokenEndpointAuthMethod selects how the client authenticates at the OAuth token
        // endpoint. When empty, credentials are sent in the request body (the historical
        // client_secret_post-shaped default). Set this to client_secret_basic explicitly
        // for providers that require HTTP Basic auth. Public clients without a secret use
        // the "none" method.
        TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty" yaml:"token_endpoint_auth_method,omitempty"`

        // RedirectURI is the callback URL where the upstream IDP will redirect after authentication.
        // For a DCR upstream (dcr_config set) an empty value is resolved to
        // `{issuer}/oauth/callback`; a pre-provisioned client (client_id set) must
        // specify it. Its hostname must match the browser-facing authorize URL
        // (issuer, or authorization_endpoint_base_url when set): /oauth/callback is
        // bound to the browser that started the login by a host-only cookie, so a
        // different host rejects every browser login (Config.Validate warns).
        RedirectURI string `json:"redirect_uri,omitempty" yaml:"redirect_uri,omitempty"`

        // Scopes are the OAuth scopes to request from the upstream IDP.
        Scopes []string `json:"scopes,omitempty" yaml:"scopes,omitempty"`

        // UserInfo contains configuration for fetching user information.
        // Optional: when nil, the upstream OAuth2 provider derives a deterministic
        // subject by SHA-256-hashing the access token (with a "tk-" prefix) instead
        // of calling a userinfo endpoint. OIDC providers always derive Subject from
        // the ID token and are unaffected.
        UserInfo *UserInfoRunConfig `json:"userinfo,omitempty" yaml:"userinfo,omitempty"`

        // TokenResponseMapping configures custom field extraction from non-standard token responses.
        // When set, the token exchange bypasses golang.org/x/oauth2 and extracts fields using
        // the configured dot-notation paths.
        //nolint:lll // field tags require full JSON+YAML names
        TokenResponseMapping *TokenResponseMappingRunConfig `json:"token_response_mapping,omitempty" yaml:"token_response_mapping,omitempty"`

        // IdentityFromToken extracts user identity (subject, name, email) directly from the
        // OAuth2 token-endpoint response body using gjson dot-notation paths. When set, the
        // embedded auth server skips the userinfo HTTP call entirely. Mirrors the CRD type
        // (cmd/thv-operator/api/v1beta1.IdentityFromTokenConfig) — the authoritative
        // trust-model and uniqueness documentation lives there.
        //nolint:lll // field tags require full JSON+YAML names
        IdentityFromToken *IdentityFromTokenRunConfig `json:"identity_from_token,omitempty" yaml:"identity_from_token,omitempty"`

        // AdditionalAuthorizationParams are extra query parameters to include in
        // authorization requests. Useful for provider-specific parameters like
        // Google's access_type=offline.
        //nolint:lll // field tags require full JSON+YAML names
        AdditionalAuthorizationParams map[string]string `json:"additional_authorization_params,omitempty" yaml:"additional_authorization_params,omitempty"`

        // AdditionalTokenParams are extra form-body parameters to include in
        // token requests (authorization code exchange and refresh). Useful for
        // providers that enforce RFC 8707 resource indicators on token requests.
        //nolint:lll // field tags require full JSON+YAML names
        AdditionalTokenParams map[string]string `json:"additional_token_params,omitempty" yaml:"additional_token_params,omitempty"`

        // OmitScopeOnRefresh stops the `scope` parameter from being sent on
        // refresh_token grant requests. Scopes are still sent on the authorization
        // request. Useful for providers that reject `scope` on refresh, such as
        // Salesforce. See upstream.CommonOAuthConfig.
        OmitScopeOnRefresh bool `json:"omit_scope_on_refresh,omitempty" yaml:"omit_scope_on_refresh,omitempty"`

        // DCRConfig enables RFC 7591 Dynamic Client Registration against the
        // upstream authorization server. When set, the client credentials are
        // obtained at runtime rather than being pre-provisioned via ClientID /
        // ClientSecretFile / ClientSecretEnvVar, and ClientID must be left empty.
        // Mutually exclusive with ClientID.
        DCRConfig *DCRUpstreamConfig `json:"dcr_config,omitempty" yaml:"dcr_config,omitempty"`

        // CAFilePath is the path to a PEM CA bundle added to the system roots.
        CAFilePath string `json:"ca_file_path,omitempty" yaml:"ca_file_path,omitempty"`

        // AllowPrivateIPs permits the upstream provider's HTTP client to connect to
        // private IP ranges (RFC-1918, link-local). When DCRConfig is set, this
        // also gates the DCR discovery and registration calls made on this
        // upstream's behalf (see pkg/authserver/runner/dcr_adapter.go), so a
        // single flag covers the whole upstream rather than needing a separate
        // DCR-specific setting. Use only when the upstream is hosted inside the
        // same cluster and has no public endpoint. HTTP-scheme restrictions are
        // unchanged — HTTPS is still required for non-localhost hosts. Defaults
        // to false.
        AllowPrivateIPs bool `json:"allow_private_ips,omitempty" yaml:"allow_private_ips,omitempty"`

        // InsecureAllowHTTP permits plain-HTTP authorization and token endpoint URLs
        // for this upstream. Only for in-cluster development environments (e.g. an
        // OAuth2 provider served over HTTP in a kind cluster) where TLS is not
        // available. Never set this in production.
        //nolint:lll // field tags require full JSON+YAML names
        InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"`
}

// DCRUpstreamConfig configures RFC 7591 Dynamic Client Registration for an
// OAuth2 or OIDC upstream. When present, the authserver performs registration
// at runtime to obtain client credentials, replacing the need to pre-provision
// a ClientID.
//
// OAuth2 upstreams must set exactly one of DiscoveryURL or
// RegistrationEndpoint. OIDC upstreams may omit both and derive DiscoveryURL
// from IssuerURL.
// points at RFC 8414 / OIDC Discovery metadata from which the registration
// endpoint is resolved; RegistrationEndpoint is used directly when the upstream
// does not publish discovery metadata.
//
// Trust assumption: DiscoveryURL and RegistrationEndpoint are operator-supplied
// URLs validated only for HTTPS-or-loopback. The DCR resolver will issue
// outbound HTTP requests — possibly carrying the RFC 7591 initial access token
// as a bearer header — to whatever address those URLs resolve to. There is
// currently no allowlist or RFC1918 / link-local / cloud-metadata-service
// guard, because the operator role is fully trusted today. If the trust
// boundary ever changes (e.g. a multi-tenant operator deployment, or a less-
// privileged role gains write access to this struct via a CRD or YAML
// surface), this field becomes a confused-deputy SSRF vector. Hardening is
// tracked in https://github.com/stacklok/toolhive/issues/5135.
type DCRUpstreamConfig struct {
        // DiscoveryURL is the exact RFC 8414 / OIDC Discovery document URL to
        // fetch at runtime. The resolver issues a single GET against this URL
        // (no well-known-path fallback) and reads registration_endpoint,
        // authorization_endpoint, token_endpoint,
        // token_endpoint_auth_methods_supported, and scopes_supported from the
        // response. Per RFC 8414 §3.3, the document's "issuer" field must
        // exactly match the upstream issuer configured on the parent
        // run-config.
        //
        // Use this field when the upstream publishes discovery metadata at a
        // path that differs from the issuer-derived well-known paths — for
        // example a multi-tenant IdP whose metadata lives at
        // https://idp.example.com/tenants/acme/.well-known/openid-configuration.
        //
        // Mutually exclusive with RegistrationEndpoint.
        DiscoveryURL string `json:"discovery_url,omitempty" yaml:"discovery_url,omitempty"`

        // RegistrationEndpoint is the RFC 7591 registration endpoint URL used
        // directly, bypassing discovery. Because no discovery is performed,
        // server-capability fields (token_endpoint_auth_methods_supported,
        // scopes_supported) are unavailable on this code path; the caller is
        // expected to also supply AuthorizationEndpoint, TokenEndpoint, and an
        // explicit Scopes list on the parent OAuth2UpstreamRunConfig. Auth
        // method falls back to the resolver's default (client_secret_basic).
        //
        // Mutually exclusive with DiscoveryURL.
        RegistrationEndpoint string `json:"registration_endpoint,omitempty" yaml:"registration_endpoint,omitempty"`

        // RotationKey is an optional operator-chosen label (1-64 characters from
        // [A-Za-z0-9._-]) that becomes part of the cached registration's identity
        // without being sent to the upstream. Changing it makes the next start
        // register a fresh client instead of reusing the cached one. At start the
        // resolver checks the cached client at the upstream (see RegistrationProbe)
        // and replaces it by itself when the upstream rejects it. Use RotationKey
        // when the upstream has dropped the registration and the configured probe
        // cannot tell, so the rejection may not be visible to the resolver, or when
        // a registration must be rotated on request. The row written under the
        // previous value is left in place. Upstream refresh tokens bound to the
        // previous client_id stop refreshing once the client changes; affected
        // users sign in again.
        RotationKey string `json:"rotation_key,omitempty" yaml:"rotation_key,omitempty"`

        // RegistrationProbe selects how the resolver checks a cached registration
        // against the upstream at start. "tokenEndpoint" (the default) redeems a
        // random authorization code with the cached client and reads
        // invalid_client as "registration gone"; it decides only for upstreams
        // that authenticate the client before validating the grant.
        // "authorizationEndpoint" sends an authorization request for the cached
        // client without following redirects and reads a direct 400, 401 or 403
        // that carries an RFC 6749 JSON error object (error invalid_client,
        // unauthorized_client or invalid_request) as "gone" (RFC 6749 §4.1.2.1),
        // for upstreams whose token endpoint answers invalid_grant for any
        // client_id; an HTML page or a bare gateway status is not a verdict. A
        // false rejection re-registers at every start and, when registration also
        // fails, fails startup. "none" skips the check; the registration is then
        // retired only through RotationKey.
        RegistrationProbe string `json:"registration_probe,omitempty" yaml:"registration_probe,omitempty"`

        // InitialAccessTokenFile is the path to a file containing the RFC 7591
        // initial access token presented to the registration endpoint. Mutually
        // exclusive with InitialAccessTokenEnvVar. Both may be omitted for open
        // registration endpoints.
        //nolint:lll // field tags require full JSON+YAML names
        InitialAccessTokenFile string `json:"initial_access_token_file,omitempty" yaml:"initial_access_token_file,omitempty"`

        // InitialAccessTokenEnvVar is the name of an environment variable
        // containing the RFC 7591 initial access token. Mutually exclusive with
        // InitialAccessTokenFile.
        //nolint:lll // field tags require full JSON+YAML names
        InitialAccessTokenEnvVar string `json:"initial_access_token_env_var,omitempty" yaml:"initial_access_token_env_var,omitempty"`

        // SoftwareID is the RFC 7591 "software_id" registration metadata value,
        // identifying the client software independent of any particular
        // registration instance.
        SoftwareID string `json:"software_id,omitempty" yaml:"software_id,omitempty"`

        // SoftwareStatement is the RFC 7591 "software_statement" JWT asserting
        // metadata about the client software, signed by a party the authorization
        // server trusts.
        SoftwareStatement string `json:"software_statement,omitempty" yaml:"software_statement,omitempty"`
}

// TokenResponseMappingRunConfig maps non-standard token response fields to standard fields.
// Paths support dot-notation for nested JSON fields (e.g., "authed_user.access_token").
type TokenResponseMappingRunConfig struct {
        // AccessTokenPath is the dot-notation path to the access token (required).
        AccessTokenPath string `json:"access_token_path" yaml:"access_token_path"`

        // ScopePath is the dot-notation path to the scope. Defaults to "scope".
        ScopePath string `json:"scope_path,omitempty" yaml:"scope_path,omitempty"`

        // RefreshTokenPath is the dot-notation path to the refresh token. Defaults to "refresh_token".
        RefreshTokenPath string `json:"refresh_token_path,omitempty" yaml:"refresh_token_path,omitempty"`

        // ExpiresInPath is the dot-notation path to the expires_in value. Defaults to "expires_in".
        ExpiresInPath string `json:"expires_in_path,omitempty" yaml:"expires_in_path,omitempty"`
}

// IdentityFromTokenRunConfig configures extracting user identity claims directly from
// the token-endpoint response body. Mirrors the CRD type
// (cmd/thv-operator/api/v1beta1.IdentityFromTokenConfig) — the authoritative
// trust-model and uniqueness documentation lives there.
type IdentityFromTokenRunConfig struct {
        // SubjectPath is the dot-notation path to the subject (user ID) field.
        // Required when IdentityFromToken is set.
        SubjectPath string `json:"subject_path" yaml:"subject_path"`

        // NamePath is the dot-notation path to the display name field.
        NamePath string `json:"name_path,omitempty" yaml:"name_path,omitempty"`

        // EmailPath is the dot-notation path to the email address field.
        EmailPath string `json:"email_path,omitempty" yaml:"email_path,omitempty"`
}

// UserInfoRunConfig contains UserInfo endpoint configuration.
// This supports both standard OIDC UserInfo endpoints and custom provider-specific endpoints.
type UserInfoRunConfig struct {
        // EndpointURL is the URL of the userinfo endpoint.
        EndpointURL string `json:"endpoint_url" yaml:"endpoint_url"`

        // HTTPMethod is the HTTP method to use for the userinfo request.
        // If not specified, defaults to GET.
        HTTPMethod string `json:"http_method,omitempty" yaml:"http_method,omitempty"`

        // AdditionalHeaders contains extra headers to include in the userinfo request.
        // Useful for providers that require specific headers (e.g., GitHub's Accept header).
        AdditionalHeaders map[string]string `json:"additional_headers,omitempty" yaml:"additional_headers,omitempty"`

        // FieldMapping contains custom field mapping configuration for non-standard providers.
        // If nil, standard OIDC field names are used ("sub", "name", "email").
        FieldMapping *UserInfoFieldMappingRunConfig `json:"field_mapping,omitempty" yaml:"field_mapping,omitempty"`
}

// UserInfoFieldMappingRunConfig maps provider-specific field names to standard UserInfo fields.
// This allows adapting non-standard provider responses to the canonical UserInfo structure.
type UserInfoFieldMappingRunConfig struct {
        // SubjectFields is an ordered list of field names to try for the user ID.
        // The first non-empty value found will be used.
        // Default: ["sub"]
        SubjectFields []string `json:"subject_fields,omitempty" yaml:"subject_fields,omitempty"`

        // NameFields is an ordered list of field names to try for the display name.
        // The first non-empty value found will be used.
        // Default: ["name"]
        NameFields []string `json:"name_fields,omitempty" yaml:"name_fields,omitempty"`

        // EmailFields is an ordered list of field names to try for the email address.
        // The first non-empty value found will be used.
        // Default: ["email"]
        EmailFields []string `json:"email_fields,omitempty" yaml:"email_fields,omitempty"`
}

// UpstreamConfig wraps an upstream IDP configuration with identifying metadata.
// It supports both OIDC providers (with discovery) and pure OAuth 2.0 providers.
type UpstreamConfig struct {
        // Name uniquely identifies this upstream.
        // Used for routing decisions and session binding in multi-upstream scenarios.
        // If empty when only one upstream is configured, defaults to "default".
        Name string `json:"name,omitempty" yaml:"name,omitempty"`

        // DisplayName is an optional human-readable label for this upstream, shown to
        // end users on the consent page in place of Name. It is display-only: never
        // compared, stored, or used to route. It is sanitized at render time and never
        // rejected by validation. It is shown as trusted text, so it must come from
        // operator-controlled configuration only.
        DisplayName string `json:"display_name,omitempty" yaml:"display_name,omitempty"`

        // Type specifies the provider type: "oidc" or "oauth2".
        Type UpstreamProviderType `json:"type" yaml:"type"`

        // OAuth2Config contains OAuth 2.0 provider configuration.
        // Used when Type is "oauth2". Must be nil when Type is "oidc".
        OAuth2Config *upstream.OAuth2Config `json:"oauth2_config,omitempty" yaml:"oauth2_config,omitempty"`

        // OIDCConfig contains OIDC provider configuration (uses discovery).
        // Used when Type is "oidc". Must be nil when Type is "oauth2".
        OIDCConfig *upstream.OIDCConfig `json:"oidc_config,omitempty" yaml:"oidc_config,omitempty"`
}

// Config is the pure configuration for the OAuth authorization server.
// All values must be fully resolved (no file paths, no env vars).
// This is the interface that consumers should use to configure the server.
type Config struct {
        // Issuer is the issuer identifier for this authorization server.
        // This will be included in the "iss" claim of issued tokens.
        Issuer string

        // AuthorizationEndpointBaseURL overrides the base URL used for the authorization_endpoint
        // in the OAuth discovery document. When empty, defaults to Issuer.
        AuthorizationEndpointBaseURL string

        // KeyProvider provides signing keys for JWT operations.
        // Supports key rotation by returning multiple public keys for JWKS.
        // If nil, an ephemeral key will be auto-generated (development only).
        //
        // Production: Use keys.NewFileProvider() or keys.NewProviderFromConfig()
        // Testing: Use a mock or keys.NewGeneratingProvider()
        KeyProvider keys.KeyProvider

        // HMACSecrets contains the symmetric secrets used for signing authorization codes
        // and refresh tokens (opaque tokens). Unlike the asymmetric SigningKey which
        // signs JWTs for distributed verification, these secrets are used internally
        // by the authorization server only.
        // Current secret must be at least 32 bytes and cryptographically random.
        // Must be consistent across all replicas in multi-instance deployments.
        // Supports secret rotation via the Rotated field.
        HMACSecrets *servercrypto.HMACSecrets

        // AccessTokenLifespan is the duration that access tokens are valid.
        // If zero, defaults to 1 hour.
        AccessTokenLifespan time.Duration

        // RefreshTokenLifespan is the duration that refresh tokens are valid.
        // If zero, defaults to 7 days.
        RefreshTokenLifespan time.Duration

        // AuthCodeLifespan is the duration that authorization codes are valid.
        // If zero, defaults to 10 minutes.
        AuthCodeLifespan time.Duration

        // DelegationTokenLifespan is the maximum lifetime for delegated tokens issued
        // via RFC 8693 token exchange. The actual lifetime is the minimum of this value
        // and the subject token's remaining lifetime. If zero, defaults to 15 minutes.
        DelegationTokenLifespan time.Duration

        // Upstreams contains configurations for connecting to upstream IDPs.
        // At least one upstream is required - the server delegates authentication to the upstream IDP.
        // Multiple upstreams form a sequential authorization chain.
        Upstreams []UpstreamConfig

        // ConsentExemptClientIDs preauthorizes listed clients to skip ToolHive's consent page.
        //
        // Listing a client ID preauthorizes whatever registration currently holds that
        // ID, for whichever user completes upstream authentication. It is not per-user
        // consent, and runtime eligibility checks do not prove where the registration
        // came from.
        //
        // The embedder must control registration writes for exempt IDs. Replacing a
        // registration or expanding its permissions transfers or broadens the
        // preauthorization. Removing an ID affects new authorization requests only; it
        // does not revoke tokens or cancel already-approved flows. Upstream
        // authentication still runs.
        ConsentExemptClientIDs []string

        // UpstreamFactory, when set, is used instead of DefaultUpstreamFactory to
        // construct each configured upstream's OAuth2Provider. Two reasons to set it:
        //
        //   - Connection reuse. DefaultUpstreamFactory builds a private HTTP client
        //     per upstream, so a process that reconstructs the server to change its
        //     upstream set creates a fresh connection pool each time. A factory that
        //     passes a shared client via upstream.WithHTTPClient /
        //     upstream.WithOAuth2HTTPClient creates none. Key the sharing by
        //     distinct trust posture rather than by host: two upstreams can share a
        //     host while differing in CAFilePath, AllowPrivateIPs or
        //     InsecureAllowHTTP. An injected client stays owned by the caller — see
        //     upstream.IdleConnectionCloser — who is then responsible for its pool,
        //     its TLS trust and its dial-time SSRF guard.
        //   - Per-upstream failure isolation. Upstream construction inside New is
        //     otherwise all-or-nothing: one unreachable or mistyped issuer_url fails
        //     the whole call. A factory owning construction can apply a per-upstream
        //     deadline and substitute a provider for a single failing upstream while
        //     the rest serve.
        //
        // The factory cannot drop an upstream: every entry in Upstreams must yield a
        // provider, and returning a nil provider with a nil error is rejected by New
        // rather than producing an upstream that panics on the first authorization
        // request. To serve without an upstream, omit it from Upstreams; to keep its
        // slot in the chain, return a substitute provider.
        //
        // SECURITY: the returned provider validates the upstream's ID tokens and
        // resolves user identity. A factory that returns a permissive provider
        // bypasses upstream authentication for that leg of the chain. Delegate to
        // DefaultUpstreamFactory for upstreams you do not need to customize.
        UpstreamFactory UpstreamProviderFactory

        // RefreshGrantHandlers are extra token endpoint handlers that validate a
        // refresh_token grant. They run in order and only on that grant (never on
        // authorization_code, device-code, token exchange, revocation or
        // introspection requests), after the built-in refresh handler has
        // authenticated the client, loaded the refresh session and checked the
        // refresh token's expiry, scope and client binding. The requester's session
        // is then the stored refresh session (a *session.Session whose
        // GetIDPSessionID returns the token session ID, tsid). A handler whose
        // CanHandleTokenEndpointRequest returns false is skipped.
        //
        // Returning an error from HandleTokenEndpointRequest refuses the grant
        // before any tokens are issued and before the refresh token is rotated;
        // later handlers do not run. Return a fosite RFC 6749 error:
        // fosite.ErrInvalidGrant answers HTTP 400 invalid_grant, which sends the
        // client to a full authorization instead of looping on refresh. Another
        // fosite error answers with its own status and code. A plain Go error
        // answers HTTP 500 with fosite's generic "error" code, which clients treat
        // as a server failure and may retry. fosite.ErrUnknownRequest is treated as no objection,
        // not a refusal, so returning it lets the grant through. The built-in
        // upstream refresh guard already refuses a refresh whose upstream
        // credential is dead; use the seam for checks of your own. When deciding
        // from upstream.ErrGrantRejected, note its doc: under a cross-replica
        // refresh race it can occasionally be reported for a live credential.
        //
        // The seam only validates. PopulateTokenEndpointResponse and
        // CanSkipClientAuth are never called, so client authentication still
        // applies. Handlers must treat the requester as read-only: they receive the
        // live request the built-in handler has populated, and changing it would
        // change what is issued. Handlers are called concurrently and must be safe
        // for concurrent use.
        //
        // Handlers are fixed when New is called, before the server exists. A
        // handler that needs the server's upstream state receives it after New
        // returns, for example by wrapping UpstreamTokenRefresher and
        // IDPTokenStorage in a service it holds a pointer to. A nil entry,
        // including a nil pointer of a concrete type, is rejected by Validate.
        //
        // These handlers run after the built-in upstream refresh guard (see
        // DisableUpstreamRefreshGuard), so they see only grants it accepted.
        RefreshGrantHandlers []fosite.TokenEndpointHandler

        // DisableUpstreamRefreshGuard turns off the built-in refresh-grant guard.
        // When on (the default, with at least one upstream configured), a
        // refresh_token grant is refused with invalid_grant when the token
        // session's upstream credential is definitively dead, and its refresh
        // token revoked when the stored rows prove the credential can never be
        // used again. docs/arch/11-auth-server-storage.md ("Upstream refresh
        // guard") lists every cause and whether it revokes. The guard refreshes
        // expired rows, one provider at a time and within a 10-second budget,
        // before deciding, so an idle session whose upstream refresh succeeds is
        // granted. A transient failure (network, 5xx, 429, storage, the budget
        // running out) grants the refresh.
        //
        // Without the guard, a refreshed access token keeps the token session and
        // therefore its dead upstream credential, so a client that answers the
        // bearer validator's 401 by refreshing loops. The refusal sends it to a
        // full authorization instead. With the guard on, a refresh token stays
        // useful only while the upstream credential does.
        //
        // Only Go embedders can set this: the RunConfig, CLI and operator CRDs do
        // not expose it.
        DisableUpstreamRefreshGuard bool

        // UpstreamFilter, when set, narrows the upstream authorization chain after the
        // first leg resolves (see handlers.WithUpstreamFilter). When nil, all
        // configured upstreams are walked — the current behavior. Pass nil itself,
        // not a nil-valued concrete pointer implementing UpstreamFilter — a typed-nil
        // interface value is non-nil and will still be wired in. Has no effect with
        // fewer than 2 configured upstreams; Validate rejects that combination.
        UpstreamFilter handlers.UpstreamFilter

        // ScopesSupported lists the OAuth 2.0 scope values advertised in discovery documents.
        // If nil or empty, defaults to registration.DefaultScopes (["openid", "profile", "email", "offline_access"]).
        // This is advertised in /.well-known/openid-configuration and
        // /.well-known/oauth-authorization-server discovery endpoints.
        ScopesSupported []string

        // BaselineClientScopes is a baseline set of OAuth 2.0 scopes the embedded
        // DCR handler unions into every newly registered client's scope set. Empty
        // means current behavior is preserved (DCR registers exactly what the client
        // requested, or the intersection of registration.DefaultScopes with
        // ScopesSupported if the client requested none).
        // All entries must also be present in ScopesSupported. When ScopesSupported
        // is empty, the validation gate uses registration.DefaultScopes as the
        // superset — so standard OIDC scopes (e.g. "offline_access") work without
        // enumerating ScopesSupported explicitly.
        BaselineClientScopes []string

        // AllowedAudiences is the list of valid resource URIs that tokens can be issued for.
        // Per RFC 8707, the "resource" parameter in authorization and token requests is
        // validated against this list. MCP clients are required to include the resource
        // parameter, so this should be configured with the canonical URIs of all MCP servers
        // this authorization server issues tokens for.
        //
        // Security: An empty list means NO audiences are permitted (secure default).
        // When empty, any request with a "resource" parameter will be rejected with
        // "invalid_target". Configure this for proper MCP specification compliance.
        AllowedAudiences []string

        // SubjectIdentifierAudiences controls disclosure of the RFC 9493 sub_id claim and the
        // https://toolhive.dev/upstream_subject claim on login-derived access tokens.
        // Empty disables both claims. Entries must be exact members of
        // AllowedAudiences. It never grants an audience or client access.
        //
        // A login records the upstream identity only when this list is non-empty. Going
        // from an empty list to a non-empty one adds the claims only for sessions whose
        // login happened while the list was non-empty; other users get them from their
        // next login. Adding an audience to a list that is already non-empty applies at
        // the next refresh. Emptying the list stops the claims at the next refresh.
        SubjectIdentifierAudiences []string

        // AccessTokenType sets the JOSE "typ" header on every JWT access token this server
        // issues: authorization code, refresh, device code, token exchange, JWT-bearer and
        // ID-JAG redemption grants. Empty or "JWT" keeps fosite's default typ: JWT.
        // "at+jwt" types access tokens as RFC 9068 §2.1 requires. Some resource servers
        // reject at+jwt by default, such as Java ones built on Nimbus JOSE+JWT's default
        // type verifier, including Spring Security before 6.5. Enable it only when every
        // resource server accepts it. It is reachable from Go only.
        AccessTokenType string

        // CIMDEnabled enables the CIMD storage decorator so the authorization server
        // accepts HTTPS URLs as client_id values without prior DCR registration.
        //
        // A resolved CIMD client is write-through persisted into the underlying
        // storage (see CIMDStorageDecorator.fetch) so that token-endpoint session
        // rehydration finds it. That row is marked DCR-issued and therefore
        // expires (DefaultDCRClientTTL, currently 30 days), but disabling
        // CIMDEnabled does not itself evict or invalidate any row a prior enabled
        // period already persisted — GetClient is instead wrapped to refuse any
        // URL-shaped client_id outright while CIMDEnabled is false (see
        // decorateStorageForCIMD / storage.NewCIMDShapeGuardStorage), so such a
        // row simply cannot be resolved for as long as CIMD stays disabled.
        // Re-enabling CIMDEnabled before that TTL expires makes the stale
        // snapshot resolvable again without a fresh document fetch, until the row
        // expires or a new fetch overwrites it. Document rotation and revocation
        // generally follow the same rule: an existing persisted row keeps
        // authenticating from its stored snapshot, with no re-validation of
        // scopes, redirect URIs, or auth method, until it is naturally re-fetched
        // or its TTL lapses.
        CIMDEnabled bool

        // CIMDCacheMaxSize is the maximum number of CIMD documents held in the LRU
        // cache. Zero is replaced by a default (256) in applyDefaults when CIMDEnabled
        // is true.
        CIMDCacheMaxSize int

        // CIMDCacheFallbackTTL is the fixed TTL applied to all cached CIMD documents
        // (Cache-Control header parsing is not yet implemented). Zero is replaced by
        // a default (5 minutes) in applyDefaults when CIMDEnabled is true.
        CIMDCacheFallbackTTL time.Duration

        // InsecureAllowHTTP permits an http:// issuer URL for non-localhost hosts.
        // Only set this for in-cluster Kubernetes deployments on a trusted network.
        // Production deployments reachable outside the cluster MUST use https://.
        InsecureAllowHTTP bool

        // TrustedIssuers lists external OIDC issuers whose tokens are accepted as
        // RFC 8693 subject tokens or RFC 7523 JWT-bearer assertions. See the
        // identically named field on RunConfig for the full doc comment, including
        // the fail-closed AllowedActors semantics and the audience/scope constraints
        // operators must account for.
        TrustedIssuers []tokenexchange.TrustedIssuer

        // AuthorizationDetailsRegistry validates RFC 9396 authorization_details on
        // RFC 8693 token-exchange requests. It is Go-only and is not part of RunConfig.
        AuthorizationDetailsRegistry *authorizationdetails.Registry

        // AllowConfidentialClientRegistration permits DCR of confidential clients
        // (client_secret_basic / client_secret_post). See RunConfig for the full
        // semantics; disabling it does not revoke already-minted secrets.
        AllowConfidentialClientRegistration bool

        // AllowPrivateKeyJWTRegistration permits DCR of clients using
        // private_key_jwt authentication. See RunConfig for the full semantics.
        AllowPrivateKeyJWTRegistration bool

        // DeviceFlowEnabled enables the RFC 8628 device authorization grant. See
        // RunConfig.DeviceFlowEnabled for the full semantics.
        DeviceFlowEnabled bool

        // DeviceCodeInterval is the minimum time a device-flow client must wait
        // between polls of the token endpoint. If zero, defaults to
        // oauthserver.DefaultDeviceCodeInterval.
        //
        // Not yet exposed via RunConfig or set by embeddedauthserver.go's Config
        // construction -- only DeviceFlowEnabled is wired through today, so this
        // field is currently reachable only from tests that construct Config
        // directly. A future increment may add a RunConfig override.
        DeviceCodeInterval time.Duration

        // ForceConfidentialRedirectURIs lists redirect URIs that are always
        // registered as confidential clients, even when the DCR request declares
        // "none". See the identically named field on RunConfig for the full
        // semantics and rationale.
        ForceConfidentialRedirectURIs []string

        // InsecureAllowConfidentialOverLoopbackHTTP opts in to confidential clients
        // when Issuer is a plain-HTTP loopback URL. See the identically named field
        // on RunConfig for the full semantics and rationale.
        InsecureAllowConfidentialOverLoopbackHTTP bool

        // DelegateClients are pre-provisioned confidential OAuth clients in their
        // resolved runtime form. ClientSecret has already been read from its file
        // or environment-variable reference. See RunConfig.DelegateClients for the
        // serialized configuration.
        DelegateClients []DelegateClient

        // DisableTokenExchange prevents registration and advertisement of the RFC
        // 8693 grant. The zero value preserves the released behavior. It is set only
        // when canonical inbound_grants explicitly omits token_exchange.
        DisableTokenExchange bool

        // SPIFFETrust is the validated, immutable runtime SPIFFE trust model. It
        // must be constructed with NewSPIFFETrustConfig; a nil value means no
        // SPIFFE associations are configured. The serialized declarations live on
        // RunConfig and are converted at the RunConfig-to-Config boundary.
        SPIFFETrust *SPIFFETrustConfig

        // RegistrationRateLimit bounds /oauth/register. Nil keeps the default of
        // 1 registration per second with a burst of 5. See the identically named
        // field on RunConfig.
        RegistrationRateLimit *RegistrationRateLimitRunConfig
}

// DelegateClient is the resolved form of DelegateClientRunConfig: the secret
// has already been read from its file or environment variable reference.
type DelegateClient struct {
        // ClientID is the OAuth client_id this client presents at the token endpoint.
        ClientID string

        // ClientSecret is the resolved (plaintext) client secret.
        ClientSecret string //nolint:gosec // G117: field legitimately holds sensitive data

        // Scopes are the OAuth scopes this client may request.
        Scopes []string

        // Audiences are the RFC 8707 resource values this client may request a token for.
        Audiences []string
}

// Validate checks that the Config is valid.
func (c *Config) Validate() error {
        slog.Debug("validating authserver config", "issuer", c.Issuer)

        if err := c.validateIssuerAndConsentConfig(); err != nil {
                return err
        }

        // KeyProvider is optional - if nil, applyDefaults() will create a GeneratingProvider

        if c.HMACSecrets == nil {
                return fmt.Errorf("HMAC secrets are required")
        }
        if len(c.HMACSecrets.Current) < servercrypto.MinSecretLength {
                return fmt.Errorf("HMAC secret must be at least %d bytes", servercrypto.MinSecretLength)
        }

        if err := c.validateUpstreams(); err != nil {
                return err
        }
        c.warnOnAuthorizeCallbackHostMismatch()

        if err := c.validateUpstreamFilter(); err != nil {
                return err
        }

        for i, h := range c.RefreshGrantHandlers {
                if isNilHandler(h) {
                        return fmt.Errorf("refresh grant handler %d is nil", i)
                }
        }

        if err := c.validateAudienceConfiguration(); err != nil {
                return err
        }

        if err := oauthserver.ValidateAccessTokenType(c.AccessTokenType); err != nil {
                return err
        }

        // BaselineClientScopes must be a subset of ScopesSupported. RunConfig.Validate
        // catches this for the YAML-loaded path, but a caller that constructs Config
        // directly bypasses that; failing here gives them a clearer call stack than
        // the inner validateParams in the provider layer.
        // When ScopesSupported is empty, use DefaultScopes as the superset (matching
        // what applyDefaults substitutes at startup).
        if err := c.validateBaselineClientScopes(); err != nil {
                return err
        }

        if err := c.validateClientAdmission(); err != nil {
                return err
        }

        if err := c.validateDelegationTokenLifespan(); err != nil {
                return err
        }

        return c.validateDelegationAndTrustConfig()
}

// validateClientAdmission checks the bounds on how unauthenticated callers
// can create clients: CIMD resolution and the /oauth/register rate.
func (c *Config) validateClientAdmission() error {
        if err := c.validateCIMDBounds(); err != nil {
                return err
        }
        return c.RegistrationRateLimit.Validate()
}

// validateIssuerAndConsentConfig validates the issuer, client transport, and
// consent-exemption settings before the remaining runtime configuration checks.
func (c *Config) validateIssuerAndConsentConfig() error {
        if err := validateIssuerURL(c.Issuer, c.InsecureAllowHTTP); err != nil {
                return fmt.Errorf("issuer: %w", err)
        }
        if err := c.validateConfidentialClientConfig(); err != nil {
                return err
        }
        if err := c.validatePrivateKeyJWTRequiresTokenExchange(); err != nil {
                return err
        }
        if err := c.validateAuthorizationDetailsRequiresTokenExchange(); err != nil {
                return err
        }
        if c.AuthorizationEndpointBaseURL != "" {
                if err := validateIssuerURL(c.AuthorizationEndpointBaseURL, c.InsecureAllowHTTP); err != nil {
                        return fmt.Errorf("authorization_endpoint_base_url: %w", err)
                }
        }
        return consentpolicy.Validate(
                c.ConsentExemptClientIDs, c.AuthorizationEndpointBaseURL, c.Issuer, len(c.Upstreams))
}

// validateDelegationAndTrustConfig groups the delegate-client and SPIFFE
// trust checks. RunConfig.Validate() also runs these checks (see the
// comment there for why), while a caller that constructs Config directly
// bypasses them.
func (c *Config) validateDelegationAndTrustConfig() error {
        if err := c.validateDelegationConfig(); err != nil {
                return err
        }
        if err := validateConfigSPIFFENotYetEnforced(c.SPIFFETrust); err != nil {
                return err
        }
        c.warnTrustedIssuerAudiences()

        slog.Debug("authserver config validation passed",
                "issuer", c.Issuer,
                "upstream_count", len(c.Upstreams),
        )
        return nil
}

// validateAudienceConfiguration validates resource audiences and the independent
// opt-in policy for including upstream subject identifiers in access tokens.
func (c *Config) validateAudienceConfiguration() error {
        if len(c.AllowedAudiences) == 0 {
                return fmt.Errorf("at least one allowed audience is required for MCP compliance (RFC 8707 resource parameter validation)")
        }
        if err := validateAllowedAudiences(c.AllowedAudiences); err != nil {
                return err
        }
        return oauthserver.ValidateSubjectIdentifierAudiences(c.SubjectIdentifierAudiences, c.AllowedAudiences)
}

// validateBaselineClientScopes ensures every baseline scope is advertised by
// ScopesSupported. When it is empty, applyDefaults supplies DefaultScopes.
func (c *Config) validateBaselineClientScopes() error {
        effective := c.ScopesSupported
        if len(effective) == 0 {
                effective = registration.DefaultScopes
        }
        return registration.ValidateScopeSubset(c.BaselineClientScopes, effective, "baseline_client_scopes")
}

// validateDelegationConfig validates delegate clients and trusted issuers for
// callers that construct a runtime Config directly.
func (c *Config) validateDelegationConfig() error {
        if err := validateResolvedDelegateClients(c.DelegateClients, c.ScopesSupported, c.AllowedAudiences); err != nil {
                return err
        }
        return validateTrustedIssuers(c.TrustedIssuers, c.Issuer, c.AllowedAudiences)
}

// validateConfidentialClientConfig groups cleartext-transport validation for
// all confidential clients and the force-confidential-redirect-uris override
// behind a single call site, keeping Config.Validate's own cyclomatic
// complexity down.
func (c *Config) validateConfidentialClientConfig() error {
        if err := ValidateConfidentialClientTransport(
                c.AllowConfidentialClientRegistration || len(c.DelegateClients) > 0, c.InsecureAllowHTTP,
                c.Issuer, c.InsecureAllowConfidentialOverLoopbackHTTP); err != nil {
                return err
        }
        return ValidateForceConfidentialRedirectURIs(c.ForceConfidentialRedirectURIs, c.AllowConfidentialClientRegistration)
}

// validatePrivateKeyJWTRequiresTokenExchange rejects
// AllowPrivateKeyJWTRegistration combined with a disabled token-exchange
// grant. private_key_jwt registrations can only ever request the RFC 8693
// token-exchange grant (see validateGrantTypes/validateResponseTypes in
// pkg/authserver/server/registration/dcr.go): there is no independent use
// of this auth method outside token exchange. Allowing registration while
// token exchange is disabled would admit clients that can never
// successfully authenticate, so reject the combination outright rather
// than let it surface later as a confusing runtime rejection.
func (c *Config) validatePrivateKeyJWTRequiresTokenExchange() error {
        if c.AllowPrivateKeyJWTRegistration && c.DisableTokenExchange {
                return fmt.Errorf(
                        "allow_private_key_jwt_registration requires token exchange to be enabled: " +
                                "private_key_jwt registrations are token-exchange-only (RFC 8693)")
        }
        return nil
}

// validateAuthorizationDetailsRequiresTokenExchange rejects a registry combined
// with a disabled token-exchange grant: authorization_details is only handled
// on token exchange, so the registry could never be reached.
func (c *Config) validateAuthorizationDetailsRequiresTokenExchange() error {
        if c.AuthorizationDetailsRegistry != nil && c.DisableTokenExchange {
                return fmt.Errorf("authorization details registry requires token exchange to be enabled")
        }
        return nil
}

// validateCIMDBounds rejects invalid CIMD cache bounds when CIMD is enabled.
// When CIMD is disabled the cache fields are ignored.
func (c *Config) validateCIMDBounds() error {
        if !c.CIMDEnabled {
                return nil
        }
        if c.CIMDCacheMaxSize < 1 {
                return fmt.Errorf("cimd.cache_max_size must be >= 1 when CIMD is enabled")
        }
        if c.CIMDCacheFallbackTTL < 0 {
                return fmt.Errorf("cimd.cache_fallback_ttl must be non-negative when CIMD is enabled")
        }
        return nil
}

// validateDelegationTokenLifespan rejects negative or excessively long delegation
// token lifespans. Capped at oauthserver.MaxAccessTokenLifespan (the same ceiling
// the token-exchange Factory enforces) so validation and construction agree on a
// single source of truth — delegated tokens should be short-lived. Zero is
// accepted; applyDefaults substitutes the default.
func (c *Config) validateDelegationTokenLifespan() error {
        if c.DelegationTokenLifespan < 0 {
                return fmt.Errorf("delegation token lifespan must not be negative")
        }
        if c.DelegationTokenLifespan > oauthserver.MaxAccessTokenLifespan {
                return fmt.Errorf("delegation token lifespan must not exceed %v", oauthserver.MaxAccessTokenLifespan)
        }
        return nil
}

// validateTrustedIssuers checks every configured TrustedIssuer as early as
// RunConfig.Validate can catch it, so a bad actor_claim, actor_matcher, or
// duplicate issuer_url fails before buildUpstreamConfigs performs live RFC 7591
// registration against upstream IdPs — not after it, on a crash loop that
// orphans an upstream registration on every restart. Shared by
// RunConfig.Validate() and Config.Validate() (mirrors
// validateBaselineClientScopes); NewMultiIssuerTokenValidator repeats these
// checks again at server startup as defence in depth.
//
// issuer_url/jwks_url endpoint shape is validated by
// tokenexchange.ValidateTrustedIssuers itself (via validateTrustedIssuer),
// so both this path and MCPExternalAuthConfig's admission-time
// ValidateInboundGrants get it from the one place, rather than each
// re-implementing the same check.
func validateTrustedIssuers(issuers []tokenexchange.TrustedIssuer, selfIssuer string, allowedAudiences []string) error {
        if err := tokenexchange.ValidateTrustedIssuers(issuers, selfIssuer, allowedAudiences); err != nil {
                return fmt.Errorf("trusted_issuers: %w", err)
        }
        for _, issuer := range issuers {
                if issuer.JWTBearerGrant == nil {
                        continue
                }
                for _, binding := range issuer.JWTBearerGrant.SubjectBindings {
                        for _, resource := range binding.AllowedResources {
                                if !slices.Contains(allowedAudiences, resource) {
                                        return fmt.Errorf(
                                                "trusted_issuers: issuer_url %q: jwt_bearer_grant subject %q resource %q is not in allowed_audiences",
                                                issuer.IssuerURL, binding.Subject, resource)
                                }
                        }
                }
        }
        return nil
}

// warnTrustedIssuerAudiences logs a warning for each TrustedIssuer whose
// ExpectedAudience is absent from AllowedAudiences. This is not a hard
// error: a subject token may carry additional audiences beyond
// ExpectedAudience, so the mismatch is sufficient-but-not-necessary for
// every exchange from that issuer to fail — an operator may know a specific
// subject token will present a matching aud even though ExpectedAudience
// itself is not in AllowedAudiences. But when it's not intentional, this is
// the invalid_target footgun documented on the TrustedIssuers field: warn so
// it surfaces at startup instead of at the first exchange attempt.
func (c *Config) warnTrustedIssuerAudiences() {
        for _, ti := range c.TrustedIssuers {
                if !slices.Contains(c.AllowedAudiences, ti.ExpectedAudience) {
                        slog.Warn("trusted issuer's expected_audience is not in allowed_audiences; "+
                                "token exchange will fail with invalid_target unless subject tokens from it "+
                                "carry an additional audience matching one",
                                "issuer", ti.IssuerURL, "expected_audience", ti.ExpectedAudience)
                }
        }
}

// Validate checks that the OIDCUpstreamRunConfig has either a pre-provisioned
// client ID or a DCR configuration, but not both. OIDC DCR may omit both endpoint
// selectors because discovery is derived from IssuerURL.
func (c *OIDCUpstreamRunConfig) Validate() error {
        hasClientID := c.ClientID != ""
        hasDCR := c.DCRConfig != nil
        if hasClientID == hasDCR {
                return fmt.Errorf("oidc upstream: exactly one of client_id or dcr_config must be set")
        }
        if !hasDCR {
                return nil
        }
        if c.DCRConfig.DiscoveryURL != "" && c.DCRConfig.RegistrationEndpoint != "" {
                return fmt.Errorf("oidc upstream: dcr_config discovery_url and registration_endpoint are mutually exclusive")
        }
        if c.DCRConfig.DiscoveryURL == "" && c.DCRConfig.RegistrationEndpoint == "" {
                if c.DCRConfig.InitialAccessTokenFile != "" && c.DCRConfig.InitialAccessTokenEnvVar != "" {
                        return fmt.Errorf(
                                "oidc upstream: dcr_config initial_access_token_file and initial_access_token_env_var are mutually exclusive")
                }
                if c.IssuerURL == "" {
                        return fmt.Errorf("oidc upstream: issuer_url is required when dcr_config omits discovery_url and registration_endpoint")
                }
                return nil
        }
        if err := c.DCRConfig.Validate(); err != nil {
                return fmt.Errorf("oidc upstream: invalid dcr_config: %w", err)
        }
        return nil
}

// Validate checks that the OAuth2UpstreamRunConfig is internally consistent.
// It enforces the mutual exclusivity of ClientID and DCRConfig: exactly one must
// be set. A ClientID is required for pre-provisioned clients; a DCRConfig is
// required when client credentials are obtained at runtime via RFC 7591
// Dynamic Client Registration. When DCRConfig is present, its own validity is
// also checked via DCRUpstreamConfig.Validate.
//
// Validate intentionally does not verify fields handled by the shared
// CommonOAuthConfig or upstream.OAuth2Config validators — it only covers the
// run-config surface area unique to OAuth2UpstreamRunConfig.
//
// Called from buildPureOAuth2Config at the RunConfig → upstream.OAuth2Config
// conversion boundary so that DCR-specific fields are validated before they
// are dropped during conversion.
func (c *OAuth2UpstreamRunConfig) Validate() error {
        hasClientID := c.ClientID != ""
        hasDCR := c.DCRConfig != nil
        switch {
        case !hasClientID && !hasDCR:
                return fmt.Errorf("oauth2 upstream: either client_id or dcr_config is required")
        case hasClientID && hasDCR:
                return fmt.Errorf("oauth2 upstream: client_id and dcr_config are mutually exclusive")
        }

        if hasDCR {
                if err := c.DCRConfig.Validate(); err != nil {
                        return fmt.Errorf("oauth2 upstream: invalid dcr_config: %w", err)
                }

                // When the operator configures DCRConfig.RegistrationEndpoint, the
                // resolver bypasses discovery and therefore cannot populate
                // AuthorizationEndpoint or TokenEndpoint from server metadata. The
                // run-config must supply both explicitly or the upstream is
                // unusable: registration would succeed and the first authorize or
                // token-exchange call would silently fail with empty endpoints.
                // Discovery flow (DCRConfig.DiscoveryURL) is unaffected — those
                // fields populate from metadata.
                if c.DCRConfig.RegistrationEndpoint != "" {
                        if c.AuthorizationEndpoint == "" || c.TokenEndpoint == "" {
                                return fmt.Errorf(
                                        "oauth2 upstream: authorization_endpoint and token_endpoint are required " +
                                                "when dcr_config.registration_endpoint is set (no discovery to populate them)")
                        }
                }
        }

        if c.IdentityFromToken != nil && c.IdentityFromToken.SubjectPath == "" {
                return fmt.Errorf("oauth2 upstream: identity_from_token.subject_path must not be empty when identity_from_token is configured")
        }

        return c.validateTokenEndpointAuthMethod()
}

// validateTokenEndpointAuthMethod checks TokenEndpointAuthMethod against the
// set of methods buildPureOAuth2Config/authStyleFromMethod support, and that
// the method is consistent with whether a client secret source is configured.
// Split out of Validate to keep that method's cyclomatic complexity down.
func (c *OAuth2UpstreamRunConfig) validateTokenEndpointAuthMethod() error {
        hasSecretSource := c.ClientSecretFile != "" || c.ClientSecretEnvVar != ""

        switch c.TokenEndpointAuthMethod {
        case "":
                // Empty means the historical POST-body default; buildPureOAuth2Config
                // passes it through unchanged, it does not resolve to a confidential method.
        case oauthproto.TokenEndpointAuthMethodNone:
                if hasSecretSource {
                        return fmt.Errorf("oauth2 upstream: token_endpoint_auth_method none cannot be used with a client secret")
                }
        case oauthproto.TokenEndpointAuthMethodClientSecretBasic, oauthproto.TokenEndpointAuthMethodClientSecretPost:
                if !hasSecretSource {
                        return fmt.Errorf(
                                "oauth2 upstream: token_endpoint_auth_method %q requires client_secret_file or client_secret_env_var",
                                c.TokenEndpointAuthMethod)
                }
        default:
                return fmt.Errorf("oauth2 upstream: unsupported token_endpoint_auth_method %q", c.TokenEndpointAuthMethod)
        }

        return nil
}

// Validate checks that the DCRUpstreamConfig specifies exactly one of
// DiscoveryURL or RegistrationEndpoint, that the configured URL is well-formed
// and uses HTTPS (or http on a loopback host for local development), and that
// the two initial-access-token sources (InitialAccessTokenFile and
// InitialAccessTokenEnvVar) are not both set.
//
// DiscoveryURL triggers runtime resolution of the registration endpoint via
// RFC 8414 / OIDC Discovery; RegistrationEndpoint bypasses discovery for
// providers that do not publish metadata. Requiring exactly one prevents
// ambiguity about which URL the authserver should contact for registration.
//
// URL well-formedness and HTTPS are enforced here at the schema-validation
// boundary so misconfiguration fails fast at startup rather than at first DCR
// attempt; the runtime callers (pkg/oauthproto/discovery.go and
// pkg/oauthproto/dcr.go) defend in depth, but this is the natural fail-fast
// point.
//
// Rejecting a config that supplies both an InitialAccessTokenFile and an
// InitialAccessTokenEnvVar prevents a credential-rotation footgun: if both
// were accepted, an operator updating the env-var value would not realize
// the file source still wins (or vice versa) and would silently keep
// presenting a stale token at registration.
func (c *DCRUpstreamConfig) Validate() error {
        hasDiscovery := c.DiscoveryURL != ""
        hasRegistration := c.RegistrationEndpoint != ""
        switch {
        case !hasDiscovery && !hasRegistration:
                return fmt.Errorf("dcr_config: either discovery_url or registration_endpoint is required")
        case hasDiscovery && hasRegistration:
                return fmt.Errorf("dcr_config: discovery_url and registration_endpoint are mutually exclusive")
        case hasDiscovery:
                if err := networking.ValidateEndpointURL(c.DiscoveryURL); err != nil {
                        return fmt.Errorf("dcr_config: invalid discovery_url: %w", err)
                }
        case hasRegistration:
                if err := networking.ValidateEndpointURL(c.RegistrationEndpoint); err != nil {
                        return fmt.Errorf("dcr_config: invalid registration_endpoint: %w", err)
                }
        }

        if c.InitialAccessTokenFile != "" && c.InitialAccessTokenEnvVar != "" {
                return fmt.Errorf(
                        "dcr_config: initial_access_token_file and initial_access_token_env_var are mutually exclusive")
        }

        // Mirrors the CRD's Pattern/MaxLength markers on rotationKey for the
        // YAML and CLI run-config paths, which have no admission-time schema.
        if c.RotationKey != "" && !dcrRotationKeyPattern.MatchString(c.RotationKey) {
                return fmt.Errorf("dcr_config: rotation_key must be 1-64 characters from [A-Za-z0-9._-]")
        }

        // Mirrors the CRD's Enum marker on registrationProbe; the resolver owns
        // the set of strategies.
        if err := dcr.RegistrationProbe(c.RegistrationProbe).Validate(); err != nil {
                return fmt.Errorf("dcr_config: %w", err)
        }

        return nil
}

// dcrRotationKeyPattern bounds DCRUpstreamConfig.RotationKey to a short,
// log- and key-safe alphabet; the CRD enforces the same shape at admission.
var dcrRotationKeyPattern = regexp.MustCompile(`^[A-Za-z0-9._-]{1,64}$`)

func validateZeroUpstreamMode(upstreamCount, delegateClientCount int, issuers []tokenexchange.TrustedIssuer) error {
        if upstreamCount > 0 || delegateClientCount > 0 || JWTBearerGrantEnabled(issuers) {
                return nil
        }
        return fmt.Errorf("at least one upstream is required unless delegate clients or a " +
                "trusted issuer with JWT bearer grant is configured")
}

// validateUpstreams validates the upstream configurations.
func (c *Config) validateUpstreams() error {
        // Track names for uniqueness checking
        seenNames := make(map[string]bool)

        for i := range c.Upstreams {
                up := &c.Upstreams[i]

                if err := c.validateUpstreamName(i, up); err != nil {
                        return err
                }

                // Check for duplicate names
                if seenNames[up.Name] {
                        return fmt.Errorf("duplicate upstream name: %q", up.Name)
                }
                seenNames[up.Name] = true

                if err := validateUpstreamType(up, c.InsecureAllowHTTP); err != nil {
                        return err
                }
        }

        return validateZeroUpstreamMode(len(c.Upstreams), len(c.DelegateClients), c.TrustedIssuers)
}

// warnOnAuthorizeCallbackHostMismatch logs a WARN for every upstream whose
// redirect_uri the browser-binding cookie cannot reach from the browser-facing
// authorize URL (AuthorizationEndpointBaseURL, else Issuer). /oauth/callback
// only completes a login for the browser that started it at /oauth/authorize
// or /oauth/device, and that binding is a host-only cookie that is Secure when
// the authorize URL is https (see handlers/browser_binding.go). It therefore
// never reaches a callback on another hostname, nor a plain-http callback on a
// non-loopback host when the cookie is Secure, and every browser login through
// that upstream is rejected at the callback. Ports are ignored, as cookies
// ignore them, and browsers do send Secure cookies to http://localhost.
//
// This is a WARN rather than a validation error: the auth server shares a
// process with the MCP proxy, and refusing to start would also take down token
// refresh and token exchange over a misconfiguration that only affects
// interactive login.
func (c *Config) warnOnAuthorizeCallbackHostMismatch() {
        authorizeBase := c.AuthorizationEndpointBaseURL
        if authorizeBase == "" {
                authorizeBase = c.Issuer
        }
        authorizeURL, err := url.Parse(authorizeBase)
        if err != nil {
                // Already rejected by validateIssuerURL above.
                return
        }
        for i := range c.Upstreams {
                up := &c.Upstreams[i]
                callbackURL, err := url.Parse(up.redirectURI())
                if err != nil || callbackURL.Hostname() == "" {
                        // validateUpstreamType has already rejected an unusable redirect_uri.
                        continue
                }
                authorizeHost, callbackHost := authorizeURL.Hostname(), callbackURL.Hostname()
                switch {
                case !strings.EqualFold(authorizeHost, callbackHost):
                        slog.Warn("upstream redirect_uri host differs from the browser-facing authorize host; "+
                                "/oauth/callback binds each login to the browser that started it with a host-only cookie, "+
                                "so browser logins through this upstream will be rejected until the two hosts match",
                                "upstream", up.Name,
                                "authorize_host", authorizeHost,
                                "callback_host", callbackHost,
                        )
                case strings.EqualFold(authorizeURL.Scheme, "https") &&
                        strings.EqualFold(callbackURL.Scheme, "http") &&
                        !networking.IsLocalhost(callbackHost):
                        slog.Warn("upstream redirect_uri is plain http while the browser-facing authorize URL is https; "+
                                "the browser-binding cookie is Secure and browsers will not send it to the callback, "+
                                "so browser logins through this upstream will be rejected until the callback uses https",
                                "upstream", up.Name,
                                "authorize_url", authorizeBase,
                                "callback_url", up.redirectURI(),
                        )
                }
        }
}

// redirectURI returns the upstream's configured callback URL for either
// provider type, or "" when no provider config is set.
func (u *UpstreamConfig) redirectURI() string {
        switch {
        case u.OAuth2Config != nil:
                return u.OAuth2Config.RedirectURI
        case u.OIDCConfig != nil:
                return u.OIDCConfig.RedirectURI
        }
        return ""
}

// validateUpstreamFilter rejects an UpstreamFilter configured with fewer than
// 2 upstreams. handlers.computeChain consults the filter only when there is a
// non-first upstream to narrow, so with a single upstream the filter would
// silently never be invoked; this fails loudly instead of letting it no-op
// without any indication to the caller.
func (c *Config) validateUpstreamFilter() error {
        if c.UpstreamFilter != nil && len(c.Upstreams) < 2 {
                return fmt.Errorf("upstream_filter is configured but has no effect with fewer than 2 upstreams")
        }
        return nil
}

// validateUpstreamName validates and defaults the upstream name.
// For single upstream, empty names default to "default".
// For multi-upstream, explicit non-"default" names are required.
func (c *Config) validateUpstreamName(i int, up *UpstreamConfig) error {
        if len(c.Upstreams) == 1 {
                if up.Name == "" {
                        up.Name = DefaultUpstreamName
                }
        } else {
                if up.Name == "" {
                        return fmt.Errorf(
                                "upstream[%d]: name must be explicitly set when multiple upstreams are configured", i)
                }
                if up.Name == DefaultUpstreamName {
                        return fmt.Errorf(
                                "upstream[%d]: name %q is reserved for single-upstream configs; use a descriptive name",
                                i, up.Name)
                }
        }

        // Validate name format (DNS-label-like) to prevent storage key injection
        if !upstreamNameRegex.MatchString(up.Name) {
                return fmt.Errorf(
                        "upstream[%d]: name %q must match %s (lowercase alphanumeric and hyphens)",
                        i, up.Name, upstreamNameRegex.String())
        }

        return nil
}

// validateUpstreamType validates the provider type and its type-specific config.
func validateUpstreamType(up *UpstreamConfig, insecureAllowHTTP bool) error {
        switch up.Type {
        case UpstreamProviderTypeOIDC:
                if up.OIDCConfig == nil {
                        return fmt.Errorf("upstream %q: oidc_config is required for OIDC provider", up.Name)
                }
                if up.OAuth2Config != nil {
                        return fmt.Errorf("upstream %q: oauth2_config must not be set when type is %q", up.Name, up.Type)
                }
                if err := up.OIDCConfig.ValidateWithInsecure(insecureAllowHTTP || up.OIDCConfig.InsecureAllowHTTP); err != nil {
                        return fmt.Errorf("upstream %q: %w", up.Name, err)
                }
        case UpstreamProviderTypeOAuth2:
                if up.OAuth2Config == nil {
                        return fmt.Errorf("upstream %q: oauth2_config is required for OAuth2 provider", up.Name)
                }
                if up.OIDCConfig != nil {
                        return fmt.Errorf("upstream %q: oidc_config must not be set when type is %q", up.Name, up.Type)
                }
                if err := up.OAuth2Config.ValidateWithInsecure(insecureAllowHTTP || up.OAuth2Config.InsecureAllowHTTP); err != nil {
                        return fmt.Errorf("upstream %q: %w", up.Name, err)
                }
        default:
                return fmt.Errorf("upstream %q: unsupported provider type: %q", up.Name, up.Type)
        }
        return nil
}

// applyDefaults applies default values to the config where not set.
func (c *Config) applyDefaults() error {
        slog.Debug("applying default values to authserver config")

        if c.AccessTokenLifespan == 0 {
                c.AccessTokenLifespan = time.Hour
                slog.Debug("applied default access token lifespan", "duration", c.AccessTokenLifespan)
        }
        if c.RefreshTokenLifespan == 0 {
                c.RefreshTokenLifespan = 24 * time.Hour * 7 // 7 days
                slog.Debug("applied default refresh token lifespan", "duration", c.RefreshTokenLifespan)
        }
        if c.AuthCodeLifespan == 0 {
                c.AuthCodeLifespan = 10 * time.Minute
                slog.Debug("applied default auth code lifespan", "duration", c.AuthCodeLifespan)
        }
        if c.DelegationTokenLifespan == 0 {
                c.DelegationTokenLifespan = 15 * time.Minute
                slog.Debug("applied default delegation token lifespan", "duration", c.DelegationTokenLifespan)
        }
        if c.HMACSecrets == nil {
                secret := make([]byte, servercrypto.MinSecretLength)
                if _, err := rand.Read(secret); err != nil {
                        return fmt.Errorf("failed to generate HMAC secret: %w", err)
                }
                c.HMACSecrets = &servercrypto.HMACSecrets{Current: secret}
                slog.Warn("no HMAC secrets configured, generating ephemeral secret",
                        "warning", "auth codes and refresh tokens will be invalid after restart")
        }
        if c.KeyProvider == nil {
                c.KeyProvider = keys.NewGeneratingProvider(keys.DefaultAlgorithm)
                slog.Warn("no key provider configured, using ephemeral signing key",
                        "warning", "JWTs will be invalid after restart")
        }
        if len(c.ScopesSupported) == 0 {
                c.ScopesSupported = registration.DefaultScopes
                slog.Debug("applied default scopes_supported", "scopes", c.ScopesSupported)
        }
        // One-time operator-facing signal for the omitted-scope registration
        // behavior (see registration.ValidateScopes and issue #6186). The outcome
        // is fully determined by ScopesSupported, so it is logged once here
        // rather than on every unauthenticated /oauth/register call or CIMD
        // resolution — those record it at Debug.
        if intersection, dropped, dcrErr := registration.ValidateScopes(nil, c.ScopesSupported); dcrErr != nil {
                slog.Warn("scopes_supported shares no scopes with the default set; "+
                        "clients that omit scope will be rejected and must declare scope explicitly",
                        "scopes_supported", c.ScopesSupported,
                        "default_scopes", registration.DefaultScopes)
        } else if len(dropped) > 0 {
                slog.Info("scopes_supported does not cover the default scope set; "+
                        "clients that omit scope will register with the intersection",
                        "intersection", intersection,
                        "dropped_defaults", dropped)
        }
        if c.CIMDEnabled && c.CIMDCacheMaxSize == 0 {
                c.CIMDCacheMaxSize = 256
                slog.Debug("applied default cimd cache_max_size", "size", c.CIMDCacheMaxSize)
        }
        if c.CIMDEnabled && c.CIMDCacheFallbackTTL == 0 {
                c.CIMDCacheFallbackTTL = 5 * time.Minute
                slog.Debug("applied default cimd cache_fallback_ttl", "ttl", c.CIMDCacheFallbackTTL)
        }
        return nil
}

// ValidateConfidentialClientTransport rejects cleartext HTTP configurations
// when any confidential client is enabled. It delegates to the server-layer
// validator so direct AuthorizationServerParams construction cannot bypass the
// same transport policy.
func ValidateConfidentialClientTransport(
        allowConfidential, insecureAllowHTTP bool,
        issuer string, insecureAllowConfidentialOverLoopbackHTTP bool,
) error {
        return oauthserver.ValidateConfidentialClientTransport(
                allowConfidential, insecureAllowHTTP, issuer, insecureAllowConfidentialOverLoopbackHTTP)
}

// ValidateForceConfidentialRedirectURIs rejects a misconfigured
// force_confidential_redirect_uris list: a non-empty list requires
// allowConfidential (there is no confidential-client path to force a
// registration onto otherwise), and every entry must be a valid https
// non-loopback redirect URI. The https-non-loopback requirement mirrors the
// restriction validateAuthMethod (pkg/authserver/server/registration/dcr.go)
// already applies to ordinary confidential DCR: a redirect URI reachable on
// loopback or a private scheme is by construction a public client (OAuth 2.1
// §2.1), and this override must not be a way to punch through that
// restriction and hand a distributed native app a secret.
func ValidateForceConfidentialRedirectURIs(uris []string, allowConfidential bool) error {
        if len(uris) == 0 {
                return nil
        }
        if !allowConfidential {
                return fmt.Errorf(
                        "force_confidential_redirect_uris requires allow_confidential_client_registration to be true")
        }
        for _, uri := range uris {
                if err := oauthproto.ValidateRedirectURI(uri, oauthproto.RedirectURIPolicyStrict); err != nil {
                        return fmt.Errorf("force_confidential_redirect_uris: %q: %w", uri, err)
                }
                parsed, err := url.Parse(uri)
                if err != nil {
                        return fmt.Errorf("force_confidential_redirect_uris: %q: invalid URL: %w", uri, err)
                }
                if networking.IsLocalhost(parsed.Hostname()) {
                        return fmt.Errorf(
                                "force_confidential_redirect_uris: %q must not be a loopback redirect URI; "+
                                        "a loopback client is a public client by construction and must not be issued a secret", uri)
                }
        }
        return nil
}

// validateIssuerURL validates that the issuer is a valid URL.
// Per OIDC Core Section 3.1.2.1 and RFC 8414 Section 2, the issuer
// MUST use the "https" scheme, except for localhost during development.
// When insecureAllowHTTP is true, http:// is also permitted for non-localhost
// hosts (for in-cluster Kubernetes deployments on trusted networks).
//
// This server's own issuer is additionally held to a no-trailing-slash rule
// that OIDC itself does not require; we control this value, unlike trusted
// external issuers.
func validateIssuerURL(issuer string, insecureAllowHTTP bool) error {
        return validateIssuerURLCore(issuer, insecureAllowHTTP)
}

// validateIssuerURLCore validates this authorization server's issuer.
func validateIssuerURLCore(issuer string, insecureAllowHTTP bool) error {
        if issuer == "" {
                return fmt.Errorf("issuer is required")
        }

        parsed, err := url.Parse(issuer)
        if err != nil {
                return errors.New("invalid URL")
        }

        if parsed.Scheme == "" {
                return fmt.Errorf("scheme is required")
        }

        if parsed.Hostname() == "" {
                return fmt.Errorf("host is required")
        }

        // Per RFC 8414 Section 2, the issuer identifier has no query or fragment components
        if parsed.RawQuery != "" {
                return fmt.Errorf("must not contain query component")
        }
        if parsed.Fragment != "" {
                return fmt.Errorf("must not contain fragment component")
        }
        // Userinfo is rejected on two independent grounds. It cannot work: OIDC
        // Discovery 1.0 Section 4.3 compares the discovery document's "issuer"
        // against this value by exact string match, and no provider echoes back
        // embedded credentials, so such an issuer always fails discovery. And it
        // must not be stored: a password here would sit in the RunConfig and could
        // be exposed by callers that log the configured endpoint. Rejecting it
        // outright beats relying on every caller to redact it.
        // Note that parsed.User is non-nil even for "https://user@host" with no
        // password, which is equally unusable as an issuer identifier.
        if parsed.User != nil {
                return fmt.Errorf("must not contain userinfo (credentials in the URL)")
        }

        // HTTPS is required unless it's a loopback address (for development, and
        // only when localhostExempt) or insecureAllowHTTP is explicitly set.
        if parsed.Scheme != "https" {
                if parsed.Scheme != "http" {
                        return fmt.Errorf("scheme must be https (or http for localhost)")
                }
                if !insecureAllowHTTP && !networking.IsLocalhost(parsed.Host) {
                        return fmt.Errorf("http scheme is only allowed for localhost, use https")
                }
        }

        // ToolHive's own issuer is held to this stricter, self-imposed rule
        // since we control the value.
        if strings.HasSuffix(issuer, "/") {
                return fmt.Errorf("must not have trailing slash")
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package identity defines the authserver's credential-free canonical-user context contract.
package identity

import "context"

// CanonicalUserContextKey identifies the canonical-user value in a context.
// It is distinct from any authenticated-principal context key.
type CanonicalUserContextKey struct{}

// WithCanonicalUser stores a canonical user identifier in ctx.
// An empty userID leaves ctx unchanged.
func WithCanonicalUser(ctx context.Context, userID string) context.Context {
        if userID == "" {
                return ctx
        }
        return context.WithValue(ctx, CanonicalUserContextKey{}, userID)
}

// CanonicalUserFromContext retrieves a canonical user identifier directly from
// the canonical-user context key. It does not infer a user from an identity.
// A directly stored empty string is present and returns ("", true).
func CanonicalUserFromContext(ctx context.Context) (string, bool) {
        userID, ok := ctx.Value(CanonicalUserContextKey{}).(string)
        return userID, ok
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "fmt"
        "slices"

        tx "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange"
)

// TokenExchangeInboundGrantRunConfig configures RFC 8693 inbound clients and issuer policies.
type TokenExchangeInboundGrantRunConfig struct {
        DelegateClients []DelegateClientRunConfig            `json:"delegate_clients,omitempty" yaml:"delegate_clients,omitempty"`
        IssuerPolicies  []TokenExchangeIssuerPolicyRunConfig `json:"issuer_policies,omitempty" yaml:"issuer_policies,omitempty"`
}

// TokenExchangeIssuerPolicyRunConfig binds RFC 8693 policy to a trusted issuer declaration.
type TokenExchangeIssuerPolicyRunConfig struct {
        IssuerRef              string   `json:"issuer_ref" yaml:"issuer_ref"`
        ExpectedAudience       string   `json:"expected_audience" yaml:"expected_audience"`
        ActorClaim             string   `json:"actor_claim,omitempty" yaml:"actor_claim,omitempty"`
        AllowedActors          []string `json:"allowed_actors,omitempty" yaml:"allowed_actors,omitempty"`
        ActorMatcher           string   `json:"actor_matcher,omitempty" yaml:"actor_matcher,omitempty"`
        AllowedDelegateClients []string `json:"allowed_delegate_clients" yaml:"allowed_delegate_clients"`
        AllowMayAct            bool     `json:"allow_may_act,omitempty" yaml:"allow_may_act,omitempty"`
}

// JWTBearerInboundGrantRunConfig configures RFC 7523 issuer policies.
type JWTBearerInboundGrantRunConfig struct {
        IssuerPolicies []JWTBearerIssuerPolicyRunConfig `json:"issuer_policies,omitempty" yaml:"issuer_policies,omitempty"`
}

// JWTBearerIssuerPolicyRunConfig binds RFC 7523 policy to a trusted issuer declaration.
type JWTBearerIssuerPolicyRunConfig struct {
        IssuerRef         string                       `json:"issuer_ref" yaml:"issuer_ref"`
        MaxAssertionAge   string                       `json:"max_assertion_age" yaml:"max_assertion_age"`
        SubjectBindings   []tx.JWTBearerSubjectBinding `json:"subject_bindings" yaml:"subject_bindings"`
        AcceptedAudiences []string                     `json:"accepted_audiences,omitempty" yaml:"accepted_audiences,omitempty"`
}

// InboundGrantCapabilities reports the effective grant families after normalization.
type InboundGrantCapabilities struct {
        TokenExchange bool
        JWTBearer     bool
}

// DeprecatedFieldPath identifies a populated legacy field and its canonical replacement.
type DeprecatedFieldPath struct {
        Path        string
        Replacement string
}

// NormalizedInboundGrants contains copied effective inputs consumed by existing runtime validators and registration.
type NormalizedInboundGrants struct {
        DelegateClients  []DelegateClientRunConfig
        TrustedIssuers   []tx.TrustedIssuer
        Capabilities     InboundGrantCapabilities
        DeprecatedFields []DeprecatedFieldPath
}

// NormalizeInboundGrants converts legacy and canonical declarations to the existing effective runtime structures.
// It never mutates cfg or any nested caller-owned slice.
func NormalizeInboundGrants(cfg *RunConfig) (*NormalizedInboundGrants, error) {
        if cfg == nil {
                return nil, fmt.Errorf("config is required")
        }

        result := &NormalizedInboundGrants{
                DelegateClients: cloneDelegateClients(cfg.DelegateClients),
                TrustedIssuers:  cloneTrustedIssuers(cfg.TrustedIssuers),
                Capabilities: InboundGrantCapabilities{
                        TokenExchange: cfg.InboundGrants == nil,
                },
        }
        issuerByName, err := indexTrustedIssuers(result.TrustedIssuers)
        if err != nil {
                return nil, err
        }

        legacyTokenExchange := len(cfg.DelegateClients) > 0
        for i := range result.TrustedIssuers {
                issuer := &result.TrustedIssuers[i]
                if hasLegacyTokenExchangePolicy(*issuer) {
                        legacyTokenExchange = true
                        result.DeprecatedFields = append(result.DeprecatedFields, DeprecatedFieldPath{
                                Path: fmt.Sprintf("trusted_issuers[%d]", i), Replacement: "inbound_grants.token_exchange.issuer_policies",
                        })
                }
                if issuer.JWTBearerGrant != nil {
                        result.Capabilities.JWTBearer = true
                        result.DeprecatedFields = append(result.DeprecatedFields, DeprecatedFieldPath{
                                Path: fmt.Sprintf("trusted_issuers[%d].jwt_bearer_grant", i), Replacement: "inbound_grants.jwt_bearer.issuer_policies",
                        })
                }
        }
        if len(cfg.DelegateClients) > 0 {
                result.DeprecatedFields = append(result.DeprecatedFields, DeprecatedFieldPath{
                        Path: "delegate_clients", Replacement: "inbound_grants.token_exchange.delegate_clients",
                })
        }
        if cfg.InboundGrants == nil {
                return result, nil
        }

        if cfg.InboundGrants.TokenExchange != nil {
                if legacyTokenExchange {
                        return nil, fmt.Errorf(
                                "inbound_grants.token_exchange conflicts with legacy delegate_clients " +
                                        "or RFC 8693 policy in trusted_issuers",
                        )
                }
                result.Capabilities.TokenExchange = true
                result.DelegateClients = cloneDelegateClients(cfg.InboundGrants.TokenExchange.DelegateClients)
                if err := applyTokenExchangePolicies(result.TrustedIssuers, issuerByName,
                        cfg.InboundGrants.TokenExchange.IssuerPolicies); err != nil {
                        return nil, err
                }
        } else {
                result.Capabilities.TokenExchange = legacyTokenExchange
        }
        // SPIFFE client-auth associations always require the token-exchange grant
        // (SPIFFEGrantTypeTokenExchange is the only grant type they may declare —
        // see validateSPIFFEGrants), independent of the legacy/canonical
        // token-exchange projection above. A SPIFFE-only configuration must not
        // leave Capabilities.TokenExchange false: that would disable the RFC 8693
        // grant handler server-wide and reject every SPIFFE client's own token
        // requests before authentication is even checked.
        if len(cfg.InboundGrants.SPIFFEClientAuth) > 0 {
                result.Capabilities.TokenExchange = true
        }

        if cfg.InboundGrants.JWTBearer != nil {
                if result.Capabilities.JWTBearer {
                        return nil, fmt.Errorf("inbound_grants.jwt_bearer conflicts with legacy trusted_issuers[*].jwt_bearer_grant")
                }
                result.Capabilities.JWTBearer = true
                if err := applyJWTBearerPolicies(result.TrustedIssuers, issuerByName,
                        cfg.InboundGrants.JWTBearer.IssuerPolicies); err != nil {
                        return nil, err
                }
        }
        return result, nil
}

func indexTrustedIssuers(issuers []tx.TrustedIssuer) (map[string]int, error) {
        byName := make(map[string]int, len(issuers))
        byURL := make(map[string]int, len(issuers))
        for i, issuer := range issuers {
                if previous, ok := byURL[issuer.IssuerURL]; ok {
                        return nil, fmt.Errorf(
                                "trusted_issuers[%d].issuer_url duplicates trusted_issuers[%d].issuer_url (configured more than once)",
                                i,
                                previous,
                        )
                }
                byURL[issuer.IssuerURL] = i
                if issuer.Name == "" {
                        continue
                }
                if previous, ok := byName[issuer.Name]; ok {
                        return nil, fmt.Errorf("trusted_issuers[%d].name duplicates trusted_issuers[%d].name %q", i, previous, issuer.Name)
                }
                byName[issuer.Name] = i
        }
        return byName, nil
}

func applyTokenExchangePolicies(
        issuers []tx.TrustedIssuer, byName map[string]int, policies []TokenExchangeIssuerPolicyRunConfig,
) error {
        seen := make(map[string]int, len(policies))
        for i, policy := range policies {
                issuerIndex, err := resolveIssuerRef(byName, seen, policy.IssuerRef,
                        fmt.Sprintf("inbound_grants.token_exchange.issuer_policies[%d]", i))
                if err != nil {
                        return err
                }
                seen[policy.IssuerRef] = i
                issuer := &issuers[issuerIndex]
                issuer.ExpectedAudience = policy.ExpectedAudience
                issuer.ActorClaim = policy.ActorClaim
                issuer.AllowedActors = slices.Clone(policy.AllowedActors)
                issuer.ActorMatcher = policy.ActorMatcher
                issuer.AllowedDelegateClients = slices.Clone(policy.AllowedDelegateClients)
                issuer.AllowMayAct = policy.AllowMayAct
        }
        return nil
}

func applyJWTBearerPolicies(
        issuers []tx.TrustedIssuer, byName map[string]int, policies []JWTBearerIssuerPolicyRunConfig,
) error {
        seen := make(map[string]int, len(policies))
        for i, policy := range policies {
                issuerIndex, err := resolveIssuerRef(byName, seen, policy.IssuerRef,
                        fmt.Sprintf("inbound_grants.jwt_bearer.issuer_policies[%d]", i))
                if err != nil {
                        return err
                }
                seen[policy.IssuerRef] = i
                issuers[issuerIndex].JWTBearerGrant = &tx.JWTBearerGrantPolicy{
                        MaxAssertionAge:   policy.MaxAssertionAge,
                        SubjectBindings:   cloneSubjectBindings(policy.SubjectBindings),
                        AcceptedAudiences: slices.Clone(policy.AcceptedAudiences),
                }
        }
        return nil
}

func resolveIssuerRef(byName map[string]int, seen map[string]int, ref, path string) (int, error) {
        if ref == "" {
                return 0, fmt.Errorf("%s.issuer_ref is required", path)
        }
        if previous, ok := seen[ref]; ok {
                return 0, fmt.Errorf("%s.issuer_ref duplicates issuer policy [%d] for %q", path, previous, ref)
        }
        index, ok := byName[ref]
        if !ok {
                return 0, fmt.Errorf("%s.issuer_ref references unknown or unnamed trusted issuer %q", path, ref)
        }
        return index, nil
}

func hasLegacyTokenExchangePolicy(issuer tx.TrustedIssuer) bool {
        return issuer.ExpectedAudience != "" || issuer.ActorClaim != "" || len(issuer.AllowedActors) > 0 ||
                issuer.ActorMatcher != "" || len(issuer.AllowedDelegateClients) > 0 || issuer.AllowMayAct
}

func cloneDelegateClients(clients []DelegateClientRunConfig) []DelegateClientRunConfig {
        cloned := slices.Clone(clients)
        for i := range cloned {
                cloned[i].Scopes = slices.Clone(cloned[i].Scopes)
                cloned[i].Audiences = slices.Clone(cloned[i].Audiences)
        }
        return cloned
}

func cloneTrustedIssuers(issuers []tx.TrustedIssuer) []tx.TrustedIssuer {
        cloned := slices.Clone(issuers)
        for i := range cloned {
                cloned[i].AllowedActors = slices.Clone(cloned[i].AllowedActors)
                cloned[i].AllowedDelegateClients = slices.Clone(cloned[i].AllowedDelegateClients)
                if cloned[i].JWTBearerGrant != nil {
                        policy := *cloned[i].JWTBearerGrant
                        policy.SubjectBindings = cloneSubjectBindings(policy.SubjectBindings)
                        policy.AcceptedAudiences = slices.Clone(policy.AcceptedAudiences)
                        cloned[i].JWTBearerGrant = &policy
                }
        }
        return cloned
}

func cloneSubjectBindings(bindings []tx.JWTBearerSubjectBinding) []tx.JWTBearerSubjectBinding {
        cloned := slices.Clone(bindings)
        for i := range cloned {
                cloned[i].AllowedResources = slices.Clone(cloned[i].AllowedResources)
        }
        return cloned
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package consentpolicy validates the syntax, registration eligibility predicates,
// and deployment prerequisites for administrator consent preauthorization.
package consentpolicy

import (
        "fmt"
        "net/url"
        "slices"
        "strings"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// Validate checks the configured client IDs and the prerequisites for enabling
// consent preauthorization. An empty list preserves existing behavior.
func Validate(clientIDs []string, authorizationBase, issuer string, upstreamCount int) error {
        if len(clientIDs) == 0 {
                return nil
        }

        seen := make(map[string]struct{}, len(clientIDs))
        for _, clientID := range clientIDs {
                if clientID == "" {
                        return fmt.Errorf("consent_exempt_client_ids must not contain an empty ID")
                }
                if _, ok := seen[clientID]; ok {
                        return fmt.Errorf("consent_exempt_client_ids contains duplicate ID %q", clientID)
                }
                seen[clientID] = struct{}{}
                if isURLShaped(clientID) {
                        return fmt.Errorf("consent_exempt_client_ids must not contain URL-shaped ID %q", clientID)
                }
        }

        if upstreamCount == 0 {
                return fmt.Errorf("consent_exempt_client_ids requires at least one configured upstream")
        }

        if authorizationBase == "" {
                authorizationBase = issuer
        }
        parsedBase, err := url.Parse(authorizationBase)
        if err != nil || !strings.EqualFold(parsedBase.Scheme, "https") || parsedBase.Host == "" {
                return fmt.Errorf("consent_exempt_client_ids requires an HTTPS authorization endpoint base URL")
        }
        return nil
}

// IsBackChannelOnlyClient reports whether client has no interactive /authorize
// flow, using the same explicit marker and metadata-shape exclusions at startup
// and runtime.
func IsBackChannelOnlyClient(client fosite.Client) bool {
        return registration.BackChannelOnly(client) ||
                len(client.GetResponseTypes()) == 0 ||
                client.GetGrantTypes().ExactOne(oauthproto.GrantTypeTokenExchange)
}

// IsEligibleClient reports whether client supports ToolHive's interactive
// authorization-code flow and can be preauthorized to skip its consent page.
func IsEligibleClient(client fosite.Client) bool {
        return client != nil &&
                !client.IsPublic() &&
                !IsBackChannelOnlyClient(client) &&
                slices.Contains(client.GetGrantTypes(), oauthproto.GrantTypeAuthorizationCode) &&
                slices.Contains(client.GetResponseTypes(), oauthproto.ResponseTypeCode) &&
                len(client.GetRedirectURIs()) > 0
}

// isURLShaped reports whether clientID is an absolute URI, a network-path URL
// with an authority, or malformed as a URL (fail closed).
func isURLShaped(clientID string) bool {
        parsed, err := url.Parse(clientID)
        return err != nil || parsed.Scheme != "" || parsed.Host != ""
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package oauthparams provides shared definitions for reserved OAuth2
// authorization parameters that are managed by the framework.
package oauthparams

import (
        "fmt"
        "net/url"
        "strings"
)

// ReservedAuthorizationParams are OAuth2 parameters managed by the framework
// that must not be set via AdditionalAuthorizationParams.
var ReservedAuthorizationParams = map[string]bool{
        "response_type":         true,
        "client_id":             true,
        "redirect_uri":          true,
        "scope":                 true,
        "state":                 true,
        "code_challenge":        true,
        "code_challenge_method": true,
        "nonce":                 true,
}

// ReservedTokenParams are OAuth2 parameters managed by the framework that
// must not be set via AdditionalTokenParams. They cover both token-endpoint
// grant types the framework issues (authorization_code and refresh_token),
// plus the RFC 7523 client-authentication credentials: letting those through
// would combine an assertion with the configured client_secret or HTTP Basic
// credentials into an invalid multi-method client-auth request, and would put
// a credential in plain configuration instead of a secret-backed path.
var ReservedTokenParams = map[string]bool{
        "grant_type":            true,
        "code":                  true,
        "redirect_uri":          true,
        "client_id":             true,
        "client_secret":         true,
        "code_verifier":         true,
        "refresh_token":         true,
        "scope":                 true,
        "client_assertion":      true,
        "client_assertion_type": true,
}

// Validate checks that no key in params is a reserved OAuth2 authorization
// parameter. Reserved parameters are managed by the framework and cannot be
// overridden via additional authorization params.
func Validate(params map[string]string) error {
        for k := range params {
                if ReservedAuthorizationParams[k] {
                        return fmt.Errorf("reserved parameter %q is managed by the framework and cannot be overridden", k)
                }
        }
        return nil
}

// ValidateTokenParams checks that no key in params is a reserved OAuth2
// token-request parameter. Reserved parameters are managed by the framework
// and cannot be overridden via additional token params. A "resource" entry is
// additionally checked against the RFC 8707 shape.
func ValidateTokenParams(params map[string]string) error {
        for k, v := range params {
                if ReservedTokenParams[k] {
                        return fmt.Errorf("reserved token parameter %q is managed by the framework and cannot be overridden", k)
                }
                if k == resourceParam {
                        if err := validateResourceIndicator(v); err != nil {
                                return err
                        }
                }
        }
        return nil
}

// resourceParam is the RFC 8707 resource indicator parameter name.
const resourceParam = "resource"

// validateResourceIndicator enforces the RFC 8707 shape for a resource
// indicator: an absolute URI carrying no fragment component. Rejecting it here
// turns an authorization-server error at first login into a config-time error.
func validateResourceIndicator(value string) error {
        if value == "" {
                return fmt.Errorf("token parameter %q must not be empty", resourceParam)
        }
        u, err := url.Parse(value)
        if err != nil {
                return fmt.Errorf("token parameter %q must be an absolute URI: %w", resourceParam, err)
        }
        if !u.IsAbs() {
                return fmt.Errorf("token parameter %q must be an absolute URI, got %q", resourceParam, value)
        }
        if u.Fragment != "" || strings.Contains(value, "#") {
                return fmt.Errorf("token parameter %q must not contain a fragment, got %q", resourceParam, value)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "context"
        "errors"
        "log/slog"
        "maps"
        "slices"
        "time"

        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"

        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
)

// rejectionRecheckDelay is how long the guard waits after an upstream
// invalid_grant before re-reading the row, giving a replica that won a
// cross-replica refresh race time to land its write. The refresher's own
// re-read runs immediately and misses a write still in flight.
const rejectionRecheckDelay = 500 * time.Millisecond

// guardCheckTimeout bounds the guard's whole check, including the upstream
// refreshes of expired rows. A check that runs out of time is uncertain, so
// the grant goes through. A refresh still in flight when it does keeps
// running on the refresher's own detached context and stores its result.
const guardCheckTimeout = 10 * time.Second

// guardRevokeTimeout bounds the revocation of a refused refresh token and the
// deletion of a session's rows after a confirmed upstream rejection. Both run
// on a context detached from the request so a client that disconnects after
// the refusal does not skip them.
const guardRevokeTimeout = 5 * time.Second

// upstreamRefreshGuard is the built-in refresh-grant validator. It refuses a
// refresh_token grant when the token session's upstream credential is
// definitively dead. A refreshed access token keeps the token session (tsid),
// so granting it would hand the client a token the bearer validator answers
// with 401 again, and a client that only ever refreshes would loop. Refusing
// with invalid_grant sends the client to a full authorization, which creates a
// new tsid with fresh upstream rows.
//
// See Config.DisableUpstreamRefreshGuard for when a credential is dead, and
// docs/arch/11-auth-server-storage.md for the full behavior. Any other failure
// (network, 5xx, storage, the check running out of time) is uncertain and the
// grant goes through: the bearer validator answers those with 503 or a
// retryable 401 anyway.
//
// Expired rows are refreshed, one provider at a time, before deciding, because
// upstream refresh is lazy (it runs only on a resource request). An idle
// client's stored token is often expired while its upstream refresh token
// still works, and refusing on stored state alone would force every idle
// client to sign in again.
//
// The guard revokes the refused refresh token only when the stored rows prove
// the credential can never be used again. It refuses without revoking when the
// cause could be wrong and later correct itself: the IdP rejected a refresh
// (a replica that loses a cross-replica refresh race to a rotating IdP sees
// the same rejection), no row could be read (a storage backend may skip rows
// it cannot parse), or a row's provider is not configured (a decision from
// this replica's config, which can lag a rolling deploy that adds the
// provider). A refusal that turns out wrong then costs the client one
// re-authorization, not its grant.
//
// Before refusing on an IdP rejection, the guard re-reads the row once more
// after rejectionRecheckDelay and treats a row now holding a different refresh
// token as live. A recheck that cannot complete (the check is cancelled or
// times out, or the read fails) leaves the outcome uncertain, so the grant
// goes through. Once the rejection is confirmed, the guard deletes all of the
// session's upstream rows: the client is sent to a new tsid, so nothing reads
// them again, and a replay of the kept refresh token is then refused as having
// no rows, without another IdP call or recheck, on every replica.
//
// A refusal is logged at WARN only when it revokes or follows a confirmed IdP
// rejection; the row deletion makes the latter happen once per session. The
// other refusals repeat on every replay of a kept refresh token, so they are
// logged at DEBUG.
type upstreamRefreshGuard struct {
        tokens       storage.UpstreamTokenStorage
        refresher    storage.UpstreamTokenRefresher
        providers    map[string]struct{}
        revoker      oauth2.TokenRevocationStorage
        recheckDelay time.Duration
        checkTimeout time.Duration
}

var _ fosite.TokenEndpointHandler = (*upstreamRefreshGuard)(nil)

// newUpstreamRefreshGuard returns the guard for the given storage, the
// server's shared refresher and the names of the configured upstreams.
func newUpstreamRefreshGuard(
        stor storage.Storage, refresher storage.UpstreamTokenRefresher, providerNames []string,
) (*upstreamRefreshGuard, error) {
        if stor == nil {
                return nil, errors.New("upstream refresh guard requires storage")
        }
        if refresher == nil {
                return nil, errors.New("upstream refresh guard requires an upstream token refresher")
        }
        providers := make(map[string]struct{}, len(providerNames))
        for _, name := range providerNames {
                providers[name] = struct{}{}
        }
        return &upstreamRefreshGuard{
                tokens:       stor,
                refresher:    refresher,
                providers:    providers,
                revoker:      stor,
                recheckDelay: rejectionRecheckDelay,
                checkTimeout: guardCheckTimeout,
        }, nil
}

// CanSkipClientAuth always returns false. refreshGrantValidators never calls
// it; it is part of fosite.TokenEndpointHandler.
func (*upstreamRefreshGuard) CanSkipClientAuth(context.Context, fosite.AccessRequester) bool {
        return false
}

// CanHandleTokenEndpointRequest reports whether the request is a refresh_token
// grant.
func (*upstreamRefreshGuard) CanHandleTokenEndpointRequest(_ context.Context, requester fosite.AccessRequester) bool {
        return isRefreshTokenGrant(requester)
}

// HandleTokenEndpointRequest refuses the grant with invalid_grant when the
// session's upstream credential is definitively dead, revoking the refresh
// token when the refusal is final. A session without a tsid (a session-less
// grant) is not checked.
func (g *upstreamRefreshGuard) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error {
        if !g.CanHandleTokenEndpointRequest(ctx, requester) {
                return fosite.ErrUnknownRequest
        }
        sess, ok := requester.GetSession().(session.UpstreamSession)
        if !ok || sess.GetIDPSessionID() == "" {
                return nil
        }
        tsid := sess.GetIDPSessionID()

        refused := g.deadUpstream(ctx, tsid)
        if refused == nil {
                return nil
        }
        clientID := requester.GetClient().GetID()
        revoked := refused.revoke && g.revoke(ctx, requester.GetID(), clientID)
        level := slog.LevelDebug
        if refused.revoke || refused.rejected {
                level = slog.LevelWarn
        }
        slog.Log(ctx, level, "refusing refresh grant: upstream credential is no longer usable; re-authorization required",
                "client_id", clientID,
                "provider_id", refused.provider,
                "reason", refused.reason,
                "revoked", revoked,
        )
        // Log tsid at DEBUG only; it is credential-adjacent.
        slog.DebugContext(ctx, "refused refresh grant for token session", "tsid", tsid)
        return fosite.ErrInvalidGrant.WithHint(
                "The upstream identity provider session is no longer valid. Re-authorization is required.")
}

// PopulateTokenEndpointResponse contributes nothing; the guard only refuses.
// refreshGrantValidators never calls it; it is part of
// fosite.TokenEndpointHandler.
func (*upstreamRefreshGuard) PopulateTokenEndpointResponse(
        context.Context, fosite.AccessRequester, fosite.AccessResponder,
) error {
        return fosite.ErrUnknownRequest
}

// refusal is why the guard refuses a grant.
type refusal struct {
        // provider is the upstream the refusal applies to, or empty when it
        // applies to the session as a whole.
        provider string
        reason   string
        // revoke is set when the stored rows prove the credential can never be
        // used again, so the refused refresh token is revoked too.
        revoke bool
        // rejected is set when the IdP's rejection of the refresh was confirmed
        // and the session's rows were deleted.
        rejected bool
}

// deadUpstream returns why the session's upstream credential is definitively
// dead, or nil when some row is usable or its state is uncertain. Rows are
// checked in provider-name order and checking stops at the first dead one, so
// a refusal costs no further upstream refreshes. The bearer validator answers
// a session without any upstream access token with 401, so a session with no
// usable row is dead too.
func (g *upstreamRefreshGuard) deadUpstream(ctx context.Context, tsid string) *refusal {
        ctx, cancel := context.WithTimeout(ctx, g.checkTimeout)
        defer cancel()

        rows, err := g.tokens.GetAllUpstreamTokens(ctx, tsid)
        if err != nil {
                slog.DebugContext(ctx, "refresh grant guard could not read upstream tokens; granting", "error", err)
                return nil
        }
        if len(rows) == 0 {
                // Not revoked: a storage backend may skip rows it cannot parse and
                // report none, which reads the same as an absent session.
                return &refusal{reason: "no upstream tokens stored"}
        }
        anyUsable := false
        for _, name := range slices.Sorted(maps.Keys(rows)) {
                usable, refused := g.checkRow(ctx, tsid, name, rows[name])
                if refused != nil {
                        return refused
                }
                anyUsable = anyUsable || usable
        }
        if !anyUsable {
                return &refusal{reason: "no upstream access token stored", revoke: true}
        }
        return nil
}

// checkRow classifies one upstream row, refreshing it first when it is
// expired. It returns a refusal when the row is dead. Otherwise usable reports
// whether the row holds an access token or its state is uncertain (a transient
// failure), which the guard does not refuse; a row that is neither is not
// dead on its own.
func (g *upstreamRefreshGuard) checkRow(
        ctx context.Context, tsid, name string, row *storage.UpstreamTokens,
) (usable bool, refused *refusal) {
        if row == nil {
                return false, nil
        }
        if !row.IsExpired(time.Now()) {
                return row.AccessToken != "", nil
        }
        if _, ok := g.providers[name]; !ok {
                // Not revoked: this replica's config decides it, and another replica
                // may have the provider configured during a rolling deploy.
                return false, &refusal{
                        provider: name,
                        reason:   "upstream access token expired and its provider is no longer configured",
                }
        }
        if row.RefreshToken == "" {
                return false, &refusal{
                        provider: name, revoke: true,
                        reason: "upstream access token expired and no refresh token is stored",
                }
        }

        refreshed, err := g.refresh(ctx, tsid, row)
        if err == nil {
                return refreshed != nil && refreshed.AccessToken != "", nil
        }
        if errors.Is(err, storage.ErrNotFound) {
                // The row was deleted after the guard read it (logout, expiry, another
                // replica's cleanup). Not revoked, for the same reason as no rows.
                return false, &refusal{provider: name, reason: "upstream token row was deleted during refresh"}
        }
        if !errors.Is(err, upstream.ErrGrantRejected) {
                slog.DebugContext(ctx, "upstream refresh failed transiently; refresh grant guard is not refusing",
                        "provider_id", name, "error", err)
                return true, nil
        }
        if !g.rejectionConfirmed(ctx, tsid, name, row.RefreshToken) {
                return true, nil
        }
        g.dropSession(ctx, tsid)
        return false, &refusal{provider: name, reason: "upstream refresh was rejected with invalid_grant", rejected: true}
}

// refresh runs the refresher, returning ctx's error without calling the IdP
// when ctx has already ended, and early when ctx ends first. The refresher
// detaches from ctx, so without this a hung IdP would hold the token endpoint
// for the refresher's full timeout. ctx's error is not ErrGrantRejected, so
// the caller treats it as uncertain.
func (g *upstreamRefreshGuard) refresh(
        ctx context.Context, tsid string, row *storage.UpstreamTokens,
) (*storage.UpstreamTokens, error) {
        if err := ctx.Err(); err != nil {
                return nil, err
        }
        type result struct {
                tokens *storage.UpstreamTokens
                err    error
        }
        // Buffered so the goroutine never blocks once the guard stops waiting.
        done := make(chan result, 1)
        go func() {
                tokens, err := g.refresher.RefreshAndStore(ctx, tsid, row)
                done <- result{tokens: tokens, err: err}
        }()
        select {
        case r := <-done:
                return r.tokens, r.err
        case <-ctx.Done():
                return nil, ctx.Err()
        }
}

// rejectionConfirmed waits recheckDelay, re-reads the row and reports whether
// the upstream rejection of rejectedRT is confirmed: the row still holds
// rejectedRT, or is gone. A row holding any other refresh token means another
// replica won a refresh race and its write has landed, so the session may be
// live even if that row's access token has already expired. rejectedRT is the
// guard's own snapshot, which the refresher may have replaced with a newer
// re-read before redeeming; a mismatch then reads as a successor too, which
// errs toward granting. The outcome is not known when the wait is cancelled or
// the read fails with anything but ErrNotFound; that reports false too, so a
// client disconnect or a storage fault never refuses a session that may be
// live.
func (g *upstreamRefreshGuard) rejectionConfirmed(ctx context.Context, tsid, provider, rejectedRT string) bool {
        if g.recheckDelay > 0 {
                timer := time.NewTimer(g.recheckDelay)
                defer timer.Stop()
                select {
                case <-ctx.Done():
                        slog.DebugContext(ctx, "refresh grant guard recheck cancelled; not refusing", "provider_id", provider)
                        return false
                case <-timer.C:
                }
        }
        stored, err := g.tokens.GetUpstreamTokens(ctx, tsid, provider)
        switch {
        case errors.Is(err, storage.ErrNotFound):
                return true
        case errors.Is(err, storage.ErrExpired) && stored != nil:
                // Storage returns an expired row alongside ErrExpired; compare it.
        case err != nil:
                slog.DebugContext(ctx, "refresh grant guard recheck read failed; not refusing",
                        "provider_id", provider, "error", err)
                return false
        case stored == nil:
                return true
        }
        return stored.RefreshToken == rejectedRT
}

// dropSession deletes all of the session's upstream rows after a confirmed
// upstream rejection. Every row goes, not only the rejected provider's: with
// the rejected row gone and its siblings kept, a replay would see only usable
// rows and be granted a token that never gets a credential for the rejected
// provider. A replica that wins a refresh race after the recheck then fails
// its compare-and-swap, which costs nothing more, since this client is
// already refused. A failure is logged at DEBUG and the grant is still
// refused; a replay then repeats the check.
func (g *upstreamRefreshGuard) dropSession(ctx context.Context, tsid string) {
        delCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), guardRevokeTimeout)
        defer cancel()
        if err := g.tokens.DeleteUpstreamTokens(delCtx, tsid); err != nil && !errors.Is(err, storage.ErrNotFound) {
                slog.DebugContext(ctx, "refresh grant guard could not delete rejected upstream rows", "error", err)
        }
}

// revoke revokes the refused refresh token and the access tokens of the same
// grant, as an RFC 7009 revocation of the refresh token would, and reports
// whether it succeeded. A failure is logged and the grant is still refused:
// the client is sent to a full authorization either way, and the guard
// refuses the token again if it is presented.
func (g *upstreamRefreshGuard) revoke(ctx context.Context, requestID, clientID string) bool {
        revokeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), guardRevokeTimeout)
        defer cancel()
        if err := errors.Join(
                g.revoker.RevokeRefreshToken(revokeCtx, requestID),
                g.revoker.RevokeAccessToken(revokeCtx, requestID),
        ); err != nil {
                slog.WarnContext(ctx, "failed to revoke refresh token refused by the upstream refresh guard",
                        "client_id", clientID, "error", err)
                return false
        }
        return true
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "context"
        "errors"
        "reflect"
        "slices"

        "github.com/ory/fosite"

        oauthserver "github.com/stacklok/toolhive/pkg/authserver/server"
)

// refreshGrantValidators runs the refresh-grant validators as validators of
// the refresh_token grant: the built-in upstream refresh guard, when enabled,
// followed by the embedder-supplied token endpoint handlers
// (Config.RefreshGrantHandlers).
//
// All handlers live in one registered value because fosite's
// TokenEndpointHandlers.Append drops a handler whose concrete type is already
// registered: wrapping each handler separately would register only the first.
//
// Only the handlers' request phase (HandleTokenEndpointRequest) runs. The
// response phase is never delegated: by then the built-in refresh handler has
// already rotated the refresh token, so a handler error there would leave the
// client with no tokens and a dead refresh token, and a handler could rewrite
// the issued response. The type implements only fosite.TokenEndpointHandler,
// so the handlers are never registered for revocation, introspection or the
// authorize endpoint.
type refreshGrantValidators struct {
        handlers []fosite.TokenEndpointHandler
}

var _ fosite.TokenEndpointHandler = (*refreshGrantValidators)(nil)

// CanSkipClientAuth always returns false: an embedder's handler validates a
// grant the built-in refresh handler has already accepted and must never
// relax the client authentication that grant requires.
func (*refreshGrantValidators) CanSkipClientAuth(context.Context, fosite.AccessRequester) bool {
        return false
}

// CanHandleTokenEndpointRequest reports whether the request is a refresh_token
// grant.
func (*refreshGrantValidators) CanHandleTokenEndpointRequest(_ context.Context, requester fosite.AccessRequester) bool {
        return isRefreshTokenGrant(requester)
}

// HandleTokenEndpointRequest runs each handler that accepts the request, in
// order, and returns the first refusal. It runs after the built-in refresh
// handler has authenticated the client and loaded the refresh session, and
// before any handler populates the response, so a refusal issues no tokens
// and leaves the refresh token unrotated. A handler returning
// fosite.ErrUnknownRequest is treated as having nothing to say, matching how
// fosite itself treats that error.
func (v *refreshGrantValidators) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error {
        if !v.CanHandleTokenEndpointRequest(ctx, requester) {
                return fosite.ErrUnknownRequest
        }
        for _, h := range v.handlers {
                if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                        continue
                }
                if err := h.HandleTokenEndpointRequest(ctx, requester); err != nil && !errors.Is(err, fosite.ErrUnknownRequest) {
                        return err
                }
        }
        return nil
}

// PopulateTokenEndpointResponse never delegates to the handlers, so the seam
// can only refuse a grant, never alter or break an issued response. fosite
// calls it on every registered handler; ErrUnknownRequest tells fosite this
// handler contributes nothing.
func (*refreshGrantValidators) PopulateTokenEndpointResponse(
        context.Context, fosite.AccessRequester, fosite.AccessResponder,
) error {
        return fosite.ErrUnknownRequest
}

// isRefreshTokenGrant reports whether the request's sole grant type is
// refresh_token, matching the check fosite's built-in refresh handler makes.
func isRefreshTokenGrant(requester fosite.AccessRequester) bool {
        return requester.GetGrantTypes().ExactOne(string(fosite.GrantTypeRefreshToken))
}

// isNilHandler reports whether h is nil, including a nil pointer (or other
// nilable value) of a concrete type, which a plain h == nil misses and which
// would panic on the first refresh grant.
func isNilHandler(h fosite.TokenEndpointHandler) bool {
        if h == nil {
                return true
        }
        v := reflect.ValueOf(h)
        nilable := []reflect.Kind{reflect.Pointer, reflect.Map, reflect.Slice, reflect.Func, reflect.Chan, reflect.Interface}
        return slices.Contains(nilable, v.Kind()) && v.IsNil()
}

// refreshGrantHandlerFactories returns the factory that registers the
// refresh-grant validators (see refreshGrantValidatorChain) with
// createProvider, or none when there are no validators.
func refreshGrantHandlerFactories(handlers []fosite.TokenEndpointHandler) []oauthserver.Factory {
        if len(handlers) == 0 {
                return nil
        }
        validators := &refreshGrantValidators{handlers: slices.Clone(handlers)}
        return []oauthserver.Factory{
                func(*oauthserver.AuthorizationServerConfig, fosite.Storage, any) (any, error) {
                        return validators, nil
                },
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "time"

        "golang.org/x/sync/singleflight"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
)

// refreshTimeout bounds how long a singleflight-deduplicated token refresh
// may take before being cancelled. It is deliberately detached from the
// triggering request's context so that waiting callers are not abandoned.
const refreshTimeout = 30 * time.Second

// upstreamStoreMaxAttempts is the maximum number of attempts to store refreshed
// upstream tokens before giving up.
const upstreamStoreMaxAttempts = 3

// upstreamStoreRetryBackoff is the fixed delay between store retry attempts.
const upstreamStoreRetryBackoff = 50 * time.Millisecond

// upstreamDeleteTimeout bounds how long the best-effort per-provider delete may
// take. A fresh detached context is used so a ctx deadline/cancel on the store
// attempts does not also kill the delete, which would leave the stale row behind.
const upstreamDeleteTimeout = 5 * time.Second

// upstreamConflictReadTimeout bounds the single re-read resolveConcurrentRefreshConflict
// performs after a CAS conflict. It is deliberately much shorter than refreshTimeout:
// the caller of RefreshAndStore is already blocked on refreshTimeout (30s) via the
// detached refreshCtx passed in here, so reusing refreshTimeout for this read would
// stack a second independent 30s window on top — doubling the worst-case wait for no
// benefit, since a single GetUpstreamTokens call needs nowhere near that budget.
const upstreamConflictReadTimeout = 5 * time.Second

// upstreamTokenRefresher implements storage.UpstreamTokenRefresher by wrapping
// a set of upstream OAuth2Providers (keyed by provider name) and
// UpstreamTokenStorage (for persisting the refreshed tokens). On each refresh
// call it dispatches to the correct provider based on the expired token's
// ProviderID.
//
// Two layers cooperate to coordinate redeeming a single-use, rotating upstream
// refresh token under concurrency:
//   - sfGroup deduplicates concurrent refreshes by the opaque storage-row
//     identity returned by UpstreamTokenStorage, but only WITHIN THIS PROCESS
//     — it saves redundant upstream calls and redis round-trips when this
//     process alone receives several concurrent requests for the same row.
//   - CompareAndSwapUpstreamTokens (see refreshAndStore) makes the STORED row
//     deterministic ACROSS PROCESSES: every write is conditioned on the refresh
//     token this call redeemed still being the one stored, so a replica whose
//     write loses a cross-process race fails instead of clobbering a winning
//     replica's rotated token with a stale one. This is a storage-ordering
//     guarantee, not a guarantee that the redemption itself was safe: both
//     replicas still call the upstream provider before either writes, so a
//     strict single-use-rotation provider (RFC 9700 §4.14.2 replay detection)
//     can still see this as two redemptions of the same refresh token and
//     revoke the grant regardless of which write wins here — see
//     docs/arch/11-auth-server-storage.md's "Refresh Coordination Scope" for
//     the provider-behavior-dependent discussion. singleflight is an
//     optimization layered on top of the CAS guarantee, not a substitute for it.
type upstreamTokenRefresher struct {
        providers            map[string]upstream.OAuth2Provider
        storage              storage.UpstreamTokenStorage
        refreshTokenLifespan time.Duration
        sfGroup              singleflight.Group
}

// Compile-time check that upstreamTokenRefresher implements storage.UpstreamTokenRefresher.
var _ storage.UpstreamTokenRefresher = (*upstreamTokenRefresher)(nil)

// RefreshAndStore validates deterministic input errors before resolving the
// opaque storage-row identity used to deduplicate a refresh within this process.
// The leader re-reads that row under a detached timeout so it never redeems a
// stale refresh token supplied by a caller. It only short-circuits on that
// re-read when the row is unexpired AND differs from the caller's row, i.e.
// another caller already refreshed it. A storage backend that returns tokens
// without checking expiry does not fool it into skipping the refresh, and a
// caller that deliberately refreshes a still-current row (e.g. because its ID
// token expired before its access token) gets a real upstream refresh.
func (r *upstreamTokenRefresher) RefreshAndStore(
        ctx context.Context,
        sessionID string,
        expired *storage.UpstreamTokens,
) (*storage.UpstreamTokens, error) {
        if expired == nil {
                return nil, errors.New("expired tokens are required")
        }
        if expired.RefreshToken == "" {
                return nil, errors.New("no refresh token available for upstream token refresh")
        }
        if _, ok := r.providers[expired.ProviderID]; !ok {
                return nil, fmt.Errorf("no upstream provider configured for %q", expired.ProviderID)
        }

        rowID, err := r.storage.ResolveUpstreamTokenRowID(ctx, sessionID, expired.ProviderID)
        if err != nil {
                return nil, err
        }
        if rowID == "" {
                return nil, errors.New("upstream token row ID cannot be empty")
        }

        result, err, _ := r.sfGroup.Do(string(rowID), func() (any, error) {
                refreshCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), refreshTimeout)
                defer cancel()

                authoritative, readErr := r.storage.GetUpstreamTokens(refreshCtx, sessionID, expired.ProviderID)
                if readErr == nil {
                        // Some storage implementations may return tokens without checking
                        // access-token expiry (the interface does not require it; see
                        // storage.UpstreamTokenStorage.GetUpstreamTokens). Only short-circuit
                        // when the re-read row is actually unexpired; otherwise fall through
                        // and refresh using it as the authoritative row.
                        if authoritative == nil {
                                return nil, fmt.Errorf(
                                        "upstream token row missing on re-read for provider %q", expired.ProviderID,
                                )
                        }
                        if !authoritative.IsExpired(time.Now()) && rowChanged(expired, authoritative) {
                                return authoritative, nil
                        }
                        return r.refreshAndStore(refreshCtx, sessionID, authoritative)
                }
                if !errors.Is(readErr, storage.ErrExpired) || authoritative == nil {
                        return nil, readErr
                }
                return r.refreshAndStore(refreshCtx, sessionID, authoritative)
        })
        if err != nil {
                return nil, err
        }
        refreshed, ok := result.(*storage.UpstreamTokens)
        if !ok || refreshed == nil {
                return nil, errors.New("unexpected nil result from upstream token refresh")
        }
        refreshedCopy := *refreshed
        return &refreshedCopy, nil
}

// refreshAndStore performs the actual token refresh and storage write.
func (r *upstreamTokenRefresher) refreshAndStore(
        ctx context.Context,
        sessionID string,
        expired *storage.UpstreamTokens,
) (*storage.UpstreamTokens, error) {
        if expired.RefreshToken == "" {
                return nil, errors.New("no refresh token available for upstream token refresh")
        }

        slog.Debug("attempting upstream token refresh",
                "session_id", sessionID,
                "provider_id", expired.ProviderID,
        )

        // Look up the provider that issued this token
        provider, ok := r.providers[expired.ProviderID]
        if !ok {
                return nil, fmt.Errorf("no upstream provider configured for %q", expired.ProviderID)
        }

        // Refresh tokens via the upstream provider
        newTokens, err := provider.RefreshTokens(ctx, expired.RefreshToken, expired.UpstreamSubject)
        if err != nil {
                if errors.Is(err, upstream.ErrGrantRejected) {
                        if winner := r.rotatedElsewhere(ctx, sessionID, expired); winner != nil {
                                return winner, nil
                        }
                }
                return nil, fmt.Errorf("upstream token refresh failed: %w", err)
        }

        // Defensive re-anchor of SessionExpiresAt: the post-PR callback write sets
        // SessionExpiresAt unconditionally so it can be carried forward here as a
        // storage TTL bound. Pre-PR rows persisted without that field decode as
        // zero. If such a legacy row is refreshed and the upstream rotation drops
        // expires_in, both ExpiresAt and SessionExpiresAt would be zero, the row
        // would be stored without any TTL bound, and the Memory backend would
        // retain it indefinitely. Re-anchor to now+RefreshTokenLifespan to restore
        // the invariant. The Redis 30-day per-key TTL also caps the legacy
        // behavior, but Memory has no such backstop.
        sessionExpiresAt := expired.SessionExpiresAt
        if sessionExpiresAt.IsZero() && newTokens.ExpiresAt.IsZero() {
                sessionExpiresAt = time.Now().Add(r.refreshTokenLifespan)
                slog.Debug("re-anchored zero SessionExpiresAt on refresh of legacy upstream token row",
                        "session_id", sessionID,
                        "provider_id", expired.ProviderID,
                        "refresh_token_lifespan", r.refreshTokenLifespan,
                )
        }

        // Build updated storage tokens preserving binding fields from the original
        updated := &storage.UpstreamTokens{
                ProviderID:       expired.ProviderID,
                AccessToken:      newTokens.AccessToken,
                RefreshToken:     newTokens.RefreshToken,
                IDToken:          newTokens.IDToken,
                ExpiresAt:        newTokens.ExpiresAt,
                SessionExpiresAt: sessionExpiresAt,
                UserID:           expired.UserID,
                UpstreamSubject:  expired.UpstreamSubject,
                ClientID:         expired.ClientID,
        }

        // Detect rotation BEFORE the old-RT backfill below so the original
        // (provider-issued) value is compared, not the backfilled fallback.
        rotated := updated.RefreshToken != "" && updated.RefreshToken != expired.RefreshToken

        // If the provider didn't rotate the refresh token, keep the original.
        if updated.RefreshToken == "" {
                updated.RefreshToken = expired.RefreshToken
        }

        // OIDC Core 1.0 §12.2 permits but does not require a new id_token on refresh.
        // When the provider omits one, keep the ID token captured at the initial login
        // so it is not erased from storage. The write below replaces the whole row,
        // so without this the persisted IDToken would be overwritten with "" and the
        // original login ID token would be lost for the remainder of the session.
        // Mirrors the RefreshToken carry-forward above.
        if updated.IDToken == "" {
                updated.IDToken = expired.IDToken
        }

        // expectedRefreshToken is the value this call actually redeemed with the
        // upstream provider. Writing through CompareAndSwapUpstreamTokens (rather
        // than an unconditional StoreUpstreamTokens) makes the STORED row
        // deterministic when another process — a different replica of this auth
        // server sharing this storage — is redeeming the same refresh token
        // concurrently: only the replica whose expected value still matches the
        // stored row may write, so a losing replica cannot clobber the winner's
        // rotated token with its own now-stale redemption. This orders the writes;
        // it does not by itself guarantee the upstream provider accepts both
        // redemptions (see the type-level doc comment above).
        expectedRefreshToken := expired.RefreshToken

        if err := r.compareAndSwapWithRetry(ctx, sessionID, expired.ProviderID, expectedRefreshToken, updated); err != nil {
                if errors.Is(err, storage.ErrConcurrentRefresh) {
                        return r.resolveConcurrentRefreshConflict(ctx, sessionID, expired.ProviderID)
                }

                if !rotated {
                        // The old refresh token is still valid in storage; the caller can
                        // proceed with the refreshed access token for this request.
                        slog.Warn("failed to persist refreshed upstream tokens; old refresh token still valid",
                                "provider_id", expired.ProviderID,
                                "error", err,
                        )
                        // Log session_id at DEBUG only; it is credential-adjacent.
                        slog.Debug("upstream token persist failure session", "session_id", sessionID)
                        return updated, nil
                }

                // The IdP rotated the refresh token: the new RT was redeemed but could
                // not be persisted. The old RT is now dead in storage and will trigger
                // reuse-detection lockout on the next refresh attempt.
                //
                // Residual window: if the process crashes here, the dead old RT stays in
                // storage. ErrNotFound on the next refresh attempt forces re-auth, which
                // is the backstop. Store-intent-before-redeem is out of scope for this fix.
                deleteCtx, deleteCancel := context.WithTimeout(context.WithoutCancel(ctx), upstreamDeleteTimeout)
                defer deleteCancel()
                if delErr := r.storage.DeleteUpstreamTokensForProvider(deleteCtx, sessionID, expired.ProviderID); delErr != nil {
                        slog.Error("failed to delete stale upstream token row after rotation persist failure",
                                "provider_id", expired.ProviderID,
                                "error", delErr,
                        )
                        // Log session_id at DEBUG only; it is credential-adjacent.
                        slog.Debug("stale upstream token row delete failure session", "session_id", sessionID)
                }
                return nil, fmt.Errorf(
                        "failed to persist rotated upstream refresh token for provider %q: %w",
                        expired.ProviderID, err,
                )
        }

        slog.Debug("upstream tokens refreshed successfully",
                "session_id", sessionID,
                "provider_id", expired.ProviderID,
        )

        return updated, nil
}

// resolveConcurrentRefreshConflict is reached after CompareAndSwapUpstreamTokens
// returns storage.ErrConcurrentRefresh: the stored refresh token no longer
// matched the value this call redeemed with. That mismatch has two distinct
// causes this function must tell apart on re-read, not just to pick the right
// return value but to log something a human can act on:
//   - Another process (typically a different replica of this auth server)
//     redeemed and persisted a rotation of the same row first. This IS a
//     genuine lost race, worth surfacing loudly.
//   - The row no longer exists at all — deleted by a logout, evicted by TTL,
//     or never created. This is NOT a race: nothing else needed to "win"
//     against this call, the row was simply gone before the write. Logging it
//     the same as the case above would send whoever is on call chasing a
//     concurrency bug that didn't happen.
//
// The re-read result decides which happened: an unexpired row means the other
// process's write is the correct result to hand back to this call — this
// call's own redemption produced a token pair the IdP (or the winning
// replica) already superseded, so it must not be used or written. A read that
// comes back storage.ErrNotFound means the row is gone; CompareAndSwapUpstreamTokens
// correctly refused to resurrect it (StoreUpstreamTokens would not have). Any
// other outcome (still expired, or a genuine read error) is ambiguous and
// treated as the loss it most likely is. In every non-recoverable case the
// refresh token this call redeemed is unusable going forward — dead at the
// IdP if it enforces single-use, and either superseded or deliberately absent
// in storage — so the caller must surface an error rather than retry the same
// redemption.
func (r *upstreamTokenRefresher) resolveConcurrentRefreshConflict(
        ctx context.Context, sessionID, providerID string,
) (*storage.UpstreamTokens, error) {
        readCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), upstreamConflictReadTimeout)
        defer cancel()

        winner, err := r.storage.GetUpstreamTokens(readCtx, sessionID, providerID)
        if err == nil && winner != nil && !winner.IsExpired(time.Now()) {
                slog.Debug("lost concurrent upstream token refresh race; using winning replica's tokens",
                        "session_id", sessionID,
                        "provider_id", providerID,
                )
                return winner, nil
        }

        if errors.Is(err, storage.ErrNotFound) {
                slog.Warn("upstream token row no longer exists after a concurrent-refresh conflict; "+
                        "this is expected on logout or TTL eviction and is not necessarily a lost race",
                        "provider_id", providerID,
                )
        } else {
                slog.Error("lost concurrent upstream token refresh race and re-read found no usable winner; "+
                        "the redeemed refresh token is unrecoverable for this attempt",
                        "provider_id", providerID,
                        "error", err,
                )
        }
        // Log session_id at DEBUG only; it is credential-adjacent.
        slog.Debug("concurrent upstream token refresh conflict session", "session_id", sessionID)
        if err != nil {
                return nil, fmt.Errorf(
                        "concurrent upstream token refresh for provider %q: %w: %w",
                        providerID, storage.ErrConcurrentRefresh, err,
                )
        }
        return nil, fmt.Errorf(
                "concurrent upstream token refresh for provider %q: %w",
                providerID, storage.ErrConcurrentRefresh,
        )
}

// rotatedElsewhere is consulted when the upstream provider rejected this
// call's refresh token with invalid_grant. A rotating provider answers that
// way to the replica that loses a cross-process refresh race (singleflight
// only deduplicates within this process), even though the winner has already
// stored a valid row. If the stored row now holds a different refresh token
// and an unexpired access token, that row is the result to use, so a lost race
// is not reported as a dead upstream credential.
//
// "Another process refreshed" is tested by refresh-token inequality, unlike
// rowChanged, which also compares the access token, ID token and expiry. The
// rejection is of the refresh token itself, so only a row holding a different
// refresh token can be a successor; a row that changed in other fields but
// still holds the rejected refresh token is just as dead.
//
// This narrows the race but does not close it: when the winner's storage write
// has not landed by the time of the re-read, the rejection is still returned.
// Any read failure also returns nil, leaving the caller's original error.
func (r *upstreamTokenRefresher) rotatedElsewhere(
        ctx context.Context, sessionID string, expired *storage.UpstreamTokens,
) *storage.UpstreamTokens {
        readCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), upstreamConflictReadTimeout)
        defer cancel()

        stored, err := r.storage.GetUpstreamTokens(readCtx, sessionID, expired.ProviderID)
        if err != nil {
                // ErrExpired and ErrNotFound are ordinary outcomes here: the row is
                // still the rejected one, or it is gone. Anything else is a storage
                // fault worth a diagnostic, though the caller's rejection still stands.
                if !errors.Is(err, storage.ErrExpired) && !errors.Is(err, storage.ErrNotFound) {
                        slog.Debug("re-read after upstream invalid_grant failed; returning the rejection",
                                "session_id", sessionID,
                                "provider_id", expired.ProviderID,
                                "error", err,
                        )
                }
                return nil
        }
        if stored == nil || stored.RefreshToken == expired.RefreshToken || stored.IsExpired(time.Now()) {
                return nil
        }
        slog.Debug("upstream rejected a refresh token another process already rotated; using the stored row",
                "session_id", sessionID,
                "provider_id", expired.ProviderID,
        )
        return stored
}

// compareAndSwapWithRetry attempts CompareAndSwapUpstreamTokens up to
// upstreamStoreMaxAttempts times, waiting upstreamStoreRetryBackoff between
// attempts. Returns nil on first success. A storage.ErrConcurrentRefresh
// result is returned immediately without retrying: the comparison failed
// because expectedRefreshToken is now definitively stale, and retrying the
// identical compare-and-swap cannot change that outcome. For any other error,
// returns the last error after all attempts are exhausted; ctx-cancellation
// short-circuits between attempts and returns the last store error (not
// ctx.Err()).
func (r *upstreamTokenRefresher) compareAndSwapWithRetry(
        ctx context.Context,
        sessionID, providerID, expectedRefreshToken string,
        updated *storage.UpstreamTokens,
) error {
        var lastErr error
        for attempt := 1; attempt <= upstreamStoreMaxAttempts; attempt++ {
                casErr := r.storage.CompareAndSwapUpstreamTokens(ctx, sessionID, providerID, expectedRefreshToken, updated)
                if casErr == nil {
                        return nil
                }
                if errors.Is(casErr, storage.ErrConcurrentRefresh) {
                        return casErr
                }
                lastErr = casErr
                slog.Debug("failed to store refreshed upstream tokens",
                        "session_id", sessionID,
                        "provider_id", providerID,
                        "attempt", attempt,
                        "error", lastErr,
                )
                if attempt < upstreamStoreMaxAttempts {
                        select {
                        case <-ctx.Done():
                                return lastErr
                        case <-time.After(upstreamStoreRetryBackoff):
                        }
                }
        }
        return lastErr
}

// rowChanged reports whether the stored row was rewritten since the caller read
// it, i.e. another caller already refreshed it. Expiry is compared as well as
// the token values so a provider that re-issues the same access token string
// with a new expiry (and does not rotate the refresh token) still counts as a
// change, and the caller does not redeem again.
func rowChanged(callerRow, stored *storage.UpstreamTokens) bool {
        return callerRow.AccessToken != stored.AccessToken ||
                callerRow.RefreshToken != stored.RefreshToken ||
                callerRow.IDToken != stored.IDToken ||
                !callerRow.ExpiresAt.Equal(stored.ExpiresAt)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "fmt"
        "strings"

        "github.com/stacklok/toolhive/pkg/auth/dcr"
        "github.com/stacklok/toolhive/pkg/authserver"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
)

// This file contains the embedded-authserver-specific adapter layer between
// pkg/auth/dcr's profile-neutral API and the authserver's domain types
// (authserver.UpstreamRunConfig, upstream.OAuth2Config, upstream.OIDCConfig). Each helper
// here:
//
//   - newDCRRequest converts either typed upstream run-config into the
//     profile-neutral dcr.Request, including resolving the file-or-env token.
//   - consumeResolution / consumeOIDCResolution and the apply helpers fold the
//     returned dcr.Resolution back into run-config and built provider configs.
//     Both are value-in / value-out: the caller's original is never observably
//     mutated.
//
// Keeping these helpers in the runner package (not pkg/auth/dcr) is the
// architectural split the dcr package's profile-agnostic charter requires:
// dcr cannot import authserver, but the embedded authserver freely depends
// on dcr.

// newDCRRequest builds the profile-neutral dcr.Request for an upstream that
// requires runtime Dynamic Client Registration. It returns nil when the
// upstream has a pre-provisioned client.
func newDCRRequest(rc *authserver.UpstreamRunConfig, localIssuer string) (*dcr.Request, error) {
        if rc == nil {
                return nil, fmt.Errorf("upstream run-config is required")
        }

        var (
                clientID              string
                dcrConfig             *authserver.DCRUpstreamConfig
                redirectURI           string
                scopes                []string
                caFilePath            string
                allowPrivateIPs       bool
                authorizationEndpoint string
                tokenEndpoint         string
        )
        switch rc.Type {
        case authserver.UpstreamProviderTypeOIDC:
                if rc.OIDCConfig == nil {
                        return nil, fmt.Errorf("oidc upstream run-config is required")
                }
                clientID = rc.OIDCConfig.ClientID
                dcrConfig = rc.OIDCConfig.DCRConfig
                redirectURI = rc.OIDCConfig.RedirectURI
                scopes = rc.OIDCConfig.Scopes
                caFilePath = rc.OIDCConfig.CAFilePath
                allowPrivateIPs = rc.OIDCConfig.AllowPrivateIPs
        case authserver.UpstreamProviderTypeOAuth2:
                if rc.OAuth2Config == nil {
                        return nil, fmt.Errorf("oauth2 upstream run-config is required")
                }
                clientID = rc.OAuth2Config.ClientID
                dcrConfig = rc.OAuth2Config.DCRConfig
                redirectURI = rc.OAuth2Config.RedirectURI
                scopes = rc.OAuth2Config.Scopes
                caFilePath = rc.OAuth2Config.CAFilePath
                allowPrivateIPs = rc.OAuth2Config.AllowPrivateIPs
                authorizationEndpoint, tokenEndpoint = rc.OAuth2Config.AuthorizationEndpoint, rc.OAuth2Config.TokenEndpoint
        default:
                return nil, fmt.Errorf("unsupported upstream type: %s", rc.Type)
        }
        if clientID != "" || dcrConfig == nil {
                return nil, nil
        }

        discoveryURL := dcrConfig.DiscoveryURL
        if rc.Type == authserver.UpstreamProviderTypeOIDC && discoveryURL == "" && dcrConfig.RegistrationEndpoint == "" {
                discoveryURL = strings.TrimRight(rc.OIDCConfig.IssuerURL, "/") + "/.well-known/openid-configuration"
        }
        initialAccessToken, err := resolveSecret(dcrConfig.InitialAccessTokenFile, dcrConfig.InitialAccessTokenEnvVar)
        if err != nil {
                return nil, fmt.Errorf("dcr: resolve initial access token: %w", err)
        }
        return &dcr.Request{
                Issuer: localIssuer, RedirectURI: redirectURI, Scopes: scopes,
                DiscoveryURL: discoveryURL, RegistrationEndpoint: dcrConfig.RegistrationEndpoint,
                RotationKey:           dcrConfig.RotationKey,
                RegistrationProbe:     dcr.RegistrationProbe(dcrConfig.RegistrationProbe),
                AuthorizationEndpoint: authorizationEndpoint, TokenEndpoint: tokenEndpoint,
                InitialAccessToken: initialAccessToken, CAFilePath: caFilePath, AllowPrivateIPs: allowPrivateIPs,
        }, nil
}

// consumeResolution returns a copy of rc with the resolved credentials and
// endpoints from res copied in and DCRConfig consumed (set to nil),
// transitioning the run-config from "DCR-pending" (ClientID == "" &&
// DCRConfig != nil) to "DCR-resolved" (ClientID populated && DCRConfig
// == nil). The "consume" name is deliberate: a second call on the
// returned value is a no-op only because the first cleared DCRConfig —
// this is a one-shot state transition, not an idempotent default-fill.
//
// rc is taken by value and the modified copy is returned. The caller's
// original is never observably mutated; the value-in / value-out shape
// makes the no-mutation contract compile-time enforced rather than a
// prose discipline the caller is required to remember. Pointer-typed
// fields (DCRConfig) share storage with the caller's copy via the struct
// shallow-copy, but the only mutation here is to assign nil to the
// copy's DCRConfig — nil-assignment to the local field does not reach
// back through the original pointer.
//
// Why DCRConfig is cleared: OAuth2UpstreamRunConfig.Validate enforces
// ClientID xor DCRConfig — a resolved copy that left DCRConfig set would
// fail the validator that runs downstream in buildPureOAuth2Config.
//
// ClientID, the endpoints, and RedirectURI are written only when rc leaves
// them empty — explicit caller configuration always wins. The conditional
// ClientID write is defence-in-depth against future call sites that bypass
// the needsDCR precondition; an unconditional overwrite would silently
// clobber a pre-provisioned ClientID with no error.
//
// The defaulted RedirectURI write closes the loop on resolver-side defaulting:
// when the caller's run-config left RedirectURI empty, the resolver derived
// issuer + /oauth/callback and persisted it on the resolution; copying it
// back here means the downstream upstream.OAuth2Config has a non-empty
// RedirectURI, which authserver.Config validation requires.
//
// Note on ClientSecret: consumeResolution does NOT write the resolved
// secret because OAuth2UpstreamRunConfig models secrets as file-or-env
// references only. To propagate the DCR-resolved secret into the final
// upstream.OAuth2Config, callers must pair this call with
// applyResolutionToOAuth2Config once the config has been built. Keeping
// the two helpers side-by-side localises the DCR-specific application
// logic.
func consumeResolution(rc authserver.OAuth2UpstreamRunConfig, res *dcr.Resolution) authserver.OAuth2UpstreamRunConfig {
        if res == nil {
                return rc
        }
        if rc.ClientID == "" {
                rc.ClientID = res.ClientID
        }
        rc.DCRConfig = nil
        if rc.AuthorizationEndpoint == "" {
                rc.AuthorizationEndpoint = res.AuthorizationEndpoint
        }
        if rc.TokenEndpoint == "" {
                rc.TokenEndpoint = res.TokenEndpoint
        }
        if rc.RedirectURI == "" {
                rc.RedirectURI = res.RedirectURI
        }
        return rc
}

func consumeOIDCResolution(cfg authserver.OIDCUpstreamRunConfig, res *dcr.Resolution) authserver.OIDCUpstreamRunConfig {
        if res == nil {
                return cfg
        }
        if cfg.ClientID == "" {
                cfg.ClientID = res.ClientID
        }
        if cfg.RedirectURI == "" {
                cfg.RedirectURI = res.RedirectURI
        }
        cfg.DCRConfig = nil
        return cfg
}

func applyResolutionToOIDCConfig(cfg upstream.OIDCConfig, res *dcr.Resolution) upstream.OIDCConfig {
        if res != nil {
                cfg.ClientSecret = res.ClientSecret
        }
        return cfg
}

// resolved ClientSecret and TokenEndpointAuthMethod overlaid onto it. This
// is the companion to consumeResolution: where that function writes fields
// representable in the file-or-env run-config model, this one writes the
// inline-only fields directly on the runtime config.
//
// cfg is taken by value and the modified copy is returned, mirroring
// consumeResolution. The no-mutation contract is compile-time enforced
// by the signature rather than a prose discipline.
//
// TokenEndpointAuthMethod carries the method the upstream negotiated during
// registration (RFC 7591); see upstream.authStyleFromMethod for why it matters.
//
// The split exists because buildPureOAuth2Config intentionally retains a
// narrow file-or-env contract (no DCR awareness) and because OAuth2's
// ClientSecret on the run-config is modelled as a reference rather than
// an inline string. Any future output path from OAuth2UpstreamRunConfig
// to upstream.OAuth2Config must call BOTH consumeResolution (run-config
// side) AND applyResolutionToOAuth2Config (built-config side) to get a
// fully-resolved DCR client. Forgetting the second call leaves both
// ClientSecret and TokenEndpointAuthMethod empty, producing silent auth
// failures at request time — a client that negotiated client_secret_basic
// reverts to POST-body credentials and is rejected. The type system does
// not enforce the pair, so the invariant lives here.
func applyResolutionToOAuth2Config(cfg upstream.OAuth2Config, res *dcr.Resolution) upstream.OAuth2Config {
        if res == nil {
                return cfg
        }
        cfg.ClientSecret = res.ClientSecret
        cfg.TokenEndpointAuthMethod = res.TokenEndpointAuthMethod
        return cfg
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package runner provides integration between the proxy runner and the auth server.
package runner

import (
        "bytes"
        "context"
        "crypto/x509"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "net/url"
        "os"
        "slices"
        "strings"
        "sync"
        "time"

        "github.com/stacklok/toolhive-core/redisconn"
        "github.com/stacklok/toolhive/pkg/auth/dcr"
        "github.com/stacklok/toolhive/pkg/authserver"
        "github.com/stacklok/toolhive/pkg/authserver/authorizationdetails"
        "github.com/stacklok/toolhive/pkg/authserver/internal/consentpolicy"
        servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
        "github.com/stacklok/toolhive/pkg/authserver/server/handlers"
        "github.com/stacklok/toolhive/pkg/authserver/server/keys"
        "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
        "github.com/stacklok/toolhive/pkg/authserver/upstreamfilter"
        "github.com/stacklok/toolhive/pkg/bodylimit"
        "github.com/stacklok/toolhive/pkg/oauthproto"
        "github.com/stacklok/toolhive/pkg/redisconfig"
)

const (
        // TLSListenerPort is the fixed port of the auth server TLS listener. It is
        // part of the issuer URL, so the listener never falls back to another port.
        TLSListenerPort        = 8443
        tlsListenerStopTimeout = 5 * time.Second
)

// Redis ACL credential environment variable names.
// These are set by the operator when Redis storage is configured.
const (
        // RedisUsernameEnvVar is the environment variable for the Redis ACL username.
        // #nosec G101 -- This is an environment variable name, not a hardcoded credential
        RedisUsernameEnvVar = "TOOLHIVE_AUTH_SERVER_REDIS_USERNAME"

        // RedisPasswordEnvVar is the environment variable for the Redis ACL password.
        // #nosec G101 -- This is an environment variable name, not a hardcoded credential
        RedisPasswordEnvVar = "TOOLHIVE_AUTH_SERVER_REDIS_PASSWORD"

        // maxAuthServerBodySize caps auth-server request bodies; shared with the DCR
        // endpoint via handlers.MaxDCRBodySize so the two bounds cannot drift.
        maxAuthServerBodySize = handlers.MaxDCRBodySize
)

// EmbeddedAuthServerOption configures an EmbeddedAuthServer.
type EmbeddedAuthServerOption func(*embeddedAuthServerOptions)

type embeddedAuthServerOptions struct {
        consentExemptClientIDs       []string
        subjectIdentifierAudiences   []string
        accessTokenType              string
        listenerHost                 string
        authorizationDetailsRegistry *authorizationdetails.Registry
}

// WithAuthorizationDetailsRegistry configures the embedded authorization server
// to pass the registry to its RFC 8693 token-exchange handler.
func WithAuthorizationDetailsRegistry(registry *authorizationdetails.Registry) EmbeddedAuthServerOption {
        return func(options *embeddedAuthServerOptions) { options.authorizationDetailsRegistry = registry }
}

// WithListenerHost sets the address the TLS listener binds. Callers pass the
// same host as their MCP listener. There is no default.
func WithListenerHost(host string) EmbeddedAuthServerOption {
        return func(options *embeddedAuthServerOptions) { options.listenerHost = host }
}

// WithConsentExemptClientIDs preauthorizes listed clients to skip ToolHive's
// consent page. Listing a client ID preauthorizes whatever registration currently
// holds that ID, for whichever user completes upstream authentication. It is not
// per-user consent, and runtime eligibility checks do not prove where the
// registration came from. The caller must control registration writes for exempt
// IDs; replacing a registration or expanding its permissions transfers or broadens
// preauthorization. Removing an ID affects new requests only and does not revoke
// tokens or cancel already-approved flows. Upstream authentication still runs.
// The input is copied when the option is created, and repeated options replace the
// list.
func WithConsentExemptClientIDs(ids []string) EmbeddedAuthServerOption {
        copiedIDs := slices.Clone(ids)
        return func(options *embeddedAuthServerOptions) {
                options.consentExemptClientIDs = slices.Clone(copiedIDs)
        }
}

// WithSubjectIdentifierAudiences enables the RFC 9493 sub_id claim and the
// https://toolhive.dev/upstream_subject claim on access tokens from authorization-code
// and device-flow logins, and their refreshes, when every audience of the token is in
// this list. sub_id carries the issuer and subject of the validated upstream OIDC ID
// token. The second claim carries the issuer and the value of the claim the upstream
// keys users on (its subject_claim, or sub). With subject_claim: oid on Entra ID it
// discloses the tenant-wide oid, so list only first-party resource servers. Each
// entry must be an exact member of the RunConfig's allowed_audiences, or server
// construction fails. An empty list disables both claims, and logins then do not
// record the upstream identity. Going from an empty list to a non-empty one adds the
// claims only for sessions whose login happened while the list was non-empty; other
// users get them from their next login. Adding an audience to a non-empty list
// applies at the next refresh. The setting never grants an audience. It is
// deliberately not a RunConfig field yet, so config files and operator CRDs cannot
// set it. The input is copied when the option is created, and repeated options
// replace the list.
func WithSubjectIdentifierAudiences(audiences []string) EmbeddedAuthServerOption {
        copied := slices.Clone(audiences)
        return func(options *embeddedAuthServerOptions) {
                options.subjectIdentifierAudiences = slices.Clone(copied)
        }
}

// WithAccessTokenType sets the JOSE "typ" header on the JWT access tokens the
// embedded server issues. The value may be empty or "JWT" (the default), or
// "at+jwt" (RFC 9068). Some resource servers reject at+jwt by default, such as
// Java ones built on Nimbus JOSE+JWT's default type verifier, including Spring
// Security before 6.5, so enable it only when every resource server accepts it.
// Any other value makes server construction fail. It is deliberately not a
// RunConfig field yet, so config files and operator CRDs cannot set it. Repeated
// options replace the value.
func WithAccessTokenType(typ string) EmbeddedAuthServerOption {
        return func(options *embeddedAuthServerOptions) { options.accessTokenType = typ }
}

// EmbeddedAuthServer wraps the authorization server for integration with the proxy runner.
// It handles configuration transformation from authserver.RunConfig to authserver.Config,
// manages resource lifecycle, and provides HTTP handlers for OAuth/OIDC endpoints.
//
// The DCR credential store is owned by the underlying authserver.Server and
// reached via DCRStore(); see that accessor's doc for SECURITY and lifecycle
// notes. Storing it twice on this struct would create a drift window with
// the server's copy, so we delegate through e.server.DCRStore() instead.
type EmbeddedAuthServer struct {
        server       authserver.Server
        keyProvider  keys.KeyProvider
        tlsListener  *tlsListener
        listenerHost string
        closeOnce    sync.Once
        closeErr     error
}

func resolveEmbeddedAuthServerOptions(opts []EmbeddedAuthServerOption) embeddedAuthServerOptions {
        var options embeddedAuthServerOptions
        for _, opt := range opts {
                opt(&options)
        }
        return options
}

// NewEmbeddedAuthServer creates an EmbeddedAuthServer from authserver.RunConfig.
// It loads signing keys from files, reads HMAC secrets from files, resolves
// upstream and delegate-client secret references from files or environment
// variables, and initializes all auth server components.
//
// The cfg parameter contains file paths and environment variable names that are
// resolved at runtime to build the underlying authserver.Config.
func NewEmbeddedAuthServer(
        ctx context.Context,
        cfg *authserver.RunConfig,
        opts ...EmbeddedAuthServerOption,
) (*EmbeddedAuthServer, error) {
        if cfg == nil {
                return nil, fmt.Errorf("config is required")
        }
        options := resolveEmbeddedAuthServerOptions(opts)
        if err := validateConsentPolicyForRunConfig(cfg, options); err != nil {
                return nil, fmt.Errorf("invalid consent policy: %w", err)
        }

        // Register gjson modifiers used by IdentityFromToken configs (e.g. @upstreamjwt).
        // Without this, modifier-bearing paths silently fail to resolve.
        upstream.RegisterModifiers()

        // Fail loudly on operator-supplied misconfiguration (e.g. a baseline
        // scope absent from scopes_supported) BEFORE touching storage or any
        // other side-effecting work, so a bad config never reaches the network
        // or filesystem. NewEmbeddedAuthServerWithStorage re-validates for callers
        // that invoke it directly; this earlier check keeps the failure ahead of
        // createStorage on this path.
        if err := cfg.Validate(); err != nil {
                return nil, fmt.Errorf("invalid run config: %w", err)
        }
        if err := validateTLSListenerHost(cfg, options.listenerHost); err != nil {
                return nil, err
        }
        // A pure config check, like the validations above, so a misconfigured
        // registry never reaches storage creation.
        if err := validateAuthorizationDetailsRegistryForRunConfig(cfg, options); err != nil {
                return nil, err
        }
        // Create the storage backend FIRST so the DCR resolver and the auth
        // server share the same persistence. Both MemoryStorage and RedisStorage
        // satisfy storage.DCRCredentialStore (verified by package-level var _
        // checks in pkg/authserver/storage), so an explicit type assertion at
        // the boundary is provably safe and keeps the wider Storage interface
        // from advertising secret-bearing DCR methods to every consumer. This
        // is the wiring change that lets a Redis-backed authserver reuse RFC
        // 7591 client registrations across replicas and restarts.
        stor, err := createStorage(ctx, cfg.Storage)
        if err != nil {
                return nil, fmt.Errorf("failed to create storage: %w", err)
        }
        return newEmbeddedAuthServerWithStorage(ctx, cfg, stor, nil, options)
}

// NewEmbeddedAuthServerWithStorage is the exported core constructor that
// builds an EmbeddedAuthServer around a caller-supplied storage backend. It
// lets external composition (e.g. an enterprise build) inject a decorated
// storage.Storage aggregate.
//
// What the injection does NOT solve: the supplied storage is the sole
// persistence boundary for this server instance. Injecting a shared backend
// (e.g. Redis) lets replicas share persisted state — DCR registrations,
// pending authorizations, upstream tokens — but it does not provide
// cross-replica message delivery or fan-out, and it does not establish session
// affinity. A request must still reach a replica that can resolve its session
// from the shared store; pinning a session to a replica (or ensuring all
// session state is in the shared store) remains the caller's responsibility,
// typically at the load balancer.
//
// The supplied storage MUST also implement storage.DCRCredentialStore (both
// OSS MemoryStorage and RedisStorage do); the constructor returns an error if
// it does not. It also validates cfg, so direct callers get the same
// fail-loud config check NewEmbeddedAuthServer performs before dispatch. Optional
// Go-only settings, such as the token-exchange authorization-details registry,
// are supplied through EmbeddedAuthServerOption and are not part of RunConfig.
//
// NewEmbeddedAuthServer dispatches into this helper after running
// createStorage; tests dispatch into it directly so they can supply a
// closeTrackingStorage wrapper to verify the deferred-cleanup contract.
//
// Resource ownership: on success, the returned EmbeddedAuthServer takes
// ownership of stor (its Close releases the backend). On any error path
// after entry, the deferred cleanup closes stor before returning so a
// crash-looping caller (typical when DCR's network I/O fails) does not
// leak the Redis client connection pool / MemoryStorage cleanup goroutine
// on every restart. The named return retErr is the gate.
func NewEmbeddedAuthServerWithStorage(
        ctx context.Context,
        cfg *authserver.RunConfig,
        stor storage.Storage,
        opts ...EmbeddedAuthServerOption,
) (*EmbeddedAuthServer, error) {
        return newEmbeddedAuthServerWithStorage(ctx, cfg, stor, nil, resolveEmbeddedAuthServerOptions(opts))
}

func warnDeprecatedInboundGrantFields(fields []authserver.DeprecatedFieldPath) {
        if len(fields) == 0 {
                return
        }
        paths := make([]string, len(fields))
        for i, field := range fields {
                paths[i] = field.Path + " -> " + field.Replacement
        }
        slog.Warn("deprecated inbound grant configuration; migrate to canonical fields", "paths", strings.Join(paths, ", "))
}

func prepareInboundGrantConfiguration(
        cfg *authserver.RunConfig,
        delegateClients []authserver.DelegateClient,
) (*authserver.NormalizedInboundGrants, []authserver.DelegateClient, *authserver.SPIFFETrustConfig, error) {
        normalized, err := authserver.NormalizeInboundGrants(cfg)
        if err != nil {
                return nil, nil, nil, fmt.Errorf("normalize inbound grants: %w", err)
        }
        warnDeprecatedInboundGrantFields(normalized.DeprecatedFields)

        if delegateClients == nil && len(normalized.DelegateClients) > 0 {
                delegateClients, err = resolveDelegateClients(normalized.DelegateClients)
                if err != nil {
                        return nil, nil, nil, err
                }
        }
        // SPIFFE client authentication is independent of the legacy/canonical
        // token-exchange projection above: it is read straight from cfg.InboundGrants,
        // never through NormalizedInboundGrants, so authentication method and
        // grant-family enablement stay separately configurable.
        spiffeTrust, err := authserver.NewSPIFFETrustConfig(
                cfg.SPIFFETrustDomains, cfg.InboundGrants, cfg.ScopesSupported, cfg.AllowedAudiences,
        )
        if err != nil {
                return nil, nil, nil, fmt.Errorf("build SPIFFE trust config: %w", err)
        }
        // SPIFFE client authentication exclusively uses the token-exchange grant
        // (validateSPIFFEGrants enforces this), so configuring any SPIFFE
        // association implies token-exchange capability independent of
        // legacy/canonical token-exchange enablement.
        normalized.Capabilities.TokenExchange = normalized.Capabilities.TokenExchange || hasSPIFFEClientAuth(cfg)
        return normalized, delegateClients, spiffeTrust, nil
}

// hasSPIFFEClientAuth reports whether cfg declares any SPIFFE client-auth
// association.
func hasSPIFFEClientAuth(cfg *authserver.RunConfig) bool {
        return cfg.InboundGrants != nil && len(cfg.InboundGrants.SPIFFEClientAuth) > 0
}

// buildUpstreamFilter constructs a handlers.UpstreamFilter from the RunConfig,
// if one is configured. It returns an untyped nil (not a typed-nil interface
// value) when cfg.UpstreamFilter is unset so authserver.Config.UpstreamFilter
// remains truly nil and the handler walks all configured upstreams —
// preserving behavior from before the filter hook existed.
//
// RunConfig.Validate has already applied the same static-invariant checks the
// filter constructor runs here (unknown provider references, mandatory-first
// references, fewer-than-two upstreams, etc.), so a Validate-passing cfg will
// not fail here in practice. The second construction is cheap and acts as a
// belt-and-suspenders guard for consumers that bypassed Validate.
func buildUpstreamFilter(cfg *authserver.RunConfig) (handlers.UpstreamFilter, error) {
        if cfg.UpstreamFilter == nil {
                return nil, nil
        }
        configured := make([]string, len(cfg.Upstreams))
        for i := range cfg.Upstreams {
                configured[i] = cfg.Upstreams[i].Name
        }
        f, err := upstreamfilter.NewGroupBasedFilter(*cfg.UpstreamFilter, configured)
        if err != nil {
                return nil, err
        }
        return f, nil
}

// newEmbeddedAuthServerWithStorage assembles the auth server from a RunConfig
// in a strict ordered sequence: validate inputs, derive crypto bits, resolve
// upstream providers and inbound grants, build the handler Config, construct
// the server. Each step is a required stage with its own error-return and
// cannot be reordered, so the function's cyclomatic complexity sits above the
// package lint ceiling by design rather than from avoidable branching. The
// shape mirrors the "sequential boot" exceptions elsewhere in the repository
// (CLI command handlers, reconcile loops) and is preferred over splitting the
// chain into helpers whose only purpose would be to spread the error checks
// across functions without clarifying the boot order.
//
//nolint:gocyclo // Sequential auth-server bootstrap; see function docstring.
func newEmbeddedAuthServerWithStorage(
        ctx context.Context,
        cfg *authserver.RunConfig,
        stor storage.Storage,
        delegateClients []authserver.DelegateClient,
        options embeddedAuthServerOptions,
) (retEAS *EmbeddedAuthServer, retErr error) {
        // Validate required inputs before the deferred cleanup is installed: cfg is
        // dereferenced during validation and stor is closed by that cleanup.
        if err := validateEmbeddedAuthServerInputs(cfg, stor); err != nil {
                return nil, err
        }

        // From here on, any error must close stor before returning.
        defer closeStorageOnError(stor, &retErr)

        // Validate consent prerequisites before upstream resolution, which may
        // perform discovery and DCR network I/O. This remains inside the cleanup
        // gate so caller-supplied storage is closed on constructor failure.
        if err := validateConsentPolicyForRunConfig(cfg, options); err != nil {
                return nil, fmt.Errorf("invalid consent policy: %w", err)
        }

        // Validate cfg here too. NewEmbeddedAuthServer validates before
        // createStorage, but direct callers of this exported constructor would
        // otherwise skip the check. Placed inside the deferred-cleanup gate above so
        // a validation failure still closes the caller-supplied storage per the
        // resource-ownership contract.
        if err := cfg.Validate(); err != nil {
                return nil, fmt.Errorf("invalid run config: %w", err)
        }
        if err := validateTLSListenerHost(cfg, options.listenerHost); err != nil {
                return nil, err
        }
        normalized, delegateClients, spiffeTrust, err := prepareInboundGrantConfiguration(cfg, delegateClients)
        if err != nil {
                return nil, err
        }
        if options.authorizationDetailsRegistry != nil && !normalized.Capabilities.TokenExchange {
                return nil, fmt.Errorf("authorization details registry requires token exchange to be enabled")
        }

        if err := authserver.PreflightSPIFFEStaticClientCollisions(ctx, stor, spiffeTrust); err != nil {
                return nil, fmt.Errorf("preflight SPIFFE static client collisions: %w", err)
        }

        // 1. Create key provider from RunConfig.SigningKeyConfig
        keyProvider, err := createKeyProvider(cfg.SigningKeyConfig)
        if err != nil {
                return nil, fmt.Errorf("failed to create key provider: %w", err)
        }

        // 2. Load HMAC secrets from files
        hmacSecrets, err := loadHMACSecrets(cfg.HMACSecretFiles)
        if err != nil {
                return nil, fmt.Errorf("failed to load HMAC secrets: %w", err)
        }

        // 3. Parse token lifespans
        accessLifespan, refreshLifespan, authCodeLifespan, err := parseTokenLifespans(cfg.TokenLifespans)
        if err != nil {
                return nil, fmt.Errorf("failed to parse token lifespans: %w", err)
        }

        // 4. Type-assert to the DCR-capable handle for the resolver. The
        // per-backend `var _ DCRCredentialStore = (*MemoryStorage)(nil)` /
        // `(*RedisStorage)(nil)` checks make this provably safe for production
        // backends; surfacing a non-DCR backend as a constructor error keeps
        // misconfiguration fail-loud at boot rather than at first DCR resolve.
        dcrStore, ok := stor.(storage.DCRCredentialStore)
        if !ok {
                return nil, fmt.Errorf("storage backend %T does not implement storage.DCRCredentialStore", stor)
        }

        // 5. Build upstream configurations. The DCR resolver caches RFC 7591
        // resolutions in dcrStore so re-entrant boot/reload paths reuse
        // previously-registered upstream clients instead of re-registering.
        upstreams, err := buildUpstreamConfigs(
                ctx, cfg.Upstreams, cfg.Issuer, dcr.NewStorageBackedStore(dcrStore), cfg.InsecureAllowHTTP,
        )
        if err != nil {
                return nil, fmt.Errorf("failed to build upstream configs: %w", err)
        }

        // 6. Parse delegation token lifespan if configured.
        delegationLifespan, err := parseOptionalDuration(cfg.DelegationTokenLifespan, "delegation token lifespan")
        if err != nil {
                return nil, err
        }

        // 7. Build the resolved Config.
        //
        // Defensive copies of the scope/audience slices: cfg is operator-supplied
        // input that may be retained or mutated by the caller (e.g. tests, a
        // future hot-reload path). The DCR handler reads these slices on every
        // request, so a mid-request mutation of the original would race. Cloning
        // here once at the boundary lets all downstream stages share by reference
        // safely. Cost is negligible — each slice is bounded by validation (≤10
        // for BaselineClientScopes, low cardinality in practice for the others).
        cimdEnabled, cimdCacheMaxSize, cimdCacheFallbackTTL := resolveCIMDConfig(cfg.CIMD)

        trustedIssuers, err := tokenexchange.ResolveJWTBearerGrantPolicies(normalized.TrustedIssuers)
        if err != nil {
                return nil, fmt.Errorf("failed to resolve JWT-bearer grant policies: %w", err)
        }

        upstreamFilter, err := buildUpstreamFilter(cfg)
        if err != nil {
                return nil, fmt.Errorf("failed to build upstream filter: %w", err)
        }

        resolvedCfg := authserver.Config{
                Issuer:                                    cfg.Issuer,
                AuthorizationEndpointBaseURL:              cfg.AuthorizationEndpointBaseURL,
                KeyProvider:                               keyProvider,
                HMACSecrets:                               hmacSecrets,
                AccessTokenLifespan:                       accessLifespan,
                RefreshTokenLifespan:                      refreshLifespan,
                AuthCodeLifespan:                          authCodeLifespan,
                DelegationTokenLifespan:                   delegationLifespan,
                Upstreams:                                 upstreams,
                ConsentExemptClientIDs:                    slices.Clone(options.consentExemptClientIDs),
                ScopesSupported:                           slices.Clone(cfg.ScopesSupported),
                BaselineClientScopes:                      slices.Clone(cfg.BaselineClientScopes),
                AllowedAudiences:                          slices.Clone(cfg.AllowedAudiences),
                SubjectIdentifierAudiences:                slices.Clone(options.subjectIdentifierAudiences),
                AccessTokenType:                           options.accessTokenType,
                CIMDEnabled:                               cimdEnabled,
                CIMDCacheMaxSize:                          cimdCacheMaxSize,
                CIMDCacheFallbackTTL:                      cimdCacheFallbackTTL,
                InsecureAllowHTTP:                         cfg.InsecureAllowHTTP,
                AllowConfidentialClientRegistration:       cfg.AllowConfidentialClientRegistration,
                AllowPrivateKeyJWTRegistration:            cfg.AllowPrivateKeyJWTRegistration,
                ForceConfidentialRedirectURIs:             slices.Clone(cfg.ForceConfidentialRedirectURIs),
                InsecureAllowConfidentialOverLoopbackHTTP: cfg.InsecureAllowConfidentialOverLoopbackHTTP,
                RegistrationRateLimit:                     cfg.RegistrationRateLimit,
                // slices.Clone is shallow: each TrustedIssuer's own AllowedActors
                // slice is still shared with cfg. NewMultiIssuerTokenValidator's
                // constructor clones AllowedActors per issuer before use, so the
                // authorization-critical data is protected without a deep copy here.
                TrustedIssuers:  trustedIssuers,
                DelegateClients: delegateClients,
                // SPIFFE client authentication factors into Capabilities.TokenExchange
                // already (see prepareInboundGrantConfiguration): it exclusively uses
                // the token-exchange grant, independent of legacy/canonical enablement.
                DisableTokenExchange:         !normalized.Capabilities.TokenExchange,
                AuthorizationDetailsRegistry: options.authorizationDetailsRegistry,
                SPIFFETrust:                  spiffeTrust,
                DeviceFlowEnabled:            cfg.DeviceFlowEnabled,
                UpstreamFilter:               upstreamFilter,
        }

        listenerKeyPair, err := loadTLSListenerKeyPair(cfg)
        if err != nil {
                return nil, err
        }

        // 8. Create the auth server. authserver.New also asserts the DCR
        // capability internally so its DCRStore() accessor returns the same
        // asserted handle this constructor used for buildUpstreamConfigs.
        server, err := authserver.New(ctx, resolvedCfg, stor)
        if err != nil {
                return nil, fmt.Errorf("failed to create auth server: %w", err)
        }

        e := &EmbeddedAuthServer{
                server:       server,
                keyProvider:  keyProvider,
                listenerHost: options.listenerHost,
        }
        e.attachTLSListener(cfg, listenerKeyPair)
        return e, nil
}

// closeStorageOnError closes stor when the constructor is returning an error.
//
// Both errors are passed through dcr.SanitizeErrorForLog before being
// recorded: closeErr for symmetry with retErr, retErr because the
// most common cause of reaching this gate is a wrapped DCR failure
// whose error chain may inline several KiB of the upstream's raw
// /register response body — that body is attacker-influenced and may
// contain URL components that carry credentials (userinfo, query,
// fragment). The existing dcr.LogStepError boundary log routes
// through the same sanitiser; keep the two log paths consistent so
// the cleanup log cannot regress to a less-defended state. The
// "cause" key matches the package-wide vocabulary for the
// triggering error.
func closeStorageOnError(stor storage.Storage, retErr *error) {
        if *retErr == nil {
                return
        }
        if closeErr := stor.Close(); closeErr != nil {
                slog.Warn("failed to close storage on NewEmbeddedAuthServer error path",
                        "error", dcr.SanitizeErrorForLog(closeErr),
                        "cause", dcr.SanitizeErrorForLog(*retErr),
                )
        }
}

// loadTLSListenerKeyPair loads and validates the TLS listener certificate. It
// returns nil when no listener is configured.
func loadTLSListenerKeyPair(cfg *authserver.RunConfig) (*tlsKeyPairCache, error) {
        if cfg.TLSListener == nil {
                return nil, nil
        }
        keyPair, err := newTLSKeyPairCache(cfg.TLSListener.CertFile, cfg.TLSListener.KeyFile)
        if err != nil {
                return nil, fmt.Errorf("load auth server TLS listener certificate: %w", err)
        }
        if err := validateTLSListenerIssuer(cfg, keyPair.leaf); err != nil {
                return nil, err
        }
        return keyPair, nil
}

// attachTLSListener builds the TLS listener when one is configured.
func (e *EmbeddedAuthServer) attachTLSListener(cfg *authserver.RunConfig, keyPair *tlsKeyPairCache) {
        if cfg.TLSListener == nil {
                return
        }
        authorities := func() []*x509.Certificate {
                certs, _ := authserver.SPIFFEX509Authorities(e.server)
                return certs
        }
        e.tlsListener = newTLSListenerWithCache(
                e.listenerHost, keyPair, e.Routes(), authorities, cfg.HasSPIFFEX509ClientAuth(),
        )
}

// Handler returns the HTTP handler for OAuth/OIDC endpoints.
// The handler uses internal chi routing and serves all endpoints:
//   - /oauth/authorize, /oauth/callback, /oauth/token, /oauth/register
//   - /.well-known/jwks.json, /.well-known/oauth-authorization-server, /.well-known/openid-configuration
//
// All auth-server endpoints are body-size-limited to handlers.MaxDCRBodySize
// (64KB) and reject oversized requests with HTTP 413 Request Entity Too Large.
func (e *EmbeddedAuthServer) Handler() http.Handler {
        // Cap request bodies on all auth-server endpoints (e.g. POST /oauth/token,
        // /oauth/register) so they cannot be used for memory-exhaustion DoS. These
        // routes are mounted outside the MCP middleware chain (the proxies and vMCP
        // mount them via Routes()/RegisterHandlers, which derive from this handler),
        // so the cap is applied here to cover every consumer.
        return bodylimit.Middleware(maxAuthServerBodySize)(e.server.Handler())
}

// Start binds Host:TLSListenerPort synchronously and serves in a goroutine.
// No-op when cfg.TLSListener is nil. Returns an error when cfg.TLSListener is
// set but no host was given with WithListenerHost. A bind failure is returned,
// never retried on another port.
func (e *EmbeddedAuthServer) Start() error {
        if e.tlsListener == nil {
                return nil
        }
        if e.listenerHost == "" {
                return fmt.Errorf("auth server TLS listener host is required")
        }
        return e.tlsListener.start()
}

// Close releases resources held by the EmbeddedAuthServer.
// This method is idempotent - subsequent calls after the first will return
// the same error (if any) without attempting to close resources again.
// Should be called during runner shutdown.
func (e *EmbeddedAuthServer) Close() error {
        e.closeOnce.Do(func() {
                var listenerErr error
                if e.tlsListener != nil {
                        ctx, cancel := context.WithTimeout(context.Background(), tlsListenerStopTimeout)
                        listenerErr = e.tlsListener.shutdown(ctx)
                        cancel()
                }
                e.closeErr = errors.Join(listenerErr, e.server.Close())
        })
        return e.closeErr
}

// CloseIdleConnections releases the idle keep-alive connections pooled by the
// underlying server's upstream IDP providers, without touching storage. See
// authserver.CloseIdleConnections for when to use this rather than Close.
func (e *EmbeddedAuthServer) CloseIdleConnections() {
        authserver.CloseIdleConnections(e.server)
}

// IDPTokenStorage returns storage for upstream IDP tokens.
// Returns nil if no upstream IDP is configured.
// This is used by the upstream swap middleware to exchange ToolHive JWTs
// for upstream IDP tokens.
func (e *EmbeddedAuthServer) IDPTokenStorage() storage.UpstreamTokenStorage {
        return e.server.IDPTokenStorage()
}

// UpstreamTokenRefresher returns a refresher that can refresh expired upstream
// tokens using the upstream provider's refresh token grant.
func (e *EmbeddedAuthServer) UpstreamTokenRefresher() storage.UpstreamTokenRefresher {
        return e.server.UpstreamTokenRefresher()
}

// KeyProvider returns the signing key provider used by the authorization server.
// This enables in-process JWKS key lookups, eliminating the need for
// self-referential HTTP calls when the token validator runs in the same process.
func (e *EmbeddedAuthServer) KeyProvider() keys.KeyProvider {
        return e.keyProvider
}

// DCRStore returns the persistent DCR credential store the authorization
// server is wired against. This delegates to the underlying authserver.Server
// so this struct does not hold a redundant copy that could drift if the
// server ever swaps backends. See authserver.Server.DCRStore for SECURITY
// and lifecycle notes — the returned interface surfaces raw client_secret
// and registration_access_token values and MUST NOT be logged or rendered.
func (e *EmbeddedAuthServer) DCRStore() storage.DCRCredentialStore {
        return e.server.DCRStore()
}

// Routes returns the authorization server's HTTP route map.
//
// The /.well-known/ paths are registered explicitly because that namespace is shared:
// the vMCP server owns /.well-known/oauth-protected-resource (RFC 9728) on the same
// mux. Adding a new AS /.well-known/ endpoint therefore requires an explicit entry here.
//
// Discovery paths are registered with both exact and trailing-slash (prefix) patterns.
// The trailing-slash variants support RFC 8414 Section 3.1 path-based issuers, where
// the client constructs /.well-known/oauth-authorization-server/{issuer-path}.
//
// The /oauth/ subtree is registered as a prefix, so new /oauth/* endpoints added to
// the chi router are picked up automatically without changes to this method.
func (e *EmbeddedAuthServer) Routes() map[string]http.Handler {
        handler := e.Handler()
        return map[string]http.Handler{
                "/.well-known/openid-configuration":        handler,
                "/.well-known/openid-configuration/":       handler,
                "/.well-known/oauth-authorization-server":  handler,
                "/.well-known/oauth-authorization-server/": handler,
                "/.well-known/jwks.json":                   handler,
                "/oauth/":                                  handler,
        }
}

// RegisterHandlers registers the authorization server's HTTP routes on the given mux.
func (e *EmbeddedAuthServer) RegisterHandlers(mux *http.ServeMux) {
        for pattern, handler := range e.Routes() {
                mux.Handle(pattern, handler)
        }
}

func validateConsentPolicyForRunConfig(cfg *authserver.RunConfig, options embeddedAuthServerOptions) error {
        return consentpolicy.Validate(
                options.consentExemptClientIDs,
                cfg.AuthorizationEndpointBaseURL,
                cfg.Issuer,
                len(cfg.Upstreams),
        )
}

func validateAuthorizationDetailsRegistryForRunConfig(cfg *authserver.RunConfig, options embeddedAuthServerOptions) error {
        if options.authorizationDetailsRegistry == nil {
                return nil
        }
        normalized, err := authserver.NormalizeInboundGrants(cfg)
        if err != nil {
                return fmt.Errorf("normalize inbound grants: %w", err)
        }
        if !normalized.Capabilities.TokenExchange {
                return fmt.Errorf("authorization details registry requires token exchange to be enabled")
        }
        return nil
}

func validateTLSListenerIssuer(cfg *authserver.RunConfig, leaf *x509.Certificate) error {
        if !cfg.HasSPIFFEX509ClientAuth() {
                return nil
        }
        issuer, err := url.Parse(cfg.Issuer)
        if err != nil {
                return fmt.Errorf("parse auth server issuer: %w", err)
        }
        if issuer.Scheme != "https" {
                return fmt.Errorf("SPIFFE X.509 authentication requires an https issuer")
        }
        port := issuer.Port()
        if port == "" {
                port = "443"
        }
        if port != fmt.Sprintf("%d", TLSListenerPort) {
                slog.Warn("auth server issuer port differs from TLS listener port", "issuer", cfg.Issuer, "port", port)
        }
        if leaf != nil {
                if err := leaf.VerifyHostname(issuer.Hostname()); err != nil {
                        slog.Warn("auth server TLS listener certificate does not match issuer hostname",
                                "issuer", cfg.Issuer, "hostname", issuer.Hostname())
                }
        }
        return nil
}

func validateTLSListenerHost(cfg *authserver.RunConfig, listenerHost string) error {
        if cfg.TLSListener != nil && listenerHost == "" {
                return fmt.Errorf("auth server TLS listener host is required; configure it with WithListenerHost")
        }
        return nil
}

// validateEmbeddedAuthServerInputs rejects required inputs before construction
// can dereference cfg or install cleanup that closes stor.
func validateEmbeddedAuthServerInputs(cfg *authserver.RunConfig, stor storage.Storage) error {
        if cfg == nil {
                return fmt.Errorf("config is required")
        }
        if stor == nil {
                return fmt.Errorf("storage is required")
        }
        return nil
}

// createKeyProvider creates a KeyProvider from SigningKeyRunConfig.
// Returns a GeneratingProvider if config is nil or empty (development mode).
func createKeyProvider(cfg *authserver.SigningKeyRunConfig) (keys.KeyProvider, error) {
        if cfg == nil || cfg.SigningKeyFile == "" {
                // Development mode: use ephemeral key
                return keys.NewGeneratingProvider(keys.DefaultAlgorithm), nil
        }

        keyCfg := keys.Config{
                KeyDir:           cfg.KeyDir,
                SigningKeyFile:   cfg.SigningKeyFile,
                FallbackKeyFiles: cfg.FallbackKeyFiles,
        }

        return keys.NewFileProvider(keyCfg)
}

// loadHMACSecrets reads HMAC secrets from files.
// Returns nil if no files are configured (development mode - authserver will generate ephemeral secret).
func loadHMACSecrets(files []string) (*servercrypto.HMACSecrets, error) {
        if len(files) == 0 {
                // Development mode: let authserver generate ephemeral secret
                return nil, nil
        }

        // Read current (first) secret
        // #nosec G304 - file path is from configuration, not user input
        current, err := os.ReadFile(files[0])
        if err != nil {
                return nil, fmt.Errorf("failed to read HMAC secret from %s: %w", files[0], err)
        }

        // HMAC secrets are raw binary (see MCPExternalAuthConfig.HMACSecretRefs doc) — do not
        // trim; any byte, including whitespace-valued ones, may be genuine key material.
        if len(current) < servercrypto.MinSecretLength {
                return nil, fmt.Errorf("HMAC secret in %s is %d bytes, but must be at least %d bytes",
                        files[0], len(current), servercrypto.MinSecretLength)
        }

        secrets := &servercrypto.HMACSecrets{
                Current: current,
        }

        // Read rotated secrets (remaining files)
        for _, file := range files[1:] {
                if file == "" {
                        continue // Skip empty paths
                }
                // #nosec G304 - file path is from configuration, not user input
                secret, err := os.ReadFile(file)
                if err != nil {
                        return nil, fmt.Errorf("failed to read rotated HMAC secret from %s: %w", file, err)
                }
                if len(secret) < servercrypto.MinSecretLength {
                        return nil, fmt.Errorf("rotated HMAC secret in %s is %d bytes, but must be at least %d bytes",
                                file, len(secret), servercrypto.MinSecretLength)
                }
                secrets.Rotated = append(secrets.Rotated, secret)
        }

        return secrets, nil
}

func parseOptionalDuration(value, name string) (time.Duration, error) {
        if value == "" {
                return 0, nil
        }
        duration, err := time.ParseDuration(value)
        if err != nil {
                return 0, fmt.Errorf("invalid %s: %w", name, err)
        }
        return duration, nil
}

// parseTokenLifespans parses duration strings from TokenLifespanRunConfig.
// Returns zero values for unset durations (defaults applied by authserver).
func parseTokenLifespans(cfg *authserver.TokenLifespanRunConfig) (access, refresh, authCode time.Duration, err error) {
        if cfg == nil {
                return 0, 0, 0, nil
        }

        if cfg.AccessTokenLifespan != "" {
                access, err = time.ParseDuration(cfg.AccessTokenLifespan)
                if err != nil {
                        return 0, 0, 0, fmt.Errorf("invalid access token lifespan: %w", err)
                }
        }

        if cfg.RefreshTokenLifespan != "" {
                refresh, err = time.ParseDuration(cfg.RefreshTokenLifespan)
                if err != nil {
                        return 0, 0, 0, fmt.Errorf("invalid refresh token lifespan: %w", err)
                }
        }

        if cfg.AuthCodeLifespan != "" {
                authCode, err = time.ParseDuration(cfg.AuthCodeLifespan)
                if err != nil {
                        return 0, 0, 0, fmt.Errorf("invalid auth code lifespan: %w", err)
                }
        }

        return access, refresh, authCode, nil
}

// buildUpstreamConfigs converts UpstreamRunConfig slice to UpstreamConfig slice.
// It preserves the provider type so the factory can create the correct provider
// (OIDCProviderImpl for OIDC, BaseOAuth2Provider for OAuth2).
//
// For OAuth2 upstreams configured with DCRConfig, buildUpstreamConfigs performs
// RFC 7591 Dynamic Client Registration against the upstream authorization
// server (hitting the network on first call, using dcrStore on subsequent
// calls) and overlays the resulting ClientID / ClientSecret onto the output
// config via consumeResolution + applyResolutionToOAuth2Config (see
// dcr_adapter.go). The caller's runConfigs slice is not mutated: in-place
// mutation of caller-provided values surprises callers and can cause data
// races, so each element is cloned before applying DCR resolution.
//
// Error logging: this function is the boundary for DCR errors — on any
// failure from dcr.ResolveCredentials it emits exactly one structured
// slog.Error via dcr.LogStepError and returns the wrapped error to the
// caller without logging further. The resolver itself does not log
// errors, which avoids the log-and-return double-reporting pattern.
func buildUpstreamConfigs(
        ctx context.Context,
        runConfigs []authserver.UpstreamRunConfig,
        issuer string,
        dcrStore dcr.CredentialStore,
        insecureAllowHTTP bool,
) ([]authserver.UpstreamConfig, error) {
        configs := make([]authserver.UpstreamConfig, 0, len(runConfigs))

        for _, rc := range runConfigs {
                cfg, err := buildOneUpstreamConfig(ctx, rc, issuer, dcrStore, insecureAllowHTTP)
                if err != nil {
                        return nil, err
                }
                configs = append(configs, *cfg)
        }

        return configs, nil
}

// buildOneUpstreamConfig validates rc, resolves DCR credentials if required,
// and builds the resulting authserver.UpstreamConfig. Split out of
// buildUpstreamConfigs to keep the per-upstream steps (validate, resolve,
// build, apply) independently testable and under the complexity limit.
func buildOneUpstreamConfig(
        ctx context.Context,
        rc authserver.UpstreamRunConfig,
        issuer string,
        dcrStore dcr.CredentialStore,
        insecureAllowHTTP bool,
) (*authserver.UpstreamConfig, error) {
        // Shallow copy of the outer UpstreamRunConfig so DCR resolution never
        // mutates the caller's slice element.
        rcCopy := rc

        if err := validateUpstreamRunConfig(&rcCopy); err != nil {
                return nil, fmt.Errorf("upstream %q: %w", rc.Name, err)
        }

        resolution, err := resolveUpstreamDCR(ctx, &rcCopy, issuer, dcrStore)
        if err != nil {
                return nil, fmt.Errorf("upstream %q: %w", rc.Name, err)
        }

        cfg, err := buildUpstreamConfig(&rcCopy, insecureAllowHTTP)
        if err != nil {
                return nil, fmt.Errorf("upstream %q: %w", rc.Name, err)
        }

        // Apply the DCR-resolved ClientSecret to the built provider config. The
        // split between consumeResolution/consumeOIDCResolution (run-config
        // fields) and applyResolutionTo*Config (inline-only ClientSecret) is
        // documented in dcr_adapter.go — both calls must be paired to produce a
        // fully-resolved DCR client.
        applyDCRResolutionToConfig(cfg, rcCopy.Type, resolution)

        return cfg, nil
}

// validateUpstreamRunConfig enforces that rc carries the sub-config matching
// its declared Type and that the sub-config's own invariants hold.
func validateUpstreamRunConfig(rc *authserver.UpstreamRunConfig) error {
        switch rc.Type {
        case authserver.UpstreamProviderTypeOIDC:
                if rc.OIDCConfig == nil {
                        return fmt.Errorf("oidc_config required")
                }
                return rc.OIDCConfig.Validate()
        case authserver.UpstreamProviderTypeOAuth2:
                if rc.OAuth2Config == nil {
                        return fmt.Errorf("oauth2_config required")
                }
                return rc.OAuth2Config.Validate()
        }
        return nil
}

// resolveUpstreamDCR performs Dynamic Client Registration for rc when
// required, folding the result back into rc's sub-config in place. It
// returns a nil Resolution when rc has a pre-provisioned client and no DCR
// is needed.
func resolveUpstreamDCR(
        ctx context.Context,
        rc *authserver.UpstreamRunConfig,
        issuer string,
        dcrStore dcr.CredentialStore,
) (*dcr.Resolution, error) {
        req, err := newDCRRequest(rc, issuer)
        if err != nil {
                return nil, err
        }
        if req == nil {
                return nil, nil
        }

        resolution, err := dcr.ResolveCredentials(ctx, req, dcrStore)
        if err != nil {
                dcr.LogStepError(rc.Name, err)
                return nil, err
        }

        switch rc.Type {
        case authserver.UpstreamProviderTypeOIDC:
                resolved := consumeOIDCResolution(*rc.OIDCConfig, resolution)
                rc.OIDCConfig = &resolved
        case authserver.UpstreamProviderTypeOAuth2:
                resolved := consumeResolution(*rc.OAuth2Config, resolution)
                rc.OAuth2Config = &resolved
        }
        return resolution, nil
}

// applyDCRResolutionToConfig overlays the DCR-resolved ClientSecret onto the
// built provider config. No-op when resolution is nil (no DCR was performed).
func applyDCRResolutionToConfig(
        cfg *authserver.UpstreamConfig,
        upstreamType authserver.UpstreamProviderType,
        resolution *dcr.Resolution,
) {
        if resolution == nil {
                return
        }
        switch upstreamType {
        case authserver.UpstreamProviderTypeOIDC:
                if cfg.OIDCConfig != nil {
                        applied := applyResolutionToOIDCConfig(*cfg.OIDCConfig, resolution)
                        cfg.OIDCConfig = &applied
                }
        case authserver.UpstreamProviderTypeOAuth2:
                if cfg.OAuth2Config != nil {
                        applied := applyResolutionToOAuth2Config(*cfg.OAuth2Config, resolution)
                        cfg.OAuth2Config = &applied
                }
        }
}

// buildUpstreamConfig builds an authserver.UpstreamConfig from UpstreamRunConfig.
// It preserves the provider type and builds the appropriate config.
func buildUpstreamConfig(rc *authserver.UpstreamRunConfig, insecureAllowHTTP bool) (*authserver.UpstreamConfig, error) {
        switch rc.Type {
        case authserver.UpstreamProviderTypeOIDC:
                oidcCfg, err := buildOIDCConfig(rc, insecureAllowHTTP)
                if err != nil {
                        return nil, err
                }
                return &authserver.UpstreamConfig{
                        Name:        rc.Name,
                        DisplayName: rc.DisplayName,
                        Type:        authserver.UpstreamProviderTypeOIDC,
                        OIDCConfig:  oidcCfg,
                }, nil

        case authserver.UpstreamProviderTypeOAuth2:
                oauth2Cfg, err := buildPureOAuth2Config(rc, insecureAllowHTTP)
                if err != nil {
                        return nil, err
                }
                return &authserver.UpstreamConfig{
                        Name:         rc.Name,
                        DisplayName:  rc.DisplayName,
                        Type:         authserver.UpstreamProviderTypeOAuth2,
                        OAuth2Config: oauth2Cfg,
                }, nil

        default:
                return nil, fmt.Errorf("unsupported upstream type: %s", rc.Type)
        }
}

// buildOIDCConfig builds an upstream.OIDCConfig for an OIDC provider.
// Discovery is deferred to the provider factory - we only resolve secrets here.
//
// Note: OIDCUpstreamRunConfig.UserInfoOverride is intentionally NOT propagated.
// OIDC providers resolve user identity from the ID token's "sub" claim (validated
// by OIDCProviderImpl.ExchangeCodeForIdentity), not from the UserInfo endpoint.
// The UserInfo endpoint may still be discovered via OIDC discovery for other
// purposes, but it is not used for identity resolution.
func buildOIDCConfig(rc *authserver.UpstreamRunConfig, insecureAllowHTTP bool) (*upstream.OIDCConfig, error) {
        if rc.OIDCConfig == nil {
                return nil, fmt.Errorf("oidc_config required for OIDC provider")
        }

        oidc := rc.OIDCConfig
        if err := oidc.Validate(); err != nil {
                return nil, err
        }

        // Warn if UserInfoOverride is configured but won't be used
        if oidc.UserInfoOverride != nil {
                slog.Warn("userinfo_override is configured for OIDC provider but will not be used; "+
                        "OIDC providers resolve identity from the ID token, not the UserInfo endpoint",
                        "upstream", rc.Name,
                )
        }

        clientSecret, err := resolveSecret(oidc.ClientSecretFile, oidc.ClientSecretEnvVar)
        if err != nil {
                return nil, fmt.Errorf("failed to resolve OIDC client secret: %w", err)
        }

        // Default scopes if not specified. The default includes offline_access
        // (standard OIDC mechanism for refresh tokens). Providers like Google that
        // use access_type=offline instead should specify explicit scopes in their
        // config to avoid sending both mechanisms.
        scopes := oidc.Scopes
        if len(scopes) == 0 {
                scopes = []string{"openid", "offline_access"}
        }

        return &upstream.OIDCConfig{
                CommonOAuthConfig: upstream.CommonOAuthConfig{
                        ClientID:                      oidc.ClientID,
                        ClientSecret:                  clientSecret,
                        RedirectURI:                   oidc.RedirectURI,
                        Scopes:                        scopes,
                        AdditionalAuthorizationParams: oidc.AdditionalAuthorizationParams,
                        AdditionalTokenParams:         oidc.AdditionalTokenParams,
                        OmitScopeOnRefresh:            oidc.OmitScopeOnRefresh,
                },
                Issuer:            oidc.IssuerURL,
                SubjectClaim:      oidc.SubjectClaim,
                AllowPrivateIPs:   oidc.AllowPrivateIPs,
                CAFilePath:        oidc.CAFilePath,
                InsecureAllowHTTP: insecureAllowHTTP || oidc.InsecureAllowHTTP,
        }, nil
}

// buildPureOAuth2Config builds an upstream.OAuth2Config for a pure OAuth2 provider.
//
// Run-config-specific invariants (e.g. ClientID/DCRConfig mutual exclusion) are
// enforced here via OAuth2UpstreamRunConfig.Validate before secrets are
// resolved, since the downstream upstream.OAuth2Config validator only sees the
// flattened runtime shape and cannot observe DCR fields.
func buildPureOAuth2Config(rc *authserver.UpstreamRunConfig, insecureAllowHTTP bool) (*upstream.OAuth2Config, error) {
        if rc.OAuth2Config == nil {
                return nil, fmt.Errorf("oauth2_config required for OAuth2 provider")
        }

        oauth2 := rc.OAuth2Config
        if err := oauth2.Validate(); err != nil {
                return nil, err
        }
        clientSecret, err := resolveSecret(oauth2.ClientSecretFile, oauth2.ClientSecretEnvVar)
        if err != nil {
                return nil, fmt.Errorf("failed to resolve OAuth2 client secret: %w", err)
        }

        authMethod := oauth2.TokenEndpointAuthMethod
        if isConfidentialAuthMethod(authMethod) && clientSecret == "" {
                return nil, fmt.Errorf(
                        "oauth2 upstream: token_endpoint_auth_method %q requires a non-empty client secret, "+
                                "but the configured secret resolved to an empty value", authMethod)
        }

        cfg := &upstream.OAuth2Config{
                CommonOAuthConfig: upstream.CommonOAuthConfig{
                        ClientID:                      oauth2.ClientID,
                        ClientSecret:                  clientSecret,
                        RedirectURI:                   oauth2.RedirectURI,
                        Scopes:                        oauth2.Scopes,
                        AdditionalAuthorizationParams: oauth2.AdditionalAuthorizationParams,
                        AdditionalTokenParams:         oauth2.AdditionalTokenParams,
                        OmitScopeOnRefresh:            oauth2.OmitScopeOnRefresh,
                },
                AuthorizationEndpoint:   oauth2.AuthorizationEndpoint,
                TokenEndpoint:           oauth2.TokenEndpoint,
                TokenEndpointAuthMethod: authMethod,
                UserInfo:                convertUserInfoConfig(oauth2.UserInfo),
                CAFilePath:              oauth2.CAFilePath,
                AllowPrivateIPs:         oauth2.AllowPrivateIPs,
                InsecureAllowHTTP:       insecureAllowHTTP || oauth2.InsecureAllowHTTP,
        }

        if oauth2.TokenResponseMapping != nil {
                cfg.TokenResponseMapping = &upstream.TokenResponseMapping{
                        AccessTokenPath:  oauth2.TokenResponseMapping.AccessTokenPath,
                        ScopePath:        oauth2.TokenResponseMapping.ScopePath,
                        RefreshTokenPath: oauth2.TokenResponseMapping.RefreshTokenPath,
                        ExpiresInPath:    oauth2.TokenResponseMapping.ExpiresInPath,
                }
        }

        if oauth2.IdentityFromToken != nil {
                cfg.IdentityFromToken = &upstream.IdentityFromTokenConfig{
                        SubjectPath: oauth2.IdentityFromToken.SubjectPath,
                        NamePath:    oauth2.IdentityFromToken.NamePath,
                        EmailPath:   oauth2.IdentityFromToken.EmailPath,
                }
        }

        return cfg, nil
}

// isConfidentialAuthMethod reports whether method requires a client secret to
// be presented at the token endpoint. Used to catch a secret file that reads
// successfully but is empty after trimming -- a case OAuth2UpstreamRunConfig.Validate
// cannot see, since it only knows whether a secret source is configured, not
// what that source resolves to.
func isConfidentialAuthMethod(method string) bool {
        switch method {
        case oauthproto.TokenEndpointAuthMethodClientSecretBasic, oauthproto.TokenEndpointAuthMethodClientSecretPost:
                return true
        default:
                return false
        }
}

// resolveSecret reads a secret from file or environment variable.
// File takes precedence over env var. Returns an error if file is specified but
// unreadable, or if envVar is specified but not set. Returns empty string with
// no error if neither file nor envVar is specified.
func resolveSecret(file, envVar string) (string, error) {
        if file != "" {
                // #nosec G304 - file path is from configuration, not user input
                data, err := os.ReadFile(file)
                if err != nil {
                        return "", fmt.Errorf("failed to read secret file %q: %w", file, err)
                }
                return string(bytes.TrimSpace(data)), nil
        }
        if envVar != "" {
                value := os.Getenv(envVar)
                if value == "" {
                        return "", fmt.Errorf("environment variable %q is not set", envVar)
                }
                return value, nil
        }
        slog.Debug("no client secret configured (neither file nor env var specified)")
        return "", nil
}

// resolveDelegateClients resolves secret references and copies authorization
// permissions before configuration crosses into the running authorization server.
func resolveDelegateClients(clients []authserver.DelegateClientRunConfig) ([]authserver.DelegateClient, error) {
        resolved := make([]authserver.DelegateClient, len(clients))
        for i, client := range clients {
                secret, err := resolveSecret(client.ClientSecretFile, client.ClientSecretEnvVar)
                if err != nil {
                        return nil, fmt.Errorf("delegate client %q: %w", client.ClientID, err)
                }
                if secret == "" {
                        return nil, fmt.Errorf("delegate client %q: resolved client secret is empty", client.ClientID)
                }
                resolved[i] = authserver.DelegateClient{
                        ClientID:     client.ClientID,
                        ClientSecret: secret,
                        Scopes:       slices.Clone(client.Scopes),
                        Audiences:    slices.Clone(client.Audiences),
                }
        }
        return resolved, nil
}

// convertUserInfoConfig converts UserInfoRunConfig to upstream.UserInfoConfig.
func convertUserInfoConfig(rc *authserver.UserInfoRunConfig) *upstream.UserInfoConfig {
        if rc == nil {
                return nil
        }
        return &upstream.UserInfoConfig{
                EndpointURL:       rc.EndpointURL,
                HTTPMethod:        rc.HTTPMethod,
                AdditionalHeaders: rc.AdditionalHeaders,
                FieldMapping:      convertFieldMapping(rc.FieldMapping),
        }
}

// convertFieldMapping converts UserInfoFieldMappingRunConfig to upstream.UserInfoFieldMapping.
func convertFieldMapping(rc *authserver.UserInfoFieldMappingRunConfig) *upstream.UserInfoFieldMapping {
        if rc == nil {
                return nil
        }
        return &upstream.UserInfoFieldMapping{
                SubjectFields: rc.SubjectFields,
                NameFields:    rc.NameFields,
                EmailFields:   rc.EmailFields,
        }
}

// createStorage creates the appropriate storage backend based on configuration.
func createStorage(ctx context.Context, cfg *storage.RunConfig) (storage.Storage, error) {
        maxDCRClients, err := cfg.ResolvedMaxDCRClients()
        if err != nil {
                return nil, err
        }
        if cfg == nil || cfg.Type == "" || cfg.Type == string(storage.TypeMemory) {
                return storage.NewMemoryStorage(storage.WithMaxDCRClients(maxDCRClients)), nil
        }
        if cfg.Type == string(storage.TypeRedis) {
                redisCfg, err := convertRedisRunConfig(cfg.RedisConfig)
                if err != nil {
                        return nil, fmt.Errorf("invalid Redis config: %w", err)
                }
                return storage.NewRedisStorage(ctx, redisCfg, cfg.RedisConfig.KeyPrefix,
                        storage.WithRedisMaxDCRClients(maxDCRClients))
        }
        return nil, fmt.Errorf("unsupported storage type: %s", cfg.Type)
}

// convertRedisRunConfig converts a serializable RedisRunConfig to a runtime
// redisconn.Config. It resolves ACL credentials from environment variables and
// parses duration strings. Connection-mode topology and defaulting are handled
// by the shared toolhive-core redisconn package when the client is constructed.
//
// An ACL config is not required: a nil rc.ACLUserConfig produces a no-auth
// connection (for a Redis/Valkey instance running without authentication) and
// emits one startup WARN naming the store. A populated ACL config whose
// password resolves to empty is still rejected — see convertRedisACLConfig.
func convertRedisRunConfig(rc *storage.RedisRunConfig) (redisconn.Config, error) {
        if rc == nil {
                return redisconn.Config{}, fmt.Errorf("redis config is required when storage type is redis")
        }

        // AuthType declares authenticated intent. If it selects ACL-user auth, a
        // nil ACLUserConfig is a misconfiguration, not a no-auth request: fail
        // loudly rather than silently downgrading to an unauthenticated connection.
        // An empty AuthType with a nil ACLUserConfig remains a valid no-auth config.
        if rc.AuthType == storage.AuthTypeACLUser && rc.ACLUserConfig == nil {
                return redisconn.Config{}, fmt.Errorf(
                        "auth_type %q requires acl_user_config; omit auth_type for a no-auth connection", rc.AuthType)
        }

        cfg := redisconn.Config{
                Addr:        rc.Addr,
                ClusterMode: rc.ClusterMode,
        }

        if rc.SentinelConfig != nil {
                cfg.SentinelConfig = &redisconn.SentinelConfig{
                        MasterName:    rc.SentinelConfig.MasterName,
                        SentinelAddrs: rc.SentinelConfig.SentinelAddrs,
                }
                cfg.DB = rc.SentinelConfig.DB
        }

        acl, err := convertRedisACLConfig(rc.ACLUserConfig)
        if err != nil {
                return redisconn.Config{}, fmt.Errorf("failed to convert ACL config: %w", err)
        }
        cfg.Username = acl.username
        cfg.Password = acl.password

        if err := applyRedisTimeouts(rc, &cfg); err != nil {
                return redisconn.Config{}, fmt.Errorf("failed to apply redis timeouts: %w", err)
        }

        tlsCfg, err := convertRedisTLSRunConfig(rc.TLS)
        if err != nil {
                return redisconn.Config{}, fmt.Errorf("master TLS config: %w", err)
        }
        cfg.TLS = tlsCfg

        // SentinelTLS only applies in Sentinel mode
        if rc.SentinelConfig != nil {
                sentinelTLSCfg, err := convertRedisTLSRunConfig(rc.SentinelTLS)
                if err != nil {
                        return redisconn.Config{}, fmt.Errorf("sentinel TLS config: %w", err)
                }
                cfg.SentinelTLS = sentinelTLSCfg
        }

        // A nil ACL config resolves to a no-auth connection. Emit exactly one
        // startup WARN naming the store so an unintended downgrade (e.g. a caller
        // that constructed a RedisRunConfig without an ACL config) is visible in
        // logs rather than silent. Logged only after the fallible timeout/TLS steps
        // succeed, so the message describes a config that will actually be used to
        // connect. A populated-but-empty config already returned an error in
        // convertRedisACLConfig, so reaching here without credentials means no
        // credential was intended.
        if rc.ACLUserConfig == nil {
                slog.Warn("Redis storage connecting without authentication (no acl_user_config configured)",
                        "store", redisStoreName(rc), "key_prefix", rc.KeyPrefix)
        }

        return cfg, nil
}

// redisACLCredentials carries resolved Redis ACL credentials between
// convertRedisACLConfig and its caller. Named fields prevent positional
// swaps of two same-typed strings at the call site.
type redisACLCredentials struct {
        username string
        password string
}

// convertRedisACLConfig resolves ACL user credentials from environment variables.
//
// A nil rc is a valid no-auth configuration: it returns empty credentials with
// no error, and go-redis then connects without AUTH. The split between no-auth
// and misconfiguration is on presence of the config block, not on whether the
// resolved password happens to be empty — a populated block whose PasswordEnvVar
// resolves to empty is a misconfiguration (mis-keyed secret, wrong-namespace
// secret, unsynced external-secrets store) and is rejected, because silently
// booting unauthenticated against a store holding sensitive rows (OAuth clients,
// DCR registrations, session data) would be a downgrade, not a choice.
//
// When UsernameEnvVar is empty, no username is resolved; go-redis then sends
// HELLO with "default" as the username (or falls back to legacy AUTH <password>
// for servers that do not support HELLO). This is required for managed Redis
// tiers without ACL users (e.g. GCP Memorystore Basic/Standard HA, Azure Cache
// for Redis).
func convertRedisACLConfig(rc *storage.ACLUserRunConfig) (redisACLCredentials, error) {
        if rc == nil {
                return redisACLCredentials{}, nil
        }
        var username string
        if rc.UsernameEnvVar != "" {
                var err error
                username, err = resolveEnvVar(rc.UsernameEnvVar)
                if err != nil {
                        return redisACLCredentials{}, fmt.Errorf("failed to resolve Redis username: %w", err)
                }
        }
        // A populated ACL config with no password_env_var is a misconfiguration.
        // Guard it explicitly so the message matches the empty-resolved-password
        // case below, rather than resolveEnvVar's generic "name is empty" error.
        if rc.PasswordEnvVar == "" {
                return redisACLCredentials{}, fmt.Errorf(
                        "populated ACL user config has no password_env_var; " +
                                "omit acl_user_config for a no-auth connection")
        }
        password, err := resolveEnvVar(rc.PasswordEnvVar)
        if err != nil {
                return redisACLCredentials{}, fmt.Errorf("failed to resolve Redis password: %w", err)
        }
        // A populated ACL config that resolves to an empty password is a
        // misconfiguration, not a request for no-auth. Fail loudly rather than
        // downgrading silently; omit acl_user_config entirely for a no-auth
        // connection.
        if password == "" {
                return redisACLCredentials{}, fmt.Errorf(
                        "resolved Redis password is empty for a populated ACL user config; " +
                                "omit acl_user_config for a no-auth connection")
        }
        return redisACLCredentials{username: username, password: password}, nil
}

// redisStoreName produces a human-readable identifier for the Redis store a
// config points at, for use in startup logs. It prefers the Sentinel master
// name when Sentinel mode is configured (Addr is empty in that mode) and
// otherwise returns the standalone/cluster address.
func redisStoreName(rc *storage.RedisRunConfig) string {
        if rc.SentinelConfig != nil {
                return "sentinel:" + rc.SentinelConfig.MasterName
        }
        return rc.Addr
}

// applyRedisTimeouts parses and applies optional timeout duration strings to cfg.
func applyRedisTimeouts(rc *storage.RedisRunConfig, cfg *redisconn.Config) error {
        if rc.DialTimeout != "" {
                d, err := time.ParseDuration(rc.DialTimeout)
                if err != nil {
                        return fmt.Errorf("invalid dial timeout: %w", err)
                }
                cfg.DialTimeout = d
        }
        if rc.ReadTimeout != "" {
                d, err := time.ParseDuration(rc.ReadTimeout)
                if err != nil {
                        return fmt.Errorf("invalid read timeout: %w", err)
                }
                cfg.ReadTimeout = d
        }
        if rc.WriteTimeout != "" {
                d, err := time.ParseDuration(rc.WriteTimeout)
                if err != nil {
                        return fmt.Errorf("invalid write timeout: %w", err)
                }
                cfg.WriteTimeout = d
        }
        return nil
}

// convertRedisTLSRunConfig converts a RedisTLSRunConfig to a runtime
// redisconn.TLSConfig. A configured CA certificate file must be readable and
// contain valid certificates; errors never fall back to system roots.
func convertRedisTLSRunConfig(rc *storage.RedisTLSRunConfig) (*redisconn.TLSConfig, error) {
        if rc == nil {
                return nil, nil
        }
        // Preserve the auth server's existing snake_case serialization while sharing
        // CA loading and validation with the session storage configuration.
        return (&redisconfig.TLSConfig{
                InsecureSkipVerify: rc.InsecureSkipVerify,
                CACertFile:         rc.CACertFile,
        }).Load()
}

// resolveCIMDConfig extracts CIMD settings from a CIMDRunConfig.
// Returns zero values when cfg is nil (CIMD disabled).
// The CacheFallbackTTL string is parsed to time.Duration; callers must ensure
// CIMDRunConfig.Validate() has already been called so the string is well-formed.
func resolveCIMDConfig(cfg *authserver.CIMDRunConfig) (enabled bool, cacheMaxSize int, cacheFallbackTTL time.Duration) {
        if cfg == nil {
                return false, 0, 0
        }
        var ttl time.Duration
        if cfg.CacheFallbackTTL != "" {
                var err error
                ttl, err = time.ParseDuration(cfg.CacheFallbackTTL)
                if err != nil {
                        // Should not happen when called after CIMDRunConfig.Validate().
                        slog.Warn("invalid cimd cache_fallback_ttl, zero will be replaced by default",
                                "value", cfg.CacheFallbackTTL, "err", err)
                }
        }
        return cfg.Enabled, cfg.CacheMaxSize, ttl
}

// resolveEnvVar reads a value from the named environment variable.
// An empty value is returned without error — empty credentials are valid for
// unauthenticated backends (e.g. a no-auth Redis where the operator injects a
// blank password from a Kubernetes Secret).
func resolveEnvVar(envVar string) (string, error) {
        if envVar == "" {
                return "", fmt.Errorf("environment variable name is empty")
        }
        value, ok := os.LookupEnv(envVar)
        if !ok {
                return "", fmt.Errorf("environment variable %q is not set", envVar)
        }
        return value, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "context"
        "crypto/sha256"
        "crypto/tls"
        "crypto/x509"
        "encoding/hex"
        "errors"
        "fmt"
        "log/slog"
        "net"
        "net/http"
        "os"
        "sync"
        "time"

        "github.com/stacklok/toolhive/pkg/authserver"
        spiffeauth "github.com/stacklok/toolhive/pkg/authserver/spiffe"
)

const (
        certificateCheckInterval  = 30 * time.Second
        listenerReadHeaderTimeout = 10 * time.Second
        listenerReadTimeout       = 30 * time.Second
        listenerWriteTimeout      = 30 * time.Second
        listenerIdleTimeout       = 60 * time.Second
)

type tlsKeyPairCache struct {
        mu            sync.Mutex
        certFile      string
        keyFile       string
        pair          tls.Certificate
        leaf          *x509.Certificate
        loadedHash    string
        failedHash    string
        lastCheck     time.Time
        checkInterval time.Duration
}

type caHintCache struct {
        mu   sync.Mutex
        hash string
        pool *x509.CertPool
}

type tlsListener struct {
        host              string
        port              int
        server            *http.Server
        listener          net.Listener
        keyPair           *tlsKeyPairCache
        caHints           caHintCache
        authorities       func() []*x509.Certificate
        requestClientCert bool
        mu                sync.Mutex
}

func newTLSListener(
        host string, cfg *authserver.TLSListenerRunConfig, routes map[string]http.Handler,
        authorities func() []*x509.Certificate, requestClientCert bool,
) (*tlsListener, error) {
        cache, err := newTLSKeyPairCache(cfg.CertFile, cfg.KeyFile)
        if err != nil {
                return nil, err
        }
        return newTLSListenerWithCache(host, cache, routes, authorities, requestClientCert), nil
}

func newTLSListenerWithCache(
        host string, cache *tlsKeyPairCache, routes map[string]http.Handler,
        authorities func() []*x509.Certificate, requestClientCert bool,
) *tlsListener {
        mux := http.NewServeMux()
        for pattern, handler := range routes {
                mux.Handle(pattern, handler)
        }
        base := &tlsListener{
                host: host, port: TLSListenerPort, keyPair: cache,
                authorities: authorities, requestClientCert: requestClientCert,
        }
        clientAuth := tls.NoClientCert
        if requestClientCert {
                clientAuth = tls.RequestClientCert
        }
        tlsConfig := &tls.Config{
                MinVersion:     tls.VersionTLS12,
                ClientAuth:     clientAuth,
                GetCertificate: cache.getCertificate,
        }
        if requestClientCert {
                tlsConfig.GetConfigForClient = func(*tls.ClientHelloInfo) (*tls.Config, error) {
                        clone := tlsConfig.Clone()
                        clone.ClientCAs = base.caPool()
                        return clone, nil
                }
        }
        base.server = &http.Server{
                Handler:           spiffeauth.Middleware(mux),
                ReadHeaderTimeout: listenerReadHeaderTimeout,
                ReadTimeout:       listenerReadTimeout,
                WriteTimeout:      listenerWriteTimeout,
                IdleTimeout:       listenerIdleTimeout,
                ErrorLog:          slog.NewLogLogger(slog.Default().Handler(), slog.LevelDebug),
                TLSConfig:         tlsConfig,
        }
        return base
}

func (l *tlsListener) start() error {
        l.mu.Lock()
        defer l.mu.Unlock()
        if l.listener != nil {
                return errors.New("auth server TLS listener already started")
        }
        address := net.JoinHostPort(l.host, fmt.Sprintf("%d", l.port))
        listener, err := net.Listen("tcp", address)
        if err != nil {
                return fmt.Errorf("bind auth server TLS listener on %s: %w", address, err)
        }
        l.listener = tls.NewListener(listener, l.server.TLSConfig)
        l.server.Addr = address
        go func() {
                if err := l.server.Serve(l.listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
                        slog.Error("auth server TLS listener stopped unexpectedly", "error", err)
                }
        }()
        return nil
}

func (l *tlsListener) shutdown(ctx context.Context) error {
        l.mu.Lock()
        server, listener := l.server, l.listener
        l.mu.Unlock()
        err := server.Shutdown(ctx)
        if errors.Is(err, context.DeadlineExceeded) {
                slog.Warn("auth server TLS listener graceful shutdown timed out, force-closing connections")
                if closeErr := server.Close(); closeErr != nil {
                        err = errors.Join(err, closeErr)
                }
        }
        // Shutdown only closes listeners Serve has already registered; close ours
        // directly so the port is freed even if the serve goroutine has not run yet.
        if listener != nil {
                if closeErr := listener.Close(); closeErr != nil && !errors.Is(closeErr, net.ErrClosed) {
                        err = errors.Join(err, closeErr)
                }
        }
        return err
}

func (l *tlsListener) caPool() *x509.CertPool {
        if l.authorities == nil {
                return nil
        }
        authorities := l.authorities()
        hash := sha256.New()
        for _, cert := range authorities {
                if cert != nil {
                        _, _ = hash.Write(cert.Raw)
                }
        }
        hashValue := hex.EncodeToString(hash.Sum(nil))
        l.caHints.mu.Lock()
        defer l.caHints.mu.Unlock()
        if hashValue == l.caHints.hash {
                return l.caHints.pool
        }
        if len(authorities) == 0 {
                l.caHints.hash, l.caHints.pool = hashValue, nil
                return nil
        }
        pool := x509.NewCertPool()
        for _, cert := range authorities {
                if cert != nil {
                        pool.AddCert(cert)
                }
        }
        l.caHints.hash, l.caHints.pool = hashValue, pool
        return pool
}

func newTLSKeyPairCache(certFile, keyFile string) (*tlsKeyPairCache, error) {
        cache := &tlsKeyPairCache{certFile: certFile, keyFile: keyFile, checkInterval: certificateCheckInterval}
        if err := cache.load(); err != nil {
                return nil, err
        }
        return cache, nil
}

func (c *tlsKeyPairCache) getCertificate(*tls.ClientHelloInfo) (*tls.Certificate, error) {
        c.mu.Lock()
        defer c.mu.Unlock()
        if time.Since(c.lastCheck) >= c.checkInterval {
                c.lastCheck = time.Now()
                certBytes, certErr := os.ReadFile(c.certFile)
                keyBytes, keyErr := os.ReadFile(c.keyFile)
                if certErr != nil || keyErr != nil {
                        slog.Warn("failed to reload auth server TLS listener certificate", "error", errors.Join(certErr, keyErr))
                } else {
                        hash := pairHash(certBytes, keyBytes)
                        if hash != c.loadedHash && hash != c.failedHash {
                                pair, err := tls.X509KeyPair(certBytes, keyBytes)
                                if err != nil {
                                        c.failedHash = hash
                                        slog.Warn("failed to reload auth server TLS listener certificate", "error", err)
                                } else if err := c.setPair(pair, hash); err != nil {
                                        c.failedHash = hash
                                        slog.Warn("failed to parse auth server TLS listener certificate", "error", err)
                                }
                        }
                }
        }
        pair := c.pair
        return &pair, nil
}

func (c *tlsKeyPairCache) load() error {
        certBytes, err := os.ReadFile(c.certFile)
        if err != nil {
                return fmt.Errorf("read certificate file: %w", err)
        }
        keyBytes, err := os.ReadFile(c.keyFile)
        if err != nil {
                return fmt.Errorf("read key file: %w", err)
        }
        pair, err := tls.X509KeyPair(certBytes, keyBytes)
        if err != nil {
                return err
        }
        return c.setPair(pair, pairHash(certBytes, keyBytes))
}

func (c *tlsKeyPairCache) setPair(pair tls.Certificate, hash string) error {
        leaf, err := x509.ParseCertificate(pair.Certificate[0])
        if err != nil {
                return err
        }
        c.pair, c.leaf, c.loadedHash = pair, leaf, hash
        c.failedHash = ""
        return nil
}

func pairHash(certBytes, keyBytes []byte) string {
        h := sha256.New()
        _, _ = h.Write(certBytes)
        _, _ = h.Write([]byte{0})
        _, _ = h.Write(keyBytes)
        return hex.EncodeToString(h.Sum(nil))
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "context"
        "crypto/x509"
        "log/slog"
        "net/http"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
)

// UpstreamProviderFactory constructs the OAuth2Provider for one configured
// upstream IDP. Set it on Config.UpstreamFactory to own upstream construction;
// DefaultUpstreamFactory is the built-in implementation and can be delegated to.
type UpstreamProviderFactory func(ctx context.Context, cfg *UpstreamConfig) (upstream.OAuth2Provider, error)

// Server is the OAuth authorization server.
// It provides HTTP handlers that serve all OAuth/OIDC endpoints.
type Server interface {
        // Handler returns an http.Handler that serves all OAuth/OIDC endpoints:
        //   - /.well-known/openid-configuration (OIDC Discovery)
        //   - /.well-known/oauth-authorization-server (RFC 8414 OAuth AS Metadata)
        //   - /.well-known/jwks.json (JSON Web Key Set)
        //   - /oauth/authorize (Authorization endpoint)
        //   - /oauth/token (Token endpoint)
        //   - /oauth/callback (Upstream IDP callback)
        //   - /oauth/register (Dynamic Client Registration, RFC 7591)
        //
        // The handler uses internal routing - the consumer doesn't need to know
        // about the endpoint structure.
        Handler() http.Handler

        // IDPTokenStorage returns storage for upstream IDP tokens.
        // Returns nil if no upstream IDP is configured.
        IDPTokenStorage() storage.UpstreamTokenStorage

        // UpstreamTokenRefresher returns a refresher that can refresh expired upstream
        // tokens using the upstream provider's refresh token grant.
        // Returns nil if no upstream IDP is configured.
        UpstreamTokenRefresher() storage.UpstreamTokenRefresher

        // DCRStore returns the persistent DCR credential store the server is wired
        // against. This is the same DCRCredentialStore used by the upstream-DCR
        // resolver at boot, so callers can read RFC 7591 client registrations
        // without bypassing the storage backend the server itself reads from.
        //
        // SECURITY: the returned interface surfaces raw `client_secret` and
        // `registration_access_token` values. Callers MUST NOT log or render the
        // returned values; treat the handle the same way you would treat a
        // secrets manager client. Intended for admin / diagnostic code paths and
        // integration tests, not for general consumers.
        //
        // Lifecycle: the returned handle's lifetime is bound to Server.Close —
        // methods invoked after Close have backend-specific behavior (a
        // MemoryStorage continues to serve reads; a RedisStorage will error on
        // its closed connection pool).
        DCRStore() storage.DCRCredentialStore

        // Close releases resources held by the server. It drains upstream idle
        // connections, stops SPIFFE bundle and trusted-issuer refresh workers, and
        // closes storage. It joins shutdown errors. Do not call it on a server whose
        // storage is shared with another live server; retire that one with
        // CloseIdleConnections instead.
        Close() error
}

// idleConnectionCloser is the capability CloseIdleConnections detects. It is
// deliberately not a method on Server: adding one would be a source-incompatible
// change for out-of-tree implementations and test doubles, and the only in-tree
// consumer needs a single call. Kept unexported so callers go through the
// CloseIdleConnections function rather than asserting themselves — the
// `var _ idleConnectionCloser = (*server)(nil)` check in server_impl.go is what
// keeps the real implementation compile-time guaranteed to satisfy it.
type idleConnectionCloser interface {
        CloseIdleConnections()
}

// spiffeX509AuthorityLister is the capability SPIFFEX509Authorities detects.
// It is deliberately not a method on Server so out-of-tree implementations and
// test doubles are not forced to implement the optional operation.
type spiffeX509AuthorityLister interface {
        SPIFFEX509Authorities() []*x509.Certificate
}

// SPIFFEX509Authorities returns the current X.509 authorities of every trust
// domain that enables spiffe_x509, and whether s supports the operation.
// The result reflects live bundle rotation. It is used only as a TLS
// certificate_authorities hint, never for verification.
func SPIFFEX509Authorities(s Server) ([]*x509.Certificate, bool) {
        lister, ok := s.(spiffeX509AuthorityLister)
        if !ok {
                return nil, false
        }
        return lister.SPIFFEX509Authorities(), true
}

// CloseIdleConnections releases the idle keep-alive connections pooled by s's
// upstream IDP providers, without touching storage, and reports whether s
// supported the operation. Only an implementation of Server that does not come
// from New can report false; the production type is checked at compile time.
//
// An embedder that reconstructs the server to change its upstream set — passing
// the same storage.Storage and keys.KeyProvider so token and JWKS continuity is
// preserved — must retire the superseded server with this function rather than
// Close, because Close would also close the storage the new server is now
// serving through. That split is the whole point: without it there is no correct
// call to make.
//
// Safe to call on a server that is still serving: only idle connections are
// closed and in-flight requests are unaffected. See upstream.IdleConnectionCloser
// for the per-provider capability this delegates to, and for the
// caller-owned-client exemption.
//
// Scope: upstream HTTP pools only. A server configured with TrustedIssuers also
// holds one JWKS refresh worker pool per issuer; those are released by Close,
// not by this function — draining them would stop the background key refresh a
// still-serving server depends on, and this function is safe to call on a live
// server. An embedder retiring a superseded server with this function therefore
// keeps its JWKS workers until the storage it shares with the replacement can
// be closed via Close.
func CloseIdleConnections(s Server) bool {
        closer, ok := s.(idleConnectionCloser)
        if ok {
                closer.CloseIdleConnections()
        }
        return ok
}

// New creates a new OAuth authorization server.
// The storage parameter is required and determines where OAuth state is persisted.
// Use storage.NewMemoryStorage() for single-instance deployments or provide
// a distributed storage backend for production deployments.
func New(ctx context.Context, cfg Config, stor storage.Storage) (Server, error) {
        slog.Debug("creating new OAuth authorization server", "issuer", cfg.Issuer)
        return newServer(ctx, cfg, stor)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package server

import (
        "context"
        "maps"

        "github.com/ory/fosite/token/jwt"
)

var _ jwt.Signer = (*atJWTAccessTokenSigner)(nil)

// atJWTAccessTokenSigner types the tokens it signs as at+jwt (RFC 9068 §2.1)
// unless the caller has already set a typ.
type atJWTAccessTokenSigner struct {
        jwt.Signer
}

// NewATJWTAccessTokenSigner wraps signer so that tokens it generates carry the
// JOSE header typ: at+jwt, unless the caller's header already sets a typ. All
// other Signer methods delegate to signer.
func NewATJWTAccessTokenSigner(signer jwt.Signer) jwt.Signer {
        return &atJWTAccessTokenSigner{Signer: signer}
}

// Generate signs claims with a copy of header. The caller's header is never
// changed; the copy gets typ: at+jwt only if it has no typ key. A typ key is
// kept even when its value is empty.
func (s *atJWTAccessTokenSigner) Generate(ctx context.Context, claims jwt.MapClaims, header jwt.Mapper) (string, string, error) {
        // Copy the header: the session's stored header must not change.
        hdr := func() map[string]any {
                if header != nil {
                        if m := maps.Clone(header.ToMap()); m != nil {
                                return m
                        }
                }
                return map[string]any{}
        }()
        if _, ok := hdr["typ"]; !ok {
                hdr["typ"] = AccessTokenTypeATJWT
        }
        return s.Signer.Generate(ctx, claims, &jwt.Headers{Extra: hdr})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package server

import (
        "fmt"
        "net/http"
        "net/url"
        "slices"

        "github.com/ory/fosite"
)

// ErrInvalidTarget is the RFC 8707 error for invalid or unauthorized resource parameters.
// This error is returned when:
// - The resource URI format is invalid (not absolute, has fragment, wrong scheme)
// - The resource is not in the server's allowed audiences list
var ErrInvalidTarget = &fosite.RFC6749Error{
        ErrorField:       "invalid_target",
        DescriptionField: "The requested resource is invalid, unknown, or malformed.",
        CodeField:        http.StatusBadRequest,
}

// ValidateAudienceURI validates that a resource URI conforms to RFC 8707 requirements.
// According to RFC 8707, a valid resource parameter must be:
// - An absolute URI (has scheme and host)
// - No fragment component
// - Use http or https scheme
func ValidateAudienceURI(resource string) error {
        if resource == "" {
                return nil // Empty resource is valid (means no audience binding requested)
        }

        parsed, err := url.Parse(resource)
        if err != nil {
                return ErrInvalidTarget.WithHintf("Resource parameter is not a valid URI: %s", err.Error())
        }

        // Must be absolute (have a scheme)
        if !parsed.IsAbs() {
                return ErrInvalidTarget.WithHint("Resource must be an absolute URI")
        }

        // Must have a host
        if parsed.Host == "" {
                return ErrInvalidTarget.WithHint("Resource must include a host")
        }

        // Must not have a fragment (RFC 8707 Section 2)
        if parsed.Fragment != "" {
                return ErrInvalidTarget.WithHint("Resource must not contain a fragment")
        }

        // Only allow http or https schemes for security
        if parsed.Scheme != "http" && parsed.Scheme != "https" {
                return ErrInvalidTarget.WithHint("Resource must use http or https scheme")
        }

        return nil
}

// ValidateSubjectIdentifierAudiences validates the runtime sub_id disclosure allowlist.
// Every entry must be an exact member of the independently validated
// allowed-audience list; an empty list disables disclosure.
func ValidateSubjectIdentifierAudiences(subjectAudiences, allowedAudiences []string) error {
        if len(subjectAudiences) == 0 {
                return nil
        }
        for _, audience := range subjectAudiences {
                if !slices.Contains(allowedAudiences, audience) {
                        return fmt.Errorf("subject identifier audience %q is not in allowed audiences", audience)
                }
        }
        return nil
}

// ValidateAudienceAllowed checks if the resource is in the allowed audiences list.
// Returns nil if allowed, or ErrInvalidTarget if not.
//
// Security: An empty allowedAudiences list means NO audiences are permitted (secure default).
func ValidateAudienceAllowed(resource string, allowedAudiences []string) error {
        if resource == "" {
                return nil // No resource requested, nothing to validate
        }

        // Secure default: empty allowlist means reject all
        if len(allowedAudiences) == 0 {
                return ErrInvalidTarget.WithHint("No resource audiences are configured on this server")
        }

        // Exact string matching
        if slices.Contains(allowedAudiences, resource) {
                return nil
        }

        return ErrInvalidTarget.WithHintf("Resource %q is not a registered audience", resource)
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package crypto provides cryptographic utilities for the OAuth authorization server.
package crypto

import (
        "crypto"
        "crypto/ecdsa"
        "crypto/ed25519"
        "crypto/elliptic"
        "crypto/rsa"
        "crypto/x509"
        "encoding/base64"
        "encoding/pem"
        "fmt"
        "os"
        "slices"

        "github.com/go-jose/go-jose/v4"
)

// MinSecretLength is the minimum required length for the HMAC secret in bytes.
// Using 256 bits (32 bytes) provides adequate security for HMAC-SHA256.
const MinSecretLength = 32

// MinRSAKeyBits is the minimum required RSA key size in bits.
// NIST recommends at least 2048 bits for RSA keys.
const MinRSAKeyBits = 2048

// MaxRSAKeyBits is the maximum accepted RSA public modulus size in bits.
const MaxRSAKeyBits = 8192

// MaxRSAPublicExponent is the largest accepted public exponent for client RSA keys.
const MaxRSAPublicExponent = 65537

// LoadSigningKey loads a private key from a PEM file.
// Supports both RSA (PKCS1 and PKCS8) and ECDSA (PKCS8) formats.
// Returns a crypto.Signer that can be used for JWT signing.
func LoadSigningKey(keyPath string) (crypto.Signer, error) {
        keyPEM, err := os.ReadFile(keyPath) // #nosec G304 - keyPath is provided by user via CLI flag or config
        if err != nil {
                return nil, fmt.Errorf("failed to read signing key: %w", err)
        }

        block, _ := pem.Decode(keyPEM)
        if block == nil {
                return nil, fmt.Errorf("failed to decode PEM block from signing key")
        }

        // Try PKCS1 first (RSA only)
        if rsaKey, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
                if keyBits := rsaKey.N.BitLen(); keyBits < MinRSAKeyBits {
                        return nil, fmt.Errorf("RSA key size %d bits is below minimum required %d bits", keyBits, MinRSAKeyBits)
                }
                return rsaKey, nil
        }

        // Try EC private key (SEC 1, ASN.1 DER form)
        if ecKey, err := x509.ParseECPrivateKey(block.Bytes); err == nil {
                return ecKey, nil
        }

        // Try PKCS8 (supports both RSA and EC)
        key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
        if err != nil {
                return nil, fmt.Errorf("failed to parse signing key: %w", err)
        }

        // Validate RSA key size if parsed key is RSA
        if rsaKey, ok := key.(*rsa.PrivateKey); ok {
                if keyBits := rsaKey.N.BitLen(); keyBits < MinRSAKeyBits {
                        return nil, fmt.Errorf("RSA key size %d bits is below minimum required %d bits", keyBits, MinRSAKeyBits)
                }
        }

        signer, ok := key.(crypto.Signer)
        if !ok {
                return nil, fmt.Errorf("signing key does not implement crypto.Signer")
        }

        return signer, nil
}

// DeriveKeyID computes a key ID from the public key using RFC 7638 JWK Thumbprint.
// The thumbprint is computed as base64url(SHA-256(JWK canonical form)).
func DeriveKeyID(key crypto.Signer) (string, error) {
        // Create a JWK from the public key
        jwk := jose.JSONWebKey{
                Key: key.Public(),
        }

        // Compute the thumbprint using go-jose's built-in RFC 7638 implementation
        thumbprint, err := jwk.Thumbprint(crypto.SHA256)
        if err != nil {
                return "", fmt.Errorf("failed to compute key thumbprint: %w", err)
        }

        // Base64url encode without padding (RFC 7638 standard)
        return base64.RawURLEncoding.EncodeToString(thumbprint), nil
}

// DeriveAlgorithm determines the appropriate JWT signing algorithm for the given key.
// Returns the algorithm string (e.g., "RS256", "ES256", "EdDSA") based on key type and parameters.
func DeriveAlgorithm(key crypto.Signer) (string, error) {
        switch k := key.(type) {
        case *rsa.PrivateKey:
                return "RS256", nil
        case *ecdsa.PrivateKey:
                return deriveECAlgorithm(k.Curve)
        case ed25519.PrivateKey:
                return "EdDSA", nil
        default:
                return "", fmt.Errorf("unsupported key type: %T", key)
        }
}

// deriveECAlgorithm determines the ECDSA algorithm based on the curve.
func deriveECAlgorithm(curve elliptic.Curve) (string, error) {
        switch curve {
        case elliptic.P256():
                return "ES256", nil
        case elliptic.P384():
                return "ES384", nil
        case elliptic.P521():
                return "ES512", nil
        default:
                return "", fmt.Errorf("unsupported EC curve: %s", curve.Params().Name)
        }
}

// rsaAlgorithmsForClientKeys are the RSA JWS algorithms accepted for a client's public key.
var rsaAlgorithmsForClientKeys = []string{"RS256", "RS384", "RS512", "PS256", "PS384", "PS512"}

// rsaAlgorithmsForSigningKeys are the RSA JWS algorithms accepted for the server's own
// signing key. Deliberately narrower than rsaAlgorithmsForClientKeys: PS* is excluded.
var rsaAlgorithmsForSigningKeys = []string{"RS256", "RS384", "RS512"}

// SupportedClientKeyAlgorithms returns the allowlist of JWS algorithms
// accepted for a private_key_jwt client's registered key: rsaAlgorithmsForClientKeys
// plus the three EC algorithms deriveECAlgorithm maps curves to. This is the
// single source of truth for both key-compatibility validation here and DCR's
// signing-algorithm allowlist, so the two lists cannot silently drift apart
// (as happened when EdDSA was previously missing from a hand-maintained copy
// of this list).
//
// EdDSA is deliberately excluded even though validateAlgorithmForPublicKeyValue
// and ValidateAlgorithmForKey both support Ed25519 keys structurally: fosite
// v0.49.0's client-assertion verification (client_authentication.go) only
// switches on the RS*/ES*/PS*/HS* families, so an EdDSA client could register
// successfully and then never be able to authenticate. Revisit when fosite
// adds EdDSA support or this pins a version that has it.
func SupportedClientKeyAlgorithms() []string {
        algs := make([]string, 0, len(rsaAlgorithmsForClientKeys)+3)
        algs = append(algs, rsaAlgorithmsForClientKeys...)
        return append(algs, "ES256", "ES384", "ES512")
}

// validateAlgorithmForPublicKeyValue checks whether alg is compatible with a public key,
// accepting exactly the RSA algorithms in allowedRSAAlgorithms. boundRSA
// additionally caps the RSA modulus and public exponent; it is set for
// untrusted client keys only, not for operator-provisioned signing keys.
func validateAlgorithmForPublicKeyValue(
        alg string, key crypto.PublicKey, allowedRSAAlgorithms []string, boundRSA bool,
) error {
        switch k := key.(type) {
        case *rsa.PublicKey:
                return validateAlgorithmForRSAPublicKey(alg, k, allowedRSAAlgorithms, boundRSA)
        case *ecdsa.PublicKey:
                expectedAlg, err := deriveECAlgorithm(k.Curve)
                if err != nil {
                        return err
                }
                if alg != expectedAlg {
                        return fmt.Errorf("algorithm %s is not compatible with EC key using curve %s (expected %s)",
                                alg, k.Curve.Params().Name, expectedAlg)
                }
                return nil
        case ed25519.PublicKey:
                if alg != "EdDSA" {
                        return fmt.Errorf("algorithm %s is not compatible with Ed25519 key (expected EdDSA)", alg)
                }
                return nil
        default:
                return fmt.Errorf("unsupported key type: %T", key)
        }
}

// validateAlgorithmForRSAPublicKey checks an RSA public key's modulus and
// exponent and whether alg is in allowedRSAAlgorithms. boundRSA caps the
// modulus size and public exponent, as for validateAlgorithmForPublicKeyValue.
func validateAlgorithmForRSAPublicKey(
        alg string, k *rsa.PublicKey, allowedRSAAlgorithms []string, boundRSA bool,
) error {
        if k == nil || k.N == nil {
                return fmt.Errorf("RSA public key is missing a modulus")
        }
        keyBits := k.N.BitLen()
        if keyBits < MinRSAKeyBits {
                return fmt.Errorf("RSA key size %d bits is below minimum required %d bits", keyBits, MinRSAKeyBits)
        }
        if boundRSA && keyBits > MaxRSAKeyBits {
                return fmt.Errorf("RSA key size %d bits exceeds maximum allowed %d bits", keyBits, MaxRSAKeyBits)
        }
        // Reject a degenerate or unsafe public exponent (e.g. e=1 makes RSA
        // the identity function; e must be odd for the exponentiation to be
        // well-defined).
        if k.E < 3 || k.E%2 == 0 {
                return fmt.Errorf("RSA public exponent %d is invalid or unsafe (must be odd and at least 3)", k.E)
        }
        if boundRSA && k.E > MaxRSAPublicExponent {
                return fmt.Errorf("RSA public exponent %d exceeds maximum allowed %d", k.E, MaxRSAPublicExponent)
        }
        if slices.Contains(allowedRSAAlgorithms, alg) {
                return nil
        }
        return fmt.Errorf("algorithm %s is not compatible with RSA key", alg)
}

// ValidateAlgorithmForPublicKey checks whether alg is compatible with a public key.
// It accepts the public-key values stored in jose.JSONWebKey.Key.
func ValidateAlgorithmForPublicKey(alg string, key crypto.PublicKey) error {
        return validateAlgorithmForPublicKeyValue(alg, key, rsaAlgorithmsForClientKeys, true)
}

// ValidateAlgorithmForKey checks if the provided algorithm is compatible with the key type.
// Returns an error if the algorithm doesn't match the key type.
func ValidateAlgorithmForKey(alg string, key crypto.Signer) error {
        switch k := key.(type) {
        case *rsa.PrivateKey:
                if k == nil {
                        return fmt.Errorf("signing key is nil")
                }
                return validateAlgorithmForPublicKeyValue(alg, &k.PublicKey, rsaAlgorithmsForSigningKeys, false)
        case *ecdsa.PrivateKey:
                if k == nil {
                        return fmt.Errorf("signing key is nil")
                }
                return validateAlgorithmForPublicKeyValue(alg, &k.PublicKey, rsaAlgorithmsForSigningKeys, false)
        case ed25519.PrivateKey:
                if k == nil {
                        return fmt.Errorf("signing key is nil")
                }
                return validateAlgorithmForPublicKeyValue(alg, k.Public(), rsaAlgorithmsForSigningKeys, false)
        default:
                return fmt.Errorf("unsupported key type: %T", key)
        }
}

// SigningKeyParams contains the derived or configured parameters for a signing key.
type SigningKeyParams struct {
        // Key is the private key used for signing.
        Key crypto.Signer
        // KeyID is the key identifier (either derived from thumbprint or configured).
        KeyID string
        // Algorithm is the signing algorithm (either derived from key type or configured).
        Algorithm string
}

// HMACSecrets holds the current secret and any rotated (previous) secrets.
// This supports zero-downtime secret rotation for OAuth token signing.
type HMACSecrets struct {
        // Current is the active secret used for signing new tokens.
        Current []byte
        // Rotated contains previously-used secrets for verifying existing tokens.
        Rotated [][]byte
}

// NewHMACSecrets creates an HMACSecrets with just a current secret (no rotation).
// This is a convenience function for cases where secret rotation is not needed.
func NewHMACSecrets(current []byte) *HMACSecrets {
        return &HMACSecrets{
                Current: current,
                Rotated: nil,
        }
}

// DeriveSigningKeyParams derives or validates signing key parameters.
// If keyID or algorithm are empty, they are derived from the key.
// If they are provided, they are validated against the key type.
func DeriveSigningKeyParams(key crypto.Signer, keyID, algorithm string) (*SigningKeyParams, error) {
        params := &SigningKeyParams{Key: key}

        // Derive or use provided KeyID
        if keyID == "" {
                derivedID, err := DeriveKeyID(key)
                if err != nil {
                        return nil, fmt.Errorf("failed to derive key ID: %w", err)
                }
                params.KeyID = derivedID
        } else {
                params.KeyID = keyID
        }

        // Derive or validate Algorithm
        if algorithm == "" {
                derivedAlg, err := DeriveAlgorithm(key)
                if err != nil {
                        return nil, fmt.Errorf("failed to derive algorithm: %w", err)
                }
                params.Algorithm = derivedAlg
        } else {
                // Validate that provided algorithm matches key type
                if err := ValidateAlgorithmForKey(algorithm, key); err != nil {
                        return nil, err
                }
                params.Algorithm = algorithm
        }

        return params, nil
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package crypto

import (
        "golang.org/x/oauth2"
)

// PKCEChallengeMethodS256 is the PKCE challenge method using SHA-256 (RFC 7636).
const PKCEChallengeMethodS256 = "S256"

// GeneratePKCEVerifier generates a cryptographically random code_verifier
// per RFC 7636 Section 4.1.
// The verifier is 43 characters (32 bytes base64url encoded without padding),
// using characters from the base64url alphabet: [A-Z], [a-z], [0-9], "-", "_".
//
// This function delegates to oauth2.GenerateVerifier() from golang.org/x/oauth2.
// It will panic on crypto/rand read failure (which is appropriate for this case).
func GeneratePKCEVerifier() string {
        return oauth2.GenerateVerifier()
}

// ComputePKCEChallenge computes the code_challenge from a code_verifier
// using the S256 method per RFC 7636 Section 4.2.
// code_challenge = BASE64URL(SHA256(code_verifier))
//
// This function delegates to oauth2.S256ChallengeFromVerifier() from golang.org/x/oauth2.
func ComputePKCEChallenge(verifier string) string {
        return oauth2.S256ChallengeFromVerifier(verifier)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package deviceflow

import (
        "fmt"
        "time"

        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"

        "github.com/stacklok/toolhive/pkg/authserver/server"
        authstorage "github.com/stacklok/toolhive/pkg/authserver/storage"
)

// Factory returns a server.Factory that registers the RFC 8628 device-code
// grant, mirroring how the tokenexchange and jwtbearer factories are
// constructed in server_impl.go's buildProvider.
func Factory(deviceStorage authstorage.DeviceCodeStorage, minInterval time.Duration) server.Factory {
        return func(config *server.AuthorizationServerConfig, stor fosite.Storage, strategy any) (any, error) {
                coreStorage, ok := stor.(oauth2.CoreStorage)
                if !ok {
                        return nil, fmt.Errorf("deviceflow: storage backend %T does not implement oauth2.CoreStorage", stor)
                }
                coreStrategy, ok := strategy.(oauth2.CoreStrategy)
                if !ok {
                        return nil, fmt.Errorf("deviceflow: strategy %T does not implement oauth2.CoreStrategy", strategy)
                }
                return &Handler{
                        DeviceStorage: deviceStorage,
                        CoreStorage:   coreStorage,
                        Strategy:      coreStrategy,
                        // The embedded *fosite.Config, not config itself:
                        // AuthorizationServerConfig's own no-context adapter methods of the
                        // same name shadow the ctx-taking ones fosite's provider interfaces
                        // require (see tokenexchange.Factory's identical concern).
                        Config:      config.Config,
                        MinInterval: minInterval,
                }, nil
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package deviceflow implements the RFC 8628 OAuth 2.0 Device Authorization
// Grant's token-endpoint handler.
package deviceflow

import (
        "context"
        "errors"
        "time"

        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"
        "github.com/ory/x/errorsx"

        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        authstorage "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// Compile-time check that Handler implements fosite.TokenEndpointHandler.
var _ fosite.TokenEndpointHandler = (*Handler)(nil)

// deviceFlowConfig is the narrow config surface this handler needs. The
// embedded *fosite.Config on server.AuthorizationServerConfig satisfies this
// (its own no-context adapter methods of the same name would otherwise
// shadow the ctx-taking ones fosite's provider interfaces require — see
// tokenexchange.Factory's identical concern).
type deviceFlowConfig interface {
        fosite.AccessTokenLifespanProvider
        fosite.RefreshTokenLifespanProvider
}

// Handler implements fosite's TokenEndpointHandler for RFC 8628's
// urn:ietf:params:oauth:grant-type:device_code grant.
//
// A device_code is single-use: HandleTokenEndpointRequest deletes it from
// DeviceStorage as soon as it observes DeviceRequestStatusAuthorized, before
// PopulateTokenEndpointResponse ever issues a token, so a device_code can
// never be redeemed twice.
type Handler struct {
        DeviceStorage authstorage.DeviceCodeStorage
        CoreStorage   oauth2.CoreStorage
        Strategy      oauth2.CoreStrategy
        Config        deviceFlowConfig
        // MinInterval is the minimum time a client must wait between polls
        // (RFC 8628 Section 3.5). A poll that arrives sooner is rejected with
        // slow_down.
        MinInterval time.Duration
}

// CanHandleTokenEndpointRequest returns true if the request's grant_type is
// the RFC 8628 device_code grant type.
func (*Handler) CanHandleTokenEndpointRequest(_ context.Context, requester fosite.AccessRequester) bool {
        return requester.GetGrantTypes().ExactOne(oauthproto.GrantTypeDeviceCode)
}

// CanSkipClientAuth always returns false: the device_code grant does not
// exempt the client from standard authentication. The client authenticates
// with whatever method it registered with, exactly as for every other grant;
// only grants that attach a synthetic, unregistered client (see
// storage.NewSyntheticClient) skip authentication, and this is not one of
// them.
func (*Handler) CanSkipClientAuth(_ context.Context, _ fosite.AccessRequester) bool {
        return false
}

// HandleTokenEndpointRequest validates the device_code, enforces the RFC
// 8628 polling contract (slow_down, authorization_pending, access_denied,
// expired_token), and — once the device request has been authorized —
// attaches a session and consumes (deletes) the device_code so it cannot be
// redeemed twice.
func (h *Handler) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }

        client := requester.GetClient()
        if !client.GetGrantTypes().Has(oauthproto.GrantTypeDeviceCode) {
                return errorsx.WithStack(fosite.ErrUnauthorizedClient.WithHint(
                        "The OAuth 2.0 Client is not allowed to use authorization grant 'urn:ietf:params:oauth:grant-type:device_code'."))
        }

        deviceCode := requester.GetRequestForm().Get("device_code")
        if deviceCode == "" {
                return errorsx.WithStack(fosite.ErrInvalidRequest.WithHint("The device_code parameter is missing."))
        }

        device, err := h.DeviceStorage.LoadDeviceRequestByDeviceCode(ctx, deviceCode)
        switch {
        case errors.Is(err, authstorage.ErrNotFound):
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                        "The device_code is unknown or has already been redeemed."))
        case errors.Is(err, authstorage.ErrExpired):
                return errorsx.WithStack(ErrExpiredToken)
        case err != nil:
                return errorsx.WithStack(fosite.ErrServerError.WithWrap(err))
        }

        // A device_code presented by a client other than the one it was issued to
        // is treated identically to an unknown code: revealing that the code
        // exists but belongs to someone else would leak information to a client
        // guessing at codes.
        if device.ClientID != client.GetID() {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                        "The device_code is unknown or has already been redeemed."))
        }

        if err := h.enforcePollInterval(ctx, device, deviceCode); err != nil {
                return err
        }

        switch device.Status {
        case authstorage.DeviceRequestStatusPending:
                return errorsx.WithStack(ErrAuthorizationPending)
        case authstorage.DeviceRequestStatusDenied:
                return errorsx.WithStack(fosite.ErrAccessDenied)
        case authstorage.DeviceRequestStatusAuthorized:
                // Proceed below.
        default:
                return errorsx.WithStack(fosite.ErrServerError.WithHintf("unrecognized device request status %q", device.Status))
        }

        h.attachSession(ctx, requester, client, device)

        // Consume the device_code now, before any token is issued, so a second
        // concurrent or replayed request for the same code always fails at the
        // LoadDeviceRequestByDeviceCode lookup above rather than racing
        // PopulateTokenEndpointResponse into issuing two tokens for one grant.
        if err := h.DeviceStorage.DeleteDeviceRequest(ctx, deviceCode); err != nil {
                return errorsx.WithStack(fosite.ErrServerError.WithWrap(err))
        }

        return nil
}

// attachSession builds and attaches the session for an authorized device
// request: identity, granted scopes/audience, and access/refresh token
// expiry (mirroring fosite's own authorization-code grant handler).
func (h *Handler) attachSession(
        ctx context.Context, requester fosite.AccessRequester, client fosite.Client, device *authstorage.DeviceRequest,
) {
        sess := session.New(device.ResolvedUserID, device.SessionID, device.ClientID, session.UserClaims{
                Name:  device.ResolvedUserName,
                Email: device.ResolvedUserEmail,
        })
        sess.VerifiedUpstream = device.VerifiedUpstream
        requester.SetSession(sess)

        for _, scope := range device.Scopes {
                requester.GrantScope(scope)
        }
        for _, aud := range device.Audience {
                requester.GrantAudience(aud)
        }

        deviceGrant := fosite.GrantType(oauthproto.GrantTypeDeviceCode)
        atLifespan := fosite.GetEffectiveLifespan(client, deviceGrant, fosite.AccessToken, h.Config.GetAccessTokenLifespan(ctx))
        sess.SetExpiresAt(fosite.AccessToken, time.Now().UTC().Add(atLifespan).Round(time.Second))

        if client.GetGrantTypes().Has(oauthproto.GrantTypeRefreshToken) {
                rtLifespan := fosite.GetEffectiveLifespan(client, deviceGrant, fosite.RefreshToken, h.Config.GetRefreshTokenLifespan(ctx))
                sess.SetExpiresAt(fosite.RefreshToken, time.Now().UTC().Add(rtLifespan).Round(time.Second))
        }
}

// enforcePollInterval applies RFC 8628 Section 3.5's minimum polling
// interval. It records this poll's timestamp unconditionally — including on
// the slow_down path — so a client polling faster than MinInterval cannot
// reset its own window by polling again before the interval elapses.
func (h *Handler) enforcePollInterval(ctx context.Context, device *authstorage.DeviceRequest, deviceCode string) error {
        tooSoon := !device.LastPolledAt.IsZero() && time.Since(device.LastPolledAt) < h.MinInterval
        if err := h.DeviceStorage.UpdateDeviceRequestLastPolledAt(ctx, deviceCode, time.Now()); err != nil {
                return errorsx.WithStack(fosite.ErrServerError.WithWrap(err))
        }
        if tooSoon {
                return errorsx.WithStack(ErrSlowDown)
        }
        return nil
}

// PopulateTokenEndpointResponse issues the access token and, when the client
// is registered for the refresh_token grant, a refresh token.
func (h *Handler) PopulateTokenEndpointResponse(
        ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder,
) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }

        access, accessSignature, err := h.Strategy.GenerateAccessToken(ctx, requester)
        if err != nil {
                return errorsx.WithStack(fosite.ErrServerError.WithWrap(err))
        }

        var refresh, refreshSignature string
        if requester.GetClient().GetGrantTypes().Has(oauthproto.GrantTypeRefreshToken) {
                refresh, refreshSignature, err = h.Strategy.GenerateRefreshToken(ctx, requester)
                if err != nil {
                        return errorsx.WithStack(fosite.ErrServerError.WithWrap(err))
                }
        }

        if err := h.CoreStorage.CreateAccessTokenSession(ctx, accessSignature, requester.Sanitize([]string{})); err != nil {
                return errorsx.WithStack(fosite.ErrServerError.WithWrap(err))
        }
        if refreshSignature != "" {
                refreshReq := requester.Sanitize([]string{})
                if err := h.CoreStorage.CreateRefreshTokenSession(ctx, refreshSignature, accessSignature, refreshReq); err != nil {
                        return errorsx.WithStack(fosite.ErrServerError.WithWrap(err))
                }
        }

        deviceGrant := fosite.GrantType(oauthproto.GrantTypeDeviceCode)
        atLifespan := fosite.GetEffectiveLifespan(
                requester.GetClient(), deviceGrant, fosite.AccessToken, h.Config.GetAccessTokenLifespan(ctx))
        responder.SetAccessToken(access)
        responder.SetTokenType("bearer")
        responder.SetExpiresIn(expiresIn(requester, fosite.AccessToken, atLifespan))
        responder.SetScopes(requester.GetGrantedScopes())
        if refresh != "" {
                responder.SetExtra("refresh_token", refresh)
        }

        return nil
}

// expiresIn mirrors fosite's own unexported helper of the same purpose
// (handler/oauth2/helper.go): it reports the session's actual expiry when
// one was set (HandleTokenEndpointRequest always sets one for AccessToken),
// falling back to defaultLifespan otherwise.
func expiresIn(r fosite.Requester, tokenType fosite.TokenType, defaultLifespan time.Duration) time.Duration {
        expiresAt := r.GetSession().GetExpiresAt(tokenType)
        if expiresAt.IsZero() {
                return defaultLifespan
        }
        return time.Until(expiresAt)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "crypto/rand"
        "errors"
        "log/slog"
        "net/http"
        "net/url"
        "slices"
        "strings"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/internal/consentpolicy"
        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// upstreamAuthSecrets holds cryptographic values needed for upstream IDP authorization.
type upstreamAuthSecrets struct {
        // State is the internal state for correlating the upstream callback.
        State string
        // PKCEVerifier is the code_verifier for upstream PKCE (RFC 7636).
        PKCEVerifier string
        // PKCEChallenge is the code_challenge derived from PKCEVerifier.
        PKCEChallenge string
        // Nonce is the OIDC nonce for ID token replay protection.
        Nonce string
}

// newUpstreamAuthSecrets generates all cryptographic secrets needed for upstream authorization.
func newUpstreamAuthSecrets() *upstreamAuthSecrets {
        verifier := crypto.GeneratePKCEVerifier()
        return &upstreamAuthSecrets{
                State:         rand.Text(),
                PKCEVerifier:  verifier,
                PKCEChallenge: crypto.ComputePKCEChallenge(verifier),
                Nonce:         rand.Text(),
        }
}

// AuthorizeHandler handles GET /oauth/authorize requests.
// It validates the client's request, then displays ToolHive consent when required
// or redirects to an upstream provider for exempt or remembered requests.
func (h *Handler) AuthorizeHandler(w http.ResponseWriter, req *http.Request) {
        ctx := req.Context()

        // Back-channel-only clients must be indistinguishable from unknown clients
        // at this unauthenticated endpoint. This runs before redirect URI matching
        // (including the loopback matcher below), which would otherwise expose that
        // the client is configured through an invalid_request response.
        if h.rejectBackChannelOnlyAuthorizeClient(ctx, w, req) {
                return
        }

        // See rewriteLoopbackRedirectURI's doc comment for what this does and why.
        // Keep the parsed form before any rewrite so exemption checks can distinguish
        // an explicitly registered redirect from a rewritten or inferred value.
        var originalForm url.Values
        if err := req.ParseForm(); err == nil {
                originalForm = make(url.Values, len(req.Form))
                for key, values := range req.Form {
                        originalForm[key] = slices.Clone(values)
                }
        }
        rewrittenFrom := h.rewriteLoopbackRedirectURI(ctx, req)

        // Let fosite validate everything: client_id, redirect_uri, response_type, PKCE, scopes
        ar, err := h.provider.NewAuthorizeRequest(ctx, req)
        // If the rewrite above fired, fosite validated (and stored) the
        // registered portless literal, not the dynamic port the client actually
        // requested. Wrapping ar restores the real listener as the error-redirect
        // target for every error path below -- see loopbackAuthorizeRequester's
        // doc comment for how -- not just a NewAuthorizeRequest failure.
        errAr := wrapLoopbackErrorRequester(ar, rewrittenFrom)
        if err != nil {
                h.writeAuthorizeRequestError(ctx, w, errAr, err)
                return
        }

        // Extract validated data from the authorize request
        clientID := ar.GetClient().GetID()
        // Use the original requested redirect_uri (before the loopback rewrite above)
        // so the dynamic port survives into PendingAuthorization and everything
        // downstream. rewrittenFrom is only non-empty when a rewrite actually
        // happened; every other case falls through to ar.GetRedirectURI() unchanged.
        redirectURI := ar.GetRedirectURI().String()
        if rewrittenFrom != "" {
                redirectURI = rewrittenFrom
        }
        state := ar.GetState()
        codeChallenge := ar.GetRequestForm().Get("code_challenge")
        codeChallengeMethod := ar.GetRequestForm().Get("code_challenge_method")
        scopes := []string(ar.GetRequestedScopes())

        // Token-only servers have no interactive authorization flow. Validation above
        // deliberately still runs first, so client and redirect errors retain their
        // normal OAuth semantics.
        if h.tokenOnly {
                h.provider.WriteAuthorizeError(ctx, w, errAr, fosite.ErrUnsupportedResponseType)
                return
        }

        // Bind the resource while the client request is still being authorized.
        // Multi-audience deployments require an explicit choice before upstream login.
        resource, ok := h.validateAuthorizationResource(ctx, w, errAr, ar)
        if !ok {
                return
        }

        slog.Debug("authorize request received",
                "client_id", clientID,
                "redirect_uri", redirectURI,
                "scope_count", len(scopes),
        )

        // Prepare the pending request before choosing the consent or upstream path;
        // beginUpstream persists its secret-bearing copy before redirecting.
        binding := newBrowserBinding()
        handle := rand.Text()
        pending := &storage.PendingAuthorization{
                ClientID: clientID, RedirectURI: redirectURI, State: state,
                PKCEChallenge: codeChallenge, PKCEMethod: codeChallengeMethod,
                Scopes: scopes, Resource: resource, ConsentStage: storage.ConsentStageAwaiting,
                BrowserBindingHash: binding.hash, UpstreamProviderName: h.upstreams[0].Name,
                SessionID: rand.Text(), CreatedAt: time.Now(),
        }
        promptConsent := promptRequiresConsent(ar.GetRequestForm())
        exempt := !promptConsent && h.consentExempt(ar, originalForm, rewrittenFrom)
        if len(h.consentExemptClientIDSet) > 0 {
                slog.Debug("consent exemption decision",
                        "client_id", clientID,
                        "exempt", exempt,
                        "prompt_requires_consent", promptConsent,
                )
        }
        if !promptConsent {
                if exempt {
                        pending.ConsentStage = storage.ConsentStageApproved
                        h.beginUpstream(ctx, w, req, errAr, pending)
                        return
                }
                if h.beginIfRemembered(ctx, w, req, errAr, pending, redirectURI, scopes, resource) {
                        return
                }
        }
        if err := h.storage.StorePendingAuthorization(ctx, handle, pending); err != nil {
                slog.Error("failed to store pending authorization", "error", err)
                h.provider.WriteAuthorizeError(ctx, w, errAr, fosite.ErrServerError.WithHint("failed to store authorization request"))
                return
        }
        h.setBrowserBindingCookie(w, handle, binding.value)
        h.renderConsent(w, handle, pending, ar.GetClient())
}

// validateAuthorizationResource resolves and validates the requested resource
// before authorization continues to consent or upstream authentication.
func (h *Handler) validateAuthorizationResource(
        ctx context.Context, w http.ResponseWriter, errAr fosite.AuthorizeRequester, ar fosite.AuthorizeRequester,
) (string, bool) {
        resources := ar.GetRequestForm()["resource"]
        if len(resources) > 1 || (len(resources) == 1 && resources[0] == "") {
                h.provider.WriteAuthorizeError(ctx, w, errAr, server.ErrInvalidTarget.WithHint("Exactly one non-empty resource is required"))
                return "", false
        }
        resource := ""
        if len(resources) == 1 {
                resource = resources[0]
        } else if len(h.config.AllowedAudiences) == 1 {
                resource = h.config.AllowedAudiences[0]
        } else {
                h.provider.WriteAuthorizeError(ctx, w, errAr, server.ErrInvalidTarget.WithHint("Specify one resource for this authorization"))
                return "", false
        }
        if resource == "" {
                h.provider.WriteAuthorizeError(ctx, w, errAr, server.ErrInvalidTarget.WithHint("No resource audience is configured"))
                return "", false
        }
        if err := server.ValidateAudienceURI(resource); err != nil {
                h.provider.WriteAuthorizeError(ctx, w, errAr, err)
                return "", false
        }
        if err := server.ValidateAudienceAllowed(resource, h.config.AllowedAudiences); err != nil {
                h.provider.WriteAuthorizeError(ctx, w, errAr, err)
                return "", false
        }
        return resource, true
}

// beginIfRemembered skips the consent page when a covering server-side approval
// exists and reports whether it answered the request. Callers must apply
// promptRequiresConsent before reaching this remembered-consent lookup.
func (h *Handler) beginIfRemembered(ctx context.Context, w http.ResponseWriter, req *http.Request,
        errAr fosite.AuthorizeRequester, pending *storage.PendingAuthorization,
        redirectURI string, scopes []string, resource string) bool {
        // Only a covering server-side approval can bypass the consent page.
        remembered, digest, err := h.coveringConsent(ctx, req, errAr.GetClient(), redirectURI, scopes, resource)
        if err != nil {
                slog.Error("failed to load remembered consent", "error", err)
                h.provider.WriteAuthorizeError(ctx, w, errAr, fosite.ErrServerError)
                return true
        }
        if remembered == nil {
                return false
        }
        pending.ConsentStage = storage.ConsentStageApproved
        pending.ExpectedUserID = remembered.UserID
        pending.ExpectedProviderSubject = remembered.ProviderSubject
        pending.ConsentSessionDigest = digest
        h.beginUpstream(ctx, w, req, errAr, pending)
        return true
}

// promptRequiresConsent reports whether any prompt parameter contains consent.
func promptRequiresConsent(form url.Values) bool {
        for _, prompt := range form["prompt"] {
                for _, token := range strings.Fields(prompt) {
                        if token == "consent" {
                                return true
                        }
                }
        }
        return false
}

// consentExempt reports whether this authorization request meets every runtime
// exemption condition: the client is allowlisted, confidential, interactive,
// supports the authorization-code grant and code response type, and has a
// registered redirect; the request asks for exactly code, supplies exactly one
// non-empty redirect explicitly, that redirect is registered and equals Fosite's
// validated final redirect, and no loopback redirect rewrite occurred.
func (h *Handler) consentExempt(
        ar fosite.AuthorizeRequester, originalForm url.Values, rewrittenFrom string,
) bool {
        client := ar.GetClient()
        if client == nil || !ar.GetResponseTypes().ExactOne(oauthproto.ResponseTypeCode) {
                return false
        }
        if _, ok := h.consentExemptClientIDSet[client.GetID()]; !ok || !consentpolicy.IsEligibleClient(client) {
                return false
        }
        redirectURIs := originalForm["redirect_uri"]
        return rewrittenFrom == "" && len(redirectURIs) == 1 && redirectURIs[0] != "" &&
                slices.Contains(client.GetRedirectURIs(), redirectURIs[0]) &&
                ar.GetRedirectURI().String() == redirectURIs[0]
}

// beginUpstream starts an approved request using fresh per-leg secrets.
func (h *Handler) beginUpstream(ctx context.Context, w http.ResponseWriter, req *http.Request,
        ar fosite.AuthorizeRequester, pending *storage.PendingAuthorization,
) {
        provider, ok := h.upstreamByName(pending.UpstreamProviderName)
        if !ok {
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("upstream provider not configured"))
                return
        }
        secrets := newUpstreamAuthSecrets()
        binding := newBrowserBinding()
        next := *pending
        next.InternalState = secrets.State
        next.UpstreamPKCEVerifier = secrets.PKCEVerifier
        next.UpstreamNonce = secrets.Nonce
        next.BrowserBindingHash = binding.hash
        if err := h.storage.StorePendingAuthorization(ctx, secrets.State, &next); err != nil {
                slog.Error("failed to store pending authorization", "error", err)
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to store authorization request"))
                return
        }
        var authOpts []upstream.AuthorizationOption
        if secrets.Nonce != "" {
                authOpts = append(authOpts, upstream.WithAdditionalParams(map[string]string{"nonce": secrets.Nonce}))
        }
        upstreamURL, err := provider.AuthorizationURL(secrets.State, secrets.PKCEChallenge, authOpts...)
        if err != nil {
                slog.Error("failed to build upstream authorization URL", "error", err)
                _ = h.storage.DeletePendingAuthorization(ctx, secrets.State)
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to build authorization URL"))
                return
        }
        h.setBrowserBindingCookie(w, secrets.State, binding.value)
        http.Redirect(w, req, upstreamURL, http.StatusSeeOther)
}

// writeAuthorizeRequestError answers a failed NewAuthorizeRequest: 429 when the
// CIMD fetch limiter refused the client lookup, fosite's error response
// otherwise.
func (h *Handler) writeAuthorizeRequestError(
        ctx context.Context, w http.ResponseWriter, ar fosite.AuthorizeRequester, err error,
) {
        if errors.Is(err, storage.ErrCIMDRateLimited) {
                writeCIMDRateLimited(w)
                return
        }
        h.provider.WriteAuthorizeError(ctx, w, ar, err)
}

// rejectBackChannelOnlyAuthorizeClient rejects a configured client that has no
// authorization response types before fosite validates redirect_uri. It returns
// false for unknown clients so fosite retains responsibility for its normal
// client lookup and error handling.
func (h *Handler) rejectBackChannelOnlyAuthorizeClient(ctx context.Context, w http.ResponseWriter, req *http.Request) bool {
        if err := req.ParseForm(); err != nil {
                return false
        }
        clientID := req.Form.Get("client_id")
        if clientID == "" {
                return false
        }
        client, err := h.storage.GetClient(ctx, clientID)
        if err != nil || client == nil || !consentpolicy.IsBackChannelOnlyClient(client) {
                return false
        }
        authorizeRequest := fosite.NewAuthorizeRequest()
        authorizeRequest.Client = client
        h.provider.WriteAuthorizeError(
                ctx,
                w,
                authorizeRequest,
                fosite.ErrInvalidClient.WithHint("The requested OAuth 2.0 Client does not exist."),
        )
        return true
}

// loopbackAuthorizeRequester wraps a fosite.AuthorizeRequester to make the
// client's real, dynamic-port loopback redirect_uri the error-redirect target
// for fosite's WriteAuthorizeError, instead of the registered portless
// literal that rewriteLoopbackRedirectURI substituted for validation.
//
// WriteAuthorizeError touches the requester only through the
// fosite.AuthorizeRequester interface (GetResponseMode, IsRedirectURIValid,
// GetRedirectURI, GetState, plus a G11NContext type-assertion covered below)
// -- it never type-asserts to the concrete *fosite.AuthorizeRequest -- so an
// embedding wrapper survives it intact.
//
// Overriding GetRedirectURI is the only way to restore the dynamic port:
// fosite.AuthorizeRequester has no SetRedirectURI method, and RedirectURI is
// a plain field on the concrete *fosite.AuthorizeRequest, so it can't be
// mutated back without a concrete-type assertion.
//
// Overriding IsRedirectURIValid is required alongside it: WriteAuthorizeError
// gates the redirect on this check, and fosite's own implementation cannot
// recognize "localhost" as loopback any more than the original /authorize
// request could -- it would reject the dynamic-port URI and fall back to a
// bare JSON body. The override only ever widens fosite's own answer with the
// loopback matcher's; see the method's doc comment for why it can't narrow.
//
// Known side effect: fosite's getLangFromRequester (i18n_helper.go) type-
// asserts the requester to fosite.G11NContext to read GetLang(), which lives
// on the embedded concrete *fosite.Request, not on the AuthorizeRequester
// interface -- so the wrapper always falls back to language.English. This is
// harmless here because no MessageCatalog is configured on the provider, but
// it's a real behavioural difference from an unwrapped requester.
type loopbackAuthorizeRequester struct {
        fosite.AuthorizeRequester
        // redirectURI is the client's real, requested dynamic-port URI.
        redirectURI *url.URL
}

// wrapLoopbackErrorRequester wraps ar in loopbackAuthorizeRequester when
// rewrittenFrom is non-empty (i.e. rewriteLoopbackRedirectURI fired), so a
// later validation failure redirects the error to the client's real listener
// instead of the registered portless placeholder. Returns ar unwrapped when
// no rewrite happened, or when rewrittenFrom fails to parse.
func wrapLoopbackErrorRequester(ar fosite.AuthorizeRequester, rewrittenFrom string) fosite.AuthorizeRequester {
        if rewrittenFrom == "" {
                return ar
        }
        parsed, err := url.Parse(rewrittenFrom)
        if err != nil {
                return ar
        }
        return &loopbackAuthorizeRequester{AuthorizeRequester: ar, redirectURI: parsed}
}

// IsRedirectURIValid may only ever WIDEN what fosite considers a valid
// redirect target, never narrow it: it first defers to the embedded
// requester's own check, and falls back to the loopback matcher only when
// that check says no. Narrowing would stop WriteAuthorizeError from
// delivering an error to a client whose redirect_uri fosite's own logic
// would have accepted on its own -- degrading a proper error redirect to a
// bare JSON body.
//
// The fallback adds exactly the "localhost" dynamic-port case fosite's own
// matcher cannot recognize (see rewriteLoopbackRedirectURI). It cannot be
// the sole check: registration.RegisteredLoopbackRedirectURI is
// public-clients-only (it early-returns false for any confidential client),
// so relying on it alone would reject confidential clients that fosite's
// own check would have validated.
func (r *loopbackAuthorizeRequester) IsRedirectURIValid() bool {
        if r.AuthorizeRequester.IsRedirectURIValid() {
                return true
        }
        client := r.GetClient()
        if client == nil {
                return false
        }
        _, ok := registration.RegisteredLoopbackRedirectURI(client, r.redirectURI.String())
        return ok
}

// GetRedirectURI returns a copy of redirectURI, never the stored pointer:
// WriteAuthorizeError mutates the returned *url.URL in place (clearing
// Fragment, overwriting RawQuery), which would otherwise corrupt
// r.redirectURI on first use.
func (r *loopbackAuthorizeRequester) GetRedirectURI() *url.URL {
        u := *r.redirectURI
        return &u
}

// logClientLookupFailure logs a GetClient failure from rewriteLoopbackRedirectURI
// at a level determined by the error: Debug for a not-found client (an
// unknown client_id is ordinary traffic on this unauthenticated endpoint --
// any caller can send a random client_id alongside a localhost redirect_uri,
// so logging every one at Warn would make log flooding trivial) and Warn for
// any other error, which indicates an actual backend problem worth surfacing.
//
// Both storage.ErrNotFound and fosite.ErrNotFound are checked: an opaque
// DCR-issued client_id not-found comes back wrapping storage.ErrNotFound, but
// a CIMD client_id (an https:// URL resolved live via
// CIMDStorageDecorator.fetch) that fails to resolve wraps fosite.ErrNotFound
// instead -- the same log-flooding risk exists for a caller who sends a
// bogus or unreachable CIMD URL as client_id.
func logClientLookupFailure(ctx context.Context, clientID string, err error) {
        level := slog.LevelWarn
        if errors.Is(err, storage.ErrNotFound) || errors.Is(err, fosite.ErrNotFound) {
                level = slog.LevelDebug
        }
        slog.Log(ctx, level, "failed to look up client for loopback redirect_uri rewrite",
                "client_id", clientID,
                "error", err,
        )
}

// rewriteLoopbackRedirectURI checks whether req's redirect_uri is a genuine
// RFC 8252 §7.3 loopback dynamic-port match against a client's registered
// redirect_uri (via registration.RegisteredLoopbackRedirectURI, so this works
// regardless of which concrete fosite.Client type the storage backend
// returns) and, if so, rewrites req.Form's redirect_uri to the registered
// (portless) literal so fosite's exact-match validation accepts it. It
// returns the original requested redirect_uri if a rewrite was made, or ""
// otherwise (including when req.Form can't be parsed, the client can't be
// resolved, isn't a loopback client, or the hostname is an IP literal --
// fosite's own validation and native loopback matching handle those cases
// as they do today).
//
// Why this rewrite exists at all: fosite has no client-side hook for
// "localhost" loopback matching (see registration.RegisteredLoopbackRedirectURI)
// -- it only matches redirect_uri by exact string equality against a client's registered
// URIs (checked first) or its own IP-literal-only loopback exception, so this
// rewrite is what makes fosite's exact-match branch accept a "localhost"
// dynamic-port request. See loopbackAuthorizeRequester above for how the
// error path recovers the dynamic port that this rewrite hides from fosite.
func (h *Handler) rewriteLoopbackRedirectURI(ctx context.Context, req *http.Request) string {
        if err := req.ParseForm(); err != nil {
                return ""
        }

        requestedRedirectURI := req.Form.Get("redirect_uri")
        clientID := req.Form.Get("client_id")
        if requestedRedirectURI == "" || clientID == "" {
                return ""
        }

        // Cheap, storage-free pre-check before doing a second client lookup (on top
        // of fosite's own internal one): a rewrite can only ever be needed for an
        // http redirect_uri whose hostname is "localhost" specifically. IP
        // literals (127.0.0.1, [::1]) are intentionally excluded: fosite's own
        // isMatchingAsLoopback already matches those natively, dynamic port and
        // all, on both success and error paths, so treating them as needing the
        // rewrite below would only add risk for no benefit. This also keeps the
        // extra storage lookup below off the common case (non-loopback
        // redirect_uris), narrowing its cost/availability impact to genuine
        // "localhost" loopback requests only.
        parsed, err := url.Parse(requestedRedirectURI)
        if err != nil || parsed.Scheme != "http" || !registration.IsLocalhostHostname(parsed.Hostname()) {
                return ""
        }

        // A lookup error here is not necessarily fatal to the request: fosite's
        // own independent lookup below can still succeed and validate the
        // (unrewritten) redirect_uri normally. It only becomes user-visible as a
        // generic "redirect_uri does not match any pre-registered redirect urls"
        // when BOTH lookups fail, or when this is a transient error (e.g. a
        // storage timeout) that fosite's second lookup doesn't hit -- in which
        // case the rewrite silently never happens. Log it so that case is
        // diagnosable; client_id is not a secret.
        client, err := h.storage.GetClient(ctx, clientID)
        if err != nil {
                logClientLookupFailure(ctx, clientID, err)
                return ""
        }

        registered, ok := registration.RegisteredLoopbackRedirectURI(client, requestedRedirectURI)
        if !ok || registered == requestedRedirectURI {
                return ""
        }

        req.Form.Set("redirect_uri", registered)
        return requestedRedirectURI
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "crypto/rand"
        "crypto/sha256"
        "crypto/subtle"
        "encoding/base64"
        "encoding/hex"
        "errors"
        "net/http"
        "net/url"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

// Browser binding ties an upstream callback to the user agent that started the
// authorization flow, as RFC 6749 Section 10.12 requires of an OAuth client
// (this server is the OAuth client toward the upstream IdP).
//
// /oauth/authorize mints a random secret, stores its hash on the
// PendingAuthorization, and sets the secret as a cookie on the redirect to the
// upstream. /oauth/callback loads the record by the upstream state and refuses
// to complete the leg unless the calling browser presents a cookie whose hash
// matches. Without this, an attacker who starts a flow for their own client can
// hand the resulting upstream URL to a victim and receive an authorization code
// bound to the victim's identity. Each multi-upstream chain leg mints its own
// secret, since every leg has a fresh state and pending record. The device
// flow's verification page (POST /oauth/device) shares the callback and the
// same shape, so it binds its PendingDeviceLogin the same way.
//
// Cookie shape:
//
//   - One cookie per pending record, named from the state, so concurrent flows in
//     one browser never overwrite each other. Cookies ignore ports, so several
//     local servers on different localhost ports share one jar; a fixed name
//     would collide across them.
//   - The "__Host-" prefix is used whenever the cookie is Secure. Browsers then
//     reject the cookie unless it was set from a secure origin with Path=/ and no
//     Domain, which stops a sibling subdomain or a plain-http origin on the same
//     host from planting a cookie whose value the attacker already knows. A
//     plain-http authorize URL has no browser-enforced defense of this kind, so it
//     gets the unprefixed name.
//   - Secure follows the scheme of the browser-facing authorize base URL, not the
//     inbound connection: pods serve plain HTTP behind a TLS-terminating ingress,
//     and the browser judges Secure against the connection it used.
//   - Path=/ is required by "__Host-" and also survives a prefix-stripping proxy
//     in front of a path-based issuer. It means a browser would also send the
//     cookie on /mcp, where the transparent proxy forwards Cookie to the backend
//     when strip-auth is not installed; the value is one-shot, useless without the
//     matching state, and expires with the pending record, so this is accepted.
//   - SameSite=Lax is sent on the IdP's top-level GET redirect back to
//     /oauth/callback. The callback is GET-only; if response_mode=form_post is
//     ever supported, the cookie would need SameSite=None and Secure.
//
// Only the hash of the secret is stored, so storage never holds the value.

const (
        // browserBindingCookieBase is the cookie name stem; the state-derived suffix
        // and the optional "__Host-" prefix are added by browserBindingCookieName.
        browserBindingCookieBase = "thv_authz_"

        // hostCookiePrefix is the RFC 6265bis cookie-name prefix that browsers only
        // accept for Secure, host-only, Path=/ cookies set from a secure origin.
        hostCookiePrefix = "__Host-"

        // browserBindingNameHashLen is the number of hex characters of
        // SHA-256(state) appended to the cookie name.
        browserBindingNameHashLen = 16
)

// The device flow's verification form (GET /oauth/device renders it, POST
// /oauth/device submits it) carries a double-submit anti-forgery token: the GET
// sets a random cookie (or reuses the one the browser already holds) and embeds
// the same value in a hidden field, and the POST accepts the form only when the
// two match. Without it a cross-site page
// could auto-submit an attacker's user_code from the victim's browser; that
// browser would then be issued the binding cookie above and could complete
// the login as if the victim had typed the code. The cookie has the same shape
// as the binding cookie, so a cross-site POST neither carries it (SameSite=Lax)
// nor can plant it (host-only, "__Host-" when Secure).
const (
        // deviceFormCookieBase is the anti-forgery cookie name stem for the device
        // verification form; one name per host, since the form has no state yet.
        deviceFormCookieBase = "thv_device_form"

        // deviceFormTokenField is the hidden form field that must echo the cookie.
        deviceFormTokenField = "form_token"

        // deviceFormTokenTTL bounds how long a rendered verification form may be
        // submitted; each render refreshes the cookie's Max-Age.
        deviceFormTokenTTL = 10 * time.Minute
)

var (
        errBrowserBindingCookieMissing = errors.New("browser binding cookie missing")
        errBrowserBindingUnbound       = errors.New("pending authorization has no browser binding")
        errBrowserBindingMismatch      = errors.New("browser binding cookie does not match")
        errDeviceFormTokenInvalid      = errors.New("device verification form token missing or mismatched")
)

// browserBinding is a freshly minted binding secret and its stored form.
type browserBinding struct {
        // value goes into the cookie and is never stored.
        value string
        // hash is BASE64URL(SHA256(value)) and is stored on the pending record.
        hash string
}

// newBrowserBinding mints a browser-binding secret.
func newBrowserBinding() browserBinding {
        value := rand.Text()
        return browserBinding{value: value, hash: hashBrowserBindingSecret(value)}
}

// hashBrowserBindingSecret returns the stored form of a cookie value.
func hashBrowserBindingSecret(value string) string {
        sum := sha256.Sum256([]byte(value))
        return base64.RawURLEncoding.EncodeToString(sum[:])
}

// browserBindingSecure reports whether the browser-facing authorize URL is
// served over https, which decides both the Secure attribute and the "__Host-"
// prefix. Config validation guarantees the URL parses with an http or https
// scheme; an unparsable value falls back to the plain-http shape, which every
// browser accepts.
func (h *Handler) browserBindingSecure() bool {
        u, err := url.Parse(h.config.GetAuthorizationEndpointBaseURL())
        if err != nil {
                return false
        }
        return strings.EqualFold(u.Scheme, "https")
}

// browserBindingCookieName derives the cookie name for one pending record from
// its upstream state. secure selects the "__Host-" prefixed shape. Set and
// verify must agree, so both go through here.
func browserBindingCookieName(internalState string, secure bool) string {
        sum := sha256.Sum256([]byte(internalState))
        name := browserBindingCookieBase + hex.EncodeToString(sum[:])[:browserBindingNameHashLen]
        if secure {
                return hostCookiePrefix + name
        }
        return name
}

// setBrowserBindingCookie writes the binding cookie for internalState onto a
// response that is about to redirect the browser to an upstream IdP.
func (h *Handler) setBrowserBindingCookie(w http.ResponseWriter, internalState, value string) {
        // G124 wants a literal Secure: true; here Secure deliberately follows the
        // configured authorize URL scheme because plain-http loopback issuers are
        // supported and a browser drops a Secure cookie set over plain http.
        secure := h.browserBindingSecure()
        http.SetCookie(w, &http.Cookie{ //nolint:gosec // G124 - Secure is derived from the authorize URL scheme
                Name:     browserBindingCookieName(internalState, secure),
                Value:    value,
                Path:     "/",
                MaxAge:   int(storage.DefaultPendingAuthorizationTTL.Seconds()),
                HttpOnly: true,
                Secure:   secure,
                SameSite: http.SameSiteLaxMode,
        })
}

// verifyBrowserBinding checks that req carries the binding cookie for the
// pending record loaded under internalState, whose stored hash is storedHash.
// Every cookie with the expected name is tried, so a planted duplicate cannot
// shadow the legitimate one. A record with an empty hash (written before
// browser binding existed) is rejected outright rather than treated as
// unbound-and-therefore-valid. The cookie is not cleared after use: the
// pending record is deleted on first use, so the cookie is inert and the
// browser drops it when Max-Age elapses.
func (h *Handler) verifyBrowserBinding(req *http.Request, internalState, storedHash string) error {
        if storedHash == "" {
                return errBrowserBindingUnbound
        }
        name := browserBindingCookieName(internalState, h.browserBindingSecure())
        expected := []byte(storedHash)
        found := false
        for _, c := range req.Cookies() {
                if c.Name != name {
                        continue
                }
                found = true
                if subtle.ConstantTimeCompare([]byte(hashBrowserBindingSecret(c.Value)), expected) == 1 {
                        return nil
                }
        }
        if !found {
                return errBrowserBindingCookieMissing
        }
        return errBrowserBindingMismatch
}

// deviceFormCookieName is the anti-forgery cookie name for the device
// verification form; secure selects the "__Host-" prefixed shape.
func deviceFormCookieName(secure bool) string {
        if secure {
                return hostCookiePrefix + deviceFormCookieBase
        }
        return deviceFormCookieBase
}

// issueDeviceFormToken returns the anti-forgery token for a render of the
// device verification form and (re)sets it as a cookie. A token the browser
// already holds is reused so that two open forms, in two tabs or on two local
// servers sharing a cookie jar, do not invalidate each other; a fresh one is
// minted only when the browser has none. Reuse is safe: a cross-site page can
// neither read the cookie nor make the browser send it (SameSite=Lax), so
// knowing the token is only possible from this origin.
func (h *Handler) issueDeviceFormToken(w http.ResponseWriter, req *http.Request) string {
        secure := h.browserBindingSecure()
        name := deviceFormCookieName(secure)
        token := ""
        for _, c := range req.Cookies() {
                if c.Name == name && c.Value != "" {
                        token = c.Value
                        break
                }
        }
        if token == "" {
                token = rand.Text()
        }
        http.SetCookie(w, &http.Cookie{ //nolint:gosec // G124 - Secure is derived from the authorize URL scheme
                Name:     name,
                Value:    token,
                Path:     "/",
                MaxAge:   int(deviceFormTokenTTL.Seconds()),
                HttpOnly: true,
                Secure:   secure,
                SameSite: http.SameSiteLaxMode,
        })
        return token
}

// verifyDeviceFormToken checks that the submitted form's hidden token matches
// an anti-forgery cookie on req. The caller has already parsed the form. Every
// cookie with the expected name is tried, as in verifyBrowserBinding.
func (h *Handler) verifyDeviceFormToken(req *http.Request) error {
        token := req.PostForm.Get(deviceFormTokenField)
        if token == "" {
                return errDeviceFormTokenInvalid
        }
        name := deviceFormCookieName(h.browserBindingSecure())
        for _, c := range req.Cookies() {
                if c.Name == name && subtle.ConstantTimeCompare([]byte(c.Value), []byte(token)) == 1 {
                        return nil
                }
        }
        return errDeviceFormTokenInvalid
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "net/url"
        "time"

        "github.com/ory/fosite"

        serveridentity "github.com/stacklok/toolhive/pkg/authserver/identity"
        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
)

// isNotFoundOrExpired reports whether err is a genuine "no such record"
// result -- storage.ErrNotFound or storage.ErrExpired (or an error wrapping
// either, e.g. via fosite's notFoundRFC6749Error) -- as opposed to a
// transient backend failure (a Redis timeout, a marshal error, etc.). Callers
// use this to avoid conflating the two: a real storage error must be reported
// as a server error and logged loudly, never silently reinterpreted as "not
// found" or used to trigger a fallback lookup.
func isNotFoundOrExpired(err error) bool {
        return errors.Is(err, storage.ErrNotFound) || errors.Is(err, storage.ErrExpired)
}

// deviceLoginOutcome is tryCompleteDeviceLogin's result: whether internalState
// matched a pending device login, was genuinely not found, or the lookup hit
// a real storage error (in which case a response has already been written).
type deviceLoginOutcome int

const (
        // deviceLoginNotFound means internalState does not match a pending device
        // login either; no response has been written and the caller should fall
        // through to its own not-found handling.
        deviceLoginNotFound deviceLoginOutcome = iota
        // deviceLoginCompleted means the device login was found and completeDeviceLogin
        // has already written a response.
        deviceLoginCompleted
        // deviceLoginStorageError means the lookup failed for a reason other than
        // not-found/expired; a response has already been written.
        deviceLoginStorageError
        // deviceLoginRejected means the device login was found but the calling
        // browser did not present its binding cookie; the login has been consumed
        // and a 400 response has already been written.
        deviceLoginRejected
)

// loadPendingOrCompleteDeviceLogin loads the pending OAuth-client
// authorization for internalState, verifies its browser binding and deadline,
// then deletes it (single-use) and builds its
// AuthorizeRequester. The upstream IDP's redirect_uri is fixed per upstream
// at construction (see upstream.OAuth2Config.RedirectURI / AuthCodeURL), so
// the device flow's verification-page login (DeviceVerificationSubmitHandler)
// reuses this same /oauth/callback endpoint rather than a distinct one --
// there is no way to register a second redirect_uri per call. Accordingly,
// when internalState genuinely does not match a pending OAuth-client
// authorization (not found/expired -- never a backend error, see
// isNotFoundOrExpired), this falls back to completing a pending device-flow
// login instead of failing outright.
//
// A pending OAuth-client authorization is only released to the browser that
// started it: req must carry the browser-binding cookie /oauth/authorize (or
// the previous chain leg) set for internalState, see browser_binding.go. The
// check runs only after the record loads, so the device-flow fallback, whose
// browser never visited /oauth/authorize, is unaffected.
//
// The bool return is false whenever a response has already been written
// (either the device-login completion, a storage error, a browser-binding
// failure, a not-found/corrupted error, or a failed lookup of the pending
// authorization's client) and the caller must return immediately without doing
// anything further.
func (h *Handler) loadPendingOrCompleteDeviceLogin(
        ctx context.Context, w http.ResponseWriter, req *http.Request, internalState, code string,
) (*storage.PendingAuthorization, fosite.AuthorizeRequester, bool) {
        pending, err := h.storage.LoadPendingAuthorization(ctx, internalState)
        if err != nil {
                if !isNotFoundOrExpired(err) {
                        // A genuine backend failure must not be silently reinterpreted as
                        // "maybe this is a device login" and reported to the client as an
                        // ordinary 400 -- that would mask an infra outage as client error.
                        slog.Error("failed to load pending authorization", "error", err)
                        http.Error(w, "internal server error", http.StatusInternalServerError)
                        return nil, nil, false
                }
                switch h.tryCompleteDeviceLogin(ctx, w, req, internalState, code) {
                case deviceLoginCompleted, deviceLoginStorageError, deviceLoginRejected:
                        return nil, nil, false
                case deviceLoginNotFound:
                        slog.Warn("pending authorization not found", "error", err)
                        http.Error(w, "authorization request not found or expired", http.StatusBadRequest)
                        return nil, nil, false
                }
        }

        if !h.validateAndConsumePending(ctx, w, req, internalState, pending, storage.ConsentStageApproved) {
                return nil, nil, false
        }

        ar, err := h.buildAuthorizeRequesterFromPending(ctx, pending)
        if err != nil {
                writePendingClientError(w, pending.ClientID, err)
                return nil, nil, false
        }
        if ar == nil {
                // Stored redirect URI was corrupt - cannot redirect, show error page
                http.Error(w, "authorization request data corrupted", http.StatusInternalServerError)
                return nil, nil, false
        }
        return pending, ar, true
}

// validateAndConsumePending checks the required consent stage, browser binding
// and absolute deadline before consuming the flow. Unbound or unapproved
// requests never invalidate the owner's state.
func (h *Handler) validateAndConsumePending(
        ctx context.Context, w http.ResponseWriter, req *http.Request, state string, pending *storage.PendingAuthorization, stage string,
) bool {
        if pending.ConsentStage != stage {
                http.Error(w, "authorization request not found or expired", http.StatusBadRequest)
                return false
        }
        if bindErr := h.verifyBrowserBinding(req, state, pending.BrowserBindingHash); bindErr != nil {
                writeUnboundCallbackRejection(w, bindErr,
                        "flow", "authorization_code",
                        "client_id", pending.ClientID,
                        "upstream_provider", pending.UpstreamProviderName,
                )
                return false
        }
        if pending.CreatedAt.IsZero() || !time.Now().Before(pending.CreatedAt.Add(storage.DefaultPendingAuthorizationTTL)) {
                http.Error(w, "authorization request not found or expired", http.StatusBadRequest)
                return false
        }
        if err := h.storage.DeletePendingAuthorization(ctx, state); err != nil {
                if isNotFoundOrExpired(err) {
                        // A concurrent request consumed the record first; single use held.
                        slog.Debug("pending authorization already consumed", "error", err)
                        http.Error(w, "authorization request not found or expired", http.StatusBadRequest)
                        return false
                }
                slog.Error("failed to delete pending authorization", "error", err)
                http.Error(w, "internal server error", http.StatusInternalServerError)
                return false
        }
        return true
}

// writeUnboundCallbackRejection logs and answers a browser-binding failure for
// either flow that shares /oauth/callback. The pending record is left intact so
// the browser that started the flow can still complete it. The response is a plain 400
// rather than a redirect: the browser making this request is by definition not
// the one that started the flow, so sending it to the client's redirect URI
// would hand a victim to whoever registered that client. attrs are extra slog
// key/value pairs identifying the flow; callers must not pass the cookie
// value, the stored hash, or device or user codes.
func writeUnboundCallbackRejection(w http.ResponseWriter, bindErr error, attrs ...any) {
        logAttrs := append([]any{"reason", bindErr.Error()}, attrs...)
        slog.Warn("rejected upstream callback: browser binding check failed", //nolint:gosec // G706 - server-side identifiers
                logAttrs...)
        http.Error(w,
                "sign-in must be completed in the browser that started it; cookies for this host must be enabled",
                http.StatusBadRequest)
}

// CallbackHandler handles GET /oauth/callback requests.
// It exchanges the upstream authorization code and issues our own authorization
// code. The request must carry the browser-binding cookie set when the flow
// started; see loadPendingOrCompleteDeviceLogin and browser_binding.go.
func (h *Handler) CallbackHandler(w http.ResponseWriter, req *http.Request) {
        ctx := req.Context()

        // Do not forward the upstream identity provider as the Referer when the
        // browser follows the final redirect to a native client's loopback listener.
        w.Header().Set("Referrer-Policy", "no-referrer")

        // Parse query parameters
        code := req.URL.Query().Get("code")
        internalState := req.URL.Query().Get("state")
        errorParam := req.URL.Query().Get("error")
        errorDescription := req.URL.Query().Get("error_description")

        // Handle upstream errors
        if errorParam != "" {
                h.handleUpstreamError(ctx, w, req, internalState, errorParam, errorDescription)
                return
        }

        // Validate required parameters - use http.Error for early errors without valid pending
        if internalState == "" {
                slog.Warn("callback missing state parameter")
                http.Error(w, "missing state parameter", http.StatusBadRequest)
                return
        }

        if code == "" {
                slog.Warn("callback missing code parameter")
                http.Error(w, "missing code parameter", http.StatusBadRequest)
                return
        }

        // Load and delete pending authorization (single-use), falling back to
        // completing a device-flow login when internalState doesn't match one.
        pending, ar, ok := h.loadPendingOrCompleteDeviceLogin(ctx, w, req, internalState, code)
        if !ok {
                return
        }

        // Look up the upstream provider that was used for this authorization leg.
        // Validating against pending.UpstreamProviderName (set during authorize) provides
        // IDP mix-up defense: we only accept callbacks for the provider we redirected to.
        upstreamProvider, ok := h.upstreamByName(pending.UpstreamProviderName)
        if !ok {
                slog.Error("upstream provider not found", "provider", pending.UpstreamProviderName)
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("upstream provider not configured"))
                return
        }

        // Exchange code and resolve identity in a single atomic operation.
        // This ensures OIDC nonce validation cannot be accidentally skipped.
        result, err := upstreamProvider.ExchangeCodeForIdentity(ctx, code, pending.UpstreamPKCEVerifier, pending.UpstreamNonce)
        if err != nil {
                slog.Error("failed to exchange code or resolve identity",
                        "error", err,
                )
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to exchange authorization code"))
                return
        }

        idpTokens := result.Tokens
        providerSubject := result.Subject

        // Use the logical upstream name as the provider identifier for storage and identity lookups.
        // This ensures write-side (StoreUpstreamTokens) and read-side (GetUpstreamTokens) keys match.
        providerID := pending.UpstreamProviderName

        // Use the session ID from the pending authorization.
        // This was generated in authorize.go and will be reused across all legs of the chain.
        sessionID := pending.SessionID

        if h.rememberedIdentityRejected(ctx, w, ar, pending, result, providerID) {
                return
        }

        // Determine identity: first leg resolves from upstream, subsequent legs
        // carry from pending. Synthetic identities (see upstream.Identity.Synthetic)
        // bypass UserResolver — the synthesized subject rotates per re-auth and
        // would otherwise grow `users` monotonically. We use the synthesized value
        // directly as an ephemeral session key and skip the LastAuthenticated
        // update (no provider_identities row to bump).
        var subject, userName, userEmail string
        if pending.ResolvedUserID == "" {
                // First leg — this is the identity provider
                if result.Synthetic {
                        subject = result.Subject
                } else {
                        user, err := h.userResolver.ResolveUser(ctx, providerID, providerSubject)
                        if err != nil {
                                slog.Error("failed to resolve user", "error", err)
                                // A UserStorage that deliberately refuses to auto-provision this
                                // identity (e.g. SCIM-only provisioning) is a client-facing denial,
                                // not a server fault — every other resolution failure keeps the
                                // generic server_error mapping.
                                resolveErr := fosite.ErrServerError.WithHint("failed to resolve user")
                                if errors.Is(err, storage.ErrUserNotProvisioned) {
                                        resolveErr = fosite.ErrAccessDenied.WithHint("user not provisioned")
                                }
                                h.provider.WriteAuthorizeError(ctx, w, ar, resolveErr)
                                return
                        }
                        subject = user.ID
                        if pending.RememberConsent {
                                clone := *pending
                                clone.FirstProviderSubject = providerSubject
                                pending = &clone
                        }
                        h.userResolver.UpdateLastAuthenticated(ctx, providerID, providerSubject)
                }
                userName = result.Name
                userEmail = result.Email
        } else {
                // Subsequent leg — use identity carried from first leg
                subject = pending.ResolvedUserID
                userName = pending.ResolvedUserName
                userEmail = pending.ResolvedUserEmail
                if !result.Synthetic {
                        h.userResolver.UpdateLastAuthenticated(ctx, providerID, providerSubject)
                }
        }

        // Capture provenance only on the callback that establishes the local ToolHive user.
        pending = carryFirstUpstreamIdentity(pending, h.recordedUpstreamIdentity(result))

        // Place the resolved canonical user in the request context so callback storage
        // calls that carry no tokens argument — GetAllUpstreamTokens during chain
        // consistency, DeleteUpstreamTokensForProvider on cleanup — can resolve the user
        // from context. The authserver-owned helper does not create a ToolHive identity:
        // no bearer has been issued, so there is no authenticated identity to assert —
        // only the canonical user for storage keying. (StoreUpstreamTokens does not need
        // this; it keys off tokens.UserID below.)
        ctx = serveridentity.WithCanonicalUser(ctx, subject)

        // Convert IDP tokens to storage tokens with binding fields.
        // SessionExpiresAt is set unconditionally as the Fosite session bound. Storage
        // backends use it as a fallback storage lifetime when ExpiresAt is zero (a
        // non-expiring upstream token). Setting it on every write — even when ExpiresAt
        // is non-zero — protects the refresh path: if the upstream provider stops
        // asserting expires_in on a later refresh, the carried-forward SessionExpiresAt
        // still bounds the storage lifetime instead of leaving the row indefinitely.
        storageTokens := &storage.UpstreamTokens{
                ProviderID:       providerID,
                AccessToken:      idpTokens.AccessToken,
                RefreshToken:     idpTokens.RefreshToken,
                IDToken:          idpTokens.IDToken,
                ExpiresAt:        idpTokens.ExpiresAt,
                SessionExpiresAt: time.Now().Add(h.config.RefreshTokenLifespan),
                ClientID:         pending.ClientID,
                UserID:           subject,         // Internal ToolHive user ID
                UpstreamSubject:  providerSubject, // Upstream IDP's subject claim
        }

        h.maybeCarryForwardRefreshToken(ctx, storageTokens, subject, providerSubject, providerID, result.Synthetic)

        if err := h.storage.StoreUpstreamTokens(ctx, sessionID, providerID, storageTokens); err != nil {
                slog.Error("failed to store upstream tokens",
                        "error", err,
                )
                // Do NOT clean up upstream tokens here: the store for this attempt just failed, so
                // this attempt wrote nothing. Under a (gateway, user)-keyed storage decorator the
                // only row at this provider is a credential from an earlier successful
                // authorization (e.g. a re-connect whose fresh store failed), and deleting it would
                // disconnect that already-connected account. Earlier legs are likewise the user's
                // own valid tokens. Leave them all; the pending is single-use and errors out below.
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to store session"))
                return
        }

        // Build the credential-free principal for the optional upstream filter, keyed on
        // the identity the first leg just established. providerSubject is the claim-mapped
        // upstream subject; subject is the canonical ToolHive user ID. On subsequent legs
        // the filter is not consulted, so this only drives filtering for upstreams[0].
        // result.Claims carries the ID-token/userinfo claims (nil for synthetic upstreams).
        principal := filterPrincipal{
                subject:         providerSubject,
                canonicalUserID: subject,
                name:            userName,
                email:           userEmail,
                claims:          result.Claims,
        }

        h.continueChainOrComplete(ctx, w, req, ar, pending, sessionID, principal)
}

// rememberedIdentityRejected checks that a remembered approval still resolves to
// the same linked account, before ResolveUser or any token write. It reports
// whether it already answered the request with an error.
func (h *Handler) rememberedIdentityRejected(ctx context.Context, w http.ResponseWriter, ar fosite.AuthorizeRequester,
        pending *storage.PendingAuthorization, result *upstream.Identity, providerID string) bool {
        providerSubject := result.Subject
        if pending.ResolvedUserID == "" && pending.ConsentSessionDigest != "" && pending.ExpectedUserID == "" {
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError)
                return true
        }
        if pending.ExpectedUserID != "" && pending.ResolvedUserID == "" {
                if pending.ConsentSessionDigest == "" || h.rememberedStorage == nil {
                        h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError)
                        return true
                }
                identity, lookupErr := h.storage.GetProviderIdentity(ctx, providerID, providerSubject)
                if !result.Synthetic && providerSubject == pending.ExpectedProviderSubject &&
                        lookupErr == nil && identity != nil && identity.UserID == pending.ExpectedUserID {
                        // The remembered identity was verified without changing user or token storage.
                } else if lookupErr != nil && !isNotFoundOrExpired(lookupErr) {
                        slog.Error("failed to verify remembered identity", "error", lookupErr)
                        h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError)
                        return true
                } else {
                        h.revokeRememberedSession(ctx, w, ar, pending.ConsentSessionDigest)
                        return true
                }
        }
        return false
}

// revokeRememberedSession deletes the originating consent session and answers
// access_denied, or server_error when the revocation itself fails.
func (h *Handler) revokeRememberedSession(ctx context.Context, w http.ResponseWriter,
        ar fosite.AuthorizeRequester, digest string) {
        if err := h.rememberedStorage.DeleteConsentSession(ctx, digest); err != nil {
                slog.Error("failed to revoke mismatched consent session", "error", err)
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError)
                return
        }
        h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrAccessDenied)
}

// maybeCarryForwardRefreshToken preserves a prior refresh token when the upstream IdP
// omits refresh_token on re-authorization (a common behavior, e.g. Google without
// prompt=consent). Without this, the new row would be written with an empty RefreshToken,
// orphaning the previously-issued RT and forcing the next refresh attempt to fail.
// Mirrors the preservation pattern in upstreamTokenRefresher.RefreshAndStore.
//
// The UpstreamSubject == providerSubject guard is defense-in-depth against account-linking
// edge cases where one internal user might have two distinct upstream subjects on the
// same provider. It is skipped for synthetic providers: a synthetic identity (OAuth2 with
// no userinfo/identity config) mints a fresh rotating subject every flow, so the equality
// can never hold and would block carry-forward entirely. There is no stable upstream
// subject to link by in that case, so the guard protects nothing — gating on a non-empty
// prior refresh token is sufficient.
//
// storageTokens is mutated in-place only when a carry-forward is warranted.
func (h *Handler) maybeCarryForwardRefreshToken(
        ctx context.Context,
        storageTokens *storage.UpstreamTokens,
        subject, providerSubject, providerID string,
        synthetic bool,
) {
        if storageTokens.RefreshToken != "" {
                return
        }
        prior, err := h.storage.GetLatestUpstreamTokensForUser(ctx, subject, providerID)
        switch {
        case err == nil:
                // Defensive: the contract returns ErrNotFound (handled below) on a miss
                // rather than a nil row, but guard anyway so a (nil, nil) return — or an
                // empty prior refresh token — can never be carried forward or panic. This
                // runs before the synthetic branch precisely so "synthetic" can never imply
                // "carry forward from a non-existent prior row".
                if prior == nil || prior.RefreshToken == "" {
                        return
                }
                // Non-synthetic: defense-in-depth account-linking guard requires the prior
                // row's upstream subject to match. Synthetic providers mint a fresh rotating
                // subject every flow, so that equality can never hold and there is no stable
                // subject to link by — skip the guard and rely on the prior RT alone.
                if synthetic || prior.UpstreamSubject == providerSubject {
                        storageTokens.RefreshToken = prior.RefreshToken
                        slog.Debug("preserved upstream refresh token from prior row",
                                "user_id", subject, "provider_id", providerID, "synthetic", synthetic,
                        )
                }
        case errors.Is(err, storage.ErrNotFound):
                // First authorization for this user/provider — nothing to preserve.
        default:
                // Storage error — log and continue with empty RT. Failing the callback
                // would be a worse user experience than the (already broken) status quo.
                slog.Warn("failed to look up prior upstream tokens for RT preservation",
                        "error", err, "user_id", subject, "provider_id", providerID,
                )
        }
}

// callbackClient resolves the client for /oauth/callback, the consent step
// (which reaches it through buildAuthorizeRequesterFromPending), the
// upstream-error branch (handleUpstreamError) and device verification. If the
// CIMD fetch limiter refuses the lookup (ErrCIMDRateLimited means no cached
// entry within the decorator's MaxStale could be served, e.g. a cold cache
// after a restart or an LRU eviction), it reads the persisted row from the
// storage below the CIMD decorator (still behind client admission) instead:
// the pending authorization was already consumed, so a 429 here
// would make the login unrecoverable.
//
// This is safe because the base row was written by the same fetch that
// /oauth/authorize validated against, and the pending authorization lives at
// most DefaultPendingAuthorizationTTL. /oauth/authorize may itself have been
// served a stale cache entry, so the snapshot is bounded by that lifetime plus
// the CIMD cache TTL plus the decorator's MaxStale, and does not depend on the
// row's own TTL.
//
// Device verification callers use the lookup only as a yes/no gate, because the
// resolved client is not used to issue anything. At device submit nothing has
// been consumed yet, but at login completion and at confirm the single-use
// pending record is already deleted, so, as on the callback, a 429 there would
// make the flow unrecoverable.
//
// The device request was created only after /oauth/device_authorization resolved
// the client, so the stale-row window there is bounded by the device code
// lifetime plus the CIMD cache TTL plus MaxStale.
//
// The base-row fallback is still needed even with stale serving, which only
// covers entries already cached in this process (not a cold cache). The shape
// guard used when CIMD is disabled returns plain not-found and never
// ErrCIMDRateLimited, so it cannot be bypassed this way. The base lookup's
// error is returned as-is.
func (h *Handler) callbackClient(ctx context.Context, id string) (fosite.Client, error) {
        client, err := h.storage.GetClient(ctx, id)
        if errors.Is(err, storage.ErrCIMDRateLimited) {
                return storage.ClientStoreBelowCIMD(h.storage).GetClient(ctx, id)
        }
        return client, err
}

// carryFirstUpstreamIdentity stores provenance only when this callback establishes the local user.
// SingleLeg authorizations connect one backend and never establish the user, so the
// upstream they authenticate against must not become the token's sub_id source.
func carryFirstUpstreamIdentity(
        pending *storage.PendingAuthorization,
        verified session.UpstreamSubject,
) *storage.PendingAuthorization {
        if pending.ResolvedUserID != "" || pending.SingleLeg {
                return pending
        }
        if !verified.IsComplete() {
                return pending
        }
        clone := *pending
        clone.VerifiedUpstream = verified
        return &clone
}

// recordedUpstreamIdentity returns the identity a login stores for later sub_id
// disclosure, or the zero value when SubjectIdentifierAudiences is empty. It decides
// whether the login records the identity at all: that needs only a non-empty list,
// while token.go still requires every audience of a token to be listed before it
// discloses. Do not share a predicate with token.go's per-audience check.
func (h *Handler) recordedUpstreamIdentity(result *upstream.Identity) session.UpstreamSubject {
        if len(h.config.SubjectIdentifierAudiences) == 0 {
                return session.UpstreamSubject{}
        }
        return verifiedUpstreamIdentity(result)
}

// verifiedUpstreamIdentity returns the validated upstream identity, or the zero
// value when the identity is synthetic or any member is missing.
func verifiedUpstreamIdentity(identity *upstream.Identity) session.UpstreamSubject {
        if identity.Synthetic {
                return session.UpstreamSubject{}
        }
        verified := session.UpstreamSubject{
                Issuer:      identity.VerifiedIssuer,
                Subject:     identity.VerifiedSubject,
                StableClaim: identity.VerifiedStableClaim,
                StableValue: identity.VerifiedStableValue,
        }
        if !verified.IsComplete() {
                return session.UpstreamSubject{}
        }
        return verified
}

// writeAuthorizationResponse generates an authorization code and writes the redirect response.
// This uses fosite's WriteAuthorizeResponse for RFC 6749 compliant redirects with proper
// status codes (303 See Other) and cache headers.
func (h *Handler) writeAuthorizationResponse(
        ctx context.Context,
        w http.ResponseWriter,
        pending *storage.PendingAuthorization,
        sessionID string,
        subject string,
        name string,
        email string,
) error {
        if pending.ConsentStage != storage.ConsentStageApproved {
                return errors.New("authorization has not been approved")
        }
        // Get the client from storage
        fositeClient, err := h.callbackClient(ctx, pending.ClientID)
        if err != nil {
                return err
        }

        // Create the session with IDP session reference, client ID, and user profile claims
        sess := session.New(subject, sessionID, pending.ClientID, session.UserClaims{
                Name:  name,
                Email: email,
        })
        sess.VerifiedUpstream = pending.VerifiedUpstream

        // Set expiration times
        now := time.Now()
        sess.SetExpiresAt(fosite.AuthorizeCode, now.Add(h.config.AuthorizeCodeLifespan))
        sess.SetExpiresAt(fosite.AccessToken, now.Add(h.config.AccessTokenLifespan))
        sess.SetExpiresAt(fosite.RefreshToken, now.Add(h.config.RefreshTokenLifespan))

        // Build the authorization request form
        form := url.Values{
                "redirect_uri":          {pending.RedirectURI},
                "code_challenge":        {pending.PKCEChallenge},
                "code_challenge_method": {pending.PKCEMethod},
        }

        // Create an authorize request using fosite
        authorizeRequest := fosite.NewAuthorizeRequest()
        authorizeRequest.Form = form
        authorizeRequest.Client = fositeClient
        authorizeRequest.Session = sess
        authorizeRequest.RequestedAt = now
        authorizeRequest.ResponseTypes = fosite.Arguments{"code"}
        authorizeRequest.State = pending.State // Set state for inclusion in redirect
        if pending.Resource != "" {
                authorizeRequest.SetRequestedAudience(fosite.Arguments{pending.Resource})
                authorizeRequest.GrantAudience(pending.Resource)
        }

        // Parse the redirect URI - this was validated by fosite during authorization,
        // so a parse error here indicates storage corruption
        redirectURI, err := url.Parse(pending.RedirectURI)
        if err != nil {
                return fmt.Errorf("stored redirect URI is invalid: %w", err)
        }
        authorizeRequest.RedirectURI = redirectURI

        // Grant only scopes that the client is registered for.
        // This prevents elevation if a tampered authorize request smuggled extra scopes
        // into the pending authorization.
        clientScopes := fositeClient.GetScopes()
        for _, scope := range pending.Scopes {
                if fosite.ExactScopeStrategy(clientScopes, scope) {
                        authorizeRequest.RequestedScope = append(authorizeRequest.RequestedScope, scope)
                        authorizeRequest.GrantedScope = append(authorizeRequest.GrantedScope, scope)
                } else {
                        slog.Warn("filtered unregistered scope from authorization", //nolint:gosec // G706 - scope from server-side storage
                                "scope", scope,
                                "client_id", pending.ClientID,
                        )
                }
        }

        // Generate the authorization response using fosite
        response, err := h.provider.NewAuthorizeResponse(ctx, authorizeRequest, sess)
        if err != nil {
                return err
        }

        // Only a successfully generated response may create a remembered browser session.
        h.maybeRememberConsent(ctx, w, pending, subject)

        // Write the redirect response using fosite's RFC 6749 compliant handler
        // This handles status code (303), cache headers, and URL building
        h.provider.WriteAuthorizeResponse(ctx, w, authorizeRequest, response)
        return nil
}

// isUnresolvedClient reports whether a client lookup failed because the client
// could not be resolved, as opposed to a storage or transport failure. A
// transient CIMD fetch failure also surfaces as fosite.ErrNotFound, so this does
// not mean the client was removed.
func isUnresolvedClient(err error) bool {
        return errors.Is(err, storage.ErrNotFound) || errors.Is(err, fosite.ErrNotFound) ||
                errors.Is(err, fosite.ErrInvalidClient)
}

// writePendingClientError answers a failed client lookup for a pending
// authorization: invalid_client for an unresolved client, 500 otherwise.
// This is a browser-facing endpoint, so the body is plain text rather than the
// JSON used on the token endpoints.
func writePendingClientError(w http.ResponseWriter, clientID string, err error) {
        if isUnresolvedClient(err) {
                slog.Debug("client for pending authorization could not be resolved", "client_id", clientID, "error", err)
                http.Error(w, "invalid_client: the client could not be resolved", http.StatusBadRequest)
                return
        }
        slog.Error("failed to resolve client for pending authorization", "client_id", clientID, "error", err)
        http.Error(w, "internal server error", http.StatusInternalServerError)
}

// buildAuthorizeRequesterFromPending creates the AuthorizeRequester for a pending
// authorization, used for error responses and to continue the authorization.
// It fails with the client lookup error instead of falling back to an empty
// client, so a removed client cannot go on to an upstream exchange. A nil
// requester with a nil error means the stored redirect URI is corrupt.
func (h *Handler) buildAuthorizeRequesterFromPending(
        ctx context.Context,
        pending *storage.PendingAuthorization,
) (fosite.AuthorizeRequester, error) {
        client, err := h.callbackClient(ctx, pending.ClientID)
        if err != nil {
                return nil, err
        }
        ar := fosite.NewAuthorizeRequest()
        ar.State = pending.State

        // Parse redirect URI - this was validated during authorization,
        // so failure indicates storage corruption
        redirectURI, err := url.Parse(pending.RedirectURI)
        if err != nil {
                slog.Error("stored redirect URI is invalid", //nolint:gosec // G706 - redirect URI from server-side storage
                        "redirect_uri", pending.RedirectURI,
                        "error", err,
                )
                return nil, nil
        }
        ar.RedirectURI = redirectURI

        ar.Client = client

        // ar.RedirectURI keeps the real dynamic-port URI from pending.RedirectURI.
        // fosite's own IsRedirectURIValid can't recognize "localhost" as loopback
        // any more than the original /authorize request could (see
        // rewriteLoopbackRedirectURI in authorize.go), so wrap here. Wrapping
        // unconditionally is safe because the wrapper's override only ever widens
        // validity (see loopbackAuthorizeRequester.IsRedirectURIValid), so a
        // non-loopback or confidential client keeps exactly the validity fosite
        // would have given it.
        return wrapLoopbackErrorRequester(ar, pending.RedirectURI), nil
}

// handleUpstreamError handles error responses from the upstream IDP.
// It attempts to redirect the error to the client if possible, otherwise
// falls back to completing a pending device-flow login as denied (see
// tryDenyDeviceLoginOnUpstreamError), otherwise shows an error page.
func (h *Handler) handleUpstreamError(
        ctx context.Context,
        w http.ResponseWriter,
        req *http.Request,
        internalState string,
        errorParam string,
        errorDescription string,
) {
        slog.Warn("upstream IDP returned error", //nolint:gosec // G706: error params from upstream IDP response
                "error", errorParam,
                "error_description", errorDescription,
        )

        // Try to load pending authorization to redirect error to client
        if internalState != "" {
                pending, err := h.storage.LoadPendingAuthorization(ctx, internalState)
                switch {
                case err == nil:
                        // Delete only the single-use pending authorization. Deliberately do NOT touch
                        // upstream tokens: the upstream IdP errored before this attempt exchanged a
                        // code, so this attempt stored nothing. Under a (gateway, user)-keyed storage
                        // decorator the only row at this provider is a credential from an earlier
                        // successful authorization — e.g. canceling a re-connect (error=access_denied)
                        // would otherwise disconnect the already-connected account — and earlier legs
                        // are likewise the user's own valid tokens. Leave them all intact.
                        if !h.validateAndConsumePending(ctx, w, req, internalState, pending, storage.ConsentStageApproved) {
                                return
                        }
                        ar, clientErr := h.buildAuthorizeRequesterFromPending(ctx, pending)
                        if clientErr != nil {
                                writePendingClientError(w, pending.ClientID, clientErr)
                                return
                        }
                        if ar != nil {
                                // Use generic error hint to avoid exposing upstream IDP details to clients.
                                // Detailed error information is logged above for server-side diagnostics.
                                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrAccessDenied.WithHint("upstream authentication failed"))
                                return
                        }
                        // ar is nil means stored redirect URI was corrupt - fall through to error page
                case isNotFoundOrExpired(err):
                        // internalState doesn't match a pending OAuth-client authorization --
                        // it may belong to a device-flow login instead (e.g. the human clicked
                        // "Deny" at the upstream IDP). See tryDenyDeviceLoginOnUpstreamError's
                        // doc comment for why this fallback matters.
                        if h.tryDenyDeviceLoginOnUpstreamError(ctx, w, req, internalState) {
                                return
                        }
                default:
                        // A genuine backend failure must not be silently swallowed, even though
                        // the response below is the same generic error either way.
                        slog.Error("failed to load pending authorization on upstream error", "error", err)
                }
        }

        // Cannot redirect to client, show generic error page.
        // Detailed error information is logged above for server-side diagnostics.
        http.Error(w, "upstream authentication failed", http.StatusBadGateway)
}

// cleanupUpstreamTokens best-effort removes upstream tokens, one provider at a time. It
// is called from exactly one place: the confirmed identity-mismatch fail-closed, where
// the session's stored state is no longer trustworthy and the whole effective chain is
// torn down. Every other callback error path deliberately preserves upstream tokens — a
// transient storage/filter failure, a downstream response-write error, or a canceled
// leg does not invalidate the credentials collected so far, and under a storage
// decorator that re-keys tokens by (gateway, user) those credentials are the user's own
// live connector logins, not throwaway session state.
//
// It deletes per provider rather than calling the session-wide DeleteUpstreamTokens
// because that method carries no provider name: a (gateway, user)-keyed decorator cannot
// narrow a session-wide delete back to one chain and would delete every connector token
// the user holds on the gateway. Naming each provider lets such a decorator scope the
// delete to exactly the effective chain's providers, so a configured provider the filter
// excluded from the chain is never touched.
//
// Each name is checked against the configured upstream set before it is deleted, so a
// name that is not a configured upstream — e.g. from a corrupted or tampered pending
// row — is skipped rather than turned into a delete.
//
// providers is the effective chain whose upstream-token rows should be removed; empty
// entries and names outside the configured upstream set are skipped, and deleting an
// already-absent row is a no-op (the storage contract makes a missing-row delete
// non-fatal), so a repeated name is harmless. Errors are logged and swallowed — cleanup
// is best-effort and must not mask the original failure.
func (h *Handler) cleanupUpstreamTokens(ctx context.Context, sessionID string, providers []string) {
        for _, provider := range providers {
                if provider == "" {
                        continue
                }
                // Only a configured upstream can legitimately have a stored token, so skip
                // anything else: it keeps a tampered chain entry from naming an arbitrary delete
                // target. The chain passed here is validateChain-checked, so this is a no-op in
                // normal operation — it is defense in depth for a corrupted pending row.
                if _, ok := h.upstreamByName(provider); !ok {
                        continue
                }
                if err := h.storage.DeleteUpstreamTokensForProvider(ctx, sessionID, provider); err != nil {
                        slog.Warn("failed to clean up upstream token for provider", //nolint:gosec // G706 - provider name from server-side chain state
                                "provider", provider,
                                "error", err,
                        )
                }
        }
}

// continueChainOrComplete checks whether all upstream providers in the authorization
// chain have been satisfied. If so, it issues the authorization code and redirects
// to the client. If not, it redirects to the next upstream provider to continue
// the chain. Called after StoreUpstreamTokens succeeds for each leg.
func (h *Handler) continueChainOrComplete(
        ctx context.Context,
        w http.ResponseWriter,
        req *http.Request,
        ar fosite.AuthorizeRequester,
        pending *storage.PendingAuthorization,
        sessionID string,
        principal FilterPrincipal,
) {
        // subject is the canonical ToolHive user ID used for chain state, token keying,
        // and the cross-leg identity check. Note this is principal.GetCanonicalUserID(), NOT
        // principal.GetSubject() (which is the upstream provider's subject).
        subject := principal.GetCanonicalUserID()
        name := principal.GetName()
        email := principal.GetEmail()

        // SingleLeg authorizations intentionally bypass chain continuation: the caller
        // scoped this flow to one specific upstream (e.g. a UI-initiated "connect one
        // backend" request), so other configured-but-tokenless upstreams must not
        // hijack it into a full chain walk. Issue the authorization code immediately.
        //
        // On failure we deliberately do NOT delete the stored upstream tokens. The
        // leg's token was already fetched and stored validly before we got here; the
        // errors writeAuthorizationResponse can return are all server-side and
        // unrelated to that credential (client lookup, redirect-URI parse, or fosite
        // failing to mint the authorization code). Wiping a good upstream token would
        // force a needless re-auth on what is a retryable error, so we keep it and
        // just surface the failure to the client.
        if pending.SingleLeg {
                if err := h.writeAuthorizationResponse(ctx, w, pending, sessionID, subject, name, email); err != nil {
                        slog.Error("failed to create authorization response", "error", err)
                        h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to create authorization code"))
                }
                return
        }

        // Resolve the effective chain of upstreams to walk. The first leg computes it
        // (consulting the optional filter with this leg's request context); every
        // subsequent leg reuses the validated chain the first leg carried forward, so
        // the filter is not re-run per leg. A subsequent leg whose pending predates the
        // chain is rejected rather than recomputed against a later leg's context.
        chain, err := h.resolveChain(ctx, pending, principal)
        if err != nil {
                slog.Error("failed to resolve upstream chain", "error", err)
                // Preserve upstream tokens: a chain-resolution failure (a transient filter/Directory
                // blip, or a stale pending) does not invalidate the credentials collected so far, and
                // this is the issue's most-reachable trigger. Fail authorization without deleting.
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to determine authorization chain"))
                return
        }

        nextProvider, err := h.nextMissingUpstream(ctx, sessionID, chain)
        if err != nil {
                slog.Error("failed to determine next upstream", "error", err)
                // Preserve upstream tokens: a transient chain-state read failure does not
                // invalidate the credentials collected so far. Fail authorization without deleting.
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to check authorization chain state"))
                return
        }

        if nextProvider == "" {
                if err := h.verifyChainIdentity(ctx, sessionID, chain, subject); err != nil {
                        // verifyChainIdentity already logged the specific cause. It fails for two very
                        // different reasons, and only one warrants deleting tokens:
                        //   - a CONFIRMED identity mismatch: the session's stored state is no longer
                        //     trustworthy, so fail closed and tear down the whole chain. This is the
                        //     ONLY path in the callback that deletes upstream tokens.
                        //   - a transient GetAllUpstreamTokens read failure: nothing is proven wrong,
                        //     so preserve every credential and just fail this authorization.
                        if errors.Is(err, errChainIdentityMismatch) {
                                h.cleanupUpstreamTokens(ctx, sessionID, chain)
                        }
                        h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("identity verification failed"))
                        return
                }

                // All upstreams satisfied — issue authorization code
                if err := h.writeAuthorizationResponse(ctx, w, pending, sessionID, subject, name, email); err != nil {
                        slog.Error("failed to create authorization response", "error", err)
                        // Preserve upstream tokens: the chain is complete and identity-verified, so the
                        // credentials are valid. Only the downstream response write failed (client
                        // lookup / redirect parse / code generation), which does not invalidate them —
                        // matching the SingleLeg branch, which preserves tokens for this same class.
                        h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to create authorization code"))
                }
                return
        }

        // Chain continues — redirect to next upstream.
        // TODO: If the user abandons the flow here (closes browser), upstream tokens from
        // completed legs remain in storage until their TTL expires. Add cascading cleanup
        // when pending authorizations expire to also delete associated upstream tokens.
        secrets := newUpstreamAuthSecrets()
        binding := newBrowserBinding()
        nextPending := &storage.PendingAuthorization{
                // Carry client request fields
                ClientID:                pending.ClientID,
                RedirectURI:             pending.RedirectURI,
                State:                   pending.State,
                PKCEChallenge:           pending.PKCEChallenge,
                PKCEMethod:              pending.PKCEMethod,
                Scopes:                  pending.Scopes,
                Resource:                pending.Resource,
                ConsentStage:            pending.ConsentStage,
                RememberConsent:         pending.RememberConsent,
                FirstProviderSubject:    pending.FirstProviderSubject,
                ExpectedUserID:          pending.ExpectedUserID,
                ExpectedProviderSubject: pending.ExpectedProviderSubject,
                ConsentSessionDigest:    pending.ConsentSessionDigest,
                // Fresh per-leg secrets
                InternalState:        secrets.State,
                UpstreamPKCEVerifier: secrets.PKCEVerifier,
                UpstreamNonce:        secrets.Nonce,
                BrowserBindingHash:   binding.hash,
                // Chain state
                UpstreamProviderName: nextProvider,
                SessionID:            sessionID,
                // Carry the effective chain forward so the filter is computed once, on the
                // first leg, and reused for every subsequent leg.
                ChainUpstreams: chain,
                // Carry resolved identity from first leg
                ResolvedUserID:    subject,
                ResolvedUserName:  name,
                ResolvedUserEmail: email,
                VerifiedUpstream:  pending.VerifiedUpstream,
                CreatedAt:         pending.CreatedAt,
        }

        if err := h.storage.StorePendingAuthorization(ctx, secrets.State, nextPending); err != nil {
                slog.Error("failed to store next chain leg", "error", err)
                // Preserve upstream tokens: failing to persist the next leg's pending does not
                // invalidate the credentials collected so far. Fail authorization without deleting.
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to continue authorization chain"))
                return
        }

        // Build authorization URL for next upstream
        var authOpts []upstream.AuthorizationOption
        if secrets.Nonce != "" {
                authOpts = append(authOpts, upstream.WithAdditionalParams(map[string]string{"nonce": secrets.Nonce}))
        }
        nextUpstream, ok := h.upstreamByName(nextProvider)
        if !ok {
                slog.Error("next upstream provider not found", "provider", nextProvider)
                _ = h.storage.DeletePendingAuthorization(ctx, secrets.State)
                // Preserve upstream tokens: a next-provider configuration error does not invalidate
                // the credentials collected so far. Fail authorization without deleting.
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("upstream provider configuration error"))
                return
        }
        nextURL, err := nextUpstream.AuthorizationURL(secrets.State, secrets.PKCEChallenge, authOpts...)
        if err != nil {
                slog.Error("failed to build next upstream authorization URL", "error", err)
                _ = h.storage.DeletePendingAuthorization(ctx, secrets.State)
                // Preserve upstream tokens: failing to build the next leg's authorization URL does
                // not invalidate the credentials collected so far. Fail authorization without deleting.
                h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to build authorization URL"))
                return
        }

        // Fresh binding for the fresh state: the same browser must also complete
        // the next leg's callback. Set only once every failure path above is past,
        // so a failed leg leaves no cookie behind. The consumed leg's cookie is
        // left to expire.
        h.setBrowserBindingCookie(w, secrets.State, binding.value)
        http.Redirect(w, req, nextURL, http.StatusFound)
}

// errChainIdentityMismatch marks the defense-in-depth failure where a chain's first
// leg's stored token belongs to a different user than the one carried through the flow.
// verifyChainIdentity wraps it so the callback can distinguish a confirmed mismatch —
// the only condition that warrants deleting upstream tokens — from a transient storage
// read failure, which must preserve them.
var errChainIdentityMismatch = errors.New("chain identity mismatch")

// verifyChainIdentity is a defense-in-depth check run once every leg of the
// effective chain is satisfied. Despite the "chain" framing, it reconciles only
// the first leg: it confirms the identity provider's stored token (chain[0]) still
// belongs to the subject carried through the flow. Intermediate/later legs are
// deliberately NOT identity-checked — those are connect-this-backend flows whose
// upstream identity can legitimately differ from the first leg's user.
//
// subject MUST be the canonical ToolHive user ID resolved from the first leg's
// upstream via that provider's configured subject-claim mapping (the OIDC
// SubjectClaim, or the "sub" claim by default) and carried forward unchanged
// through PendingAuthorization.ResolvedUserID. The caller resolves it exactly once,
// on the first leg, from the claim-mapped upstream subject. This cross-checks it
// against firstTokens.UserID — the same canonical ID persisted when the first leg
// stored its tokens — so a first leg whose stored user disagrees with the carried
// subject is rejected.
//
// It gates on the effective chain rather than the raw config: chain[0] is always
// the first (required) upstream, so a chain the filter narrowed to just that
// upstream has no first-leg cross-check to run and the check is a no-op. Returns a
// non-nil error when the storage lookup fails or an identity mismatch is detected;
// the caller maps either to a server error.
func (h *Handler) verifyChainIdentity(ctx context.Context, sessionID string, chain []string, subject string) error {
        if len(chain) <= 1 {
                return nil
        }
        allTokens, err := h.storage.GetAllUpstreamTokens(ctx, sessionID)
        if err != nil {
                slog.Error("failed to load upstream tokens for chain identity check", "error", err)
                return fmt.Errorf("failed to load upstream tokens for identity check: %w", err)
        }
        firstProvider := chain[0]
        if firstTokens, ok := allTokens[firstProvider]; ok && firstTokens.UserID != subject {
                // Emit the mismatch as discrete slog fields — not folded into the returned
                // error string — so log pipelines can filter/alert on this defense-in-depth
                // check, matching the logging from before this check was extracted here.
                slog.Error("identity mismatch between chain state and stored tokens",
                        "expected", subject,
                        "got", firstTokens.UserID,
                        "provider", firstProvider,
                )
                // Wrap the sentinel so the caller can tell a confirmed mismatch (tear down the
                // chain) apart from the storage-read failure above (preserve the tokens).
                return fmt.Errorf("%w on provider %q", errChainIdentityMismatch, firstProvider)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "crypto/sha256"
        "encoding/base64"
        "errors"
        "fmt"
        "html/template"
        "log/slog"
        "net/http"
        "net/netip"
        "net/url"
        "strconv"
        "strings"
        "time"
        "unicode"

        "github.com/ory/fosite"
        "golang.org/x/net/idna"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

type consentPageData struct {
        Handle, ActionURL, Destination, Summary, RedirectURI, ClientID, ClientName, Resource string
        // Upstream is the upstream's name (its routing and storage key); UpstreamDisplay
        // is the sanitized label shown in the "You will sign in with" line.
        Upstream, UpstreamDisplay string
        Scopes                    []string
        CanRemember               bool
        RememberFor               string
}

// consentStyles extends pageStyles for the consent page only; it is served as a
// second <style> element with its own CSP hash.
const consentStyles = `
details{margin-top:1rem}
summary{cursor:pointer; color:var(--muted); font-size:0.9rem; font-weight:500}
.subtitle.summary{margin-bottom:0.6rem}
.remember{margin-top:1.25rem; padding-top:0.75rem; border-top:1px solid var(--border)}
.remember label{
  display:flex; align-items:center; gap:0.6rem; min-height:2.75rem; cursor:pointer;
  font-size:0.9rem; font-weight:500;
}
.remember input{width:1.1rem; height:1.1rem; margin:0; flex:none; accent-color:var(--brand)}
.remember input:focus-visible{outline:2px solid var(--brand); outline-offset:2px}
.remember p{margin:0 0 0 1.7rem; color:var(--muted); font-size:0.8rem; line-height:1.4}
.remember + .btn-row{margin-top:1rem}
`

// formatDays renders d as whole days, falling back to hours when it is not a whole number of days.
func formatDays(d time.Duration) string {
        unit, n := "day", int(d/(24*time.Hour))
        if d%(24*time.Hour) != 0 {
                unit, n = "hour", int(d/time.Hour)
        }
        if n != 1 {
                unit += "s"
        }
        return strconv.Itoa(n) + " " + unit
}

// maxConsentTextRunes caps the client name and upstream label shown on the consent page.
const maxConsentTextRunes = 80

var consentPageTemplate = template.Must(template.New("consent").Parse(`<!DOCTYPE html>
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1">
<title>Sign-in request &middot; ToolHive</title><style>` + pageStyles + `</style><style>` + consentStyles + `</style></head><body>
<div class="card"><div class="card-body">
<h1>Allow this app to sign you in?</h1>
<p class="subtitle summary">{{.Summary}}</p>
<p class="subtitle">You will sign in with {{.UpstreamDisplay}} next.</p>
<dl><dt>Sends your sign-in to:</dt><dd><strong>{{.Destination}}</strong></dd>
<dt>Name given by the app (not verified):</dt><dd>{{.ClientName}}</dd></dl>
<details><summary>Technical details</summary><dl>
<dt>Full redirect URI</dt><dd>{{.RedirectURI}}</dd>
<dt>Client ID</dt><dd>{{.ClientID}}</dd>
<dt>Requested access</dt><dd>{{range .Scopes}}<span class="scope-chip">{{.}}</span>{{end}}</dd>
<dt>Resource</dt><dd>{{.Resource}}</dd><dt>Identity provider (internal name)</dt><dd>{{.Upstream}}</dd></dl></details>
<form method="POST" action="{{.ActionURL}}"><input type="hidden" name="handle" value="{{.Handle}}">
{{if .CanRemember}}<div class="remember"><label>
<input type="checkbox" name="remember" value="1" aria-describedby="remember-help">
<span>Don't ask again for this app on this browser</span></label>
<p id="remember-help">Skips this page for {{.RememberFor}} when the same app asks for the same access.</p></div>{{end}}
<div class="btn-row"><button class="btn-secondary" type="submit" name="decision" value="deny">Deny</button>
<button class="btn-primary" type="submit" name="decision" value="approve">Approve</button></div></form>
</div></div></body></html>`))

// classifyRedirect describes where the sign-in is sent without presenting a
// custom-scheme authority as a trusted destination. Anything that cannot be
// classified is shown as the full URI so a human can still inspect it.
func classifyRedirect(redirectURI string) (destination, summary string) {
        fallback := func() (string, string) {
                return redirectURI, "An app wants to receive your sign-in at the address shown below. " +
                        "Only continue if you just started this from that app or site."
        }
        u, err := url.Parse(redirectURI)
        if err != nil || u.Scheme == "" {
                return fallback()
        }
        scheme := strings.ToLower(u.Scheme)
        if scheme != "http" && scheme != "https" {
                return "an app that opens " + scheme + ":// links",
                        "An app on this device that opens " + scheme + ":// links wants to receive your sign-in. " +
                                "Only continue if you just started that app."
        }
        host := u.Hostname()
        if host == "" {
                return fallback()
        }
        if ip, err := netip.ParseAddr(host); strings.EqualFold(host, "localhost") || (err == nil && ip.IsLoopback()) {
                return "this computer", "An app running on this computer wants to receive your sign-in. " +
                        "Only continue if you just started that app."
        }
        destination = host
        // Show the punycode form too, so a lookalike IDN host cannot pass as a legitimate one.
        if ascii, err := idna.Lookup.ToASCII(host); err == nil && !strings.EqualFold(ascii, host) {
                destination = host + " (" + ascii + ")"
        }
        return destination, "An app wants to receive your sign-in at " + destination + ". " +
                "Only continue if you just started this from that app or site."
}

// displayClientName drops control and format runes (bidi overrides, zero-width
// characters, newlines) from an attacker-controlled name before it is shown.
func displayClientName(name string) string {
        return strings.Map(func(r rune) rune {
                if unicode.IsControl(r) || unicode.Is(unicode.Cf, r) {
                        return -1
                }
                return r
        }, name)
}

// truncateConsentText keeps at most maxConsentTextRunes runes of s and appends an
// ellipsis when it cuts, trimming any whitespace left before the ellipsis.
func truncateConsentText(s string) string {
        if r := []rune(s); len(r) > maxConsentTextRunes {
                return strings.TrimRightFunc(string(r[:maxConsentTextRunes]), unicode.IsSpace) + "…"
        }
        return s
}

// upstreamDisplayName returns the label for the named upstream on the consent page:
// its sanitized DisplayName when the upstream is still configured and has one,
// otherwise the name itself. The label is resolved at render time from this
// handler's upstream list, so a handler rebuilt with a different upstream set shows
// the current value; it is display-only and never a key.
func (h *Handler) upstreamDisplayName(name string) string {
        for i := range h.upstreams {
                if h.upstreams[i].Name != name {
                        continue
                }
                if label := strings.TrimSpace(displayClientName(h.upstreams[i].DisplayName)); label != "" {
                        return truncateConsentText(label)
                }
                return name
        }
        return name
}

func (h *Handler) renderConsent(
        w http.ResponseWriter, handle string, pending *storage.PendingAuthorization, client fosite.Client,
) {
        destination, summary := classifyRedirect(pending.RedirectURI)
        clientName := displayClientName(registration.ClientName(client))
        if clientName == "" {
                clientName = "No name given"
        } else {
                clientName = truncateConsentText(clientName)
        }
        data := consentPageData{
                Handle: handle, ActionURL: h.config.GetAuthorizationEndpointBaseURL() + "/oauth/consent",
                Destination: destination, Summary: summary, RedirectURI: pending.RedirectURI,
                ClientID: client.GetID(), ClientName: clientName, Resource: pending.Resource,
                Upstream: pending.UpstreamProviderName, UpstreamDisplay: h.upstreamDisplayName(pending.UpstreamProviderName),
                Scopes: pending.Scopes, CanRemember: h.rememberedStorage != nil,
                RememberFor: formatDays(storage.ConsentSessionTTL),
        }
        pageHash, consentHash := sha256.Sum256([]byte(pageStyles)), sha256.Sum256([]byte(consentStyles))
        setHTMLSecurityHeaders(w)
        w.Header().Set("Referrer-Policy", "same-origin")
        // A form-action restriction also applies to the POST's 302/303 target. Approval
        // navigates to the configured IdP and denial to the validated client URI
        // (including dynamic loopback ports), so 'self' would block both. The only
        // form target is the escaped, configured same-origin URL; base-uri 'none',
        // no scripts, frame-ancestors 'none', and the POST's Origin/binding checks
        // protect the submission. No form-action directive is set because it
        // would also restrict the validated cross-origin redirect after the POST.
        w.Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'sha256-"+
                base64.StdEncoding.EncodeToString(pageHash[:])+"' 'sha256-"+
                base64.StdEncoding.EncodeToString(consentHash[:])+"'; frame-ancestors 'none'; base-uri 'none'")
        if err := consentPageTemplate.Execute(w, data); err != nil {
                slog.Error("failed to render consent page", "error", err)
        }
}

// ConsentHandler handles the browser-bound decision before redirecting upstream.
func (h *Handler) ConsentHandler(w http.ResponseWriter, req *http.Request) {
        wantOrigin, err := expectedBrowserOrigin(h.config.GetAuthorizationEndpointBaseURL())
        origins := req.Header.Values("Origin")
        if err != nil || len(origins) != 1 || origins[0] != wantOrigin {
                http.Error(w, "invalid origin", http.StatusForbidden)
                return
        }
        values, ok := parseConsentForm(w, req, h.rememberedStorage != nil)
        if !ok {
                return
        }
        handle := values.Get("handle")
        pending, err := h.storage.LoadPendingAuthorization(req.Context(), handle)
        if err != nil {
                if errors.Is(err, storage.ErrNotFound) || errors.Is(err, storage.ErrExpired) {
                        http.Error(w, "authorization request not found or expired", http.StatusBadRequest)
                } else {
                        slog.Error("failed to load pending authorization", "error", err)
                        http.Error(w, "internal server error", http.StatusInternalServerError)
                }
                return
        }
        if !h.validateAndConsumePending(req.Context(), w, req, handle, pending, storage.ConsentStageAwaiting) {
                return
        }
        ar, err := h.buildAuthorizeRequesterFromPending(req.Context(), pending)
        if err != nil {
                writePendingClientError(w, pending.ClientID, err)
                return
        }
        if ar == nil {
                http.Error(w, "authorization request data corrupted", http.StatusInternalServerError)
                return
        }
        if values.Get("decision") == "deny" {
                h.provider.WriteAuthorizeError(req.Context(), w, ar, fosite.ErrAccessDenied)
                return
        }
        approved := *pending
        approved.ConsentStage = storage.ConsentStageApproved
        approved.RememberConsent = h.rememberedStorage != nil && values.Get("remember") == "1"
        h.beginUpstream(req.Context(), w, req, ar, &approved)
}

// expectedBrowserOrigin derives the browser's serialized origin solely from the
// trusted base URL, including compressed IPv6, ASCII domain names, and numeric ports.
func expectedBrowserOrigin(baseURL string) (string, error) {
        errInvalid := errors.New("invalid origin")
        base, err := url.Parse(baseURL)
        if err != nil || base.Host == "" || (base.Scheme != "http" && base.Scheme != "https") {
                return "", errInvalid
        }
        host, err := serializeOriginHost(base.Hostname())
        if err != nil {
                return "", err
        }
        if base.Port() != "" {
                port, err := strconv.ParseUint(base.Port(), 10, 16)
                if err != nil {
                        return "", errInvalid
                }
                if (base.Scheme != "http" || port != 80) && (base.Scheme != "https" || port != 443) {
                        host += ":" + strconv.FormatUint(port, 10)
                }
        }
        return (&url.URL{Scheme: base.Scheme, Host: host}).String(), nil
}

// parseConsentForm reads the size-limited consent form, requiring a handle and an
// approve/deny decision, plus remember=1 only when remembered consent is enabled.
// It writes the error response itself.
func parseConsentForm(w http.ResponseWriter, req *http.Request, allowRemember bool) (url.Values, bool) {
        req.Body = http.MaxBytesReader(w, req.Body, 4096)
        if err := req.ParseForm(); err != nil {
                http.Error(w, "invalid form", http.StatusBadRequest)
                return nil, false
        }
        values := req.PostForm
        if len(values) < 2 || len(values) > 3 || len(values["handle"]) != 1 || len(values["decision"]) != 1 ||
                (len(values) == 3 && (len(values["remember"]) != 1 || values.Get("remember") != "1" || !allowRemember)) ||
                values.Get("handle") == "" || (values.Get("decision") != "approve" && values.Get("decision") != "deny") {
                http.Error(w, "invalid form", http.StatusBadRequest)
                return nil, false
        }
        return values, true
}

// serializeOriginHost renders a hostname the way browsers serialize it in an origin.
func serializeOriginHost(hostname string) (string, error) {
        host := strings.ToLower(hostname)
        if !strings.Contains(host, ":") {
                ascii, err := idna.Lookup.ToASCII(host)
                if err != nil || ascii == "" {
                        return "", errors.New("invalid origin")
                }
                return ascii, nil
        }
        ip, err := netip.ParseAddr(host)
        if err != nil || ip.Zone() != "" {
                return "", errors.New("invalid origin")
        }
        host = ip.String()
        if ip.Is4In6() {
                // Browsers serialize mapped IPv4 as hexadecimal IPv6, not dotted decimal.
                b := ip.As16()
                host = fmt.Sprintf("::ffff:%x:%x", uint16(b[12])<<8|uint16(b[13]), uint16(b[14])<<8|uint16(b[15]))
        }
        return "[" + host + "]", nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "slices"
        "strconv"
        "strings"
        "time"

        "github.com/go-jose/go-jose/v3"
        "github.com/google/uuid"
        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// MaxDCRBodySize is the maximum allowed size for DCR request bodies (64KB).
// This prevents DoS attacks via extremely large payloads while being generous
// enough for legitimate requests with multiple redirect URIs.
//
// It is exported to serve as the single source of truth for the auth-server
// body-size cap: the embedded auth server (pkg/authserver/runner) derives its
// own request-body limit from this constant so the two cannot drift.
const MaxDCRBodySize = 64 * 1024

// RegisterClientHandler handles POST /oauth/register requests.
// It implements RFC 7591 Dynamic Client Registration for public clients with
// loopback redirect URIs, and additionally for confidential clients with
// https non-loopback redirect URIs when AllowConfidentialClientRegistration
// is set. Token-only servers accept only private_key_jwt clients registered
// for RFC 8693 token exchange.
func (h *Handler) RegisterClientHandler(w http.ResponseWriter, req *http.Request) {
        ctx := req.Context()

        // Limit request body size to prevent DoS attacks
        req.Body = http.MaxBytesReader(w, req.Body, MaxDCRBodySize)

        // Validate Content-Type header (RFC 7591 requires application/json)
        contentType := req.Header.Get("Content-Type")
        if !strings.HasPrefix(contentType, "application/json") {
                writeDCRError(w, http.StatusBadRequest, &registration.DCRError{
                        Error:            registration.DCRErrorInvalidClientMetadata,
                        ErrorDescription: "Content-Type must be application/json",
                })
                return
        }

        // Parse request body. oauthproto.ScopeList.UnmarshalJSON handles both
        // RFC 7591 wire formats for "scope" (space-delimited string or JSON
        // array) so we accept either shape transparently here.
        var dcrReq oauthproto.DynamicClientRegistrationRequest
        if err := json.NewDecoder(req.Body).Decode(&dcrReq); err != nil {
                writeDCRError(w, http.StatusBadRequest, &registration.DCRError{
                        Error:            registration.DCRErrorInvalidClientMetadata,
                        ErrorDescription: "invalid JSON request body",
                })
                return
        }

        // Validate request. h.config.AllowConfidentialClientRegistration gates whether
        // client_secret_basic / client_secret_post registrations are accepted.
        validated, dcrErr := h.validateDCRRequest(&dcrReq)
        if dcrErr != nil {
                writeDCRError(w, http.StatusBadRequest, dcrErr)
                return
        }

        // Validate requested scopes against server's supported scopes
        scopes, droppedDefaults, dcrErr := registration.ValidateScopes(dcrReq.Scopes, h.config.ScopesSupported)
        if dcrErr != nil {
                writeDCRError(w, http.StatusBadRequest, dcrErr)
                return
        }

        // Union with the operator-configured scope baseline. RFC 7591 §3.2.1 permits
        // the AS to replace requested client metadata values during registration; we
        // use that to expand the registered scope set so a client whose DCR request
        // narrowed the scope field can still request the baseline at /oauth/authorize.
        // h.config.BaselineClientScopes is validated at startup to be a subset of
        // ScopesSupported, so the union is guaranteed to be a subset of advertised
        // scopes. Operators should keep the baseline narrow (e.g. openid,
        // offline_access) — every DCR-registered client gains the ability to request
        // these scopes at /oauth/authorize regardless of what they registered with.
        if len(h.config.BaselineClientScopes) > 0 {
                effective := registration.UnionScopes(scopes, h.config.BaselineClientScopes)
                if !slices.Equal(effective, scopes) {
                        // Baseline-driven expansion is the intended behavior whenever
                        // baseline_client_scopes is configured, so per-registration
                        // audit lives at Debug rather than Warn. Operator-visible
                        // signal that the baseline is in effect comes from a one-time
                        // Info log at server startup (NewAuthorizationServerConfig).
                        slog.Debug("DCR registered scope set expanded by baseline_client_scopes",
                                "client_name", validated.ClientName,
                                "requested", scopes,
                                "effective", effective,
                                "baseline", h.config.BaselineClientScopes,
                        )
                        scopes = effective
                }
        }

        // Generate client ID
        clientID := uuid.NewString()

        // A request that declared itself public but whose redirect_uris exactly
        // matches an operator-configured force-confidential entry is registered
        // as confidential anyway (RFC 7591 §3.2.1 permits the server to
        // substitute client metadata). See Config.ForceConfidentialRedirectURIs
        // for the rationale: some MCP clients declare "none" and then refuse to
        // proceed without a client_secret.
        forcedURI, effectiveAuthMethod, dcrErr := resolveForceConfidentialOverride(validated, h.config.ForceConfidentialRedirectURIs)
        if dcrErr != nil {
                writeDCRError(w, http.StatusBadRequest, dcrErr)
                return
        }
        if forcedURI != "" {
                slog.Warn("DCR: forcing confidential client registration for redirect_uri "+
                        "that requested token_endpoint_auth_method 'none'",
                        "client_id", clientID,
                        "redirect_uri", forcedURI,
                )
        }

        fositeClient, clientSecret, err := buildDCRClient(
                clientID, forcedURI != "", effectiveAuthMethod, validated, scopes, h.config.AllowedAudiences,
                validated.JWKS, validated.TokenEndpointAuthSigningAlg)
        if err != nil {
                slog.Error("failed to create client", "error", err)
                writeDCRError(w, http.StatusInternalServerError, &registration.DCRError{
                        Error:            "server_error",
                        ErrorDescription: "failed to create client",
                })
                return
        }

        // Register client
        if err := h.storage.RegisterClient(ctx, fositeClient); err != nil {
                if errors.Is(err, storage.ErrClientCapacity) {
                        h.dcrMetrics.recordCapacityRejection(ctx)
                        h.capacityWarner.rejected(ctx, err)
                        w.Header().Set("Retry-After", strconv.Itoa(capacityRetryAfterSeconds(err)))
                        writeDCRError(w, http.StatusServiceUnavailable, &registration.DCRError{
                                Error:            "server_error",
                                ErrorDescription: "client registration capacity reached; try again later",
                        })
                        return
                }
                slog.Error("failed to register client", "error", err)
                writeDCRError(w, http.StatusInternalServerError, &registration.DCRError{
                        Error:            "server_error",
                        ErrorDescription: "failed to register client",
                })
                return
        }

        // Successful DCR registration is a normal operational event, not a
        // long-running operation, so it logs at Debug to stay silent at INFO+.
        // client_id, software_id, token_endpoint_auth_method, and scopes are
        // public client metadata per RFC 7591 and not credentials. If audit
        // signal is desired in future, the right home is a dedicated audit-
        // log emission path rather than promoting this record to INFO.
        //
        // Note: the "issuer" attribute here identifies THIS server (the
        // ToolHive-embedded AS that is performing the registration), not the
        // upstream AS being registered against. That distinction is important
        // when correlating these logs with the resolver's logs in
        // pkg/auth/dcr/resolver.go, which use "issuer" to mean the
        // upstream AS. The two uses live at opposite ends of the DCR flow.
        // No "upstream" attribute is emitted because the /oauth/register
        // endpoint has no upstream concept.
        //nolint:gosec // G706: client_id is public metadata per RFC 7591.
        slog.Debug("registered new DCR client",
                h.dcrClientLogAttrs(clientID, validated.SoftwareID, effectiveAuthMethod, scopes, droppedDefaults)...)

        // Build response per RFC 7591 Section 3.2.1.
        // Scopes reflects the scopes actually granted to this client: the
        // client-supplied scope set was validated against ScopesSupported by
        // ValidateScopes above, then (if configured) unioned with
        // BaselineClientScopes — which is itself guaranteed by startup-time
        // validation to be a subset of ScopesSupported. The unioned set is NOT
        // re-validated here. ScopeList.MarshalJSON emits the RFC 7591 §2
        // space-delimited wire form on the way out.
        issuedAt := time.Now().Unix()
        response := oauthproto.DynamicClientRegistrationResponse{
                ClientID:                    clientID,
                ClientIDIssuedAt:            issuedAt,
                RedirectURIs:                validated.RedirectURIs,
                ClientName:                  validated.ClientName,
                TokenEndpointAuthMethod:     effectiveAuthMethod,
                GrantTypes:                  validated.GrantTypes,
                ResponseTypes:               validated.ResponseTypes,
                Scopes:                      oauthproto.ScopeList(scopes),
                JWKS:                        validated.JWKS,
                JWKSURI:                     validated.JWKSURI,
                TokenEndpointAuthSigningAlg: validated.TokenEndpointAuthSigningAlg,
        }
        if clientSecret != "" {
                // client_secret_expires_at is 0 ("does not expire", RFC 7591 §2): the
                // storage TTL on a DCR-issued client is refreshed by RenewClientTTL on
                // every token exchange, so an actively used registration never expires
                // and advertising issued_at+TTL would be false. Advertising a real
                // expiry would also make ToolHive's own DCR client (which acts on this
                // field, pkg/auth/dcr/resolver.go) re-register against a ToolHive auth
                // server every TTL while its existing registration is still live,
                // orphaning the old row — the registration bloat the TTL exists to
                // prevent. An unused or idle registration still expires (after
                // DefaultUnusedDCRClientTTL or DefaultDCRClientTTL), but
                // re-registration mints a new client_id either way, so the field
                // cannot usefully describe that case.
                response.ClientSecret = clientSecret
                response.ClientSecretExpiresAt = new(int64) // 0: does not expire
        }

        w.Header().Set("Content-Type", "application/json")
        w.Header().Set("Cache-Control", "no-store")
        w.Header().Set("Pragma", "no-cache")
        w.WriteHeader(http.StatusCreated)
        if err := json.NewEncoder(w).Encode(response); err != nil {
                slog.Error("failed to encode DCR response", "error", err)
        }
}

func (h *Handler) validateDCRRequest(
        req *oauthproto.DynamicClientRegistrationRequest,
) (*oauthproto.DynamicClientRegistrationRequest, *registration.DCRError) {
        validated, dcrErr := registration.ValidateDCRRequest(
                req, h.config.AllowConfidentialClientRegistration, h.config.AllowPrivateKeyJWTRegistration)
        if dcrErr != nil {
                return validated, dcrErr
        }
        // validated.GrantTypes is the effective, post-defaulting grant set (e.g. an
        // empty grant_types on a private_key_jwt request defaults to token exchange
        // inside registration.validateGrantTypes), so this check catches both an
        // explicit and an implicit token-exchange request. Without it, DCR would
        // register a client whose GetGrantTypes() durably includes token exchange
        // even though the RFC 8693 factory is never registered at the token
        // endpoint when disabled, and every token request against that client
        // would fail confusingly with unsupported_grant_type instead of being
        // rejected here at registration time.
        if !h.config.TokenExchangeEnabled && slices.Contains(validated.GrantTypes, oauthproto.GrantTypeTokenExchange) {
                return nil, &registration.DCRError{
                        Error:            registration.DCRErrorInvalidClientMetadata,
                        ErrorDescription: "token exchange is disabled on this authorization server",
                }
        }
        if !h.tokenOnly {
                return validated, nil
        }
        if validated.TokenEndpointAuthMethod == oauthproto.TokenEndpointAuthMethodPrivateKeyJWT &&
                slices.Contains(validated.GrantTypes, oauthproto.GrantTypeTokenExchange) {
                return validated, nil
        }
        return nil, &registration.DCRError{
                Error:            registration.DCRErrorInvalidClientMetadata,
                ErrorDescription: "token-only authorization servers only permit private_key_jwt token-exchange registrations",
        }
}

// dcrClientLogAttrs builds the attribute list for the "registered new DCR
// client" record. dropped_defaults is attached only when the request omitted
// scope and scopes_supported did not carry the full default set (see issue
// #6186): recording it here — after the baseline union and the client_id
// mint upstream — means "scopes" is the set the client actually holds and
// the drop correlates with the client's later log lines. The drop itself is
// fully determined by startup config; the one-time operator-facing signal
// lives in Config.applyDefaults.
func (h *Handler) dcrClientLogAttrs(
        clientID, softwareID, effectiveAuthMethod string, scopes, droppedDefaults []string,
) []any {
        attrs := []any{
                "client_id", clientID,
                "software_id", softwareID,
                "token_endpoint_auth_method", effectiveAuthMethod,
                "scopes", scopes,
        }
        if len(droppedDefaults) > 0 {
                attrs = append(attrs, "dropped_defaults", droppedDefaults)
        }
        if issuer := h.issuer(); issuer != "" {
                attrs = append(attrs, "issuer", issuer)
        }
        return attrs
}

// resolveForceConfidentialOverride checks whether validated's redirect_uris
// match an operator-configured force-confidential entry and the request
// declared (or defaulted to) "none". It returns the matched URI ("" if no
// override applies) and the auth method that should actually be used to
// build and report the client.
//
// The override is reported as client_secret_post, never client_secret_basic:
// the Python MCP SDK these clients are typically built on constrains
// token_endpoint_auth_method to ["none", "client_secret_post"], and
// client_secret_basic is rejected there.
//
// A matched override upgrades the WHOLE registration to confidential, so the
// entire redirect_uris list — not just the matched entry — must pass the
// same https non-loopback policy the ordinary confidential path enforces
// (validateAuthMethod). Without this, a request mixing a matched https URI
// with an unrelated loopback URI would mint a secret for a client that also
// has a loopback callback, which the ordinary path already rejects. The
// registration is rejected rather than silently downgraded to public: a
// silent downgrade would hand back a public registration with no secret,
// reproducing the exact confusing failure this override exists to fix.
func resolveForceConfidentialOverride(
        validated *oauthproto.DynamicClientRegistrationRequest, forceConfidentialURIs []string,
) (forcedURI, effectiveAuthMethod string, dcrErr *registration.DCRError) {
        effectiveAuthMethod = validated.TokenEndpointAuthMethod
        if effectiveAuthMethod != oauthproto.TokenEndpointAuthMethodNone {
                return "", effectiveAuthMethod, nil
        }
        forcedURI = matchForceConfidentialRedirectURI(validated.RedirectURIs, forceConfidentialURIs)
        if forcedURI == "" {
                return "", effectiveAuthMethod, nil
        }
        effectiveAuthMethod = oauthproto.TokenEndpointAuthMethodClientSecretPost
        if dcrErr := registration.ValidateConfidentialRedirectURIs(validated.RedirectURIs, effectiveAuthMethod); dcrErr != nil {
                return "", "", dcrErr
        }
        return forcedURI, effectiveAuthMethod, nil
}

// buildDCRClient mints a client_secret when effectiveAuthMethod requires one
// and constructs the fosite client for a DCR registration. forced selects the
// plain *fosite.DefaultClient shape (registration.NewConfidentialPlain) used
// for a force-confidential override, rather than the OIDC shape
// registration.New produces for an ordinary confidential registration: fosite
// only enforces token_endpoint_auth_method on clients implementing
// fosite.OpenIDConnectClient, and an overridden client's credential
// presentation (HTTP Basic vs. form body) is unknown in advance, so the
// plain shape — which accepts either — is required there.
//
// Any client_secret field on the incoming request is ignored — secrets are
// always server-generated. The plaintext is returned to the caller exactly
// once; only its SHA-256 hash is stored.
func buildDCRClient(
        clientID string,
        forced bool,
        effectiveAuthMethod string,
        validated *oauthproto.DynamicClientRegistrationRequest,
        scopes, allowedAudiences []string,
        jwks *jose.JSONWebKeySet, signingAlgorithm string,
) (fositeClient fosite.Client, clientSecret string, err error) {
        if effectiveAuthMethod != oauthproto.TokenEndpointAuthMethodNone &&
                effectiveAuthMethod != oauthproto.TokenEndpointAuthMethodPrivateKeyJWT {
                clientSecret, err = registration.GenerateClientSecret()
                if err != nil {
                        return nil, "", fmt.Errorf("failed to generate client secret: %w", err)
                }
        }

        if forced {
                fositeClient, err = registration.NewConfidentialPlain(registration.Config{
                        ID:            clientID,
                        ClientName:    validated.ClientName,
                        Secret:        clientSecret,
                        RedirectURIs:  validated.RedirectURIs,
                        GrantTypes:    validated.GrantTypes,
                        ResponseTypes: validated.ResponseTypes,
                        Scopes:        scopes,
                        Audience:      allowedAudiences,
                })
        } else {
                fositeClient, err = registration.New(registration.Config{
                        ID:                                clientID,
                        ClientName:                        validated.ClientName,
                        Secret:                            clientSecret,
                        RedirectURIs:                      validated.RedirectURIs,
                        TokenEndpointAuthMethod:           validated.TokenEndpointAuthMethod,
                        GrantTypes:                        validated.GrantTypes,
                        ResponseTypes:                     validated.ResponseTypes,
                        Scopes:                            scopes,
                        Audience:                          allowedAudiences,
                        JSONWebKeys:                       jwks,
                        TokenEndpointAuthSigningAlgorithm: signingAlgorithm,
                })
        }
        if err != nil {
                return nil, "", err
        }
        return fositeClient, clientSecret, nil
}

// matchForceConfidentialRedirectURI returns the first entry in redirectURIs
// that is an exact match for one of the operator-configured
// forceConfidentialURIs, or "" if none match. Matching is deliberately exact
// (not prefix or scheme-relaxed): an attacker who registers with someone
// else's callback URI is issued a secret for a client whose authorization
// codes are delivered to that someone else's redirect endpoint, so exact
// matching does not hand out a usable credential for another client.
func matchForceConfidentialRedirectURI(redirectURIs, forceConfidentialURIs []string) string {
        for _, uri := range redirectURIs {
                if slices.Contains(forceConfidentialURIs, uri) {
                        return uri
                }
        }
        return ""
}

// writeDCRError writes a DCR error response per RFC 7591 Section 3.2.2.
func writeDCRError(w http.ResponseWriter, statusCode int, dcrErr *registration.DCRError) {
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(statusCode)
        // Encoding errors are not recoverable (headers already written), log for diagnostics
        if err := json.NewEncoder(w).Encode(dcrErr); err != nil {
                slog.Debug("failed to encode DCR error response", "error", err)
        }
}

// capacityRetryAfterSeconds returns the Retry-After value for a DCR cap
// rejection: the seconds until the oldest counted client expires, which is
// the earliest moment a slot can free up, rounded up and at least 1.
func capacityRetryAfterSeconds(err error) int {
        var capErr *storage.ClientCapacityError
        if !errors.As(err, &capErr) {
                return 1
        }
        return max(int((capErr.RetryAfter+time.Second-1)/time.Second), 1)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "errors"
        "log/slog"
        "sync"
        "time"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

// capacityWarnInterval is the minimum time between two cap-reached WARN lines.
const capacityWarnInterval = 15 * time.Minute

// capacityWarner logs registrations rejected at the DCR client cap at WARN, at
// most once per interval, so an operator without telemetry still learns why
// /oauth/register answers 503. The caller of /oauth/register is
// unauthenticated, so the volume is bounded by time rather than by request
// rate: rejections inside the interval are only counted, and the next line
// reports how many were held back. It deliberately logs neither the cap
// filling nor clearing, since a caller holding the cap flips between the two
// every time a slot expires.
type capacityWarner struct {
        interval time.Duration
        now      func() time.Time
        // counter reports the live count for the WARN line. Nil when the
        // storage backend cannot count DCR clients.
        counter storage.DCRClientCounter

        mu         sync.Mutex
        last       time.Time
        suppressed int64
}

// newCapacityWarner builds a capacityWarner for stor that logs at most once
// per capacityWarnInterval.
func newCapacityWarner(stor storage.Storage) *capacityWarner {
        counter, _ := storage.Unwrap(stor).(storage.DCRClientCounter)
        return &capacityWarner{interval: capacityWarnInterval, now: time.Now, counter: counter}
}

// rejected records one registration rejected with err, a ClientCapacityError,
// and logs it when the interval since the last line has passed. It is a no-op
// on a nil receiver, for Handlers built without NewHandler.
func (w *capacityWarner) rejected(ctx context.Context, err error) {
        if w == nil {
                return
        }
        w.mu.Lock()
        now := w.now()
        if !w.last.IsZero() && now.Sub(w.last) < w.interval {
                w.suppressed++
                w.mu.Unlock()
                return
        }
        w.last = now
        suppressed := w.suppressed
        w.suppressed = 0
        w.mu.Unlock()

        attrs := []any{
                "retry_after_seconds", capacityRetryAfterSeconds(err),
                "suppressed_rejections", suppressed,
        }
        var capErr *storage.ClientCapacityError
        if errors.As(err, &capErr) {
                attrs = append(attrs, "max_dcr_clients", capErr.Cap)
        }
        // One count per interval, so this read stays off the request path in
        // all but the request that logs.
        if w.counter != nil {
                if n, countErr := w.counter.CountDCRClients(ctx); countErr == nil {
                        attrs = append(attrs, "dcr_clients", n)
                }
        }
        attrs = append(attrs, "setting", "storage.max_dcr_clients")
        slog.Warn("DCR client cap reached; rejecting new client registrations", attrs...)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "fmt"

        "go.opentelemetry.io/otel/metric"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

const dcrInstrumentationName = "github.com/stacklok/toolhive/pkg/authserver"

// dcrMetrics lets operators see the DCR client cap. It records no per-client
// attributes. Every capacity rejection is counted here; capacityWarner logs
// them only at a time-bounded rate, because the caller of /oauth/register is
// unauthenticated.
type dcrMetrics struct {
        capacityRejections metric.Int64Counter
        // gaugeRegistration is the live-client gauge's callback, or nil when the
        // storage backend cannot count DCR clients.
        gaugeRegistration metric.Registration
}

// newDCRMetrics creates the capacity rejection counter and, when the storage
// backend implements storage.DCRClientCounter, the live-client gauge. The
// gauge is read from storage at collection time, never on a request path.
func newDCRMetrics(meterProvider metric.MeterProvider, stor storage.Storage) (*dcrMetrics, error) {
        meter := meterProvider.Meter(dcrInstrumentationName)

        capacityRejections, err := meter.Int64Counter(
                "toolhive_authserver_dcr_capacity_rejections",
                metric.WithDescription("Total number of dynamic client registrations rejected because the DCR client cap was reached"),
        )
        if err != nil {
                return nil, fmt.Errorf("create DCR capacity rejection counter: %w", err)
        }

        m := &dcrMetrics{capacityRejections: capacityRejections}
        counter, ok := storage.Unwrap(stor).(storage.DCRClientCounter)
        if !ok {
                return m, nil
        }
        gauge, err := meter.Int64ObservableGauge(
                "toolhive_authserver_dcr_clients",
                metric.WithDescription("Number of live dynamically registered clients counted against the DCR client cap"),
        )
        if err != nil {
                return nil, fmt.Errorf("create DCR client gauge: %w", err)
        }
        m.gaugeRegistration, err = meter.RegisterCallback(func(ctx context.Context, o metric.Observer) error {
                n, err := counter.CountDCRClients(ctx)
                if err != nil {
                        return err
                }
                o.ObserveInt64(gauge, n)
                return nil
        }, gauge)
        if err != nil {
                return nil, fmt.Errorf("register DCR client gauge callback: %w", err)
        }
        return m, nil
}

// close unregisters the gauge callback. It is a no-op on a nil receiver or
// when no gauge was registered.
func (m *dcrMetrics) close() error {
        if m == nil || m.gaugeRegistration == nil {
                return nil
        }
        if err := m.gaugeRegistration.Unregister(); err != nil {
                return fmt.Errorf("unregister DCR client gauge callback: %w", err)
        }
        return nil
}

// recordCapacityRejection counts one registration rejected at the DCR cap. It
// is a no-op on a nil receiver, for Handlers built without NewHandler.
func (m *dcrMetrics) recordCapacityRejection(ctx context.Context) {
        if m == nil {
                return
        }
        m.capacityRejections.Add(ctx, 1)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "crypto/rand"
        "encoding/base64"
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "net/url"
        "strings"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// maxDeviceCodeGenerationAttempts bounds retries when a freshly generated
// device_code/user_code pair collides with an existing pending request.
// Collision probability is astronomically small (32 bytes / 8 chars of
// crypto/rand output each); this only guards against exhausting the request
// with retries if storage is somehow degenerate.
const maxDeviceCodeGenerationAttempts = 5

// userCodeCharset is RFC 8628's suggested character set for the human-typed
// user_code: uppercase letters and digits with the visually ambiguous
// characters I, O, 0, and 1 removed, so a user reading the code off a screen
// cannot mistype it as a different valid character.
const userCodeCharset = "BCDFGHJKLMNPQRSTVWXZ0123456789"

// userCodeGroupLength is the length of each hyphen-separated group in the
// generated user_code (RFC 8628's example format is XXXX-XXXX).
const userCodeGroupLength = 4

// deviceAuthorizationResponse is the RFC 8628 Section 3.2 response body.
type deviceAuthorizationResponse struct {
        DeviceCode              string `json:"device_code"`
        UserCode                string `json:"user_code"`
        VerificationURI         string `json:"verification_uri"`
        VerificationURIComplete string `json:"verification_uri_complete,omitempty"`
        ExpiresIn               int    `json:"expires_in"`
        Interval                int    `json:"interval,omitempty"`
}

// clientAuthenticator is satisfied by *fosite.Fosite (the concrete type
// behind every fosite.OAuth2Provider this codebase constructs). It is
// asserted out of h.provider rather than added to a wider interface because
// fosite.OAuth2Provider itself does not expose AuthenticateClient -- it is
// only called internally by NewAccessRequest, the revocation handler, and
// the PAR handler.
type clientAuthenticator interface {
        AuthenticateClient(ctx context.Context, r *http.Request, form url.Values) (fosite.Client, error)
}

// DeviceAuthorizationHandler handles POST /oauth/device_authorization
// requests (RFC 8628 Section 3.1). It is only mounted when the server
// config enables the device-code grant — see OAuthRoutes.
//
// The request body is application/x-www-form-urlencoded, per RFC 8628
// Section 3.1; the response body is JSON, per Section 3.2.
//
// RFC 8628 Section 3.1 requires client authentication "using the same
// method that the [token endpoint] supports" -- so this handler
// authenticates through h.provider's own AuthenticateClient, the exact
// mechanism NewAccessRequest invokes at the token endpoint (client_secret_basic
// / client_secret_post / private_key_jwt / SPIFFE, per how this server's
// AuthorizationServerConfig configured it), rather than a bare GetClient
// lookup that would let anyone who learns a confidential client's client_id
// mint device_code/user_code pairs on that client's behalf without its secret.
func (h *Handler) DeviceAuthorizationHandler(w http.ResponseWriter, req *http.Request) {
        ctx := req.Context()

        req.Body = http.MaxBytesReader(w, req.Body, MaxDCRBodySize)
        if err := req.ParseForm(); err != nil {
                writeOAuthError(w, http.StatusBadRequest, "invalid_request", "failed to parse request body")
                return
        }

        clientID := req.PostForm.Get("client_id")
        if clientID == "" {
                writeOAuthError(w, http.StatusBadRequest, "invalid_request", "The client_id parameter is required.")
                return
        }

        authenticator, ok := h.provider.(clientAuthenticator)
        if !ok {
                // Should never happen: every provider this codebase constructs is
                // *fosite.Fosite. A test double substituting h.provider without this
                // method would hit this path.
                slog.Error("device authorization: provider does not support client authentication", "client_id", clientID)
                writeOAuthError(w, http.StatusInternalServerError, "server_error", "failed to authenticate client")
                return
        }
        client, err := authenticator.AuthenticateClient(ctx, req, req.PostForm)
        if err != nil {
                if errors.Is(err, storage.ErrCIMDRateLimited) {
                        writeCIMDRateLimitedOAuth(w)
                        return
                }
                rfcErr := fosite.ErrorToRFC6749Error(err)
                slog.Debug("device authorization: client authentication failed", "client_id", clientID, "error", err)
                writeOAuthError(w, rfcErr.CodeField, rfcErr.ErrorField, rfcErr.GetDescription())
                return
        }
        if !client.GetGrantTypes().Has(oauthproto.GrantTypeDeviceCode) {
                writeOAuthError(w, http.StatusBadRequest, "unauthorized_client",
                        "The OAuth 2.0 Client is not allowed to use authorization grant 'urn:ietf:params:oauth:grant-type:device_code'.")
                return
        }

        requestedScopes := strings.Fields(req.PostForm.Get("scope"))
        scopes, _, dcrErr := registration.ValidateScopes(requestedScopes, h.config.ScopesSupported)
        if dcrErr != nil {
                writeOAuthError(w, http.StatusBadRequest, "invalid_scope", dcrErr.ErrorDescription)
                return
        }

        device, err := h.storeNewDeviceRequest(ctx, client.GetID(), scopes, client.GetAudience())
        if err != nil {
                slog.Error("device authorization: failed to store device request", "client_id", clientID, "error", err)
                writeOAuthError(w, http.StatusInternalServerError, "server_error", "failed to create device authorization request")
                return
        }

        // The verification page is a browser-facing endpoint like /oauth/authorize,
        // so it is advertised from the same base URL. The browser-binding cookie the
        // page sets is host-only and must reach /oauth/callback, which lives on the
        // browser-facing host; advertising the page from the issuer would strand the
        // cookie whenever the two hosts differ.
        verificationURI := h.config.GetAuthorizationEndpointBaseURL() + "/oauth/device"
        response := deviceAuthorizationResponse{
                DeviceCode:              device.DeviceCode,
                UserCode:                device.UserCode,
                VerificationURI:         verificationURI,
                VerificationURIComplete: verificationURI + "?user_code=" + url.QueryEscape(device.UserCode),
                ExpiresIn:               int(storage.DefaultDeviceRequestTTL.Seconds()),
                Interval:                int(h.deviceCodeInterval.Seconds()),
        }

        slog.Debug("issued device authorization request", "client_id", clientID)

        w.Header().Set("Content-Type", "application/json")
        w.Header().Set("Cache-Control", "no-store")
        w.Header().Set("Pragma", "no-cache")
        w.WriteHeader(http.StatusOK)
        if err := json.NewEncoder(w).Encode(response); err != nil {
                slog.Error("failed to encode device authorization response", "error", err)
        }
}

// storeNewDeviceRequest generates a fresh device_code/user_code pair and
// persists it as pending, retrying on a colliding code up to
// maxDeviceCodeGenerationAttempts times.
func (h *Handler) storeNewDeviceRequest(
        ctx context.Context, clientID string, scopes, audience []string,
) (*storage.DeviceRequest, error) {
        var lastErr error
        for attempt := 0; attempt < maxDeviceCodeGenerationAttempts; attempt++ {
                deviceCode, err := generateDeviceCode()
                if err != nil {
                        return nil, fmt.Errorf("generate device_code: %w", err)
                }
                userCode, err := generateUserCode()
                if err != nil {
                        return nil, fmt.Errorf("generate user_code: %w", err)
                }
                device := &storage.DeviceRequest{
                        DeviceCode: deviceCode,
                        UserCode:   userCode,
                        ClientID:   clientID,
                        Scopes:     scopes,
                        Audience:   audience,
                        Status:     storage.DeviceRequestStatusPending,
                        Interval:   h.deviceCodeInterval,
                        CreatedAt:  time.Now(),
                }
                err = h.deviceStorage.StoreDeviceRequest(ctx, device)
                if err == nil {
                        return device, nil
                }
                if !errors.Is(err, storage.ErrAlreadyExists) {
                        return nil, err
                }
                lastErr = err
        }
        return nil, fmt.Errorf("exhausted %d attempts generating a unique device/user code: %w",
                maxDeviceCodeGenerationAttempts, lastErr)
}

// generateDeviceCode returns a high-entropy, opaque device_code: 32 bytes of
// crypto/rand, base64url-encoded without padding.
func generateDeviceCode() (string, error) {
        buf := make([]byte, 32)
        if _, err := rand.Read(buf); err != nil {
                return "", err
        }
        return base64.RawURLEncoding.EncodeToString(buf), nil
}

// generateUserCode returns a short, human-typeable code in RFC 8628's
// example "XXXX-XXXX" format, drawn from userCodeCharset via crypto/rand so
// it carries enough entropy to resist online guessing within the device
// request's TTL.
//
// Bytes at or above maxUsableRandByte are rejected and redrawn rather than
// reduced with '%', so every retained byte maps onto the charset with equal
// probability. Without rejection, 256 % len(userCodeCharset) leftover values
// would wrap around to the charset's low indices, making those characters
// marginally more likely than the rest.
func generateUserCode() (string, error) {
        const length = userCodeGroupLength * 2
        charsetLen := byte(len(userCodeCharset))
        //nolint:gosec // G115: charsetLen is len(userCodeCharset) == 30 at compile time,
        // so 256/30*30 == 240, well within byte range -- not a runtime-controlled value.
        maxUsableRandByte := byte(256/int(charsetLen)) * charsetLen

        idx := make([]byte, length)
        buf := make([]byte, 1)
        for i := 0; i < length; {
                if _, err := rand.Read(buf); err != nil {
                        return "", err
                }
                if buf[0] >= maxUsableRandByte {
                        continue
                }
                idx[i] = userCodeCharset[buf[0]%charsetLen]
                i++
        }
        return string(idx[:userCodeGroupLength]) + "-" + string(idx[userCodeGroupLength:]), nil
}

// writeOAuthError writes a bare RFC 6749-shaped JSON error body. Unlike
// writeDCRError (registration.DCRError, RFC 7591 Section 3.2.2's dedicated
// shape), it is for responses written without a fosite.AccessRequester to hand
// to h.provider.WriteAccessError/WriteAuthorizeError (the device authorization
// endpoint, and the token endpoint's CIMD rate-limit response through
// writeCIMDRateLimitedOAuth), so it mirrors their JSON shape directly.
func writeOAuthError(w http.ResponseWriter, statusCode int, errorCode, description string) {
        w.Header().Set("Content-Type", "application/json")
        w.Header().Set("Cache-Control", "no-store")
        w.Header().Set("Pragma", "no-cache")
        w.WriteHeader(statusCode)
        body := struct {
                Error            string `json:"error"`
                ErrorDescription string `json:"error_description,omitempty"`
        }{Error: errorCode, ErrorDescription: description}
        if err := json.NewEncoder(w).Encode(body); err != nil {
                slog.Debug("failed to encode OAuth error response", "error", err)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "crypto/rand"
        "errors"
        "log/slog"
        "net/http"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
)

// normalizeUserCode uppercases and strips whitespace/dashes from a
// user-submitted code, then re-inserts the canonical "XXXX-XXXX" grouping
// (see userCodeGroupLength, device_authorization.go) when the result is the
// expected length. A malformed submission is returned as-is and simply fails
// the LoadDeviceRequestByUserCode lookup, surfacing as "invalid or expired".
func normalizeUserCode(raw string) string {
        s := strings.ToUpper(strings.TrimSpace(raw))
        s = strings.Map(func(r rune) rune {
                if r == '-' || r == ' ' {
                        return -1
                }
                return r
        }, s)
        if len(s) == userCodeGroupLength*2 {
                return s[:userCodeGroupLength] + "-" + s[userCodeGroupLength:]
        }
        return s
}

// DeviceVerificationHandler handles GET /oauth/device requests: it renders
// the form the human uses to enter/confirm a device flow user_code, issuing
// the anti-forgery token the POST requires (see browser_binding.go). See
// DeviceVerificationSubmitHandler for what happens next.
func (h *Handler) DeviceVerificationHandler(w http.ResponseWriter, req *http.Request) {
        userCode := normalizeUserCode(req.URL.Query().Get("user_code"))
        h.renderVerifyForm(w, req, http.StatusOK, userCode, "")
}

// DeviceVerificationSubmitHandler handles POST /oauth/device: it validates
// the submitted user_code against a pending DeviceRequest and, on success,
// redirects to the first configured upstream IDP to authenticate the human --
// mirroring AuthorizeHandler's redirect-to-upstream step (authorize.go), but
// without any of the OAuth-client-specific machinery (no fosite
// AuthorizeRequester, no client redirect_uri): the device flow's
// verification page has no client_id/redirect_uri of its own.
func (h *Handler) DeviceVerificationSubmitHandler(w http.ResponseWriter, req *http.Request) {
        ctx := req.Context()

        req.Body = http.MaxBytesReader(w, req.Body, MaxDCRBodySize)
        if err := req.ParseForm(); err != nil {
                h.renderVerifyForm(w, req, http.StatusBadRequest, "", "Could not read the submitted form.")
                return
        }

        userCode := normalizeUserCode(req.PostForm.Get("user_code"))

        // A submission must come from a form this server rendered to this browser:
        // otherwise a cross-site page could submit an attacker's user_code from the
        // victim's browser and have that browser bound to the attacker's login.
        if err := h.verifyDeviceFormToken(req); err != nil {
                slog.Warn("device verification: rejected form submission", "reason", err.Error())
                h.renderVerifyForm(w, req, http.StatusBadRequest, userCode,
                        "This form has expired. Please enter the code again.")
                return
        }

        if userCode == "" {
                h.renderVerifyForm(w, req, http.StatusBadRequest, "", "Please enter the code shown on your device.")
                return
        }

        device, err := h.deviceStorage.LoadDeviceRequestByUserCode(ctx, userCode)
        if err != nil && !isNotFoundOrExpired(err) {
                slog.Error("device verification: failed to look up user_code", "error", err)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }
        if err != nil || device.Status != storage.DeviceRequestStatusPending {
                slog.Debug("device verification: invalid or expired user_code", "error", err)
                h.renderVerifyForm(w, req, http.StatusBadRequest, userCode,
                        "That code is invalid or has expired. Please check your device and try again.")
                return
        }

        if h.rejectUnresolvableDeviceClient(ctx, w, device.ClientID) {
                return
        }

        // Device flow has no client_id/redirect_uri of its own to route on (see
        // this function's doc comment), so it cannot pick the upstream
        // appropriate for device.ClientID the way AuthorizeHandler's per-client
        // IDP routing does. Rather than silently defaulting to h.upstreams[0] --
        // which would authenticate against whichever upstream happens to be
        // configured first in a multi-upstream deployment, exactly the kind of
        // IDP mix-up these servers otherwise defend against -- require exactly
        // one configured upstream so that limitation is a loud startup-time-visible
        // error, not a silent routing decision made per request.
        if len(h.upstreams) != 1 {
                slog.Error("device verification: device flow requires exactly one configured upstream",
                        "upstream_count", len(h.upstreams))
                renderError(w, http.StatusInternalServerError, "This server is not configured for device-flow sign-in.")
                return
        }

        secrets := newUpstreamAuthSecrets()
        binding := newBrowserBinding()
        pending := &storage.PendingDeviceLogin{
                DeviceCode:           device.DeviceCode,
                UserCode:             device.UserCode,
                UpstreamPKCEVerifier: secrets.PKCEVerifier,
                UpstreamNonce:        secrets.Nonce,
                BrowserBindingHash:   binding.hash,
                UpstreamProviderName: h.upstreams[0].Name,
                CreatedAt:            time.Now(),
        }
        if err := h.storage.StorePendingDeviceLogin(ctx, secrets.State, pending); err != nil {
                slog.Error("device verification: failed to store pending login", "error", err)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }

        var authOpts []upstream.AuthorizationOption
        if secrets.Nonce != "" {
                authOpts = append(authOpts, upstream.WithAdditionalParams(map[string]string{"nonce": secrets.Nonce}))
        }
        upstreamURL, err := h.upstreams[0].Provider.AuthorizationURL(secrets.State, secrets.PKCEChallenge, authOpts...)
        if err != nil {
                slog.Error("device verification: failed to build upstream authorization URL", "error", err)
                _ = h.storage.DeletePendingDeviceLogin(ctx, secrets.State)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }

        // Bind the login to this browser exactly as AuthorizeHandler does: only the
        // browser that submitted the user_code may complete the upstream callback,
        // so a user_code holder cannot hand the upstream URL to someone else and
        // have that person's identity land on the device. See browser_binding.go.
        h.setBrowserBindingCookie(w, secrets.State, binding.value)
        http.Redirect(w, req, upstreamURL, http.StatusFound)
}

// tryCompleteDeviceLogin attempts to treat internalState as a device-flow
// login rather than an OAuth-client authorization: it is CallbackHandler's
// fallback when internalState does not match a PendingAuthorization, kept as
// a separate function so CallbackHandler's own branching stays simple. See
// completeDeviceLogin's doc comment for why /oauth/callback is shared between
// the two flows.
//
// Returns deviceLoginNotFound (no response written) when internalState
// genuinely does not match a pending device login either, so the caller
// falls through to its own not-found handling. Returns deviceLoginStorageError
// (response already written) when the lookup failed for a reason other than
// not-found/expired -- a real backend error must not be silently
// reinterpreted as "not a device login" and reported to the client as an
// ordinary 400. Returns deviceLoginRejected (response already written) when
// req does not carry the browser-binding cookie the verification page set
// for internalState; the login is consumed so the state cannot be retried.
func (h *Handler) tryCompleteDeviceLogin(
        ctx context.Context, w http.ResponseWriter, req *http.Request, internalState, code string,
) deviceLoginOutcome {
        devicePending, err := h.storage.LoadPendingDeviceLogin(ctx, internalState)
        if err != nil {
                if isNotFoundOrExpired(err) {
                        return deviceLoginNotFound
                }
                slog.Error("failed to load pending device login", "error", err)
                http.Error(w, "internal server error", http.StatusInternalServerError)
                return deviceLoginStorageError
        }
        if delErr := h.storage.DeletePendingDeviceLogin(ctx, internalState); delErr != nil {
                slog.Warn("failed to delete pending device login", "error", delErr)
        }
        if bindErr := h.verifyBrowserBinding(req, internalState, devicePending.BrowserBindingHash); bindErr != nil {
                writeUnboundCallbackRejection(w, bindErr, "flow", "device", "upstream_provider", devicePending.UpstreamProviderName)
                return deviceLoginRejected
        }
        h.completeDeviceLogin(ctx, w, devicePending, code)
        return deviceLoginCompleted
}

// tryDenyDeviceLoginOnUpstreamError handles an upstream IDP error (the
// "error" query parameter on /oauth/callback, e.g. the human clicked "Deny"
// at the upstream login screen, or the IDP itself failed) for a device-flow
// login: handleUpstreamError's fallback when internalState does not match a
// PendingAuthorization either. Without this, an upstream error for a
// device-flow login left the DeviceRequest Pending forever (until the
// device_code itself expired) instead of surfacing the denial immediately,
// since nothing else ever calls MarkDeviceRequestDenied for this path.
//
// Returns false (no response written) when internalState does not match a
// pending device login either, so the caller falls through to its own
// generic error page. A browser that does not present the login's binding
// cookie gets a 400 and the login is consumed, but the DeviceRequest is left
// pending: that browser did not start the login and may not decide it.
func (h *Handler) tryDenyDeviceLoginOnUpstreamError(
        ctx context.Context, w http.ResponseWriter, req *http.Request, internalState string,
) bool {
        devicePending, err := h.storage.LoadPendingDeviceLogin(ctx, internalState)
        if err != nil {
                return false
        }
        if delErr := h.storage.DeletePendingDeviceLogin(ctx, internalState); delErr != nil {
                slog.Warn("failed to delete pending device login", "error", delErr)
        }
        if bindErr := h.verifyBrowserBinding(req, internalState, devicePending.BrowserBindingHash); bindErr != nil {
                writeUnboundCallbackRejection(w, bindErr, "flow", "device", "upstream_provider", devicePending.UpstreamProviderName)
                return true
        }
        if err := h.deviceStorage.MarkDeviceRequestDenied(ctx, devicePending.DeviceCode); err != nil {
                slog.Warn("device verification: failed to mark device request denied after upstream error", "error", err)
        }
        renderDenied(w, "Sign-in failed",
                "The identity provider reported an error, so this device's request was denied. You may close this window.")
        return true
}

// completeDeviceLogin finishes a device-flow verification-page login once
// the upstream IDP has redirected back to /oauth/callback (CallbackHandler
// dispatches here after determining internalState matches a
// PendingDeviceLogin rather than a PendingAuthorization -- see that
// function's doc comment for why both flows share one callback endpoint).
// It resolves the canonical identity (mirroring CallbackHandler's first-leg
// resolution) and stores it server-side behind a fresh opaque token, then
// renders the Approve/Deny confirmation page. See PendingDeviceConfirmation's
// doc comment for why the resolved identity is never sent to the browser as
// editable form data.
func (h *Handler) completeDeviceLogin(
        ctx context.Context, w http.ResponseWriter, pending *storage.PendingDeviceLogin, code string,
) {
        upstreamProvider, ok := h.upstreamByName(pending.UpstreamProviderName)
        if !ok {
                slog.Error("device verification: upstream provider not found", "provider", pending.UpstreamProviderName)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }

        // Time has passed during login: re-check the device request is still
        // pending/unexpired, and its client still resolves, before spending the
        // upstream code.
        device, err := h.deviceStorage.LoadDeviceRequestByUserCode(ctx, pending.UserCode)
        if err != nil && !isNotFoundOrExpired(err) {
                slog.Error("device verification: failed to re-check device request", "error", err)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }
        if err != nil || device.Status != storage.DeviceRequestStatusPending {
                slog.Debug("device verification: device request no longer pending", "error", err)
                renderError(w, http.StatusBadRequest,
                        "This device request is no longer valid. Please try again from your device.")
                return
        }

        if h.rejectUnresolvableDeviceClient(ctx, w, device.ClientID) {
                return
        }

        result, err := upstreamProvider.ExchangeCodeForIdentity(ctx, code, pending.UpstreamPKCEVerifier, pending.UpstreamNonce)
        if err != nil {
                slog.Error("device verification: failed to exchange code or resolve identity", "error", err)
                renderError(w, http.StatusInternalServerError, "Sign-in failed. Please try again.")
                return
        }

        userID, userName, userEmail := result.Subject, result.Name, result.Email
        if !result.Synthetic {
                user, resolveErr := h.userResolver.ResolveUser(ctx, pending.UpstreamProviderName, result.Subject)
                if resolveErr != nil {
                        slog.Error("device verification: failed to resolve user", "error", resolveErr)
                        status, msg := http.StatusInternalServerError, "Something went wrong. Please try again."
                        if errors.Is(resolveErr, storage.ErrUserNotProvisioned) {
                                status, msg = http.StatusForbidden, "Your account is not provisioned to use this application."
                        }
                        renderError(w, status, msg)
                        return
                }
                userID = user.ID
                h.userResolver.UpdateLastAuthenticated(ctx, pending.UpstreamProviderName, result.Subject)
        }

        token := rand.Text()
        verified := h.recordedUpstreamIdentity(result)
        confirmation := &storage.PendingDeviceConfirmation{
                DeviceCode:        device.DeviceCode,
                UserCode:          device.UserCode,
                ResolvedUserID:    userID,
                ResolvedUserName:  userName,
                ResolvedUserEmail: userEmail,
                VerifiedUpstream:  verified,
                // UserID/SessionExpiresAt are left zero here and filled in by
                // DeviceVerificationConfirmHandler once a session id exists -- see
                // PendingDeviceConfirmation.UpstreamTokens's doc comment.
                UpstreamTokens: &storage.UpstreamTokens{
                        ProviderID:      pending.UpstreamProviderName,
                        AccessToken:     result.Tokens.AccessToken,
                        RefreshToken:    result.Tokens.RefreshToken,
                        IDToken:         result.Tokens.IDToken,
                        ExpiresAt:       result.Tokens.ExpiresAt,
                        UpstreamSubject: result.Subject,
                        ClientID:        device.ClientID,
                },
                Synthetic: result.Synthetic,
                CreatedAt: time.Now(),
        }
        if err := h.storage.StorePendingDeviceConfirmation(ctx, token, confirmation); err != nil {
                slog.Error("device verification: failed to store pending confirmation", "error", err)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }

        renderConfirm(w, confirmPageData{
                Name:         userName,
                Email:        userEmail,
                ClientID:     device.ClientID,
                Scopes:       device.Scopes,
                ConfirmToken: token,
        })
}

// DeviceVerificationConfirmHandler handles POST /oauth/device/confirm: the
// human's explicit Approve/Deny decision from the confirmation page rendered
// by DeviceVerificationCallbackHandler. It resolves confirm_token back to the
// server-side PendingDeviceConfirmation (single-use) and calls
// storage.MarkDeviceRequestAuthorized or MarkDeviceRequestDenied accordingly
// -- the same calls the device-flow integration test drives directly against
// storage to simulate this page.
func (h *Handler) DeviceVerificationConfirmHandler(w http.ResponseWriter, req *http.Request) {
        ctx := req.Context()

        req.Body = http.MaxBytesReader(w, req.Body, MaxDCRBodySize)
        if err := req.ParseForm(); err != nil {
                renderError(w, http.StatusBadRequest, "Could not read the submitted form.")
                return
        }

        token := req.PostForm.Get("confirm_token")
        action := req.PostForm.Get("action")
        if token == "" || (action != "approve" && action != "deny") {
                renderError(w, http.StatusBadRequest, "Invalid confirmation request.")
                return
        }

        confirmation, err := h.storage.LoadPendingDeviceConfirmation(ctx, token)
        if err != nil {
                slog.Warn("device verification: pending confirmation not found", "error", err)
                renderError(w, http.StatusBadRequest, "This confirmation request was not found or has expired.")
                return
        }
        // Delete immediately (single-use): a confirm_token authorizes exactly one decision.
        if err := h.storage.DeletePendingDeviceConfirmation(ctx, token); err != nil {
                slog.Warn("device verification: failed to delete pending confirmation", "error", err)
        }

        if action == "deny" {
                if err := h.deviceStorage.MarkDeviceRequestDenied(ctx, confirmation.DeviceCode); err != nil {
                        slog.Warn("device verification: failed to mark device request denied", "error", err)
                        renderError(w, http.StatusBadRequest,
                                "This device request is no longer valid. Please try again from your device.")
                        return
                }
                renderDenied(w, "Access denied", "You have denied this device's request. You may close this window.")
                return
        }

        // The confirmation page does not carry the client, but the upstream tokens
        // captured at login do. A record without them is corrupt, not a removed client.
        if confirmation.UpstreamTokens == nil {
                slog.Error("device verification: pending confirmation has no upstream tokens")
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }
        if h.rejectUnresolvableDeviceClient(ctx, w, confirmation.UpstreamTokens.ClientID) {
                return
        }

        sessionID := rand.Text()
        if err := h.persistDeviceSessionTokens(ctx, confirmation, sessionID); err != nil {
                slog.Error("device verification: failed to store upstream tokens", "error", err)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
                return
        }
        if err := h.deviceStorage.MarkDeviceRequestAuthorized(
                ctx, confirmation.DeviceCode,
                storage.DeviceAuthorization{
                        UserID:           confirmation.ResolvedUserID,
                        UserName:         confirmation.ResolvedUserName,
                        UserEmail:        confirmation.ResolvedUserEmail,
                        SessionID:        sessionID,
                        VerifiedUpstream: confirmation.VerifiedUpstream,
                },
        ); err != nil {
                slog.Error("device verification: failed to mark device request authorized", "error", err)
                renderError(w, http.StatusBadRequest,
                        "This device request is no longer valid. Please try again from your device.")
                return
        }

        renderResult(w, "Device authorized", "You may now close this window and return to your device.")
}

// persistDeviceSessionTokens stores the upstream tokens captured at login
// time (completeDeviceLogin) under the device grant's final session id.
// Device-flow sessions have no session id until this point -- unlike the
// OAuth-client authorization_code flow, where CallbackHandler stores tokens
// under a session id minted up front -- so completeDeviceLogin could only
// stash the tokens in the PendingDeviceConfirmation for this handler to
// persist once sessionID exists. Mirrors CallbackHandler's
// UserID/SessionExpiresAt population and refresh-token carry-forward.
func (h *Handler) persistDeviceSessionTokens(
        ctx context.Context, confirmation *storage.PendingDeviceConfirmation, sessionID string,
) error {
        if confirmation.UpstreamTokens == nil {
                return errors.New("pending device confirmation has no upstream tokens")
        }
        storageTokens := *confirmation.UpstreamTokens
        storageTokens.UserID = confirmation.ResolvedUserID
        storageTokens.SessionExpiresAt = time.Now().Add(h.config.RefreshTokenLifespan)

        h.maybeCarryForwardRefreshToken(
                ctx, &storageTokens, confirmation.ResolvedUserID, storageTokens.UpstreamSubject,
                storageTokens.ProviderID, confirmation.Synthetic,
        )

        return h.storage.StoreUpstreamTokens(ctx, sessionID, storageTokens.ProviderID, &storageTokens)
}

// rejectUnresolvableDeviceClient looks up the device request's client in the
// admitted store and, when it no longer resolves, writes the error page and
// returns true: a client removed from configuration must not get an upstream
// login, a stored identity or an authorized device request. A lookup failure
// other than not-found is a 500. Like the callback, it uses callbackClient so a
// live CIMD client behind the fetch limiter still resolves.
func (h *Handler) rejectUnresolvableDeviceClient(ctx context.Context, w http.ResponseWriter, clientID string) bool {
        _, err := h.callbackClient(ctx, clientID)
        switch {
        case err == nil:
                return false
        case isUnresolvedClient(err):
                slog.Debug("device verification: client did not resolve", "client_id", clientID, "error", err)
                renderError(w, http.StatusBadRequest,
                        "This device request is no longer valid. Please try again from your device.")
        default:
                slog.Error("device verification: failed to resolve client", "client_id", clientID, "error", err)
                renderError(w, http.StatusInternalServerError, "Something went wrong. Please try again.")
        }
        return true
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "net"
        "net/http"
        "sync"
        "time"

        "golang.org/x/time/rate"
)

// perIPLimiterEvictionThreshold is how many distinct IPs perIPLimiter tracks
// before it sweeps idle entries on the next call. Kept low deliberately:
// this limiter only ever guards a low-traffic, human-facing login page, so
// the map is never expected to hold more than a handful of concurrent
// callers, and a small threshold keeps the sweep itself cheap.
const perIPLimiterEvictionThreshold = 256

// perIPLimiterIdleTTL is how long an IP's bucket survives with no requests
// before eviction reclaims it.
const perIPLimiterIdleTTL = 10 * time.Minute

// perIPLimiter buckets golang.org/x/time/rate.Limiters by client IP, unlike
// registerLimiter/deviceAuthorizationLimiter/deviceVerificationLimiter's
// per-process bucket (see registerLimiter's doc comment in handler.go for
// why those endpoints are deliberately NOT per-IP). Device flow's
// verification page is different: it's a human-facing login page a
// legitimate user retries against (mistyped codes), so a single
// process-wide bucket lets one caller -- malicious or just persistent --
// burn the shared burst and starve every other concurrent login on the
// server. Per-IP scoping bounds that blast radius to one caller's own IP.
//
// Idle entries are swept inline (no background goroutine to manage the
// lifecycle of) once the map grows past perIPLimiterEvictionThreshold, so
// long-lived processes don't grow this map unboundedly.
type perIPLimiter struct {
        mu       sync.Mutex
        limiters map[string]*ipLimiterEntry
        limit    rate.Limit
        burst    int
}

type ipLimiterEntry struct {
        limiter    *rate.Limiter
        lastSeenAt time.Time
}

// newPerIPLimiter returns a perIPLimiter granting each distinct IP its own
// token bucket at the given rate/burst.
func newPerIPLimiter(limit rate.Limit, burst int) *perIPLimiter {
        return &perIPLimiter{
                limiters: make(map[string]*ipLimiterEntry),
                limit:    limit,
                burst:    burst,
        }
}

// allow reports whether a request from ip may proceed, consuming one token
// from that IP's bucket if so.
func (l *perIPLimiter) allow(ip string) bool {
        l.mu.Lock()
        defer l.mu.Unlock()

        now := time.Now()
        entry, ok := l.limiters[ip]
        if !ok {
                entry = &ipLimiterEntry{limiter: rate.NewLimiter(l.limit, l.burst)}
                l.limiters[ip] = entry
        }
        entry.lastSeenAt = now

        if len(l.limiters) > perIPLimiterEvictionThreshold {
                l.evictIdleLocked(now)
        }
        return entry.limiter.Allow()
}

// evictIdleLocked removes buckets untouched for longer than
// perIPLimiterIdleTTL. Callers must hold l.mu.
func (l *perIPLimiter) evictIdleLocked(now time.Time) {
        for ip, entry := range l.limiters {
                if now.Sub(entry.lastSeenAt) > perIPLimiterIdleTTL {
                        delete(l.limiters, ip)
                }
        }
}

// clientIP extracts the request's client IP from RemoteAddr, deliberately
// ignoring X-Forwarded-For and similar headers: those are attacker-controlled
// input without a trusted, correctly-configured reverse proxy in front of
// this server (see registerLimiter's doc comment in handler.go for the same
// reasoning), and trusting them here would let a client forge a distinct IP
// per request and defeat the very throttling this limiter exists to provide.
func clientIP(req *http.Request) string {
        host, _, err := net.SplitHostPort(req.RemoteAddr)
        if err != nil {
                return req.RemoteAddr
        }
        return host
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "html/template"
        "log/slog"
        "net/http"
        "strings"
)

// This is the first HTML surface in pkg/authserver -- every other handler
// writes JSON. html/template (not manual string concatenation) is used
// throughout for auto-escaping, since these pages echo back a user-submitted
// user_code and an upstream-supplied name/email.
//
// Styling follows ToolHive/Stacklok's brand palette (deep emerald green
// primary action, Inter-first font stack, pill buttons, light/dark via
// prefers-color-scheme) so this page doesn't look like an unstyled internal
// tool when a human lands on it mid device-flow login. No external
// stylesheets or fonts are loaded: ToolHive commonly runs in air-gapped/
// enterprise environments, and a login-adjacent page reaching out to a
// third-party CDN is both an availability risk and a needless referrer leak.

// pageStyles is the shared stylesheet for all three device-flow pages,
// tokenized on ToolHive/Stacklok's palette (see toolhive-cloud-ui's
// globals.css: --brand hsl(161 94% 21%), --btn-primary-hover #02543a,
// --nav-background #18442e). Duplicated as a single constant across pages
// rather than a separate stylesheet request -- these are simple, low-traffic
// pages and a second HTTP round trip buys nothing.
const pageStyles = `
:root{
  --bg:#f4f4f5; --card:#ffffff; --fg:#0a0a0b; --muted:#71717a;
  --border:#e4e4e7; --brand:#026b46; --brand-hover:#02543a; --brand-fg:#ffffff;
  --nav-bg:#18442e; --danger:#dc2626; --danger-bg:#fef2f2; --danger-border:#fecaca;
  --radius:0.75rem;
}
@media (prefers-color-scheme: dark){
  :root{
    --bg:#0f0f11; --card:#18181b; --fg:#fafafa; --muted:#a1a1aa;
    --border:#27272a; --brand:#1a8f61; --brand-hover:#22a873; --brand-fg:#ffffff;
    --nav-bg:#0f2e1e; --danger:#f87171; --danger-bg:#2a1414; --danger-border:#4c1d1d;
  }
}
*{box-sizing:border-box}
body{
  font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Inter,Roboto,Helvetica,Arial,sans-serif;
  background:var(--bg); color:var(--fg); margin:0; padding:2.5rem 1rem;
  display:flex; justify-content:center; min-height:100vh;
}
.card{
  width:100%; max-width:26rem; height:fit-content;
  background:var(--card); border:1px solid var(--border); border-radius:var(--radius);
  box-shadow:0 1px 3px rgba(0,0,0,0.08), 0 1px 2px rgba(0,0,0,0.04);
}
.card-body{padding:2rem 1.75rem}
h1{font-size:1.25rem; font-weight:600; margin:0 0 0.4rem; letter-spacing:-0.01em}
.subtitle{color:var(--muted); font-size:0.9rem; margin:0 0 1.5rem; line-height:1.4}
.error{
  background:var(--danger-bg); border:1px solid var(--danger-border); color:var(--danger);
  border-radius:0.5rem; padding:0.65rem 0.85rem; font-size:0.875rem; margin-bottom:1.25rem;
}
input[type=text]{
  font-size:1.5rem; font-weight:600; letter-spacing:0.15em; text-align:center;
  text-transform:uppercase; width:100%; padding:0.75rem 0.5rem; margin:0 0 1.5rem;
  border:1.5px solid var(--border); border-radius:0.5rem; background:var(--card); color:var(--fg);
  font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;
}
input[type=text]:focus{
  outline:none; border-color:var(--brand);
  box-shadow:0 0 0 3px color-mix(in srgb, var(--brand) 20%, transparent);
}
input[type=text]::placeholder{color:var(--muted); font-weight:400; letter-spacing:0.1em}
button{
  font-size:0.9rem; font-weight:600; padding:0.7rem 1.4rem; border-radius:999px;
  border:none; cursor:pointer; width:100%; transition:background-color 0.15s ease;
  font-family:inherit;
}
button:focus-visible{outline:2px solid var(--brand); outline-offset:2px}
.btn-primary{background:var(--brand); color:var(--brand-fg)}
.btn-primary:hover{background:var(--brand-hover)}
.btn-secondary{background:transparent; color:var(--fg); border:1.5px solid var(--border)}
.btn-secondary:hover{background:color-mix(in srgb, var(--fg) 6%, transparent)}
.btn-row{display:flex; gap:0.75rem; margin-top:1.5rem}
dl{margin:0 0 0.25rem; font-size:0.9rem}
dt{color:var(--muted); font-weight:500; margin-top:0.9rem}
dt:first-child{margin-top:0}
dd{margin:0.2rem 0 0; font-weight:500}
.scope-chip{
  display:inline-block; background:color-mix(in srgb, var(--brand) 12%, transparent);
  color:var(--brand); border-radius:999px; padding:0.15rem 0.65rem; font-size:0.8rem;
  font-weight:500; margin:0.2rem 0.35rem 0.2rem 0;
}
.identity{
  display:flex; align-items:center; gap:0.75rem; padding:0.85rem 1rem; margin-bottom:1.25rem;
  background:color-mix(in srgb, var(--fg) 4%, transparent); border-radius:0.6rem;
}
.avatar{
  width:2.25rem; height:2.25rem; border-radius:999px; background:var(--brand); color:var(--brand-fg);
  display:flex; align-items:center; justify-content:center; font-weight:600; font-size:0.95rem;
  flex-shrink:0;
}
.identity-text{min-width:0}
.identity-name{font-weight:600; font-size:0.9rem; overflow:hidden; text-overflow:ellipsis; white-space:nowrap}
.identity-email{color:var(--muted); font-size:0.8rem; overflow:hidden; text-overflow:ellipsis; white-space:nowrap}
.result-icon{
  width:3rem; height:3rem; border-radius:999px; display:flex; align-items:center; justify-content:center;
  margin:0 auto 1.25rem; font-size:1.5rem;
}
.result-icon.ok{background:color-mix(in srgb, var(--brand) 15%, transparent); color:var(--brand)}
.result-icon.denied{background:color-mix(in srgb, var(--muted) 20%, transparent); color:var(--muted)}
.result-icon.err{background:var(--danger-bg); color:var(--danger)}
.result-body{text-align:center}
`

type verifyPageData struct {
        UserCode string
        Error    string
        // FormToken is the double-submit anti-forgery token that must echo the
        // cookie issueDeviceFormToken set on this response.
        FormToken string
}

var verifyPageTemplate = template.Must(template.New("device-verify").Parse(`<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Device Login &middot; ToolHive</title>
<style>` + pageStyles + `</style>
</head>
<body>
<div class="card">
  <div class="card-body">
    <h1>Enter your device code</h1>
    <p class="subtitle">Enter the code shown on your device to continue signing in.</p>
    {{if .Error}}<div class="error">{{.Error}}</div>{{end}}
    <form method="POST" action="/oauth/device">
      <input type="hidden" name="form_token" value="{{.FormToken}}">
      <input type="text" name="user_code" value="{{.UserCode}}" placeholder="XXXX-XXXX"
        autocomplete="off" autofocus required>
      <button type="submit" class="btn-primary">Continue</button>
    </form>
  </div>
</div>
</body>
</html>`))

type confirmPageData struct {
        Name         string
        Email        string
        ClientID     string
        Scopes       []string
        ConfirmToken string
        Initial      string
}

var confirmPageTemplate = template.Must(template.New("device-confirm").Parse(`<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Authorize Device &middot; ToolHive</title>
<style>` + pageStyles + `</style>
</head>
<body>
<div class="card">
  <div class="card-body">
    <h1>Authorize this device?</h1>
    <p class="subtitle">A device is requesting access to your account. Review the details below before continuing.</p>
    <div class="identity">
      <div class="avatar">{{.Initial}}</div>
      <div class="identity-text">
        <div class="identity-name">{{if .Name}}{{.Name}}{{else}}{{.Email}}{{end}}</div>
        {{if and .Name .Email}}<div class="identity-email">{{.Email}}</div>{{end}}
      </div>
    </div>
    <dl>
      <dt>Application</dt><dd>{{.ClientID}}</dd>
      {{if .Scopes}}
      <dt>Requested access</dt>
      <dd>{{range .Scopes}}<span class="scope-chip">{{.}}</span>{{end}}</dd>
      {{end}}
    </dl>
    <form method="POST" action="/oauth/device/confirm">
      <input type="hidden" name="confirm_token" value="{{.ConfirmToken}}">
      <div class="btn-row">
        <button class="btn-secondary" type="submit" name="action" value="deny">Deny</button>
        <button class="btn-primary" type="submit" name="action" value="approve">Approve</button>
      </div>
    </form>
  </div>
</div>
</body>
</html>`))

type resultPageData struct {
        Title   string
        Message string
        Icon    string // "ok", "denied", or "err"
        Glyph   string // a single character/emoji rendered inside .result-icon
}

var resultPageTemplate = template.Must(template.New("device-result").Parse(`<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>{{.Title}} &middot; ToolHive</title>
<style>` + pageStyles + `</style>
</head>
<body>
<div class="card">
  <div class="card-body result-body">
    <div class="result-icon {{.Icon}}">{{.Glyph}}</div>
    <h1>{{.Title}}</h1>
    <p class="subtitle">{{.Message}}</p>
  </div>
</div>
</body>
</html>`))

// setHTMLSecurityHeaders sets the headers common to the interactive auth pages:
// HTML content type, no caching, and anti-framing protection.
func setHTMLSecurityHeaders(w http.ResponseWriter) {
        w.Header().Set("Content-Type", "text/html; charset=utf-8")
        w.Header().Set("Cache-Control", "no-store")
        w.Header().Set("X-Frame-Options", "DENY")
        w.Header().Set("Content-Security-Policy", "frame-ancestors 'none'")
        w.Header().Set("X-Content-Type-Options", "nosniff")
}

// initials returns up to two uppercase initials for a display name/email, for
// the confirm page's avatar badge. Falls back to "?" when neither is usable.
func initials(name, email string) string {
        source := name
        if source == "" {
                source = email
        }
        fields := strings.FieldsFunc(source, func(r rune) bool {
                return r == ' ' || r == '.' || r == '@' || r == '_' || r == '-'
        })
        if len(fields) == 0 {
                return "?"
        }
        out := firstRuneUpper(fields[0])
        if len(fields) > 1 && fields[1] != "" {
                out += firstRuneUpper(fields[1])
        }
        return out
}

// firstRuneUpper returns the uppercased first Unicode rune of s. Slicing the
// first byte instead (s[:1]) would split a multi-byte rune and render
// invalid UTF-8/replacement characters for non-ASCII names.
func firstRuneUpper(s string) string {
        r := []rune(s)
        return strings.ToUpper(string(r[:1]))
}

// renderVerifyForm renders the device verification form with its anti-forgery
// token, so a re-rendered form (after an error) is submittable.
func (h *Handler) renderVerifyForm(w http.ResponseWriter, req *http.Request, status int, userCode, errMsg string) {
        data := verifyPageData{UserCode: userCode, Error: errMsg, FormToken: h.issueDeviceFormToken(w, req)}
        setHTMLSecurityHeaders(w)
        w.WriteHeader(status)
        if err := verifyPageTemplate.Execute(w, data); err != nil {
                slog.Error("device verification: failed to render verify form", "error", err)
        }
}

func renderConfirm(w http.ResponseWriter, data confirmPageData) {
        setHTMLSecurityHeaders(w)
        data.Initial = initials(data.Name, data.Email)
        if err := confirmPageTemplate.Execute(w, data); err != nil {
                slog.Error("device verification: failed to render confirm page", "error", err)
        }
}

func renderResult(w http.ResponseWriter, title, message string) {
        setHTMLSecurityHeaders(w)
        data := resultPageData{Title: title, Message: message, Icon: "ok", Glyph: "✓"}
        if err := resultPageTemplate.Execute(w, data); err != nil {
                slog.Error("device verification: failed to render result page", "error", err)
        }
}

// renderDenied renders the result page in its "denied" visual state --
// distinct from renderResult's success checkmark, since a device denial is a
// deliberate, expected outcome rather than a completed grant.
func renderDenied(w http.ResponseWriter, title, message string) {
        setHTMLSecurityHeaders(w)
        data := resultPageData{Title: title, Message: message, Icon: "denied", Glyph: "✕"}
        if err := resultPageTemplate.Execute(w, data); err != nil {
                slog.Error("device verification: failed to render result page", "error", err)
        }
}

func renderError(w http.ResponseWriter, status int, message string) {
        setHTMLSecurityHeaders(w)
        w.WriteHeader(status)
        data := resultPageData{Title: "Something went wrong", Message: message, Icon: "err", Glyph: "!"}
        if err := resultPageTemplate.Execute(w, data); err != nil {
                slog.Error("device verification: failed to render error page", "error", err)
        }
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package handlers

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "net/http"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        sharedobauth "github.com/stacklok/toolhive/pkg/oauthproto"
)

// Cache-Control max-age values for discovery endpoints.
// These are not exposed to users but extracted as constants for documentation and maintainability.
const (
        // DefaultJWKSCacheMaxAge is the Cache-Control max-age for the JWKS endpoint (1 hour).
        // This balances caching efficiency with timely key rotation propagation.
        DefaultJWKSCacheMaxAge = 3600

        // DefaultDiscoveryCacheMaxAge is the Cache-Control max-age for the discovery endpoint (1 hour).
        // Aligned with Google's OIDC discovery cache policy.
        DefaultDiscoveryCacheMaxAge = 3600
)

// getSigningAlgorithms extracts the signing algorithms from the JWKS keys.
// If no keys are available, it falls back to RS256 per OIDC Core Section 15.1.
func (h *Handler) getSigningAlgorithms() []string {
        publicJWKS := h.config.PublicJWKS()
        if publicJWKS == nil || len(publicJWKS.Keys) == 0 {
                // Fall back to RS256 per OIDC Core Section 15.1 requirement
                return []string{"RS256"}
        }

        // Collect unique algorithms from keys
        seen := make(map[string]bool)
        var algs []string
        for _, key := range publicJWKS.Keys {
                if key.Algorithm != "" && !seen[key.Algorithm] {
                        seen[key.Algorithm] = true
                        algs = append(algs, key.Algorithm)
                }
        }

        if len(algs) == 0 {
                // No algorithms found on keys, fall back to RS256
                return []string{"RS256"}
        }

        return algs
}

// JWKSHandler handles GET /.well-known/jwks.json requests.
// It returns the public keys used for verifying JWTs.
func (h *Handler) JWKSHandler(w http.ResponseWriter, _ *http.Request) {
        publicJWKS := h.config.PublicJWKS()
        if publicJWKS == nil {
                slog.Error("no public JWKS available")
                http.Error(w, "internal server error", http.StatusInternalServerError)
                return
        }

        data, err := json.Marshal(publicJWKS)
        if err != nil {
                slog.Error("failed to encode JWKS",
                        "error", err,
                )
                http.Error(w, "internal server error", http.StatusInternalServerError)
                return
        }

        w.Header().Set("Content-Type", "application/json")
        w.Header().Set("Cache-Control", fmt.Sprintf("public, max-age=%d", DefaultJWKSCacheMaxAge))
        w.Header().Set("X-Content-Type-Options", "nosniff")
        _, _ = w.Write(data) //nolint:gosec // G705: data is JSON-marshaled from internal metadata, not user input
}

// buildOAuthMetadata constructs the authorization-server metadata shared by the
// OAuth AS metadata endpoint and the OIDC discovery endpoint. Token-only
// servers intentionally omit authorization response metadata, including
// response_types_supported, because they do not implement those response types.
// They retain configured scopes because scopes describe permissions available
// through token exchange, not interactive authorization availability. This is
// ToolHive metadata for the supported token flows, not a claim of full RFC 8414
// or OIDC Discovery conformance.
func (h *Handler) buildOAuthMetadata() sharedobauth.AuthorizationServerMetadata {
        issuer := h.config.GetAccessTokenIssuer()

        metadata := sharedobauth.AuthorizationServerMetadata{
                // REQUIRED
                Issuer: issuer,

                // RECOMMENDED
                TokenEndpoint:   issuer + "/oauth/token",
                JWKSURI:         issuer + "/.well-known/jwks.json",
                ScopesSupported: h.config.ScopesSupported,

                // OPTIONAL
                GrantTypesSupported:                        h.grantTypesSupported(),
                TokenEndpointAuthMethodsSupported:          h.tokenEndpointAuthMethodsSupported(),
                TokenEndpointAuthSigningAlgValuesSupported: h.tokenEndpointAuthSigningAlgorithms(),

                // ClientIDMetadataDocumentSupported is defined in the CIMD draft as an
                // OAuth AS metadata field (RFC 8414), not in OIDC Discovery 1.0. It is
                // included here because MCP clients (e.g. VS Code) discover the AS via
                // /.well-known/openid-configuration and need this flag there to activate
                // CIMD. Spec-compliant OIDC consumers silently ignore unknown fields.
                ClientIDMetadataDocumentSupported: h.config.CIMDEnabled,
        }
        if !h.tokenOnly {
                metadata.AuthorizationEndpoint = h.config.GetAuthorizationEndpointBaseURL() + "/oauth/authorize"
                metadata.ResponseTypesSupported = []string{sharedobauth.ResponseTypeCode}
                metadata.CodeChallengeMethodsSupported = []string{crypto.PKCEChallengeMethodS256}
        }
        if !h.tokenOnly || h.config.AllowPrivateKeyJWTRegistration {
                metadata.RegistrationEndpoint = issuer + "/oauth/register"
        }
        if h.config.DeviceFlowEnabled {
                metadata.DeviceAuthorizationEndpoint = issuer + "/oauth/device_authorization"
        }
        return metadata
}

// tokenEndpointAuthMethodsSupported returns the token_endpoint_auth_methods_supported
// list for discovery, derived from config. "none" is always first — the public-client
// default. When confidential DCR is enabled or static delegate clients are configured,
// the two client_secret_* methods are appended. Static clients need these methods even
// though DCR itself remains public-only. RFC 8414 defines no ordering semantics, so
// "none"-first is a readability convention, not a security control.
func (h *Handler) tokenEndpointAuthMethodsSupported() []string {
        methods := []string{sharedobauth.TokenEndpointAuthMethodNone}
        if h.config.AllowConfidentialClientRegistration || h.config.HasStaticDelegateClients {
                methods = append(methods,
                        sharedobauth.TokenEndpointAuthMethodClientSecretBasic,
                        sharedobauth.TokenEndpointAuthMethodClientSecretPost,
                )
        }
        if h.config.AllowPrivateKeyJWTRegistration {
                methods = append(methods, sharedobauth.TokenEndpointAuthMethodPrivateKeyJWT)
        }
        return methods
}

// tokenEndpointAuthSigningAlgorithms returns the signing algorithms advertised
// for private_key_jwt client assertions. The list is shared with DCR validation
// so discovery cannot claim support for an algorithm registration rejects.
func (h *Handler) tokenEndpointAuthSigningAlgorithms() []string {
        if !h.config.AllowPrivateKeyJWTRegistration {
                return nil
        }
        return registration.SupportedSigningAlgorithms()
}

// grantTypesSupported returns only grant families registered with fosite.
func (h *Handler) grantTypesSupported() []string {
        grantTypes := make([]string, 0, 4)
        if h.config.TokenExchangeEnabled {
                grantTypes = append(grantTypes, sharedobauth.GrantTypeTokenExchange)
        }
        if !h.tokenOnly {
                grantTypes = append(grantTypes,
                        string(fosite.GrantTypeAuthorizationCode),
                        string(fosite.GrantTypeRefreshToken),
                )
        }
        if h.config.JWTBearerGrantEnabled {
                grantTypes = append(grantTypes, sharedobauth.GrantTypeJWTBearer)
        }
        if h.config.DeviceFlowEnabled {
                grantTypes = append(grantTypes, sharedobauth.GrantTypeDeviceCode)
        }
        return grantTypes
}

// OAuthDiscoveryHandler handles GET /.well-known/oauth-authorization-server requests.
// It returns the OAuth 2.0 Authorization Server Metadata per RFC 8414.
// This endpoint is useful for non-OIDC OAuth clients.
func (h *Handler) OAuthDiscoveryHandler(w http.ResponseWriter, _ *http.Request) {
        metadata := h.buildOAuthMetadata()

        data, err := json.Marshal(metadata)
        if err != nil {
                slog.Error("failed to encode OAuth AS metadata",
                        "error", err,
                )
                http.Error(w, "internal server error", http.StatusInternalServerError)
                return
        }

        w.Header().Set("Content-Type", "application/json")
        w.Header().Set("Cache-Control", fmt.Sprintf("public, max-age=%d", DefaultDiscoveryCacheMaxAge))
        w.Header().Set("X-Content-Type-Options", "nosniff")
        _, _ = w.Write(data) //nolint:gosec // G705: data is JSON-marshaled from internal metadata, not user input
}

// OIDCDiscoveryHandler handles GET /.well-known/openid-configuration requests.
// With interactive authorization enabled, it returns OIDC discovery metadata
// describing the authorization server capabilities. In token-only mode, the
// endpoint is a ToolHive compatibility alias for key discovery, not a
// standards-conforming OIDC Discovery response: generic MCP and OIDC
// authorization-code clients cannot use it because it omits the authorization
// endpoint and PKCE metadata required to begin an interactive flow.
func (h *Handler) OIDCDiscoveryHandler(w http.ResponseWriter, _ *http.Request) {
        if h.tokenOnly {
                // Token-only servers do not implement OIDC authorization. This ToolHive
                // compatibility alias lets existing key-discovery clients find the token
                // endpoint without advertising unsupported OIDC response types; it is not
                // a standards-conforming OIDC Discovery response.
                data, err := json.Marshal(h.buildOAuthMetadata())
                if err != nil {
                        slog.Error("failed to encode token-only discovery document", "error", err)
                        http.Error(w, "internal server error", http.StatusInternalServerError)
                        return
                }
                w.Header().Set("Content-Type", "application/json")
                w.Header().Set("Cache-Control", fmt.Sprintf("public, max-age=%d", DefaultDiscoveryCacheMaxAge))
                w.Header().Set("X-Content-Type-Options", "nosniff")
                _, _ = w.Write(data)
                return
        }

        // Get signing algorithms from the actual JWKS keys
        signingAlgs := h.getSigningAlgorithms()

        discovery := sharedobauth.OIDCDiscoveryDocument{
                // Include all OAuth 2.0 AS Metadata (RFC 8414)
                AuthorizationServerMetadata: h.buildOAuthMetadata(),

                // OIDC-specific REQUIRED fields
                SubjectTypesSupported:            []string{"public"},
                IDTokenSigningAlgValuesSupported: signingAlgs,
        }

        data, err := json.Marshal(discovery)
        if err != nil {
                slog.Error("failed to encode discovery document",
                        "error", err,
                )
                http.Error(w, "internal server error", http.StatusInternalServerError)
                return
        }

        w.Header().Set("Content-Type", "application/json")
        w.Header().Set("Cache-Control", fmt.Sprintf("public, max-age=%d", DefaultDiscoveryCacheMaxAge))
        w.Header().Set("X-Content-Type-Options", "nosniff")
        _, _ = w.Write(data)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

// FilterPrincipal is a getter contract for identity and claims used in upstream
// authorization decisions; it does not constrain an implementation to contain only
// these values. The authserver callback supplies a private projection containing
// credential-free values: Subject, CanonicalUserID, Name, Email, and Claims. Subject
// is the first upstream's provider-specific mapped subject: OIDC uses the
// configured SubjectClaim (or "sub" by default), OAuth2 may extract it from configured
// userinfo fields or the token response, and a synthetic fallback may apply. It is
// distinct from CanonicalUserID, the canonical storage-user ID (normally the resolved
// internal user ID, but a synthetic identity may use the mapped subject itself, which
// need not be stable). Name and Email come from the first upstream's resolved
// identity or profile and may originate from OIDC claims, configured OAuth2 userinfo
// fields, or token-response extraction. Claims contains the upstream-provided
// ID-token/userinfo claims, or is nil when unavailable; the authorization server
// does not manufacture claims. Callers must treat the returned claims map as
// read-only.
type FilterPrincipal interface {
        GetSubject() string
        GetCanonicalUserID() string
        GetName() string
        GetEmail() string
        GetClaims() map[string]any
}

// filterPrincipal contains only the credential-free values exposed to an
// upstream filter.
type filterPrincipal struct {
        subject         string
        canonicalUserID string
        name            string
        email           string
        claims          map[string]any
}

func (p filterPrincipal) GetSubject() string         { return p.subject }
func (p filterPrincipal) GetCanonicalUserID() string { return p.canonicalUserID }
func (p filterPrincipal) GetName() string            { return p.name }
func (p filterPrincipal) GetEmail() string           { return p.email }
func (p filterPrincipal) GetClaims() map[string]any  { return p.claims }

var _ FilterPrincipal = filterPrincipal{}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package handlers

import (
        "context"
        "fmt"
        "log/slog"
        "math"
        "mime"
        "net/http"
        "net/url"
        "slices"
        "strconv"
        "strings"
        "time"

        "github.com/go-chi/chi/v5"
        "github.com/ory/fosite"
        "go.opentelemetry.io/otel"
        "go.opentelemetry.io/otel/metric"
        "golang.org/x/time/rate"

        "github.com/stacklok/toolhive/pkg/authserver/internal/consentpolicy"
        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
)

// NamedUpstream pairs a logical provider name with its OAuth2Provider implementation.
// The name is used as the storage key and must be unique within the upstream slice.
type NamedUpstream struct {
        Name string
        // DisplayName is an optional human-readable label shown on the consent page.
        // It is display-only: never compared, stored, or used to route.
        DisplayName string
        Provider    upstream.OAuth2Provider
}

// Handler provides HTTP handlers for the OAuth authorization server endpoints.
type Handler struct {
        provider  fosite.OAuth2Provider
        config    *server.AuthorizationServerConfig
        storage   storage.Storage
        upstreams []NamedUpstream
        // tokenOnly is derived at construction from the absence of upstreams and is
        // the single source of truth for token-only capability behavior.
        tokenOnly    bool
        userResolver *UserResolver
        // refresher, when set, lets nextMissingUpstream transparently refresh an
        // expired upstream leg during chain evaluation instead of re-prompting. Nil
        // when no refresher is configured; an expired leg is then treated as missing.
        refresher storage.UpstreamTokenRefresher
        // filter, when set, narrows the authorization chain to a subset of the
        // configured upstreams once the first leg resolves. Nil when no filter is
        // configured; the chain then walks all configured upstreams as before.
        filter UpstreamFilter
        // consentExemptClientIDs is populated by WithConsentExemptClientIDs until
        // constructor validation builds consentExemptClientIDSet.
        consentExemptClientIDs   []string
        consentExemptClientIDSet map[string]struct{}
        // registerLimiter bounds the unauthenticated /oauth/register endpoint,
        // which mints persisted state on every call. Per-process, not per-IP: the
        // goal is to cap storage growth and log volume, not to be fair between
        // callers (per-IP limiting behind a shared NAT or proxy would either
        // under- or over-count, and X-Forwarded-For is attacker-controlled without
        // a trusted proxy list). Because the limit lives on this per-process field,
        // running N replicas multiplies the effective rate by N — size the
        // configured rate with replica count in mind. DefaultRegisterRate and
        // DefaultRegisterBurst unless WithRegisterRateLimit overrides them. Nil in
        // tests that construct Handler directly; OAuthRoutes tolerates nil by
        // skipping the gate.
        registerLimiter *rate.Limiter
        // meterProvider supplies dcrMetrics. otel.GetMeterProvider() unless
        // WithMeterProvider overrides it.
        meterProvider metric.MeterProvider
        // dcrMetrics counts DCR cap rejections and reports the live DCR client
        // count. Nil in tests that construct Handler directly; its methods
        // tolerate a nil receiver.
        dcrMetrics *dcrMetrics
        // capacityWarner logs DCR cap rejections at WARN, bounded by time. Nil in
        // tests that construct Handler directly; its methods tolerate a nil
        // receiver.
        capacityWarner *capacityWarner
        // deviceAuthorizationLimiter bounds the unauthenticated
        // /oauth/device_authorization endpoint, which mints persisted state
        // (a device_code/user_code pair) on every call, exactly like
        // /oauth/register. Same default rate and same per-process (not per-IP)
        // reasoning as registerLimiter. Nil when device flow is disabled — in
        // that case OAuthRoutes never registers the route, so the nil limiter is
        // never dereferenced. Also nil in tests that construct Handler directly.
        deviceAuthorizationLimiter *rate.Limiter
        // deviceCodeInterval is the minimum time a device-flow client must wait
        // between polls of the token endpoint, and the value advertised in the
        // device authorization response's "interval" field. Set from
        // config.DeviceCodeInterval at construction, falling back to
        // server.DefaultDeviceCodeInterval when that is zero.
        deviceCodeInterval time.Duration
        // deviceStorage is storage's storage.DeviceCodeStorage capability,
        // resolved once via storage.Unwrap(stor).(storage.DeviceCodeStorage) at
        // construction time -- storage.Storage deliberately does NOT embed
        // DeviceCodeStorage (see that interface's contract there), and CIMD/SPIFFE
        // decorators only implement the narrower Storage interface they wrap, so
        // the unwrap is needed to reach a decorated backend's underlying
        // MemoryStorage/RedisStorage. DeviceAuthorizationHandler goes through this
        // field instead of storage directly. Nil when device flow is disabled;
        // OAuthRoutes never registers DeviceAuthorizationHandler in that case, so
        // it is never dereferenced.
        deviceStorage     storage.DeviceCodeStorage
        rememberedStorage storage.RememberedConsentStorage
        // deviceVerificationLimiter bounds the unauthenticated POST /oauth/device
        // endpoint (DeviceVerificationSubmitHandler), which looks up a
        // DeviceRequest by the human-entered user_code. RFC 8628 §5.4 requires
        // throttling this lookup: user_code is drawn from a bounded charset over a
        // fixed-length code with a 10-minute TTL, so an unthrottled endpoint is a
        // brute-force guessing oracle. Unlike registerLimiter, this IS per-IP
        // (see perIPLimiter's doc comment for why): this endpoint is a
        // human-facing login page, not a machine-driven one-shot call, so a
        // shared process-wide bucket would let one caller starve every other
        // concurrent device-flow login. Nil when device flow is disabled;
        // OAuthRoutes never registers the route in that case, so it is never
        // dereferenced.
        deviceVerificationLimiter *perIPLimiter
}

// DefaultRegisterRate and DefaultRegisterBurst bound /oauth/register when no
// rate is configured: 1 registration/sec sustained, burst of 5. Generous for
// legitimate tooling (a client registers once, then reuses its client_id),
// tight enough that an unauthenticated caller cannot grow the client keyspace
// or the log stream faster than operators can react.
const (
        DefaultRegisterRate  rate.Limit = 1
        DefaultRegisterBurst int        = 5
)

// WithRegisterRateLimit overrides the /oauth/register rate (registrations per
// second) and burst. The limit is per process: N replicas admit N times the
// rate. NewHandler rejects a non-positive rate or burst.
func WithRegisterRateLimit(r rate.Limit, burst int) Option {
        return func(h *Handler) {
                h.registerLimiter = rate.NewLimiter(r, burst)
        }
}

// WithMeterProvider sets the meter provider for the DCR metrics. Without it,
// NewHandler uses the global otel.GetMeterProvider().
func WithMeterProvider(mp metric.MeterProvider) Option {
        return func(h *Handler) {
                h.meterProvider = mp
        }
}

// UpstreamFilter narrows the authorization chain to a subset of the configured
// upstreams, keyed on the identity established by the first upstream. It is
// consulted exactly once per authorization, in the callback handler, after the
// first upstream (upstreams[0]) resolves. The first upstream is always required:
// it is never passed to the filter and cannot be removed by it.
//
// FilterUpstreams receives, in order:
//   - ctx: the request context of the first leg's callback.
//   - principal: the first upstream's credential-free identity and claims, used for
//     authorization decisions. See FilterPrincipal for where each value comes from.
//     The filter MUST treat principal (including its claims map) as read-only.
//   - configured: the names of the non-first configured upstreams, in configured
//     order.
//
// Claims contract: values are untyped and come straight from the upstream IdP, so
// implementations MUST assert types defensively (e.g. comma-ok on map/slice
// access) — a malformed IdP response must not be able to panic the filter. GetClaims()
// is nil not only for OAuth2 identityFromToken / synthetic upstreams (which resolve
// no structured claim set) but also on a transient OIDC claims-extraction failure,
// so a security-relevant filter MUST treat nil/absent claims as fail-closed. Note
// that the `aud` claim, when present, is the upstream IdP's client_id — not this
// authorization server.
//
// It returns the subset to keep. The handler preserves configured order and
// ignores any returned name that is not one of the non-first configured upstreams,
// so the filter cannot reorder the chain or introduce unknown providers. A
// returned error fails the authorization with a server error — the handler never
// falls back to walking every upstream on error.
//
// The "first-leg identity" guarantee holds even across a rolling upgrade of a
// persistent backend: a subsequent-leg pending that lacks a computed chain (e.g.
// one written before this feature existed) is rejected and forces a fresh
// authorization, rather than re-running the filter against a later leg.
type UpstreamFilter interface {
        FilterUpstreams(
                ctx context.Context,
                principal FilterPrincipal,
                configured []string,
        ) ([]string, error)
}

// Option configures optional Handler behavior at construction time.
type Option func(*Handler)

// WithUpstreamRefresher injects a refresher used by nextMissingUpstream to
// transparently refresh expired upstream tokens while evaluating the
// authorization chain. When unset, an expired leg is treated as missing and the
// user is re-prompted — the behavior before this option existed.
func WithUpstreamRefresher(r storage.UpstreamTokenRefresher) Option {
        return func(h *Handler) {
                h.refresher = r
        }
}

// WithUpstreamFilter injects a filter that narrows the authorization chain to a
// subset of the configured upstreams once the first leg resolves. When unset, the
// handler walks all configured upstreams — the behavior before this option
// existed. See UpstreamFilter for the contract.
func WithUpstreamFilter(f UpstreamFilter) Option {
        return func(h *Handler) {
                h.filter = f
        }
}

// WithConsentExemptClientIDs configures the client IDs preauthorized to skip
// ToolHive's consent page. Listing a client ID preauthorizes whatever registration
// currently holds that ID, for whichever user completes upstream authentication.
// It is not per-user consent, and runtime eligibility checks do not prove where
// the registration came from. The embedder must control registration writes for
// exempt IDs; replacement or permission expansion transfers or broadens
// preauthorization. Removing an ID affects new requests only and does not revoke
// tokens or cancel already-approved flows. Upstream authentication still runs.
// The IDs are copied when the option is created, and repeated options replace the
// list.
func WithConsentExemptClientIDs(ids []string) Option {
        clientIDs := slices.Clone(ids)
        return func(h *Handler) {
                h.consentExemptClientIDs = slices.Clone(clientIDs)
        }
}

// NewHandler creates a new Handler with the given dependencies.
// upstreams defines the ordered sequence of upstream providers consulted
// during multi-upstream authorization flows (e.g., sequential token acquisition).
//
// Returns an error if config is nil, if config's embedded *fosite.Config is
// nil, or if any entry has an empty name, a nil provider, or a duplicate name.
// Upstream names must be unique: upstreamByName returns the
// first match, tokens are keyed by name, and the authorization chain is keyed by
// name — a duplicate would silently shadow a provider. Catching misconfiguration
// here is far easier to diagnose than a nil-deref panic or a shadowed provider
// deep inside an HTTP handler at request time.
func NewHandler(
        provider fosite.OAuth2Provider,
        config *server.AuthorizationServerConfig,
        stor storage.Storage,
        upstreams []NamedUpstream,
        opts ...Option,
) (*Handler, error) {
        if config == nil || config.Config == nil {
                return nil, fmt.Errorf(
                        "handlers: AuthorizationServerConfig with embedded *fosite.Config must be non-nil")
        }
        seen := make(map[string]struct{}, len(upstreams))
        for _, u := range upstreams {
                if u.Name == "" {
                        return nil, fmt.Errorf("handlers: upstream entry has empty name")
                }
                if u.Provider == nil {
                        return nil, fmt.Errorf("handlers: upstream %q has nil provider", u.Name)
                }
                if _, dup := seen[u.Name]; dup {
                        return nil, fmt.Errorf("handlers: duplicate upstream name %q", u.Name)
                }
                seen[u.Name] = struct{}{}
        }
        h := &Handler{
                provider:           provider,
                config:             config,
                storage:            stor,
                upstreams:          upstreams,
                tokenOnly:          len(upstreams) == 0,
                userResolver:       NewUserResolver(stor),
                registerLimiter:    rate.NewLimiter(DefaultRegisterRate, DefaultRegisterBurst),
                meterProvider:      otel.GetMeterProvider(),
                deviceCodeInterval: deviceCodeInterval(config.DeviceCodeInterval),
        }
        if config.DeviceFlowEnabled {
                // Same rate as registerLimiter's default: this gate protects the same kind of
                // unauthenticated persisted-state minting, just reached through
                // /oauth/device_authorization instead of /oauth/register.
                h.deviceAuthorizationLimiter = rate.NewLimiter(rate.Limit(1), 5)
                // Same rate as deviceAuthorizationLimiter, but per-IP rather than
                // per-process -- see deviceVerificationLimiter's field doc comment
                // and perIPLimiter's doc comment for why.
                h.deviceVerificationLimiter = newPerIPLimiter(rate.Limit(1), 5)
                deviceStorage, ok := storage.Unwrap(stor).(storage.DeviceCodeStorage)
                if !ok {
                        return nil, fmt.Errorf(
                                "handlers: device flow enabled but storage backend %T does not implement storage.DeviceCodeStorage", stor)
                }
                h.deviceStorage = deviceStorage
        }
        if base, err := url.Parse(config.GetAuthorizationEndpointBaseURL()); err == nil && base.Scheme == "https" {
                h.rememberedStorage, _ = storage.Unwrap(stor).(storage.RememberedConsentStorage)
        }
        for _, o := range opts {
                o(h)
        }
        if err := validateRegisterLimiter(h.registerLimiter); err != nil {
                return nil, err
        }
        if err := h.initConsentExemptions(); err != nil {
                return nil, err
        }
        // Registered last: the gauge callback outlives a failed constructor.
        metrics, err := newDCRMetrics(h.meterProvider, stor)
        if err != nil {
                return nil, fmt.Errorf("handlers: %w", err)
        }
        h.dcrMetrics = metrics
        h.capacityWarner = newCapacityWarner(stor)
        return h, nil
}

// Close releases what NewHandler registered outside the Handler: the DCR
// client gauge's callback, which holds the storage. Call it when the server
// shuts down so a rebuilt server in the same process does not leave a
// callback reading closed storage.
func (h *Handler) Close() error {
        return h.dcrMetrics.close()
}

// deviceCodeInterval returns configured, falling back to
// server.DefaultDeviceCodeInterval when configured is zero. A zero interval
// would defeat the interval's purpose of bounding poll frequency, so it is
// never used as the enforced minimum.
func deviceCodeInterval(configured time.Duration) time.Duration {
        if configured <= 0 {
                return server.DefaultDeviceCodeInterval
        }
        return configured
}

// Routes returns a router with all OAuth/OIDC endpoints registered.
func (h *Handler) Routes() http.Handler {
        r := chi.NewRouter()
        h.OAuthRoutes(r)
        h.WellKnownRoutes(r)
        return r
}

// OAuthRoutes registers OAuth endpoints (authorize, consent, callback, token, register,
// and the device flow's device_authorization/verification routes when enabled)
// on the provided router.
func (h *Handler) OAuthRoutes(r chi.Router) {
        r.Get("/oauth/authorize", rejectMultipartAuthorizeBody(h.AuthorizeHandler))
        r.Post("/oauth/consent", h.ConsentHandler)
        // /oauth/callback is shared by the OAuth-client authorization_code flow
        // and (when device flow is enabled) the verification page's upstream
        // login -- an upstream's redirect_uri is fixed per upstream at
        // construction, so there is no way to register a second callback path
        // for the device flow. CallbackHandler dispatches between the two based
        // on which pending record the state parameter matches. An OAuth-client
        // authorization additionally requires the browser-binding cookie that
        // /oauth/authorize set, so it must be served on the same host (see
        // browser_binding.go).
        r.Get("/oauth/callback", h.CallbackHandler)
        r.Post("/oauth/token", h.TokenHandler)
        r.Post("/oauth/register", h.rateLimitRegister(h.RegisterClientHandler))
        if h.config.DeviceFlowEnabled {
                r.Post("/oauth/device_authorization", h.rateLimitDeviceAuthorization(h.DeviceAuthorizationHandler))
                r.Get("/oauth/device", h.DeviceVerificationHandler)
                r.Post("/oauth/device", h.rateLimitDeviceVerification(h.DeviceVerificationSubmitHandler))
                r.Post("/oauth/device/confirm", h.DeviceVerificationConfirmHandler)
        }
}

// rateLimitRegister gates the unauthenticated registration endpoint: over the
// limit it returns 429 with a Retry-After hint rather than minting persisted
// state. RFC 7591 has no DCR-specific rate-limit response, so this uses the
// plain HTTP semantics; well-behaved DCR clients retry.
func (h *Handler) rateLimitRegister(next http.HandlerFunc) http.HandlerFunc {
        return func(w http.ResponseWriter, req *http.Request) {
                if h.registerLimiter != nil && !h.registerLimiter.Allow() {
                        w.Header().Set("Retry-After", strconv.Itoa(limiterRetryAfterSeconds(h.registerLimiter)))
                        http.Error(w, "rate limit exceeded, retry later", http.StatusTooManyRequests)
                        return
                }
                next(w, req)
        }
}

// validateRegisterLimiter rejects a register limiter without a finite
// positive rate and a positive burst. rate.Inf (math.MaxFloat64, not +Inf)
// would silently disable limiting, and NaN passes a plain <= 0 check.
func validateRegisterLimiter(l *rate.Limiter) error {
        lim := float64(l.Limit())
        if math.IsNaN(lim) || math.IsInf(lim, 0) || l.Limit() == rate.Inf || lim <= 0 || l.Burst() <= 0 {
                return fmt.Errorf("handlers: register rate limit must have a finite positive rate and positive burst, got %v/s burst %d",
                        l.Limit(), l.Burst())
        }
        return nil
}

// initConsentExemptions validates the consent-exempt client IDs set by
// WithConsentExemptClientIDs and indexes them for lookup.
func (h *Handler) initConsentExemptions() error {
        if err := consentpolicy.Validate(
                h.consentExemptClientIDs,
                h.config.GetAuthorizationEndpointBaseURL(),
                h.config.GetAccessTokenIssuer(),
                len(h.upstreams),
        ); err != nil {
                return fmt.Errorf("handlers: %w", err)
        }
        h.consentExemptClientIDSet = make(map[string]struct{}, len(h.consentExemptClientIDs))
        for _, id := range h.consentExemptClientIDs {
                h.consentExemptClientIDSet[id] = struct{}{}
        }
        h.consentExemptClientIDs = nil
        return nil
}

// limiterRetryAfterSeconds returns the Retry-After value for a request l
// refused: 1/rate seconds, the longest the limiter can take to refill one
// token, rounded up and at least 1. A partly refilled bucket frees a token
// sooner, so this is an upper bound. A rate below 1/s therefore advertises a
// wait longer than a second, so a
// client honouring the header does not retry into further 429s. The result
// is clamped to math.MaxInt32 so a vanishingly small rate cannot overflow int.
func limiterRetryAfterSeconds(l *rate.Limiter) int {
        return int(min(max(math.Ceil(1/float64(l.Limit())), 1), math.MaxInt32))
}

// rejectMultipartAuthorizeBody prevents Fosite from accepting multipart GET
// bodies on /oauth/authorize. The endpoint's parameters must come from the query
// string; otherwise client_id could bypass the client-hiding checks that use
// ParseForm, which do not read multipart GET bodies.
func rejectMultipartAuthorizeBody(next http.HandlerFunc) http.HandlerFunc {
        return func(w http.ResponseWriter, req *http.Request) {
                mediaType, _, err := mime.ParseMediaType(req.Header.Get("Content-Type"))
                if err == nil && strings.HasPrefix(mediaType, "multipart/") {
                        // WARN on purpose although a caller can trigger it: this is a
                        // request-shape rejection and the operator-visible signal. Only the
                        // media type and path are logged, never client_id or other input.
                        slog.WarnContext(req.Context(), "rejected multipart request body on authorize endpoint",
                                "media_type", mediaType, "path", req.URL.Path)
                        http.Error(w, "multipart request bodies are not supported", http.StatusBadRequest)
                        return
                }
                next(w, req)
        }
}

// writeCIMDRateLimited writes the plain-text 429 for a client lookup refused by
// the CIMD storage decorator's fetch limiter, matching rateLimitRegister's
// response. It is used only by the authorize endpoint: the client is unresolved
// there, so its redirect_uri cannot be trusted (RFC 6749 section 4.1.2.1) and
// no OAuth error can be redirected to it.
func writeCIMDRateLimited(w http.ResponseWriter) {
        w.Header().Set("Retry-After", "1")
        http.Error(w, "rate limit exceeded, retry later", http.StatusTooManyRequests)
}

// writeCIMDRateLimitedOAuth writes the 429 for a CIMD client lookup refused by
// the fetch limiter as an RFC 6749 section 5.2 JSON error. It is used by the
// token and device authorization endpoints only. The status stays 429 and the
// code is temporarily_unavailable: section 5.2 has no rate-limit code and
// defines temporarily_unavailable only for the authorization endpoint, so this
// is an extension that clients tolerate (fosite defines the same code).
func writeCIMDRateLimitedOAuth(w http.ResponseWriter) {
        w.Header().Set("Retry-After", "1")
        writeOAuthError(w, http.StatusTooManyRequests, "temporarily_unavailable",
                "client metadata fetch rate limit exceeded, retry later")
}

// rateLimitDeviceAuthorization gates the unauthenticated
// /oauth/device_authorization endpoint: over the limit it returns 429 with a
// Retry-After hint rather than minting persisted state. Mirrors
// rateLimitRegister exactly, including its nil-tolerant shape (nil when
// device flow is disabled, in which case this wrapper is never installed by
// OAuthRoutes anyway).
func (h *Handler) rateLimitDeviceAuthorization(next http.HandlerFunc) http.HandlerFunc {
        return func(w http.ResponseWriter, req *http.Request) {
                if h.deviceAuthorizationLimiter != nil && !h.deviceAuthorizationLimiter.Allow() {
                        w.Header().Set("Retry-After", "1")
                        http.Error(w, "rate limit exceeded, retry later", http.StatusTooManyRequests)
                        return
                }
                next(w, req)
        }
}

// rateLimitDeviceVerification gates the unauthenticated POST /oauth/device
// endpoint: over the limit it returns 429 with a Retry-After hint rather than
// running another user_code lookup. Per-IP (see perIPLimiter's doc comment
// for why this endpoint differs from rateLimitDeviceAuthorization's
// per-process gate), otherwise the same shape, including tolerating a nil
// limiter when device flow is disabled.
func (h *Handler) rateLimitDeviceVerification(next http.HandlerFunc) http.HandlerFunc {
        return func(w http.ResponseWriter, req *http.Request) {
                if h.deviceVerificationLimiter != nil && !h.deviceVerificationLimiter.allow(clientIP(req)) {
                        w.Header().Set("Retry-After", "1")
                        http.Error(w, "rate limit exceeded, retry later", http.StatusTooManyRequests)
                        return
                }
                next(w, req)
        }
}

// WellKnownRoutes registers well-known endpoints (JWKS, OAuth/OIDC discovery) on the provided router.
// Both discovery endpoints are registered per the MCP specification requirement to provide
// at least one discovery mechanism, with both supported for maximum interoperability:
// - /.well-known/oauth-authorization-server (RFC 8414) for OAuth-only clients
// - /.well-known/openid-configuration (OIDC Discovery 1.0) for OIDC clients
//
// The wildcard variants (/.well-known/oauth-authorization-server/*) handle RFC 8414
// Section 3.1 path-based issuers, where clients insert /.well-known/ before the
// issuer's path component (e.g., /.well-known/oauth-authorization-server/inject-test
// for issuer https://example.com/inject-test).
func (h *Handler) WellKnownRoutes(r chi.Router) {
        r.Get("/.well-known/jwks.json", h.JWKSHandler)
        r.Get("/.well-known/oauth-authorization-server", h.OAuthDiscoveryHandler)
        r.Get("/.well-known/oauth-authorization-server/*", h.OAuthDiscoveryHandler)
        r.Get("/.well-known/openid-configuration", h.OIDCDiscoveryHandler)
        r.Get("/.well-known/openid-configuration/*", h.OIDCDiscoveryHandler)
}

// nextMissingUpstream returns the name of the next upstream provider in the
// authorization chain that the user must (re-)authenticate with for this session.
// It walks the provided chain — the effective, possibly filtered, ordered set of
// upstream names for this authorization (see computeChain) — rather than the raw
// configured list, so a leg the filter dropped is never prompted for. Returns
// empty string if all legs in the chain are satisfied, or an error if the storage
// lookup fails.
//
// A leg is satisfied when it has a stored token that is live (or asserts no
// expiry). A present-but-expired token is NOT treated as satisfied by presence
// alone: the leg is refreshed transparently (mirroring upstreamtoken
// InProcessService.GetAllUpstreamCredentials) and only counts as satisfied if the refresh
// succeeds. If refresh is impossible or fails, the leg is reported as missing so
// the user is re-prompted up front, rather than the stale token surfacing as a
// runtime auth error later at MCP-request token-swap time.
func (h *Handler) nextMissingUpstream(ctx context.Context, sessionID string, chain []string) (string, error) {
        stored, err := h.storage.GetAllUpstreamTokens(ctx, sessionID)
        if err != nil {
                return "", fmt.Errorf("failed to check upstream token state: %w", err)
        }
        for _, name := range chain {
                tokens, ok := stored[name]
                if !ok || tokens == nil {
                        // No token stored for this leg — prompt.
                        return name, nil
                }
                // A live token (or one with no asserted expiry) satisfies the leg.
                if tokens.ExpiresAt.IsZero() || !tokens.IsExpired(time.Now()) {
                        continue
                }
                // Expired — attempt a transparent refresh; prompt now if it can't be done.
                if !h.refreshExpiredLeg(ctx, sessionID, name, tokens) {
                        return name, nil
                }
        }
        return "", nil
}

// computeChain returns the ordered, effective set of upstream names this
// authorization must walk. The first configured upstream always leads the chain
// and is never filtered out. When no filter is configured, the chain is the full
// configured list in order (the behavior before the filter hook existed). When a
// filter is configured, it is consulted with the names of the non-first upstreams
// (in configured order) and the chain becomes the first upstream plus the kept
// subset. Configured order is preserved and any returned name that is not a
// non-first configured upstream is ignored, so the filter can only narrow — never
// reorder or extend — the chain.
//
// A filter error is returned to the caller so the authorization fails cleanly; it
// never silently falls back to walking every upstream.
//
// principal is the first upstream's resolved identity, passed through to the filter
// (see UpstreamFilter); it is unused when no filter is configured.
func (h *Handler) computeChain(ctx context.Context, principal FilterPrincipal) ([]string, error) {
        chain := []string{h.upstreams[0].Name}
        rest := h.upstreams[1:]
        if len(rest) == 0 {
                return chain, nil
        }

        restNames := make([]string, len(rest))
        for i := range rest {
                restNames[i] = rest[i].Name
        }

        if h.filter == nil {
                return append(chain, restNames...), nil
        }

        keep, err := h.filter.FilterUpstreams(ctx, principal, restNames)
        if err != nil {
                return nil, fmt.Errorf("upstream filter failed: %w", err)
        }

        keepSet := make(map[string]struct{}, len(keep))
        for _, name := range keep {
                keepSet[name] = struct{}{}
        }
        // Iterate configured order (not the filter's return order) so the chain
        // preserves the operator-defined sequence and silently drops any name the
        // filter returned that is not a non-first configured upstream.
        for i := range rest {
                if _, ok := keepSet[rest[i].Name]; ok {
                        chain = append(chain, rest[i].Name)
                }
        }
        return chain, nil
}

// resolveChain returns the effective upstream chain for the leg that just
// completed. A true first leg — identified by an as-yet-unresolved user
// (pending.ResolvedUserID == "") — computes the chain now, consulting the optional
// filter with this leg's request context. A subsequent leg reuses the chain the
// first leg carried forward in the pending authorization, after validating it
// against the configured upstreams.
//
// A subsequent leg whose pending carries no chain (e.g. one written before the
// ChainUpstreams field existed, a possible rolling-upgrade window on a persistent
// backend) is a hard error rather than a silent recompute: the filter's contract
// is to run once, with the first leg's context, so recomputing here with a later
// leg's context could narrow the chain against the wrong request. Failing closed
// forces a fresh authorization that starts cleanly at the first leg.
//
// principal is the identity resolved by the leg that just completed; it is passed
// to computeChain (and thus the filter) only on the first leg, where it reflects
// the identity provider. Subsequent legs reuse the stored chain and ignore it.
func (h *Handler) resolveChain(
        ctx context.Context,
        pending *storage.PendingAuthorization,
        principal FilterPrincipal,
) ([]string, error) {
        if pending.ResolvedUserID == "" {
                // True first leg — compute with this (the first) leg's request context.
                return h.computeChain(ctx, principal)
        }
        if len(pending.ChainUpstreams) == 0 {
                return nil, fmt.Errorf("subsequent chain leg is missing its computed upstream chain (stale pending authorization)")
        }
        if err := h.validateChain(pending.ChainUpstreams); err != nil {
                return nil, fmt.Errorf("stored upstream chain is invalid: %w", err)
        }
        return pending.ChainUpstreams, nil
}

// validateChain rejects an upstream chain loaded from storage that does not agree
// with the configured upstreams. A valid chain is a non-empty, in-order,
// duplicate-free subsequence of the configured upstreams led by the required
// first upstream. This guards the reuse path against a corrupt or tampered
// PendingAuthorization row shrinking or reordering an in-flight chain to skip legs
// — which would also silently disable the cross-leg identity check, since
// verifyChainIdentity is a no-op for a single-element chain.
//
// Note: this validates that each name is currently configured, not that the name
// still points at the same provider it did when the chain was computed. A name
// rebound to a different IdP mid-flight (a rolling config change) still validates;
// the frozen subject and the pending TTL bound that exposure.
func (h *Handler) validateChain(chain []string) error {
        if len(chain) == 0 {
                return fmt.Errorf("chain is empty")
        }
        if chain[0] != h.upstreams[0].Name {
                return fmt.Errorf("chain must lead with the first configured upstream %q, got %q",
                        h.upstreams[0].Name, chain[0])
        }
        // Walk both, advancing the config cursor past each match. A name that never
        // matches — unconfigured, out of order, or a duplicate of an already-consumed
        // entry — is rejected, so the stored chain is a faithful narrowing of the config.
        ci := 0
        for _, name := range chain {
                for ci < len(h.upstreams) && h.upstreams[ci].Name != name {
                        ci++
                }
                if ci >= len(h.upstreams) {
                        return fmt.Errorf("chain entry %q is unconfigured, out of order, or duplicated", name)
                }
                ci++ // consume the matched upstream so it cannot be reused
        }
        return nil
}

// refreshExpiredLeg attempts to refresh an expired upstream token for one chain
// leg. It returns true when the leg can be treated as authenticated (refresh
// succeeded) and false when the user must be re-prompted: no refresher configured,
// no refresh token on the row, or the refresh itself failed (expired/revoked
// refresh token, provider error). Mirrors the refresh-then-classify behavior in
// upstreamtoken.InProcessService.GetAllUpstreamCredentials.
func (h *Handler) refreshExpiredLeg(
        ctx context.Context,
        sessionID, providerName string,
        expired *storage.UpstreamTokens,
) bool {
        if h.refresher == nil || expired.RefreshToken == "" {
                return false
        }
        // RefreshAndStore persists the refreshed token to storage as its side effect;
        // the token-swap path re-reads it from storage at MCP-request time, so the
        // returned *UpstreamTokens is intentionally discarded here.
        if _, err := h.refresher.RefreshAndStore(ctx, sessionID, expired); err != nil {
                slog.WarnContext(ctx, "upstream token refresh failed during chain evaluation; re-prompting",
                        "session_id", sessionID,
                        "provider", providerName,
                        "error", err,
                )
                return false
        }
        slog.DebugContext(ctx, "refreshed expired upstream token during chain evaluation",
                "session_id", sessionID,
                "provider", providerName,
        )
        return true
}

// upstreamByName returns the upstream provider with the given name.
// It follows the (value, bool) convention: the second return value is false
// if no upstream with that name exists.
func (h *Handler) upstreamByName(name string) (upstream.OAuth2Provider, bool) {
        for i := range h.upstreams {
                if h.upstreams[i].Name == name {
                        return h.upstreams[i].Provider, true
                }
        }
        return nil, false
}

// issuer returns the authorization-server issuer URL. NewHandler enforces
// that h.config and h.config.Config are non-nil, so this method does not
// re-validate.
func (h *Handler) issuer() string {
        return h.config.AccessTokenIssuer
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "crypto/rand"
        "crypto/sha256"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "slices"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

func (h *Handler) consentCookieName() string {
        issuer := h.config.GetAccessTokenIssuer()
        base := h.config.GetAuthorizationEndpointBaseURL()
        return "__Host-thv_consent_" + consentDigest(fmt.Sprintf("%d:%s%d:%s", len(issuer), issuer, len(base), base))
}

func consentDigest(secret string) string {
        return fmt.Sprintf("%x", sha256.Sum256([]byte(secret)))
}

// coveringConsent returns only approvals matching the exact client request envelope.
// A storage failure is never treated as a missing approval.
func (h *Handler) coveringConsent(ctx context.Context, req *http.Request, client fosite.Client,
        redirectURI string, scopes []string, resource string) (*storage.ConsentSession, string, error) {
        session, digest, err := h.rememberedSession(ctx, req)
        if session == nil || err != nil {
                return nil, "", err
        }
        approval, err := h.rememberedApproval(ctx, session.UserID, client.GetID())
        if approval == nil || err != nil {
                return nil, "", err
        }
        if !approvalCovers(client, approval, redirectURI, scopes, resource) {
                return nil, "", nil
        }
        return session, digest, nil
}

// rememberedSession loads the live consent session named by the browser cookie.
// It returns nil without error when there is no usable session.
func (h *Handler) rememberedSession(ctx context.Context, req *http.Request) (*storage.ConsentSession, string, error) {
        if h.rememberedStorage == nil {
                return nil, "", nil
        }
        cookie, err := req.Cookie(h.consentCookieName())
        if errors.Is(err, http.ErrNoCookie) {
                return nil, "", nil
        }
        if err != nil {
                return nil, "", err
        }
        if cookie.Value == "" {
                return nil, "", nil
        }
        digest := consentDigest(cookie.Value)
        session, err := h.rememberedStorage.GetConsentSession(ctx, digest)
        if isNotFoundOrExpired(err) {
                return nil, "", nil
        }
        if err != nil {
                return nil, "", err
        }
        if session == nil || session.UserID == "" || session.ProviderSubject == "" || session.ProviderID == "" ||
                session.ExpiresAt.IsZero() {
                return nil, "", nil
        }
        if !time.Now().Before(session.ExpiresAt) {
                return nil, "", nil
        }
        if session.ProviderID != h.upstreams[0].Name {
                return nil, "", nil
        }
        return session, digest, nil
}

// rememberedApproval loads the live client approval for the user.
// It returns nil without error when there is no usable approval.
func (h *Handler) rememberedApproval(ctx context.Context, userID, clientID string) (*storage.ClientApproval, error) {
        approval, err := h.rememberedStorage.GetClientApproval(ctx, userID, clientID)
        if isNotFoundOrExpired(err) {
                return nil, nil
        }
        if err != nil {
                return nil, err
        }
        if approval == nil || approval.ExpiresAt.IsZero() {
                return nil, nil
        }
        if !time.Now().Before(approval.ExpiresAt) {
                return nil, nil
        }
        return approval, nil
}

// approvalCovers reports whether approval spans the requested redirect URI
// (including the registered loopback-port rule), resource and scopes.
func approvalCovers(client fosite.Client, approval *storage.ClientApproval,
        redirectURI string, scopes []string, resource string) bool {
        redirectMatches := approval.RedirectURI == redirectURI
        if !redirectMatches {
                approvedRegistered, approvedOK := registration.RegisteredLoopbackRedirectURI(client, approval.RedirectURI)
                requestedRegistered, requestedOK := registration.RegisteredLoopbackRedirectURI(client, redirectURI)
                redirectMatches = approvedOK && requestedOK && approvedRegistered == requestedRegistered
        }
        if !redirectMatches || approval.Resource != resource {
                return false
        }
        for _, scope := range scopes {
                if !slices.Contains(approval.Scopes, scope) {
                        return false
                }
        }
        return true
}

func (h *Handler) maybeRememberConsent(ctx context.Context, w http.ResponseWriter,
        pending *storage.PendingAuthorization, userID string) {
        if h.rememberedStorage == nil || !pending.RememberConsent || pending.FirstProviderSubject == "" || userID == "" {
                return
        }
        now := time.Now()
        approval := &storage.ClientApproval{
                RedirectURI: pending.RedirectURI, Scopes: slices.Clone(pending.Scopes), Resource: pending.Resource,
                ExpiresAt: now.Add(storage.ClientApprovalTTL),
        }
        if err := h.rememberedStorage.StoreClientApproval(ctx, userID, pending.ClientID, approval); err != nil {
                slog.Warn("failed to remember client approval", "error", err)
                return
        }
        secret := rand.Text()
        session := &storage.ConsentSession{
                UserID: userID, ProviderID: pending.UpstreamProviderName, ProviderSubject: pending.FirstProviderSubject,
                ExpiresAt: now.Add(storage.ConsentSessionTTL),
        }
        // The final leg's provider can differ from the first provider.
        if len(h.upstreams) > 0 {
                session.ProviderID = h.upstreams[0].Name
        }
        if err := h.rememberedStorage.StoreConsentSession(ctx, consentDigest(secret), session); err != nil {
                slog.Warn("failed to remember browser session", "error", err)
                return
        }
        http.SetCookie(w, &http.Cookie{
                Name: h.consentCookieName(), Value: secret, Path: "/", Secure: true, HttpOnly: true,
                SameSite: http.SameSiteLaxMode, MaxAge: int(storage.ConsentSessionTTL.Seconds()),
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "errors"
        "log/slog"
        "net/http"
        "slices"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// logAccessError logs a token-endpoint failure from fosite's access
// request/response flow at a level determined by the error's HTTP status:
// Debug for 4xx (the only class an unauthenticated caller can trigger, so
// logging every one at Error would let a caller flood the log stream) and
// Error for 5xx (a storage outage or signing-key failure, which must not be
// silent at default log level). RFC6749Error.Error() returns only the error
// code, so the diagnostic fields are read individually.
func logAccessError(ctx context.Context, msg string, err error) {
        rfcErr := fosite.ErrorToRFC6749Error(err)
        lvl := slog.LevelDebug
        if rfcErr.CodeField >= http.StatusInternalServerError {
                lvl = slog.LevelError
        }
        //nolint:gosec // G706: error is request-derived (form parsing, client auth);
        // fields are error code/hint/debug, not raw client input
        slog.Log(ctx, lvl, msg,
                "error", rfcErr.ErrorField,
                "debug", rfcErr.DebugField,
                "status", rfcErr.CodeField,
        )
}

// TokenHandler handles POST /oauth/token requests.
// It processes token requests using fosite's access request/response flow.
func (h *Handler) TokenHandler(w http.ResponseWriter, req *http.Request) {
        ctx := req.Context()

        // Create a placeholder session for the token request.
        // All parameters are empty because Fosite's NewAccessRequest will:
        // 1. Extract the authorization code from the request
        // 2. Retrieve the stored authorize session from storage (created in CallbackHandler)
        // 3. Use the stored session's claims (subject, tsid, client_id) for token generation
        // This session object is only used as a deserialization template.
        sess := session.New("", "", "", session.UserClaims{})

        // Token endpoint only: lets the CIMD decorator use the persisted client
        // row when its fetch limiter refuses.
        ctx = storage.WithPersistedClientFallback(ctx)

        // Parse and validate the access request
        accessRequest, err := h.provider.NewAccessRequest(ctx, req, sess)
        if err != nil {
                if errors.Is(err, storage.ErrCIMDRateLimited) {
                        writeCIMDRateLimitedOAuth(w)
                        return
                }
                // RFC6749Error.Error() returns only the error code, so a bare "error"
                // field carries no diagnostic detail; logAccessError reads CodeField,
                // DebugField, and DescriptionField instead. Client errors (4xx, the
                // only class an unauthenticated caller can trigger) stay at Debug so a
                // flood can't drown real errors; a genuine 5xx (storage outage, signing
                // failure) still surfaces at Error.
                logAccessError(ctx, "failed to create access request", err)
                h.provider.WriteAccessError(ctx, w, accessRequest, err)
                return
        }

        if h.applyResourceAudience(ctx, w, accessRequest) {
                return
        }
        h.applySubjectIdentifier(accessRequest)

        // Generate the access response (tokens)
        response, err := h.provider.NewAccessResponse(ctx, accessRequest)
        if err != nil {
                // Same reasoning as the NewAccessRequest failure above: split on
                // status so a 500 from a storage outage or signing-key failure isn't
                // silently swallowed at Debug level.
                logAccessError(ctx, "failed to create access response", err)
                h.provider.WriteAccessError(ctx, w, accessRequest, err)
                return
        }

        // Renew the registration TTL for DCR-issued clients (public or confidential) on a
        // successful token exchange/refresh. This is the proven-use signal — unlike the
        // unauthenticated /oauth/authorize client read — so an actively-used client does
        // not expire mid-lifecycle and is not forced to re-register. The first renewal also
        // raises a newly registered client from DefaultUnusedDCRClientTTL to
        // DefaultDCRClientTTL. Best-effort: a renewal failure must not fail the token that
        // was just issued. Storage is a no-op for pre-provisioned/static clients, which
        // carry no DCRIssued marker, and the CIMD decorator skips CIMD rows (each
        // successful fetch already renews them).
        if client := accessRequest.GetClient(); client != nil {
                if err := h.storage.RenewClientTTL(ctx, client); err != nil {
                        slog.Warn("failed to renew client registration TTL",
                                "client_id", client.GetID(),
                                "error", err,
                        )
                }
        }

        // Write the token response
        h.provider.WriteAccessResponse(ctx, w, accessRequest, response)
}

// applySubjectIdentifier removes stale sub_id and upstream subject claims from
// every grant, then adds both claims only when the current runtime policy
// permits the final login-token audience.
func (h *Handler) applySubjectIdentifier(request fosite.AccessRequester) {
        stored := request.GetSession()
        if stored == nil {
                return
        }
        // Clone is the copy-before-mutate step: the stored session must not be modified in place.
        cloned := stored.Clone()
        sess, ok := cloned.(*session.Session)
        if !ok {
                return
        }
        request.SetSession(cloned)
        sess.ClearSubjectIdentifier()

        loginGrant := request.GetGrantTypes().HasOneOf(
                oauthproto.GrantTypeAuthorizationCode, oauthproto.GrantTypeRefreshToken, oauthproto.GrantTypeDeviceCode)
        if !loginGrant || !sess.VerifiedUpstream.IsComplete() {
                return
        }

        // On authorization_code redemption fosite copies the authorized granted audience
        // only inside NewAccessResponse, after this runs, and applyResourceAudience only
        // validates a matching resource without granting it, so the granted audience is
        // usually still empty here. writeAuthorizationResponse sets requested and granted
        // audience to the same pending.Resource, so the requested audience is equivalent.
        // Only authorization_code may fall back: for device_code the requested audience is
        // the caller's unvalidated audience form field, and refresh copies it forward from
        // the original request. Those grants already carry their granted audience here.
        audiences := request.GetGrantedAudience()
        if len(audiences) == 0 && request.GetGrantTypes().ExactOne(oauthproto.GrantTypeAuthorizationCode) {
                audiences = request.GetRequestedAudience()
        }
        if len(audiences) == 0 {
                return
        }
        for _, audience := range audiences {
                if !subjectIdentifierAudienceEnabled(h.config.SubjectIdentifierAudiences, audience) {
                        return
                }
        }
        sess.SetSubjectIdentifier(sess.VerifiedUpstream)
}

func subjectIdentifierAudienceEnabled(configured []string, audience string) bool {
        return audience != "" && slices.Contains(configured, audience)
}

// applyResourceAudience applies the RFC 8707 resource parameter to the access
// request's granted audience. It reports whether it already wrote an error response.
func (h *Handler) applyResourceAudience(
        ctx context.Context, w http.ResponseWriter, accessRequest fosite.AccessRequester,
) (handled bool) {
        // RFC 8707: Handle resource parameter for audience claim.
        // The resource parameter allows clients to specify which protected resource (MCP server)
        // the token is intended for. This value becomes the "aud" claim in the JWT.
        //
        // Note: RFC 8707 allows multiple resource parameters, but we explicitly reject them
        // for security reasons (simpler audience model, clearer token scope).
        resources := accessRequest.GetRequestForm()["resource"]
        if requested := accessRequest.GetRequestedAudience(); len(requested) > 0 &&
                (accessRequest.GetGrantTypes().ExactOne("authorization_code") || accessRequest.GetGrantTypes().ExactOne("refresh_token")) {
                // Fosite restores the authorized audience on code redemption and refresh;
                // other grants (notably token exchange) may request audiences independently.
                if len(requested) != 1 || len(resources) > 1 || (len(resources) == 1 && resources[0] != requested[0]) {
                        h.provider.WriteAccessError(ctx, w, accessRequest,
                                server.ErrInvalidTarget.WithHint("Resource differs from authorized audience"))
                        return true
                }
        } else if len(resources) > 1 {
                slog.Debug("multiple resource parameters not supported", //nolint:gosec // G706: count is an integer
                        "count", len(resources),
                )
                h.provider.WriteAccessError(ctx, w, accessRequest,
                        server.ErrInvalidTarget.WithHint("Multiple resource parameters are not supported"))
                return true
        } else if len(resources) == 1 && resources[0] != "" {
                resource := resources[0]
                // Validate URI format per RFC 8707
                if err := server.ValidateAudienceURI(resource); err != nil {
                        slog.Debug("invalid resource URI format", //nolint:gosec // G706: resource URI from token request
                                "resource", resource,
                                "error", err,
                        )
                        h.provider.WriteAccessError(ctx, w, accessRequest, err)
                        return true
                }

                // Validate against allowed audiences list
                if err := server.ValidateAudienceAllowed(resource, h.config.AllowedAudiences); err != nil {
                        slog.Debug("resource not in allowed audiences", //nolint:gosec // G706: resource URI from token request
                                "resource", resource,
                                "error", err,
                        )
                        h.provider.WriteAccessError(ctx, w, accessRequest, err)
                        return true
                }

                slog.Debug("granting audience from resource parameter", //nolint:gosec // G706: resource URI from token request
                        "resource", resource,
                )
                accessRequest.GrantAudience(resource)
        } else if accessRequest.GetGrantTypes().ExactOne("authorization_code") && len(h.config.AllowedAudiences) == 1 {
                // No resource parameter provided (or provided as empty) during an authorization_code
                // exchange; default to the sole allowed audience. The len == 1 guard makes the
                // intended audience unambiguous and the index access safe. We restrict this defaulting
                // to authorization_code grants: for refresh_token grants, fosite already carries the
                // originally-granted audience forward through the session, so re-granting here would
                // conflict with fosite's audience matching strategy.
                slog.Debug("no resource parameter, defaulting to sole allowed audience",
                        "audience", h.config.AllowedAudiences[0],
                )
                accessRequest.GrantAudience(h.config.AllowedAudiences[0])
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package handlers

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "time"

        "github.com/google/uuid"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

// UserResolver handles finding or creating users based on provider identity.
// It manages the mapping between upstream provider subjects and internal user IDs.
type UserResolver struct {
        storage storage.UserStorage
}

// NewUserResolver creates a new UserResolver with the given storage.
func NewUserResolver(stor storage.UserStorage) *UserResolver {
        return &UserResolver{storage: stor}
}

// ResolveUser finds an existing user or creates a new one for the provider identity.
// Returns the user whose ID will be the "sub" claim in our JWTs.
//
// The resolution process:
// 1. Look up existing identity by (providerID, providerSubject)
// 2. If found, return the linked user
// 3. If not found, create a new user and link the identity
func (r *UserResolver) ResolveUser(
        ctx context.Context,
        providerID string,
        providerSubject string,
) (*storage.User, error) {
        if providerID == "" {
                return nil, errors.New("provider ID cannot be empty")
        }
        if providerSubject == "" {
                return nil, errors.New("provider subject cannot be empty")
        }

        // Try to find existing identity link
        identity, err := r.storage.GetProviderIdentity(ctx, providerID, providerSubject)
        if err != nil {
                if !errors.Is(err, storage.ErrNotFound) {
                        return nil, fmt.Errorf("failed to lookup provider identity: %w", err)
                }
                // No existing identity — create new user and link
                return r.createUserWithIdentity(ctx, providerID, providerSubject)
        }

        // Found existing identity, get the user
        user, err := r.storage.GetUser(ctx, identity.UserID)
        if err != nil {
                return nil, fmt.Errorf("identity exists but user not found: %w", err)
        }
        return user, nil
}

// createUserWithIdentity creates a new user and links the provider identity.
// This is called when no existing identity is found for the provider subject.
func (r *UserResolver) createUserWithIdentity(
        ctx context.Context,
        providerID string,
        providerSubject string,
) (*storage.User, error) {
        now := time.Now()

        user := &storage.User{
                ID:        uuid.New().String(),
                CreatedAt: now,
                UpdatedAt: now,
        }

        if err := r.storage.CreateUser(ctx, user); err != nil {
                return nil, fmt.Errorf("failed to create user: %w", err)
        }

        identity := &storage.ProviderIdentity{
                UserID:          user.ID,
                ProviderID:      providerID,
                ProviderSubject: providerSubject,
                LinkedAt:        now,
                LastUsedAt:      now,
        }

        if err := r.storage.CreateProviderIdentity(ctx, identity); err != nil {
                // Rollback user creation on identity link failure
                if deleteErr := r.storage.DeleteUser(ctx, user.ID); deleteErr != nil {
                        slog.Warn("failed to rollback user creation", "error", deleteErr)
                }
                return nil, fmt.Errorf("failed to link provider identity: %w", err)
        }

        slog.Debug("created new user with provider identity",
                "user_id", user.ID,
                "provider_id", providerID,
        )

        return user, nil
}

// UpdateLastAuthenticated updates the last authentication timestamp for a provider identity.
// This supports OIDC max_age parameter enforcement by tracking when users last authenticated.
// Errors are logged but not fatal - callers should continue with authorization.
func (r *UserResolver) UpdateLastAuthenticated(
        ctx context.Context,
        providerID string,
        providerSubject string,
) {
        if err := r.storage.UpdateProviderIdentityLastUsed(ctx, providerID, providerSubject, time.Now()); err != nil {
                slog.Warn("failed to update identity last used timestamp", "error", err)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package keys

// Config holds configuration for creating a KeyProvider.
// The caller is responsible for populating this from their own config source
// (environment variables, YAML files, flags, etc.).
type Config struct {
        // KeyDir is the directory containing PEM-encoded private key files.
        // All key filenames are relative to this directory.
        //
        // In Kubernetes deployments, this is typically a mounted Secret volume:
        //
        //        volumeMounts:
        //        - name: signing-keys
        //          mountPath: /etc/toolhive/keys
        KeyDir string

        // SigningKeyFile is the filename of the primary signing key (relative to KeyDir).
        // This key is used for signing new tokens.
        // If empty with KeyDir set, NewProviderFromConfig returns an error.
        // If both KeyDir and SigningKeyFile are empty, an ephemeral key is generated.
        SigningKeyFile string

        // FallbackKeyFiles are filenames of additional keys for verification (relative to KeyDir).
        // These keys are included in the JWKS endpoint for token verification but are NOT
        // used for signing new tokens.
        //
        // Key rotation (single replica): update SigningKeyFile to the new key and move
        // the old filename here. Tokens signed with old keys remain verifiable until
        // they expire.
        //
        // Key rotation (multiple replicas): to avoid a window where one replica signs
        // with a key not yet advertised by another replica's JWKS endpoint:
        //  1. Add the new key to FallbackKeyFiles and roll out to all replicas.
        //  2. Promote it to SigningKeyFile, move the old key to FallbackKeyFiles, roll out.
        //  3. Remove the old key from FallbackKeyFiles after its tokens have expired.
        FallbackKeyFiles []string
}

// NewProviderFromConfig creates a KeyProvider based on the configuration.
//
// Behavior:
//   - If KeyDir and SigningKeyFile are set: load keys from directory
//   - If both are empty: return GeneratingProvider (ephemeral key for development)
//   - If KeyDir is set but SigningKeyFile is empty: returns an error
func NewProviderFromConfig(cfg Config) (KeyProvider, error) {
        if cfg.KeyDir != "" {
                return NewFileProvider(cfg)
        }

        // Generate ephemeral key (development only)
        return NewGeneratingProvider(DefaultAlgorithm), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package keys

import (
        "context"
        "crypto"
        "crypto/ecdsa"
        "crypto/elliptic"
        "crypto/rand"
        "fmt"
        "log/slog"
        "path/filepath"
        "sync"
        "time"

        servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
)

//go:generate mockgen -destination=mocks/mock_provider.go -package=mocks -source=provider.go KeyProvider,PublicKeyProvider

// PublicKeyProvider provides public keys for JWT verification.
// Use this interface when a component only needs to verify tokens,
// not sign them, to avoid leaking private key access.
type PublicKeyProvider interface {
        // PublicKeys returns all public keys for the JWKS endpoint.
        // May return multiple keys during rotation periods.
        PublicKeys(ctx context.Context) ([]*PublicKeyData, error)
}

// KeyProvider provides signing keys for JWT operations.
// Implementations handle key sourcing (file, memory, generation).
// KeyProvider implicitly satisfies PublicKeyProvider.
type KeyProvider interface {
        // SigningKey returns the current signing key.
        // Returns ErrNoSigningKey if no key is available.
        SigningKey(ctx context.Context) (*SigningKeyData, error)

        // PublicKeys returns all public keys for the JWKS endpoint.
        // May return multiple keys during rotation periods.
        PublicKeys(ctx context.Context) ([]*PublicKeyData, error)
}

// FileProvider loads signing keys from PEM files in a directory.
// The signing key is used for signing new tokens.
// All keys (signing + fallback) are exposed via PublicKeys() for JWKS.
// Keys are loaded once at construction time; changes require restart.
type FileProvider struct {
        signingKey *SigningKeyData
        allKeys    []*SigningKeyData
}

// NewFileProvider creates a provider that loads keys from a directory.
// Config.SigningKeyFile is the primary key used for signing new tokens.
// Config.FallbackKeyFiles are loaded for JWKS verification (for key rotation).
// All keys are loaded immediately and validated.
// Supports RSA (PKCS1/PKCS8), ECDSA (SEC1/PKCS8), and Ed25519 keys.
func NewFileProvider(cfg Config) (*FileProvider, error) {
        if cfg.SigningKeyFile == "" {
                return nil, fmt.Errorf("signing key file is required")
        }

        // Load the primary signing key
        signingKeyPath := filepath.Join(cfg.KeyDir, cfg.SigningKeyFile)
        signingKey, err := loadKeyFromFile(signingKeyPath)
        if err != nil {
                return nil, fmt.Errorf("failed to load signing key: %w", err)
        }

        // Start with the signing key in allKeys
        allKeys := []*SigningKeyData{signingKey}

        // Load fallback keys for JWKS verification
        for _, filename := range cfg.FallbackKeyFiles {
                keyPath := filepath.Join(cfg.KeyDir, filename)
                key, err := loadKeyFromFile(keyPath)
                if err != nil {
                        return nil, fmt.Errorf("failed to load fallback key %s: %w", filename, err)
                }
                allKeys = append(allKeys, key)
        }

        return &FileProvider{
                signingKey: signingKey,
                allKeys:    allKeys,
        }, nil
}

// loadKeyFromFile loads a single key from a PEM file.
func loadKeyFromFile(keyPath string) (*SigningKeyData, error) {
        signer, err := servercrypto.LoadSigningKey(keyPath)
        if err != nil {
                return nil, err
        }

        params, err := servercrypto.DeriveSigningKeyParams(signer, "", "")
        if err != nil {
                return nil, fmt.Errorf("failed to derive key parameters: %w", err)
        }

        return &SigningKeyData{
                KeyID:     params.KeyID,
                Algorithm: params.Algorithm,
                Key:       params.Key,
                CreatedAt: time.Now(),
        }, nil
}

// SigningKey returns the primary signing key used for signing new tokens.
// Returns a copy to prevent external mutation of internal state.
func (p *FileProvider) SigningKey(_ context.Context) (*SigningKeyData, error) {
        return &SigningKeyData{
                KeyID:     p.signingKey.KeyID,
                Algorithm: p.signingKey.Algorithm,
                Key:       p.signingKey.Key,
                CreatedAt: p.signingKey.CreatedAt,
        }, nil
}

// PublicKeys returns public keys for all loaded keys (signing + additional).
// This enables verification of tokens signed with any of the loaded keys,
// supporting key rotation scenarios where old keys must remain valid.
func (p *FileProvider) PublicKeys(_ context.Context) ([]*PublicKeyData, error) {
        pubKeys := make([]*PublicKeyData, 0, len(p.allKeys))
        for _, key := range p.allKeys {
                pubKeys = append(pubKeys, &PublicKeyData{
                        KeyID:     key.KeyID,
                        Algorithm: key.Algorithm,
                        PublicKey: key.Key.Public(),
                        CreatedAt: key.CreatedAt,
                })
        }
        return pubKeys, nil
}

// GeneratingProvider generates an ephemeral key on first access.
// Suitable for development but NOT recommended for production.
// Generated keys are lost on restart, invalidating all issued tokens.
type GeneratingProvider struct {
        algorithm string
        mu        sync.Mutex
        key       *SigningKeyData
}

// NewGeneratingProvider creates a provider that generates an ephemeral key.
// The key is generated lazily on first SigningKey() call.
// If algorithm is empty, DefaultAlgorithm (ES256) is used.
func NewGeneratingProvider(algorithm string) *GeneratingProvider {
        if algorithm == "" {
                algorithm = DefaultAlgorithm
        }
        return &GeneratingProvider{algorithm: algorithm}
}

// SigningKey returns the signing key, generating one if needed.
// Thread-safe: uses mutex to ensure only one key is generated.
// Returns a copy to prevent external mutation of internal state.
func (p *GeneratingProvider) SigningKey(_ context.Context) (*SigningKeyData, error) {
        p.mu.Lock()
        defer p.mu.Unlock()

        if p.key != nil {
                return &SigningKeyData{
                        KeyID:     p.key.KeyID,
                        Algorithm: p.key.Algorithm,
                        Key:       p.key.Key,
                        CreatedAt: p.key.CreatedAt,
                }, nil
        }

        key, err := p.generateKey()
        if err != nil {
                return nil, err
        }

        slog.Warn("generated ephemeral signing key - tokens will be invalid after restart",
                "algorithm", key.Algorithm,
                "key_id", key.KeyID,
        )

        p.key = key
        return &SigningKeyData{
                KeyID:     p.key.KeyID,
                Algorithm: p.key.Algorithm,
                Key:       p.key.Key,
                CreatedAt: p.key.CreatedAt,
        }, nil
}

// PublicKeys returns the public key for JWKS.
// Generates the signing key if it hasn't been generated yet.
func (p *GeneratingProvider) PublicKeys(ctx context.Context) ([]*PublicKeyData, error) {
        key, err := p.SigningKey(ctx)
        if err != nil {
                return nil, err
        }
        return []*PublicKeyData{{
                KeyID:     key.KeyID,
                Algorithm: key.Algorithm,
                PublicKey: key.Key.Public(),
                CreatedAt: key.CreatedAt,
        }}, nil
}

func (p *GeneratingProvider) generateKey() (*SigningKeyData, error) {
        privateKey, err := generatePrivateKey(p.algorithm)
        if err != nil {
                return nil, fmt.Errorf("failed to generate signing key: %w", err)
        }

        keyID, err := servercrypto.DeriveKeyID(privateKey)
        if err != nil {
                return nil, fmt.Errorf("failed to derive key ID: %w", err)
        }

        return &SigningKeyData{
                KeyID:     keyID,
                Algorithm: p.algorithm,
                Key:       privateKey,
                CreatedAt: time.Now(),
        }, nil
}

// generatePrivateKey creates a new private key for the specified algorithm.
func generatePrivateKey(algorithm string) (crypto.Signer, error) {
        switch algorithm {
        case "ES256":
                return ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
        case "ES384":
                return ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
        case "ES512":
                return ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
        default:
                return nil, fmt.Errorf("unsupported algorithm for key generation: %s", algorithm)
        }
}

// Compile-time interface checks.
var (
        _ KeyProvider       = (*FileProvider)(nil)
        _ KeyProvider       = (*GeneratingProvider)(nil)
        _ PublicKeyProvider = (*FileProvider)(nil)
        _ PublicKeyProvider = (*GeneratingProvider)(nil)
)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package server

import (
        "context"

        "github.com/ory/fosite"
)

// clientAssertionAudienceConfig overrides the token URLs fosite accepts as the
// audience of a client assertion. fosite reads them only when authenticating
// clients; every other Configurator method comes from the embedded config.
type clientAssertionAudienceConfig struct {
        *fosite.Config
        tokenURLs []string
}

func (c clientAssertionAudienceConfig) GetTokenURLs(context.Context) []string {
        return c.tokenURLs
}

// newClientAssertionAudienceConfig wraps config so that client assertions may
// carry either the issuer or the token endpoint URL as audience.
func newClientAssertionAudienceConfig(config *fosite.Config) clientAssertionAudienceConfig {
        return clientAssertionAudienceConfig{
                Config:    config,
                tokenURLs: []string{config.AccessTokenIssuer, config.TokenURL},
        }
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
        "context"
        "crypto"
        "fmt"
        "log/slog"
        "net/url"
        "slices"
        "strings"
        "time"

        josev3 "github.com/go-jose/go-jose/v3"
        "github.com/go-jose/go-jose/v4"
        "github.com/ory/fosite"
        "github.com/spiffe/go-spiffe/v2/bundle/jwtbundle"
        "github.com/spiffe/go-spiffe/v2/bundle/x509bundle"

        servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
        "github.com/stacklok/toolhive/pkg/authserver/server/keys"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/networking"
)

// Token lifespan bounds for validation.
const (
        // MinAccessTokenLifespan is the minimum allowed access token lifetime.
        MinAccessTokenLifespan = 1 * time.Minute
        // MaxAccessTokenLifespan is the maximum allowed access token lifetime.
        MaxAccessTokenLifespan = 24 * time.Hour
        // MinRefreshTokenLifespan is the minimum allowed refresh token lifetime.
        MinRefreshTokenLifespan = 1 * time.Hour
        // MaxRefreshTokenLifespan is the maximum allowed refresh token lifetime (30 days).
        MaxRefreshTokenLifespan = 30 * 24 * time.Hour
        // MinAuthCodeLifespan is the minimum allowed authorization code lifetime.
        MinAuthCodeLifespan = 30 * time.Second
        // MaxAuthCodeLifespan is the maximum allowed authorization code lifetime (RFC 6749 recommends 10 min max).
        MaxAuthCodeLifespan = 10 * time.Minute
        // MaxAssertionLifespan is the maximum allowed exp-now duration for a
        // private_key_jwt client_assertion. fosite persists the replay marker
        // (Storage.SetClientAssertionJWT/ConsumeAssertionJWT) with a TTL derived
        // from the assertion's own exp claim, and neither fosite nor RFC 7523
        // bounds that value — a client could otherwise set exp decades out and
        // grow replay-tracking storage indefinitely. A client assertion
        // authenticates a single token request, so it has no legitimate reason
        // to outlive that request by more than a small clock-skew margin.
        MaxAssertionLifespan = 5 * time.Minute
        // MaxAssertionJTILength bounds the byte length of a private_key_jwt
        // client_assertion's jti before it's persisted as a replay marker.
        // Comfortably fits UUIDs/ULIDs and typical nonce formats with headroom
        // while capping per-entry storage cost.
        MaxAssertionJTILength = 256
)

// DefaultDeviceCodeInterval is RFC 8628's suggested default minimum interval
// between device-flow polls of the token endpoint, applied whenever
// AuthorizationServerConfig.DeviceCodeInterval / AuthorizationServerParams.DeviceCodeInterval
// is left at its zero value. A zero interval would defeat the interval's
// purpose entirely, so it is never used as the enforced minimum.
const DefaultDeviceCodeInterval = 5 * time.Second

// Values for AuthorizationServerConfig.AccessTokenType and AuthorizationServerParams.AccessTokenType.
const (
        // AccessTokenTypeJWT is the default JOSE "typ" header for JWT access tokens, the value fosite writes.
        AccessTokenTypeJWT = "JWT"
        // AccessTokenTypeATJWT is the RFC 9068 §2.1 media type for JWT access tokens.
        AccessTokenTypeATJWT = "at+jwt"
)

// ValidateAccessTokenType accepts the empty string (meaning AccessTokenTypeJWT),
// AccessTokenTypeJWT and AccessTokenTypeATJWT. Matching is exact and case-sensitive.
func ValidateAccessTokenType(typ string) error {
        switch typ {
        case "", AccessTokenTypeJWT, AccessTokenTypeATJWT:
                return nil
        default:
                return fmt.Errorf("unsupported access token type %q: must be %q or %q", typ, AccessTokenTypeJWT, AccessTokenTypeATJWT)
        }
}

// AuthorizationServerConfig wraps fosite.Config with additional configuration
// for JWT signing and other extensions.
type AuthorizationServerConfig struct {
        *fosite.Config
        SigningKey  *jose.JSONWebKey
        SigningJWKS *jose.JSONWebKeySet
        // AllowedAudiences is the list of valid resource URIs that tokens can be issued for.
        // Per RFC 8707, the "resource" parameter in token requests is validated against this list.
        // Security: An empty list means NO audiences are permitted (secure default).
        AllowedAudiences []string
        // SubjectIdentifierAudiences controls disclosure of the RFC 9493 sub_id claim and the
        // https://toolhive.dev/upstream_subject claim on login-derived access tokens.
        // It is not an access-control or audience-grant setting.
        SubjectIdentifierAudiences []string
        // AccessTokenType is the JOSE "typ" header set on JWT access tokens. It is always
        // AccessTokenTypeJWT or AccessTokenTypeATJWT, never empty: NewAuthorizationServerConfig
        // sets an empty value to AccessTokenTypeJWT. It applies to access tokens only.
        AccessTokenType string
        // ScopesSupported lists the OAuth 2.0 scope values advertised in discovery documents.
        // This is advertised in /.well-known/openid-configuration and
        // /.well-known/oauth-authorization-server discovery endpoints.
        ScopesSupported []string
        // BaselineClientScopes is a baseline set of OAuth 2.0 scopes the DCR handler
        // unions into every newly registered client's scope set. All entries are
        // guaranteed to be a subset of ScopesSupported.
        BaselineClientScopes []string
        // AuthorizationEndpointBaseURL overrides the base URL for the authorization_endpoint
        // in the discovery document. When empty, defaults to the issuer (AccessTokenIssuer).
        AuthorizationEndpointBaseURL string
        // CIMDEnabled indicates that the CIMD storage decorator is active. When true,
        // the discovery document advertises client_id_metadata_document_supported.
        CIMDEnabled bool
        // AllowConfidentialClientRegistration permits DCR of confidential clients. When true,
        // the DCR handler accepts client_secret_basic / client_secret_post and the
        // discovery document advertises them in token_endpoint_auth_methods_supported.
        AllowConfidentialClientRegistration bool
        // AllowPrivateKeyJWTRegistration permits DCR of clients using
        // private_key_jwt authentication. This is independent of confidential-client
        // registration and is used by DCR validation and discovery advertisement.
        AllowPrivateKeyJWTRegistration bool
        // HasStaticDelegateClients indicates whether any pre-provisioned confidential
        // delegate client is registered at startup. Discovery advertises client-secret
        // authentication methods when this is true.
        HasStaticDelegateClients bool
        // InsecureAllowHTTP permits a non-loopback HTTP issuer. It is incompatible
        // with confidential clients.
        InsecureAllowHTTP bool
        // InsecureAllowConfidentialOverLoopbackHTTP explicitly permits confidential
        // clients with a loopback HTTP issuer.
        InsecureAllowConfidentialOverLoopbackHTTP bool
        // ForceConfidentialRedirectURIs lists redirect URIs that the DCR handler
        // always registers as confidential clients, overriding a requested "none"
        // auth method. See authserver.Config.ForceConfidentialRedirectURIs for the
        // full semantics.
        ForceConfidentialRedirectURIs []string
        // TokenExchangeEnabled indicates whether RFC 8693 is registered and advertised.
        TokenExchangeEnabled bool
        // JWTBearerGrantEnabled indicates that at least one trusted issuer has the
        // RFC 7523 JWT-bearer grant configured. Discovery advertises
        // urn:ietf:params:oauth:grant-type:jwt-bearer in grant_types_supported
        // only when this is true, mirroring how the grant itself is only
        // registered with fosite when true (see buildProvider).
        JWTBearerGrantEnabled bool
        // DeviceFlowEnabled indicates that the RFC 8628 device authorization
        // grant is registered and advertised. Discovery advertises
        // urn:ietf:params:oauth:grant-type:device_code in grant_types_supported,
        // and POST /oauth/device_authorization is mounted, only when this is true.
        DeviceFlowEnabled bool
        // DeviceCodeInterval is the minimum time a device-flow client must wait
        // between polls of the token endpoint (RFC 8628 Section 3.5), and the
        // value advertised in the device authorization response's "interval"
        // field. Zero is replaced by DefaultDeviceCodeInterval wherever it would
        // otherwise be used as the enforced minimum.
        DeviceCodeInterval time.Duration
        // SPIFFEClientResolver resolves a verified SPIFFE identity to its
        // configured OAuth client for the SPIFFE client-authentication strategy.
        // Nil when no SPIFFE trust is configured; package server cannot import
        // the concrete association registry and storage that back this resolver
        // (authserver imports server), so the caller supplies it as a closure.
        SPIFFEClientResolver SPIFFEClientResolver
        // SPIFFEX509BundleSource provides X.509 bundles for verifying SPIFFE
        // X.509-SVID client certificates. Not yet read by this package: carried
        // here, copied through from AuthorizationServerParams, so the X.509 and
        // JWT SPIFFE client-authentication arms land on a shared field instead of
        // each independently extending this struct.
        SPIFFEX509BundleSource x509bundle.Source
        // SPIFFEJWTBundleSource provides JWT bundles for verifying SPIFFE
        // JWT-SVID client assertions, read by newSPIFFEClientAuthenticationStrategy.
        SPIFFEJWTBundleSource jwtbundle.Source
}

// Factory is a constructor which is used to create an OAuth2 endpoint handler.
// NewAuthorizationServer handles consuming the new struct and attaching it
// to the parts of the config that it implements.
//
// The strategy parameter is typed as any because fosite uses different strategy
// interfaces for different flows (e.g., oauth2.CoreStrategy, openid.OpenIDConnectTokenStrategy)
// that do not share a common base interface.
type Factory func(config *AuthorizationServerConfig, storage fosite.Storage, strategy any) (any, error)

// AuthorizationServerParams contains the configuration needed to create an AuthorizationServerConfig.
// This is a minimal subset of the authserver.Config fields needed for OAuth2.
type AuthorizationServerParams struct {
        Issuer               string
        AccessTokenLifespan  time.Duration
        RefreshTokenLifespan time.Duration
        AuthCodeLifespan     time.Duration
        HMACSecrets          *servercrypto.HMACSecrets
        SigningKeyID         string
        SigningKeyAlgorithm  string
        SigningKey           crypto.Signer
        // AdditionalPublicKeys are published in the JWKS alongside the primary
        // signing key, in public form only, so that tokens signed with a
        // previous key remain verifiable during rotation. Typically sourced from
        // keys.KeyProvider.PublicKeys(ctx), which includes the primary signing
        // key itself; the entry matching SigningKeyID is skipped to avoid
        // publishing it twice.
        AdditionalPublicKeys []*keys.PublicKeyData
        // AllowedAudiences is the list of valid resource URIs that tokens can be issued for.
        // Per RFC 8707, the "resource" parameter in token requests is validated against this list.
        // Security: An empty list means NO audiences are permitted (secure default).
        AllowedAudiences []string
        // SubjectIdentifierAudiences controls disclosure of the RFC 9493 sub_id claim and the
        // https://toolhive.dev/upstream_subject claim on login-derived access tokens.
        // It is not an access-control or audience-grant setting.
        SubjectIdentifierAudiences []string
        // AccessTokenType sets the JOSE "typ" header on every JWT access token this server issues;
        // one of AccessTokenTypeJWT (default) or AccessTokenTypeATJWT. An empty value means
        // AccessTokenTypeJWT. It affects access tokens only, never other JWTs the server signs.
        AccessTokenType string
        // ScopesSupported lists the OAuth 2.0 scope values advertised in discovery documents.
        ScopesSupported []string
        // BaselineClientScopes is a baseline set of OAuth 2.0 scopes the DCR handler
        // unions into every newly registered client's scope set. All entries are
        // guaranteed to be a subset of ScopesSupported.
        BaselineClientScopes []string
        // AuthorizationEndpointBaseURL overrides the base URL for the authorization_endpoint
        // in the discovery document. When empty, defaults to Issuer.
        AuthorizationEndpointBaseURL string
        // CIMDEnabled indicates that the CIMD storage decorator is active. When true,
        // the discovery document advertises client_id_metadata_document_supported.
        CIMDEnabled bool
        // AllowConfidentialClientRegistration permits DCR of confidential clients. When true,
        // the DCR handler accepts client_secret_basic / client_secret_post and the
        // discovery document advertises them in token_endpoint_auth_methods_supported.
        AllowConfidentialClientRegistration bool
        // AllowPrivateKeyJWTRegistration permits DCR of clients using
        // private_key_jwt authentication. This is independent of confidential-client
        // registration and is used by DCR validation and discovery advertisement.
        AllowPrivateKeyJWTRegistration bool
        // HasStaticDelegateClients indicates whether any pre-provisioned confidential
        // delegate client is registered at startup. Discovery advertises client-secret
        // authentication methods when this is true.
        HasStaticDelegateClients bool
        // InsecureAllowHTTP permits a non-loopback HTTP issuer. It is incompatible
        // with confidential clients.
        InsecureAllowHTTP bool
        // InsecureAllowConfidentialOverLoopbackHTTP explicitly permits confidential
        // clients with a loopback HTTP issuer.
        InsecureAllowConfidentialOverLoopbackHTTP bool
        // ForceConfidentialRedirectURIs lists redirect URIs that the DCR handler
        // always registers as confidential clients, overriding a requested "none"
        // auth method. See authserver.Config.ForceConfidentialRedirectURIs for the
        // full semantics.
        ForceConfidentialRedirectURIs []string
        // DisableTokenExchange prevents RFC 8693 registration and advertisement.
        // The zero value preserves the historical enabled behavior.
        DisableTokenExchange bool
        // JWTBearerGrantEnabled indicates that at least one trusted issuer has the
        // RFC 7523 JWT-bearer grant configured. See AuthorizationServerConfig's
        // field of the same name.
        JWTBearerGrantEnabled bool
        // DeviceFlowEnabled indicates that the RFC 8628 device authorization
        // grant should be registered and advertised. See AuthorizationServerConfig's
        // field of the same name.
        DeviceFlowEnabled bool
        // DeviceCodeInterval is the minimum time a device-flow client must wait
        // between polls of the token endpoint. See AuthorizationServerConfig's
        // field of the same name.
        DeviceCodeInterval time.Duration
        // SPIFFEClientResolver resolves a verified SPIFFE identity to its
        // configured OAuth client. Nil when no SPIFFE trust is configured.
        // See the identically named field on AuthorizationServerConfig.
        SPIFFEClientResolver SPIFFEClientResolver
        // SPIFFEX509BundleSource provides X.509 bundles for verifying SPIFFE
        // X.509-SVID client certificates. Not yet read by this package: threaded
        // through here so the X.509 and JWT SPIFFE client-authentication arms
        // land on a shared field instead of each independently extending this
        // struct.
        SPIFFEX509BundleSource x509bundle.Source
        // SPIFFEJWTBundleSource provides JWT bundles for verifying SPIFFE
        // JWT-SVID client assertions, copied through to AuthorizationServerConfig.
        SPIFFEJWTBundleSource jwtbundle.Source
}

// validateIssuerURL validates that the issuer is a valid URL with http or https scheme
// and no trailing slash. Following the pattern from pkg/config/validation.go.
func validateIssuerURL(issuer string) error {
        if issuer == "" {
                return fmt.Errorf("issuer is required")
        }

        parsedURL, err := url.Parse(issuer)
        if err != nil {
                return fmt.Errorf("issuer is not a valid URL: %w", err)
        }

        if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" {
                return fmt.Errorf("issuer must use http or https scheme")
        }

        if parsedURL.Hostname() == "" {
                return fmt.Errorf("issuer must have a host")
        }

        if strings.HasSuffix(issuer, "/") {
                return fmt.Errorf("issuer must not have a trailing slash")
        }

        return nil
}

// validateAllowedAudiences validates that all allowed audiences are valid RFC 8707 URIs.
func validateAllowedAudiences(audiences []string) error {
        for i, aud := range audiences {
                if err := ValidateAudienceURI(aud); err != nil {
                        return fmt.Errorf("allowed audience [%d] %q is invalid: %w", i, aud, err)
                }
        }
        return nil
}

// validateHMACSecrets validates that all HMAC secrets meet the minimum length requirement.
func validateHMACSecrets(secrets *servercrypto.HMACSecrets) error {
        if secrets == nil {
                return fmt.Errorf("HMAC secrets are required")
        }
        if len(secrets.Current) < servercrypto.MinSecretLength {
                return fmt.Errorf("current HMAC secret must be at least %d bytes", servercrypto.MinSecretLength)
        }
        for i, rotated := range secrets.Rotated {
                if len(rotated) < servercrypto.MinSecretLength {
                        return fmt.Errorf("rotated HMAC secret [%d] must be at least %d bytes", i, servercrypto.MinSecretLength)
                }
        }
        return nil
}

// validateTokenLifespans validates that token lifespans are within allowed bounds.
func validateTokenLifespans(cfg *AuthorizationServerParams) error {
        if cfg.AccessTokenLifespan < MinAccessTokenLifespan || cfg.AccessTokenLifespan > MaxAccessTokenLifespan {
                return fmt.Errorf("access token lifespan must be between %v and %v", MinAccessTokenLifespan, MaxAccessTokenLifespan)
        }
        if cfg.RefreshTokenLifespan < MinRefreshTokenLifespan || cfg.RefreshTokenLifespan > MaxRefreshTokenLifespan {
                return fmt.Errorf("refresh token lifespan must be between %v and %v", MinRefreshTokenLifespan, MaxRefreshTokenLifespan)
        }
        if cfg.AuthCodeLifespan < MinAuthCodeLifespan || cfg.AuthCodeLifespan > MaxAuthCodeLifespan {
                return fmt.Errorf("authorization code lifespan must be between %v and %v", MinAuthCodeLifespan, MaxAuthCodeLifespan)
        }
        return nil
}

// ValidateConfidentialClientTransport rejects cleartext HTTP configurations
// when any confidential client is enabled, whether it is admitted through DCR
// or statically declared.
func ValidateConfidentialClientTransport(
        allowConfidential, insecureAllowHTTP bool,
        issuer string, insecureAllowConfidentialOverLoopbackHTTP bool,
) error {
        if !allowConfidential {
                return nil
        }
        if insecureAllowHTTP {
                return fmt.Errorf("allow_confidential_client_registration cannot be combined with insecure_allow_http: " +
                        "confidential clients would send secrets over cleartext HTTP")
        }
        parsed, err := url.Parse(issuer)
        if err != nil {
                return fmt.Errorf("confidential clients require a valid issuer URL")
        }
        if parsed.Scheme != "http" {
                return nil
        }
        // Mirrors the query/fragment/userinfo shape validateIssuerURLCore enforces
        // (pkg/authserver/config.go) for the same reasons: an issuer identifier
        // with any of these is malformed under RFC 8414 §2 regardless of scheme,
        // and this function is the sole transport authority for the loopback
        // opt-in path once a caller reaches it directly.
        if parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
                return fmt.Errorf("confidential clients require a valid issuer URL")
        }
        if insecureAllowConfidentialOverLoopbackHTTP && !networking.IsLocalhost(parsed.Host) {
                return fmt.Errorf("allow_confidential_client_registration cannot use the loopback HTTP opt-in with a non-loopback issuer: " +
                        "confidential clients would send secrets over cleartext HTTP")
        }
        if !insecureAllowConfidentialOverLoopbackHTTP && networking.IsLocalhost(parsed.Host) {
                return fmt.Errorf("allow_confidential_client_registration cannot be combined with a plain-HTTP loopback issuer unless " +
                        "insecure_allow_confidential_over_loopback_http is set: confidential clients would send secrets over cleartext HTTP")
        }
        if !networking.IsLocalhost(parsed.Host) {
                return fmt.Errorf("allow_confidential_client_registration cannot use a plain-HTTP non-loopback issuer: " +
                        "confidential clients would send secrets over cleartext HTTP")
        }
        return nil
}

// validateParams validates all fields on AuthorizationServerParams.
func validateParams(cfg *AuthorizationServerParams) error {
        if err := validateIssuerURL(cfg.Issuer); err != nil {
                return err
        }
        if cfg.SigningKeyID == "" {
                return fmt.Errorf("signing key ID is required")
        }
        if cfg.SigningKeyAlgorithm == "" {
                return fmt.Errorf("signing key algorithm is required")
        }
        if cfg.SigningKey == nil {
                return fmt.Errorf("signing key is required")
        }
        if err := validateHMACSecrets(cfg.HMACSecrets); err != nil {
                return err
        }
        if err := servercrypto.ValidateAlgorithmForKey(cfg.SigningKeyAlgorithm, cfg.SigningKey); err != nil {
                return fmt.Errorf("invalid signing configuration: %w", err)
        }
        if err := validateTokenLifespans(cfg); err != nil {
                return err
        }
        if cfg.AuthorizationEndpointBaseURL != "" {
                if err := validateIssuerURL(cfg.AuthorizationEndpointBaseURL); err != nil {
                        return fmt.Errorf("authorization endpoint base URL: %w", err)
                }
        }
        if err := validateAllowedAudiences(cfg.AllowedAudiences); err != nil {
                return err
        }
        if err := ValidateSubjectIdentifierAudiences(cfg.SubjectIdentifierAudiences, cfg.AllowedAudiences); err != nil {
                return err
        }
        if err := ValidateAccessTokenType(cfg.AccessTokenType); err != nil {
                return err
        }
        if err := ValidateConfidentialClientTransport(
                cfg.AllowConfidentialClientRegistration || cfg.HasStaticDelegateClients,
                cfg.InsecureAllowHTTP, cfg.Issuer, cfg.InsecureAllowConfidentialOverLoopbackHTTP,
        ); err != nil {
                return err
        }
        // Defense-in-depth: re-check the baseline-⊆-scopes_supported invariant.
        // RunConfig.Validate performs the same check at the operator-supplied
        // wire-format boundary; this gate covers callers that construct
        // AuthorizationServerParams programmatically and bypass that path.
        return registration.ValidateScopeSubset(cfg.BaselineClientScopes, cfg.ScopesSupported, "baseline_client_scopes")
}

// NewAuthorizationServerConfig creates an AuthorizationServerConfig from the provided configuration.
func NewAuthorizationServerConfig(cfg *AuthorizationServerParams) (*AuthorizationServerConfig, error) {
        if cfg == nil {
                return nil, fmt.Errorf("config is required")
        }
        if err := validateParams(cfg); err != nil {
                return nil, err
        }

        if len(cfg.BaselineClientScopes) > 0 {
                slog.Info("DCR registrations will be auto-granted baseline scopes",
                        "scopes", cfg.BaselineClientScopes,
                )
        }

        // Build JWK from signing key
        jwk := jose.JSONWebKey{
                Key:       cfg.SigningKey,
                KeyID:     cfg.SigningKeyID,
                Algorithm: cfg.SigningKeyAlgorithm,
                Use:       "sig",
        }

        // The JWKS published at /.well-known/jwks.json must include every
        // fallback key still in rotation, not just the primary signing key,
        // or the documented key-rotation procedure never opens its overlap
        // window (#6451). The primary key is listed first; any additional
        // key matching its KeyID is skipped to avoid publishing it twice.
        signingJWKS := &jose.JSONWebKeySet{Keys: []jose.JSONWebKey{jwk}}
        for _, pubKey := range cfg.AdditionalPublicKeys {
                if pubKey.KeyID == cfg.SigningKeyID {
                        continue
                }
                signingJWKS.Keys = append(signingJWKS.Keys, jose.JSONWebKey{
                        Key:       pubKey.PublicKey,
                        KeyID:     pubKey.KeyID,
                        Algorithm: pubKey.Algorithm,
                        Use:       "sig",
                })
        }

        fositeConfig := &fosite.Config{
                AccessTokenIssuer:              cfg.Issuer,
                AccessTokenLifespan:            cfg.AccessTokenLifespan,
                RefreshTokenLifespan:           cfg.RefreshTokenLifespan,
                AuthorizeCodeLifespan:          cfg.AuthCodeLifespan,
                GlobalSecret:                   cfg.HMACSecrets.Current,
                RotatedGlobalSecrets:           cfg.HMACSecrets.Rotated,
                TokenURL:                       cfg.Issuer + "/oauth/token",
                EnforcePKCE:                    true,
                EnablePKCEPlainChallengeMethod: false, // Only allow S256 per MCP specification
                // ScopeStrategy validates requested scopes against client's registered scopes.
                // ExactScopeStrategy requires exact matches (no wildcards) for security.
                // This prevents clients from requesting scopes beyond what they registered with.
                ScopeStrategy: fosite.ExactScopeStrategy,
                // AudienceMatchingStrategy and JWKSFetcherStrategy are set explicitly
                // because fosite's getters lazily assign their defaults on first use,
                // an unsynchronized write that races between concurrent requests.
                AudienceMatchingStrategy: fosite.DefaultAudienceMatchingStrategy,
                JWKSFetcherStrategy:      unsupportedJWKSFetcher{},
                // ClientSecretsHasher compares client secrets at the token endpoint.
                // Plain SHA-256 is correct here: the secrets are 256 bits of CSPRNG
                // output, never client-chosen, so a password-stretching KDF protects
                // nothing and would hand unauthenticated callers a CPU-amplification
                // lever (one bcrypt verify per wrong-secret attempt). Must match the
                // hasher registration.New applies at client construction.
                ClientSecretsHasher: registration.SHA256Hasher,
        }

        accessTokenType := func() string {
                if cfg.AccessTokenType == "" {
                        return AccessTokenTypeJWT
                }
                return cfg.AccessTokenType
        }()

        return &AuthorizationServerConfig{
                Config:                                    fositeConfig,
                SigningKey:                                &jwk,
                SigningJWKS:                               signingJWKS,
                AllowedAudiences:                          cfg.AllowedAudiences,
                SubjectIdentifierAudiences:                cfg.SubjectIdentifierAudiences,
                AccessTokenType:                           accessTokenType,
                ScopesSupported:                           cfg.ScopesSupported,
                BaselineClientScopes:                      cfg.BaselineClientScopes,
                AuthorizationEndpointBaseURL:              cfg.AuthorizationEndpointBaseURL,
                CIMDEnabled:                               cfg.CIMDEnabled,
                AllowConfidentialClientRegistration:       cfg.AllowConfidentialClientRegistration,
                AllowPrivateKeyJWTRegistration:            cfg.AllowPrivateKeyJWTRegistration,
                HasStaticDelegateClients:                  cfg.HasStaticDelegateClients,
                InsecureAllowHTTP:                         cfg.InsecureAllowHTTP,
                InsecureAllowConfidentialOverLoopbackHTTP: cfg.InsecureAllowConfidentialOverLoopbackHTTP,
                ForceConfidentialRedirectURIs:             cfg.ForceConfidentialRedirectURIs,
                TokenExchangeEnabled:                      !cfg.DisableTokenExchange,
                JWTBearerGrantEnabled:                     cfg.JWTBearerGrantEnabled,
                DeviceFlowEnabled:                         cfg.DeviceFlowEnabled,
                DeviceCodeInterval:                        cfg.DeviceCodeInterval,
                SPIFFEClientResolver:                      cfg.SPIFFEClientResolver,
                SPIFFEX509BundleSource:                    cfg.SPIFFEX509BundleSource,
                SPIFFEJWTBundleSource:                     cfg.SPIFFEJWTBundleSource,
        }, nil
}

// unsupportedJWKSFetcher is the fosite JWKS fetcher for this server. Client
// keys are only accepted inline: DCR rejects jwks_uri and storage drops it, so
// fosite has nothing to fetch. Unlike fosite's default fetcher, it starts no
// cache goroutines, and if a client with a jwks_uri ever reached client
// authentication it fails closed instead of fetching an arbitrary URL.
type unsupportedJWKSFetcher struct{}

// Resolve always rejects the client, because remote client key sets are not
// supported. fosite's interface uses go-jose v3.
func (unsupportedJWKSFetcher) Resolve(context.Context, string, bool) (*josev3.JSONWebKeySet, error) {
        return nil, fosite.ErrInvalidClient.WithHint("Client JSON Web Key Sets must be registered inline; jwks_uri is not supported.")
}

// NewAuthorizationServer creates a new fosite OAuth2Provider with the given configuration,
// storage, strategy, and endpoint handler factories.
func NewAuthorizationServer(
        config *AuthorizationServerConfig,
        storage fosite.Storage,
        strategy any,
        factories ...Factory,
) (fosite.OAuth2Provider, error) {
        if err := ValidateConfidentialClientTransport(
                config.AllowConfidentialClientRegistration || config.HasStaticDelegateClients,
                config.InsecureAllowHTTP,
                config.AccessTokenIssuer,
                config.InsecureAllowConfidentialOverLoopbackHTTP,
        ); err != nil {
                return nil, err
        }
        providerConfig := *config
        fositeConfig := *config.Config
        fositeConfig.AuthorizeEndpointHandlers = slices.Clone(fositeConfig.AuthorizeEndpointHandlers)
        fositeConfig.TokenEndpointHandlers = slices.Clone(fositeConfig.TokenEndpointHandlers)
        fositeConfig.TokenIntrospectionHandlers = slices.Clone(fositeConfig.TokenIntrospectionHandlers)
        fositeConfig.RevocationHandlers = slices.Clone(fositeConfig.RevocationHandlers)
        fositeConfig.PushedAuthorizeEndpointHandlers = slices.Clone(fositeConfig.PushedAuthorizeEndpointHandlers)
        providerConfig.Config = &fositeConfig
        provider := fosite.NewOAuth2Provider(storage, newClientAssertionAudienceConfig(&fositeConfig))
        // Select the fallback after constructing the provider because fosite's
        // built-in strategy is a provider method. Preserve a caller-supplied strategy
        // when present, then install the SPIFFE dispatcher on the provider-local copy.
        defaultStrategy := func() fosite.ClientAuthenticationStrategy {
                if fositeConfig.ClientAuthenticationStrategy != nil {
                        return fositeConfig.ClientAuthenticationStrategy
                }
                return provider.DefaultClientAuthenticationStrategy
        }()
        fositeConfig.ClientAuthenticationStrategy = newSPIFFEClientAuthenticationStrategy(
                defaultStrategy, config.GetAccessTokenIssuer(),
                config.SPIFFEJWTBundleSource, providerConfig.SPIFFEClientResolver,
        )

        for _, factory := range factories {
                result, err := factory(&providerConfig, storage, strategy)
                if err != nil {
                        return nil, fmt.Errorf("authorization server factory failed: %w", err)
                }

                var matched bool

                if ah, ok := result.(fosite.AuthorizeEndpointHandler); ok {
                        fositeConfig.AuthorizeEndpointHandlers.Append(ah)
                        matched = true
                }

                if th, ok := result.(fosite.TokenEndpointHandler); ok {
                        fositeConfig.TokenEndpointHandlers.Append(th)
                        matched = true
                }

                if ti, ok := result.(fosite.TokenIntrospector); ok {
                        fositeConfig.TokenIntrospectionHandlers.Append(ti)
                        matched = true
                }

                if rh, ok := result.(fosite.RevocationHandler); ok {
                        fositeConfig.RevocationHandlers.Append(rh)
                        matched = true
                }

                if ph, ok := result.(fosite.PushedAuthorizeEndpointHandler); ok {
                        fositeConfig.PushedAuthorizeEndpointHandlers.Append(ph)
                        matched = true
                }

                if result != nil && !matched {
                        return nil, fmt.Errorf("authorization server factory returned unrecognized handler type %T", result)
                }
        }

        return provider, nil
}

// GetSigningKey returns the config's signing key.
func (c *AuthorizationServerConfig) GetSigningKey(_ context.Context) *jose.JSONWebKey {
        return c.SigningKey
}

// GetPrivateSigningJWKS returns the config's signing JWKS containing private keys.
//
// WARNING: This JWKS contains PRIVATE key material and MUST NOT be exposed publicly.
// Use PublicJWKS() for the /.well-known/jwks.json endpoint.
func (c *AuthorizationServerConfig) GetPrivateSigningJWKS(_ context.Context) *jose.JSONWebKeySet {
        return c.SigningJWKS
}

// PublicJWKS returns a copy of the JWKS containing only public keys.
func (c *AuthorizationServerConfig) PublicJWKS() *jose.JSONWebKeySet {
        if c.SigningJWKS == nil {
                return nil
        }

        publicJWKS := &jose.JSONWebKeySet{
                Keys: make([]jose.JSONWebKey, 0, len(c.SigningJWKS.Keys)),
        }

        for _, key := range c.SigningJWKS.Keys {
                publicKey := key.Public()
                publicJWKS.Keys = append(publicJWKS.Keys, publicKey)
        }

        return publicJWKS
}

// GetAccessTokenIssuer returns the issuer URL for access tokens.
// This is an adapter method that wraps the embedded fosite.Config method.
func (c *AuthorizationServerConfig) GetAccessTokenIssuer() string {
        return c.AccessTokenIssuer
}

// GetAuthorizationEndpointBaseURL returns the base URL for the authorization endpoint.
// If AuthorizationEndpointBaseURL is set, it is returned; otherwise falls back to the issuer.
func (c *AuthorizationServerConfig) GetAuthorizationEndpointBaseURL() string {
        if c.AuthorizationEndpointBaseURL != "" {
                return c.AuthorizationEndpointBaseURL
        }
        return c.GetAccessTokenIssuer()
}

// GetAuthorizeCodeLifespan returns the lifetime for authorization codes.
// This is an adapter method that wraps the embedded fosite.Config method.
func (c *AuthorizationServerConfig) GetAuthorizeCodeLifespan() time.Duration {
        return c.AuthorizeCodeLifespan
}

// GetAccessTokenLifespan returns the lifetime for access tokens.
// This is an adapter method that wraps the embedded fosite.Config method.
func (c *AuthorizationServerConfig) GetAccessTokenLifespan() time.Duration {
        return c.AccessTokenLifespan
}

// GetRefreshTokenLifespan returns the lifetime for refresh tokens.
// This is an adapter method that wraps the embedded fosite.Config method.
func (c *AuthorizationServerConfig) GetRefreshTokenLifespan() time.Duration {
        return c.RefreshTokenLifespan
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package registration provides OAuth client types and utilities, including
// RFC 8252 compliant loopback redirect URI support for native OAuth clients.
package registration

import (
        "context"
        "crypto/rand"
        "encoding/base64"
        "fmt"
        "log/slog"
        "net"
        "net/url"
        "slices"
        "strings"

        "github.com/go-jose/go-jose/v3"
        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// RegisteredLoopbackRedirectURI returns the registered redirect URI of c that
// requestedURI matches -- exactly, or under RFC 8252 Section 7.3 loopback
// dynamic-port rules -- or ("", false) if none matches.
//
// Loopback dynamic-port matching is restricted to public clients: RFC 8252
// loopback redirects are a native-app pattern and a confidential client must
// never get dynamic-port flexibility on its registered redirect_uri. Keeping
// that guard here rather than in a per-client method means no storage backend
// can reconstruct a client that silently skips it.
func RegisteredLoopbackRedirectURI(c fosite.Client, requestedURI string) (string, bool) {
        if !c.IsPublic() {
                return "", false
        }
        return matchLoopbackRedirectURI(c.GetRedirectURIs(), requestedURI)
}

// IsLocalhostHostname reports whether host (as returned by url.Hostname()) is
// the string "localhost", matched case-insensitively -- the same definition
// hostnamesMatch and isLoopbackHostname use. Exported so callers outside this
// package that need to know "is this specifically localhost, not an IP
// loopback literal" (e.g. deciding whether fosite's own IP-literal-only
// loopback matching already handles a redirect_uri, or whether it needs this
// package's help) share this package's one definition instead of
// re-implementing it and risking drift.
func IsLocalhostHostname(host string) bool {
        return strings.EqualFold(host, "localhost")
}

// matchLoopbackRedirectURI returns the registered URI (from registeredURIs)
// that requestedURI matches, either exactly or under RFC 8252 Section 7.3
// loopback dynamic-port rules, or ("", false) if none matches.
//
// Exact matches take precedence over loopback matches: a full pass over
// registeredURIs checks for an exact match first, and only then falls back to
// loopback matching. Without this, a client registered with both
// "http://localhost/callback" and "http://localhost:54321/callback" would
// have a request for the second rewritten to the first, because loopback
// matching against the first entry succeeds before the loop ever reaches the
// second entry's exact match. A client that registered a specific port must
// not have its request silently redirected to a different registered entry.
func matchLoopbackRedirectURI(registeredURIs []string, requestedURI string) (string, bool) {
        if slices.Contains(registeredURIs, requestedURI) {
                return requestedURI, true
        }
        for _, registeredURI := range registeredURIs {
                if matchesAsLoopback(requestedURI, registeredURI) {
                        return registeredURI, true
                }
        }
        return "", false
}

// DefaultScopes are the default OAuth 2.0 scopes for registered clients.
// Includes offline_access to enable refresh token issuance.
var DefaultScopes = []string{"openid", "profile", "email", "offline_access"}

// Config holds configuration for creating a new OAuth client.
type Config struct {
        // ID is the unique client identifier.
        ID string

        // ClientName is unverified, display-only metadata supplied at registration.
        ClientName string

        // Secret is the client secret for confidential clients.
        // Required for client_secret_basic / client_secret_post; ignored for "none".
        Secret string //nolint:gosec // G117: field legitimately holds sensitive data

        // RedirectURIs is the list of allowed redirect URIs.
        RedirectURIs []string

        // TokenEndpointAuthMethod is the client's registered auth method:
        // "none" (public client) or one of the client_secret_* methods
        // (confidential client). Required — there is no default, so callers must
        // choose explicitly rather than drifting into one.
        TokenEndpointAuthMethod string

        // GrantTypes overrides the default grant types.
        // If nil or empty, defaultGrantTypes is used.
        GrantTypes []string

        // ResponseTypes overrides the default response types. If nil or empty,
        // defaultResponseTypes is used except for private_key_jwt clients, which
        // deliberately support no response types.
        ResponseTypes []string

        // Scopes overrides the default scopes.
        // If nil or empty, DefaultScopes is used.
        Scopes []string

        // Audience is the list of allowed audience values for this client.
        // Per RFC 8707, the "resource" parameter in token requests is validated
        // against this list. If nil, audience validation will reject all values.
        Audience []string

        // JSONWebKeys contains the validated inline public keys for private_key_jwt.
        JSONWebKeys *jose.JSONWebKeySet

        // TokenEndpointAuthSigningAlgorithm is the validated signing algorithm for
        // private_key_jwt client authentication.
        TokenEndpointAuthSigningAlgorithm string
}

// dcrIssued is the marker identifying clients built by this package (i.e.
// issued via dynamic client registration or an equivalent explicit
// registration.New call). Storage backends use it to tell DCR-issued
// registrations — which carry an anti-bloat TTL — from pre-provisioned
// clients, which must never acquire one. The method is unexported so the
// marker cannot be implemented outside this package.
type dcrIssued interface {
        dcrIssued()
}

// DCRIssued reports whether client was issued by this package.
func DCRIssued(client fosite.Client) bool {
        _, ok := client.(dcrIssued)
        return ok
}

// ClientName returns the unverified display name of a DCR client, if supplied.
func ClientName(client fosite.Client) string {
        if named, ok := client.(interface{ GetClientName() string }); ok {
                return named.GetClientName()
        }
        return ""
}

// MarkDCRIssued wraps a client rebuilt from persisted DCR-issued form so the
// DCRIssued marker — and with it the anti-bloat TTL behaviour in storage —
// survives the storage round-trip. Callers must only mark clients they know
// were DCR-issued; pre-provisioned clients must never carry the marker.
//
// client must be one of the two concrete shapes clientFromStored produces:
// *fosite.DefaultOpenIDConnectClient (a row with a recorded
// token_endpoint_auth_method) or *fosite.DefaultClient (a row with none). The
// type switch below embeds whichever concrete type it was given rather than
// the fosite.Client interface, so every method the concrete type implements —
// now and in the future, including optional fosite interfaces like
// ClientWithSecretRotation — promotes automatically instead of being silently
// dropped, which is exactly what embedding the interface would do (fosite
// type-asserts for ClientWithSecretRotation.GetRotatedHashes during secret
// validation).
//
// Any other concrete type is a caller bug, but this sits on the GetClient
// path reachable from the unauthenticated /oauth/authorize endpoint, so it
// must not crash the process. Falling back to the client unwrapped, with an
// error log naming the type, is a bounded degradation: the row loses its
// DCRIssued marker and so stops renewing its TTL, which is recoverable —
// unlike a panic on an unauthenticated request path.
func MarkDCRIssued(client fosite.Client) fosite.Client {
        return MarkDCRIssuedWithName(client, "")
}

// MarkDCRIssuedWithName restores a DCR-issued client's unverified display name.
func MarkDCRIssuedWithName(client fosite.Client, name string) fosite.Client {
        switch c := client.(type) {
        case *fosite.DefaultOpenIDConnectClient:
                return &markedDCRIssuedOIDC{DefaultOpenIDConnectClient: c, clientName: name}
        case *fosite.DefaultClient:
                return &markedDCRIssuedDefault{DefaultClient: c, clientName: name}
        default:
                slog.Error("registration: MarkDCRIssued: unsupported concrete client type, returning unmarked",
                        "type", fmt.Sprintf("%T", client))
                return client
        }
}

type markedDCRIssuedOIDC struct {
        *fosite.DefaultOpenIDConnectClient
        clientName string
}

func (c *markedDCRIssuedOIDC) GetClientName() string { return c.clientName }
func (markedDCRIssuedOIDC) dcrIssued()               {}

type markedDCRIssuedDefault struct {
        *fosite.DefaultClient
        clientName string
}

func (c *markedDCRIssuedDefault) GetClientName() string { return c.clientName }
func (markedDCRIssuedDefault) dcrIssued()               {}

type dcrIssuedMarker struct{}

func (dcrIssuedMarker) dcrIssued() {}

// backChannelOnly marks a client with no interactive /authorize flow -- it
// must never be resolvable there, regardless of what its response types or
// grant types happen to look like. Explicit, not inferred, so a future
// client class does not get silently hidden from /authorize just because it
// happens to share metadata shape with a back-channel-only client.
type backChannelOnly interface {
        backChannelOnly()
}

// BackChannelOnly reports whether client is explicitly marked as having no
// interactive /authorize flow.
func BackChannelOnly(client fosite.Client) bool {
        _, ok := client.(backChannelOnly)
        return ok
}

// BackChannelOnlyMarker is embedded anonymously by a client type -- in this
// package or any other -- to mark it as never resolvable at /authorize. The
// backChannelOnly method it carries is unexported, but Go resolves interface
// satisfaction for a promoted method by the method's defining package, not
// the embedder's, so embedding this exported struct is sufficient for the
// embedding type to satisfy the unexported backChannelOnly interface above.
type BackChannelOnlyMarker struct{}

func (BackChannelOnlyMarker) backChannelOnly() {}

// publicClient is the DCR-issued public client shape: an OIDC client (so the
// "none" method is recorded and enforced). RFC 8252 Section 7.3 loopback
// dynamic-port matching for native apps is provided separately by
// RegisteredLoopbackRedirectURI, not by this type.
type publicClient struct {
        dcrIssuedMarker
        *fosite.DefaultOpenIDConnectClient
        clientName string
}

func (c *publicClient) GetClientName() string { return c.clientName }

// confidentialClient is the DCR-issued confidential client shape: an OIDC
// client so fosite pins and enforces the registered auth method at the token
// endpoint. RegisteredLoopbackRedirectURI refuses dynamic-port matching for
// any non-public client, so a secret-holding client never gets it.
type confidentialClient struct {
        dcrIssuedMarker
        *fosite.DefaultOpenIDConnectClient
        clientName string
}

func (c *confidentialClient) GetClientName() string { return c.clientName }

// privateKeyJWTClient is a DCR-issued client authenticated with an inline
// public key rather than a client secret. It is non-public because it must
// authenticate at the token endpoint, while retaining the DCR marker.
type privateKeyJWTClient struct {
        dcrIssuedMarker
        *fosite.DefaultOpenIDConnectClient
        clientName string
}

func (c *privateKeyJWTClient) GetClientName() string { return c.clientName }

// GetResponseTypes prevents fosite.DefaultClient's implicit ["code"] default
// from representing this token-exchange-only client as an authorization-code
// client. Embedding DefaultOpenIDConnectClient still preserves fosite's
// private_key_jwt authentication-method enforcement.
func (privateKeyJWTClient) GetResponseTypes() fosite.Arguments {
        return nil
}

// GenerateClientSecret mints a new client secret: 32 bytes of crypto/rand
// output, base64url-encoded (43 characters, no padding). RawURLEncoding is
// load-bearing, not cosmetic: fosite url.QueryUnescape's both Basic-auth
// components, so a secret containing '+', '/', or '%' would be corrupted or
// rejected at the token endpoint.
func GenerateClientSecret() (string, error) {
        buf := make([]byte, 32)
        if _, err := rand.Read(buf); err != nil {
                return "", fmt.Errorf("failed to generate client secret: %w", err)
        }
        return base64.RawURLEncoding.EncodeToString(buf), nil
}

// New creates a fosite.Client from the given configuration.
// Public clients get TokenEndpointAuthMethod "none" via DefaultOpenIDConnectClient;
// RFC 8252 Section 7.3 loopback redirect URI matching for native OAuth clients is
// provided separately by RegisteredLoopbackRedirectURI, not by the client type itself.
// Confidential clients with secrets have their Secret field SHA-256 hashed
// (see SHA256Hasher) as required by fosite for credential validation.
func New(cfg Config) (fosite.Client, error) {
        // Validate the auth method explicitly: silently defaulting an empty or
        // unknown value would reclassify the client one layer up, the same
        // public-ification bug the storage read path fails closed against.
        switch cfg.TokenEndpointAuthMethod {
        case oauthproto.TokenEndpointAuthMethodNone,
                oauthproto.TokenEndpointAuthMethodClientSecretBasic,
                oauthproto.TokenEndpointAuthMethodClientSecretPost,
                oauthproto.TokenEndpointAuthMethodPrivateKeyJWT:
        default:
                return nil, fmt.Errorf("unsupported token_endpoint_auth_method: %q", cfg.TokenEndpointAuthMethod)
        }
        public := cfg.TokenEndpointAuthMethod == oauthproto.TokenEndpointAuthMethodNone
        privateKeyJWT := cfg.TokenEndpointAuthMethod == oauthproto.TokenEndpointAuthMethodPrivateKeyJWT

        // Apply defaults for empty slices
        grantTypes := cfg.GrantTypes
        if len(grantTypes) == 0 {
                grantTypes = defaultGrantTypes
        }

        responseTypes := cfg.ResponseTypes
        // RFC 7591 §2 treats an omitted response_types value as ["code"]. A
        // private_key_jwt registration supports no response types, but RFC 7591
        // has no wire encoding for that zero-element set; preserve the validated
        // empty value rather than applying the interactive-client default.
        if len(responseTypes) == 0 && !privateKeyJWT {
                responseTypes = defaultResponseTypes
        }

        scopes := cfg.Scopes
        if len(scopes) == 0 {
                scopes = DefaultScopes
        }

        // Create the DefaultClient
        defaultClient := &fosite.DefaultClient{
                ID:            cfg.ID,
                RedirectURIs:  cfg.RedirectURIs,
                ResponseTypes: responseTypes,
                GrantTypes:    grantTypes,
                Scopes:        scopes,
                Audience:      cfg.Audience,
                Public:        public,
        }

        // Hash the secret for confidential clients. fosite compares the stored
        // hash with the presented secret using the hasher configured on
        // fosite.Config.ClientSecretsHasher, so this must use the same SHA-256
        // hasher — see SHA256Hasher for why no KDF is used.
        if !public && !privateKeyJWT {
                if cfg.Secret == "" {
                        return nil, fmt.Errorf("confidential client requires a secret")
                }
                hashedSecret, err := SHA256Hasher.Hash(context.Background(), []byte(cfg.Secret))
                if err != nil {
                        return nil, fmt.Errorf("failed to hash client secret: %w", err)
                }
                defaultClient.Secret = hashedSecret
        }

        oidcClient := &fosite.DefaultOpenIDConnectClient{
                DefaultClient:                     defaultClient,
                TokenEndpointAuthMethod:           cfg.TokenEndpointAuthMethod,
                JSONWebKeys:                       cloneJSONWebKeySet(cfg.JSONWebKeys),
                TokenEndpointAuthSigningAlgorithm: cfg.TokenEndpointAuthSigningAlgorithm,
        }

        if public {
                return &publicClient{DefaultOpenIDConnectClient: oidcClient, clientName: cfg.ClientName}, nil
        }
        if privateKeyJWT {
                return &privateKeyJWTClient{DefaultOpenIDConnectClient: oidcClient, clientName: cfg.ClientName}, nil
        }
        return &confidentialClient{DefaultOpenIDConnectClient: oidcClient, clientName: cfg.ClientName}, nil
}

func cloneJSONWebKeySet(jwks *jose.JSONWebKeySet) *jose.JSONWebKeySet {
        if jwks == nil {
                return nil
        }
        return &jose.JSONWebKeySet{Keys: append([]jose.JSONWebKey(nil), jwks.Keys...)}
}

// NewConfidentialPlain creates a DCR-issued confidential client as a plain
// *fosite.DefaultClient (Public: false, hashed secret), NOT the
// *fosite.DefaultOpenIDConnectClient shape New produces for an ordinary
// confidential registration.
//
// The difference matters: fosite only enforces token_endpoint_auth_method for
// clients implementing fosite.OpenIDConnectClient (see
// client_authentication.go in fosite v0.49.0 — AuthenticateClient type-
// switches on that interface before checking the method at all). A plain
// *fosite.DefaultClient with Public=false accepts credentials via either HTTP
// Basic or the form body and verifies whichever is presented. Use this
// constructor when the caller does not know which presentation the client
// will use — pinning the wrong one yields an invalid_client the operator
// cannot debug remotely.
//
// This is the shape ValidateDCRRequest's override path uses when a request's
// redirect_uris matches an operator-configured
// Config.ForceConfidentialRedirectURIs entry: such a client declared itself
// public but requires a secret, and the operator cannot know in advance
// whether its OAuth library presents credentials via Basic or form body.
//
// TokenEndpointAuthMethod on cfg is ignored; the returned client has no
// pinned method by construction. Ignores cfg.GrantTypes/ResponseTypes
// defaulting the same way New does. The returned client always carries the
// DCRIssued marker (see MarkDCRIssued) so storage retention applies.
func NewConfidentialPlain(cfg Config) (fosite.Client, error) {
        if cfg.Secret == "" {
                return nil, fmt.Errorf("confidential client requires a secret")
        }

        grantTypes := cfg.GrantTypes
        if len(grantTypes) == 0 {
                grantTypes = defaultGrantTypes
        }
        responseTypes := cfg.ResponseTypes
        if len(responseTypes) == 0 {
                responseTypes = defaultResponseTypes
        }
        scopes := cfg.Scopes
        if len(scopes) == 0 {
                scopes = DefaultScopes
        }

        hashedSecret, err := SHA256Hasher.Hash(context.Background(), []byte(cfg.Secret))
        if err != nil {
                return nil, fmt.Errorf("failed to hash client secret: %w", err)
        }

        defaultClient := &fosite.DefaultClient{
                ID:            cfg.ID,
                Secret:        hashedSecret,
                RedirectURIs:  cfg.RedirectURIs,
                ResponseTypes: responseTypes,
                GrantTypes:    grantTypes,
                Scopes:        scopes,
                Audience:      cfg.Audience,
                Public:        false,
        }

        return MarkDCRIssuedWithName(defaultClient, cfg.ClientName), nil
}

// NewStaticDelegateClient creates an unmarked, pre-provisioned confidential
// client for token exchange. A bare DefaultClient intentionally leaves the
// token endpoint authentication method unpinned, allowing both HTTP Basic and
// form-body client-secret authentication.
func NewStaticDelegateClient(cfg Config) (*fosite.DefaultClient, error) {
        if cfg.ID == "" {
                return nil, fmt.Errorf("delegate client requires an ID")
        }
        if cfg.Secret == "" {
                return nil, fmt.Errorf("confidential client requires a secret")
        }
        if len(cfg.GrantTypes) != 1 || cfg.GrantTypes[0] != oauthproto.GrantTypeTokenExchange {
                return nil, fmt.Errorf("delegate client grant types must be exactly [%q]", oauthproto.GrantTypeTokenExchange)
        }
        if len(cfg.Scopes) == 0 {
                return nil, fmt.Errorf("delegate client requires at least one scope")
        }
        if len(cfg.Audience) == 0 {
                return nil, fmt.Errorf("delegate client requires at least one audience")
        }

        hashedSecret, err := SHA256Hasher.Hash(context.Background(), []byte(cfg.Secret))
        if err != nil {
                return nil, fmt.Errorf("failed to hash client secret: %w", err)
        }

        return &fosite.DefaultClient{
                ID:         cfg.ID,
                Secret:     hashedSecret,
                GrantTypes: slices.Clone(cfg.GrantTypes),
                Scopes:     slices.Clone(cfg.Scopes),
                Audience:   slices.Clone(cfg.Audience),
                Public:     false,
        }, nil
}

// matchesAsLoopback checks if the requested URI matches the registered URI
// using RFC 8252 Section 7.3 loopback rules.
//
// Per RFC 8252 Section 7.3:
//   - Loopback redirect URIs use the "http" scheme
//   - The host must be 127.0.0.1, [::1], or localhost
//   - The authorization server MUST allow any port
//   - The path and query components must match exactly
func matchesAsLoopback(requestedURI, registeredURI string) bool {
        requested, err := url.Parse(requestedURI)
        if err != nil {
                return false
        }

        registered, err := url.Parse(registeredURI)
        if err != nil {
                return false
        }

        // RFC 6749 Section 3.1.2: the redirection endpoint URI MUST NOT include a
        // fragment component, and must not carry userinfo. Fosite's own
        // IsValidRedirectURI enforces this on whatever redirect_uri it actually
        // validates -- but a loopback match here is what causes the ORIGINAL
        // requested URI (not fosite's validated one) to become the effective
        // redirect target, so the same check must run here too, or an invalid
        // requested URI could reach storage/token-issuance unvalidated.
        //
        // The registered side gets the same userinfo check: registration
        // validation (oauthproto.ValidateRedirectURI, backed by fosite's
        // IsValidRedirectURI) only rejects a missing fragment, not userinfo, so a
        // registered "http://user:pass@localhost/callback" would otherwise match
        // a dynamic-port request for "http://localhost:54321/callback" -- the
        // hostname/path/query comparisons below never look at User, so dropping
        // this check would let a registered entry's userinfo silently vanish from
        // what "matching" means.
        if requested.Fragment != "" || requested.User != nil || registered.User != nil {
                return false
        }

        // RFC 8252 Section 7.3: Loopback redirect URIs use the "http" scheme.
        // Dynamic port matching only applies to http loopback URIs, not https.
        if requested.Scheme != "http" || registered.Scheme != "http" {
                return false
        }

        // Both must be loopback addresses
        if !isLoopbackHostname(requested.Hostname()) || !isLoopbackHostname(registered.Hostname()) {
                return false
        }

        // Hostnames must match (e.g., both 127.0.0.1 or both localhost)
        if !hostnamesMatch(requested.Hostname(), registered.Hostname()) {
                return false
        }

        // Path must match exactly. EscapedPath() (not Path) is compared: Path is
        // percent-decoded, so an encoded separator (e.g. registered
        // "/callback%2Fchild") would otherwise compare equal to a literal,
        // unencoded path ("/callback/child") that was never actually registered.
        if requested.EscapedPath() != registered.EscapedPath() {
                return false
        }

        // Query must match exactly, including whether a bare "?" was present at
        // all (ForceQuery): RawQuery alone can't distinguish "/callback" from
        // "/callback?", since both parse to an empty RawQuery.
        if requested.RawQuery != registered.RawQuery || requested.ForceQuery != registered.ForceQuery {
                return false
        }

        // Port can be any value (this is the key RFC 8252 requirement)
        return true
}

// isLoopbackURI reports whether uri's host is a loopback address (an IP
// loopback literal or "localhost", case-insensitively) — a hostname-only
// check, with no DNS resolution. Used by ValidateConfidentialRedirectURIs in
// dcr.go to reject a confidential registration's redirect_uri that targets
// loopback, which RFC 8252 §7.3 reserves for native/public clients.
func isLoopbackURI(uri string) bool {
        parsed, err := url.Parse(uri)
        if err != nil {
                return false
        }
        return isLoopbackHostname(parsed.Hostname())
}

// isLoopbackHostname reports whether host (as returned by url.Hostname(), so
// already stripped of brackets and port) is one of the RFC 8252 §7.3 loopback
// forms: "localhost" (case-insensitive, matching hostnamesMatch below) or an
// IP loopback literal (127.0.0.1, ::1).
//
// This is deliberately self-contained rather than delegating to
// networking.IsLocalhost: that helper (via oauthproto.IsLoopbackHost) is a
// case-SENSITIVE prefix check requiring the bracketed "[::1]" form, which
// url.Hostname() never produces -- using it here would silently make
// "LOCALHOST" and "::1" both unmatchable despite hostnamesMatch's own
// case-insensitive "localhost" contract.
func isLoopbackHostname(host string) bool {
        if IsLocalhostHostname(host) {
                return true
        }
        ip := net.ParseIP(host)
        return ip != nil && ip.IsLoopback()
}

// hostnamesMatch checks if two hostnames (as returned by url.Hostname()) should
// be considered equivalent for loopback matching purposes.
//
// The parameters are expected to be pre-parsed hostname strings from url.Hostname(),
// not raw URIs. This function is called from matchesAsLoopback which handles URL parsing.
//
// Per RFC 8252, the hostname must match exactly. We normalize localhost to
// be case-insensitive, but 127.0.0.1 and localhost are treated as different
// hostnames (a client registered with 127.0.0.1 will not match localhost requests).
func hostnamesMatch(requested, registered string) bool {
        // Case-insensitive comparison for localhost
        if strings.EqualFold(requested, "localhost") && strings.EqualFold(registered, "localhost") {
                return true
        }

        // Exact match for IP addresses
        return requested == registered
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package registration

import (
        "github.com/ory/fosite"
)

// configuredClient is implemented only by the wrappers MarkConfigured returns.
// It is deliberately separate from dcrIssued: a configured client must keep
// every behaviour that skips DCR-issued clients (audience binding, TTL
// renewal) applying to it.
type configuredClient interface {
        configured()
}

// IsConfigured reports whether client carries the configured stamp, meaning
// storage read it back from a row written by ReconcileConfiguredClient.
func IsConfigured(client fosite.Client) bool {
        _, ok := client.(configuredClient)
        return ok
}

// MarkConfigured wraps a client read back from a row written by
// ReconcileConfiguredClient so the stamp survives the storage round-trip.
//
// client must be *fosite.DefaultOpenIDConnectClient or *fosite.DefaultClient.
// The wrapper embeds the concrete pointer rather than the fosite.Client
// interface so every optional interface the concrete type implements
// (fosite.OpenIDConnectClient, ClientWithSecretRotation) is promoted.
//
// For any other type it returns (nil, false). Callers must treat that as a
// denial: handing back the client unmarked would let a configured client
// escape the check the stamp exists for.
func MarkConfigured(client fosite.Client) (fosite.Client, bool) {
        switch c := client.(type) {
        case *fosite.DefaultOpenIDConnectClient:
                return &configuredOIDC{DefaultOpenIDConnectClient: c}, true
        case *fosite.DefaultClient:
                return &configuredDefault{DefaultClient: c}, true
        default:
                return nil, false
        }
}

type configuredOIDC struct {
        *fosite.DefaultOpenIDConnectClient
}

func (configuredOIDC) configured() {}

type configuredDefault struct {
        *fosite.DefaultClient
}

func (configuredDefault) configured() {}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package registration provides OAuth 2.0 Dynamic Client Registration (DCR)
// functionality per RFC 7591, including request validation and secure redirect
// URI handling for public native clients.
package registration

import (
        "fmt"
        "slices"
        "strings"

        "github.com/go-jose/go-jose/v3"
        "github.com/ory/fosite"

        servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// ValidateScopeSubset checks that every scope in subset is also present in
// superset, returning an error that names fieldName and the offending scope.
//
// Shared across the layers that validate baseline-scope configuration so the
// error message format is identical wherever the violation is caught (a
// caller using YAML-loaded config and a caller constructing config
// programmatically both see the same wording).
//
// fieldName should be the wire-format or display name of the field being
// validated (e.g. "baseline_client_scopes"). It is embedded verbatim in the
// returned error.
func ValidateScopeSubset(subset, superset []string, fieldName string) error {
        if len(subset) == 0 {
                return nil
        }
        supported := make(map[string]bool, len(superset))
        for _, s := range superset {
                supported[s] = true
        }
        for _, s := range subset {
                if !supported[s] {
                        return fmt.Errorf("%s contains %q which is not in scopes_supported", fieldName, s)
                }
        }
        return nil
}

// DCR error codes per RFC 7591 Section 3.2.2
const (
        // DCRErrorInvalidRedirectURI indicates that the value of one or more
        // redirect_uris is invalid.
        DCRErrorInvalidRedirectURI = "invalid_redirect_uri"

        // DCRErrorInvalidClientMetadata indicates that the value of one of the
        // client metadata fields is invalid and the server has rejected this request.
        DCRErrorInvalidClientMetadata = "invalid_client_metadata"
)

// Validation limits to prevent DoS attacks via excessively large requests.
const (
        // MaxRedirectURICount is the maximum number of redirect URIs allowed per client.
        MaxRedirectURICount = 10

        // maxInlineJWKSKeys limits the number of keys validated for private_key_jwt registration.
        maxInlineJWKSKeys = 10

        // MaxClientNameLength is the maximum allowed length for a client name.
        MaxClientNameLength = 256

        // MaxSoftwareIDLength is the maximum allowed length for a software_id
        // value. RFC 7591 does not mandate an upper bound, so we reuse the
        // client_name cap for consistency — a software_id is a similar-purpose
        // human-oriented identifier and the same ballpark DoS concerns apply.
        MaxSoftwareIDLength = 256
)

// DCRError represents an OAuth 2.0 Dynamic Client Registration error
// response per RFC 7591 Section 3.2.2.
type DCRError struct {
        // Error is a single ASCII error code from the defined set.
        Error string `json:"error"`

        // ErrorDescription is a human-readable text providing additional information.
        ErrorDescription string `json:"error_description,omitempty"`
}

// defaultGrantTypes are the default grant types for registered clients.
var defaultGrantTypes = []string{oauthproto.GrantTypeAuthorizationCode, oauthproto.GrantTypeRefreshToken}

// allowedGrantTypes defines the grant types permitted for public clients.
var allowedGrantTypes = map[string]bool{
        oauthproto.GrantTypeAuthorizationCode: true,
        oauthproto.GrantTypeRefreshToken:      true,
        oauthproto.GrantTypeDeviceCode:        true,
}

// defaultResponseTypes are the default response types for registered clients.
var defaultResponseTypes = []string{"code"}

// allowedResponseTypes defines the response types permitted for public clients.
var allowedResponseTypes = map[string]bool{
        "code": true,
}

// ValidateDCRRequest validates a DCR request according to RFC 7591
// and the server's security policy. When allowConfidential is false the
// policy is unchanged from the historical default: public clients only.
// When true, token_endpoint_auth_method may also be client_secret_basic or
// client_secret_post; such confidential registrations are additionally
// restricted to https non-loopback redirect URIs, because a client on a
// loopback or private-scheme URI is by construction a public client
// (OAuth 2.1 §2.1) and minting it a secret ships that secret inside a
// distributed binary.
// Returns the validated request with defaults applied, or an error.
//
// The returned request preserves validated client metadata, including scopes and
// private_key_jwt key metadata. Scope authorization against the server's supported
// set is a separate step handled by ValidateScopes using the caller's policy inputs.
func ValidateDCRRequest(
        req *oauthproto.DynamicClientRegistrationRequest,
        allowConfidential bool,
        allowPrivateKeyJWT bool,
) (*oauthproto.DynamicClientRegistrationRequest, *DCRError) {
        // Validate the authentication method before applying method-specific metadata
        // rules. Existing registration methods may carry jwks_uri metadata even though
        // only private_key_jwt requires inline keys.
        authMethod, dcrErr := validateAuthMethod(req.TokenEndpointAuthMethod, req.RedirectURIs, allowConfidential, allowPrivateKeyJWT)
        if dcrErr != nil {
                return nil, dcrErr
        }

        if len(req.RedirectURIs) == 0 && authMethod != oauthproto.TokenEndpointAuthMethodPrivateKeyJWT &&
                !isDeviceCodeOnlyRequest(req.GrantTypes) {
                return nil, &DCRError{
                        Error:            DCRErrorInvalidRedirectURI,
                        ErrorDescription: "redirect_uris is required",
                }
        }

        // Validate redirect_uris count limit
        if len(req.RedirectURIs) > MaxRedirectURICount {
                return nil, &DCRError{
                        Error:            DCRErrorInvalidRedirectURI,
                        ErrorDescription: "too many redirect_uris (maximum 10)",
                }
        }

        // 3. Validate all redirect_uris per RFC 8252
        for _, uri := range req.RedirectURIs {
                if err := ValidateRedirectURI(uri); err != nil {
                        return nil, err
                }
        }

        // 4. Validate client_name length
        if len(req.ClientName) > MaxClientNameLength {
                return nil, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "client_name too long (maximum 256 characters)",
                }
        }

        // 4a. Validate software_id: length cap + printable-ASCII charset.
        // RFC 7591 does not mandate an upper bound or a character class for
        // software_id, but since we capture the value in audit logs we want a
        // predictable shape and a hard cap against DoS — a caller sending
        // multi-MB strings would slip past only the 64 KiB body cap otherwise.
        if dcrErr := validateSoftwareID(req.SoftwareID); dcrErr != nil {
                return nil, dcrErr
        }

        if dcrErr := validatePrivateKeyJWTMetadata(req, authMethod); dcrErr != nil {
                return nil, dcrErr
        }

        // 6. Validate/default grant_types
        grantTypes, err := validateGrantTypes(req.GrantTypes, authMethod)
        if err != nil {
                return nil, err
        }

        // 7. Validate/default response_types
        responseTypes, err := validateResponseTypes(req.ResponseTypes, authMethod, grantTypes)
        if err != nil {
                return nil, err
        }

        // Only private_key_jwt uses client key metadata. Keep unrelated DCR
        // registrations from echoing or retaining attacker-supplied key material.
        clientJWKS := cloneJSONWebKeySet(req.JWKS)
        clientJWKSURI := req.JWKSURI
        clientSigningAlgorithm := req.TokenEndpointAuthSigningAlg
        if authMethod != oauthproto.TokenEndpointAuthMethodPrivateKeyJWT {
                clientJWKS = nil
                clientJWKSURI = ""
                clientSigningAlgorithm = ""
        }

        // Return validated request with defaults applied
        return &oauthproto.DynamicClientRegistrationRequest{
                RedirectURIs:                req.RedirectURIs,
                ClientName:                  req.ClientName,
                TokenEndpointAuthMethod:     authMethod,
                GrantTypes:                  grantTypes,
                ResponseTypes:               responseTypes,
                Scopes:                      req.Scopes,
                JWKS:                        clientJWKS,
                JWKSURI:                     clientJWKSURI,
                TokenEndpointAuthSigningAlg: clientSigningAlgorithm,
                SoftwareID:                  req.SoftwareID,
        }, nil
}

// validateSoftwareID enforces the length cap and printable-ASCII charset
// for a DCR request's software_id. An empty value is accepted (the field
// is optional per RFC 7591).
func validateSoftwareID(softwareID string) *DCRError {
        if softwareID == "" {
                return nil
        }
        if len(softwareID) > MaxSoftwareIDLength {
                return &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "software_id too long (maximum 256 characters)",
                }
        }
        // Allow printable ASCII (0x20..0x7E) only. Control characters and
        // non-ASCII input are rejected: we log this value and want a
        // predictable on-disk representation regardless of handler choice.
        for i := 0; i < len(softwareID); i++ {
                c := softwareID[i]
                if c < 0x20 || c > 0x7E {
                        return &DCRError{
                                Error:            DCRErrorInvalidClientMetadata,
                                ErrorDescription: "software_id must contain only printable ASCII characters",
                        }
                }
        }
        return nil
}

// validateAuthMethod validates token_endpoint_auth_method and returns the
// effective method with the empty default applied. When allowConfidential is
// false the policy is unchanged from the historical default: public clients
// only. When true, the client_secret_* methods are accepted but restricted to
// https non-loopback redirect URIs, because a client on a loopback or
// private-scheme URI is by construction a public client (OAuth 2.1 §2.1) and
// minting it a secret ships that secret inside a distributed binary.
func validateAuthMethod(
        authMethod string,
        redirectURIs []string,
        allowConfidential bool,
        allowPrivateKeyJWT bool,
) (string, *DCRError) {
        if authMethod == "" {
                // RFC 7591 §2 says an omitted token_endpoint_auth_method defaults to
                // "client_secret_basic". We deliberately default to "none" instead.
                // That RFC default predates PKCE-based public clients becoming the
                // norm for native and CLI apps, which are most MCP clients today.
                // Following it here would silently turn any public client that
                // omits the field into a confidential one — and fosite pins the
                // auth method at registration, so that client's first token request
                // would fail with invalid_client because it has no secret and never
                // asked for one. It also wouldn't help the clients confidential
                // support was added for: the known cases send "none" explicitly
                // rather than omitting the field, and an explicit value always
                // overrides the default.
                authMethod = oauthproto.TokenEndpointAuthMethodNone
        }
        switch authMethod {
        case oauthproto.TokenEndpointAuthMethodNone:
                // Always accepted; the public-client default.
                return authMethod, nil
        case oauthproto.TokenEndpointAuthMethodPrivateKeyJWT:
                if !allowPrivateKeyJWT {
                        return "", &DCRError{Error: DCRErrorInvalidClientMetadata, ErrorDescription: "private_key_jwt registration is disabled"}
                }
                return authMethod, nil
        case oauthproto.TokenEndpointAuthMethodClientSecretBasic,
                oauthproto.TokenEndpointAuthMethodClientSecretPost:
                if !allowConfidential {
                        return "", &DCRError{
                                Error:            DCRErrorInvalidClientMetadata,
                                ErrorDescription: "this authorization server only supports token_endpoint_auth_method 'none'",
                        }
                }
                // Server policy (not a spec mandate): confidential registrations must
                // use https non-loopback redirect URIs. A client reachable on loopback
                // or a private scheme is typically a distributed native app that cannot
                // keep a secret, so this server declines to mint it one.
                if dcrErr := ValidateConfidentialRedirectURIs(redirectURIs, authMethod); dcrErr != nil {
                        return "", dcrErr
                }
                return authMethod, nil
        default:
                return "", &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "unsupported token_endpoint_auth_method: " + authMethod,
                }
        }
}

// ValidateConfidentialRedirectURIs checks that every entry in redirectURIs
// meets the confidential-client policy: https non-loopback, RFC 8252 strict
// scheme rules. authMethod is embedded in the loopback-rejection message and
// should be the auth method the client is being registered with.
//
// Shared by the ordinary confidential registration path (validateAuthMethod)
// and the force-confidential override (resolveForceConfidentialOverride in
// pkg/authserver/server/handlers/dcr.go), so a registration cannot become
// confidential by either path while carrying a loopback redirect_uri.
func ValidateConfidentialRedirectURIs(redirectURIs []string, authMethod string) *DCRError {
        for _, uri := range redirectURIs {
                if err := oauthproto.ValidateRedirectURI(uri, oauthproto.RedirectURIPolicyStrict); err != nil {
                        return &DCRError{
                                Error:            DCRErrorInvalidRedirectURI,
                                ErrorDescription: err.Error(),
                        }
                }
                // isLoopbackURI is a literal-string check (127.0.0.0/8, ::1,
                // ::ffff:127.0.0.1, localhost); it does not resolve the hostname,
                // so a registrant-controlled name that merely resolves to a
                // loopback address slips through. This is acceptable: the
                // registrant controls both its own redirect URI and its own DNS,
                // so at worst this lets an actor mint a secret for its own
                // loopback app, not attack another tenant. Resolving DNS at
                // registration time would add an unreliable (TTL/rebinding) and
                // unnecessary network dependency to a request that should stay
                // purely local.
                if isLoopbackURI(uri) {
                        return &DCRError{
                                Error: DCRErrorInvalidRedirectURI,
                                ErrorDescription: "token_endpoint_auth_method '" + authMethod +
                                        "' requires https non-loopback redirect_uris; native and loopback clients must use 'none'" +
                                        " (offending redirect_uri: " + uri + ")",
                        }
                }
        }
        return nil
}

// ValidatePrivateKeyJWTClientKeys checks persisted private_key_jwt metadata before
// fosite can use it for signature verification. This also rejects registrations
// made before the key-size and JWKS limits were introduced. It returns nil for
// clients that do not use private_key_jwt.
func ValidatePrivateKeyJWTClientKeys(client fosite.Client) *DCRError {
        c, ok := client.(fosite.OpenIDConnectClient)
        if !ok || c.GetTokenEndpointAuthMethod() != oauthproto.TokenEndpointAuthMethodPrivateKeyJWT {
                return nil
        }
        return validatePrivateKeyJWTMetadata(&oauthproto.DynamicClientRegistrationRequest{
                JWKS:                        c.GetJSONWebKeys(),
                TokenEndpointAuthSigningAlg: c.GetTokenEndpointAuthSigningAlgorithm(),
        }, oauthproto.TokenEndpointAuthMethodPrivateKeyJWT)
}

// validatePrivateKeyJWTMetadata is also run against stored clients on every
// load (ValidatePrivateKeyJWTClientKeys), so tightening it revokes existing
// private_key_jwt registrations that no longer pass.
func validatePrivateKeyJWTMetadata(req *oauthproto.DynamicClientRegistrationRequest, authMethod string) *DCRError {
        if authMethod != oauthproto.TokenEndpointAuthMethodPrivateKeyJWT {
                return nil
        }
        if req.JWKSURI != "" {
                return &DCRError{Error: DCRErrorInvalidClientMetadata, ErrorDescription: "jwks_uri is not supported; provide inline jwks"}
        }
        if req.JWKS == nil || len(req.JWKS.Keys) == 0 {
                return &DCRError{Error: DCRErrorInvalidClientMetadata, ErrorDescription: "private_key_jwt requires a non-empty inline jwks"}
        }
        if len(req.JWKS.Keys) > maxInlineJWKSKeys {
                return &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: fmt.Sprintf("too many inline jwks keys (maximum %d)", maxInlineJWKSKeys),
                }
        }
        if !supportedSigningAlgorithm(req.TokenEndpointAuthSigningAlg) {
                return &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "token_endpoint_auth_signing_alg is missing or unsupported",
                }
        }
        matching := false
        for i := range req.JWKS.Keys {
                keyMatches, dcrErr := validatePrivateKeyJWTKey(&req.JWKS.Keys[i], req.TokenEndpointAuthSigningAlg)
                if dcrErr != nil {
                        return dcrErr
                }
                matching = matching || keyMatches
        }
        if !matching {
                return &DCRError{Error: DCRErrorInvalidClientMetadata, ErrorDescription: "signing algorithm does not match a registered key"}
        }
        return nil
}

func validatePrivateKeyJWTKey(key *jose.JSONWebKey, algorithm string) (bool, *DCRError) {
        if !key.Valid() || !key.IsPublic() || key.Use != "sig" {
                return false, &DCRError{Error: DCRErrorInvalidClientMetadata, ErrorDescription: "jwks must contain valid public signing keys"}
        }
        if key.Algorithm != "" && !supportedSigningAlgorithm(key.Algorithm) {
                return false, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "jwks contains an unsupported signing algorithm",
                }
        }
        if key.Algorithm != "" && key.Algorithm != algorithm {
                return false, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "jwks contains a signing algorithm that does not match the request",
                }
        }
        // Always validate structural compatibility between the declared algorithm
        // and the actual key type/curve, even when the JWK's own alg label
        // already matched above: a client could mislabel an EC key as "RS256"
        // and have both label-only checks pass despite the key never being able
        // to actually sign or verify with that algorithm.
        if err := servercrypto.ValidateAlgorithmForPublicKey(algorithm, key.Key); err != nil {
                return false, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "jwks key is not compatible with the requested signing algorithm",
                }
        }
        return true, nil
}

// SupportedSigningAlgorithms returns the allowlist accepted for private_key_jwt
// client assertions and their registered signing keys. Delegates to
// servercrypto.SupportedClientKeyAlgorithms, the single source of truth also
// consulted by key-compatibility validation, so the two cannot drift apart.
func SupportedSigningAlgorithms() []string {
        return servercrypto.SupportedClientKeyAlgorithms()
}

func supportedSigningAlgorithm(alg string) bool {
        return slices.Contains(SupportedSigningAlgorithms(), alg)
}

// isDeviceCodeOnlyRequest reports whether the requested grant_types are
// exactly {device_code} or {device_code, refresh_token}. RFC 8628 device
// flow clients never receive a redirect and have no authorize-endpoint
// response_type, so a request matching this carves out of the
// authorization_code-shaped requirements the same way private_key_jwt
// registrations do. Mixing device_code into an authorization_code
// registration is intentionally NOT covered here -- such a request must
// still satisfy the ordinary authorization_code requirements below.
func isDeviceCodeOnlyRequest(grantTypes []string) bool {
        if !slices.Contains(grantTypes, oauthproto.GrantTypeDeviceCode) {
                return false
        }
        for _, gt := range grantTypes {
                if gt != oauthproto.GrantTypeDeviceCode && gt != oauthproto.GrantTypeRefreshToken {
                        return false
                }
        }
        return true
}

func validateGrantTypes(grantTypes []string, authMethod string) ([]string, *DCRError) {
        if authMethod == oauthproto.TokenEndpointAuthMethodPrivateKeyJWT {
                if len(grantTypes) == 0 {
                        return []string{oauthproto.GrantTypeTokenExchange}, nil
                }
                if len(grantTypes) != 1 || grantTypes[0] != oauthproto.GrantTypeTokenExchange {
                        return nil, &DCRError{
                                Error:            DCRErrorInvalidClientMetadata,
                                ErrorDescription: "private_key_jwt registrations require exactly the token exchange grant",
                        }
                }
                return append([]string(nil), grantTypes...), nil
        }
        if isDeviceCodeOnlyRequest(grantTypes) {
                for _, gt := range grantTypes {
                        if !allowedGrantTypes[gt] {
                                return nil, &DCRError{
                                        Error:            DCRErrorInvalidClientMetadata,
                                        ErrorDescription: "unsupported grant_type: " + gt,
                                }
                        }
                }
                return append([]string(nil), grantTypes...), nil
        }
        if len(grantTypes) == 0 {
                grantTypes = defaultGrantTypes
        }
        // Require authorization_code explicitly - provides a clearer error for the
        // "refresh_token only" case that would otherwise pass the allowlist.
        if !slices.Contains(grantTypes, oauthproto.GrantTypeAuthorizationCode) {
                return nil, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "grant_types must include 'authorization_code'",
                }
        }
        for _, gt := range grantTypes {
                if !allowedGrantTypes[gt] {
                        return nil, &DCRError{
                                Error:            DCRErrorInvalidClientMetadata,
                                ErrorDescription: "unsupported grant_type: " + gt,
                        }
                }
        }
        return grantTypes, nil
}

func validateResponseTypes(responseTypes []string, authMethod string, grantTypes []string) ([]string, *DCRError) {
        if authMethod == oauthproto.TokenEndpointAuthMethodPrivateKeyJWT {
                if len(responseTypes) != 0 {
                        return nil, &DCRError{
                                Error:            DCRErrorInvalidClientMetadata,
                                ErrorDescription: "private_key_jwt registrations cannot use response_types",
                        }
                }
                return nil, nil
        }
        if isDeviceCodeOnlyRequest(grantTypes) {
                if len(responseTypes) != 0 {
                        return nil, &DCRError{
                                Error:            DCRErrorInvalidClientMetadata,
                                ErrorDescription: "device_code registrations cannot use response_types",
                        }
                }
                return nil, nil
        }
        if len(responseTypes) == 0 {
                responseTypes = defaultResponseTypes
        }
        // Require "code" explicitly - purely defense-in-depth since the allowlist
        // currently only contains "code", but provides a clearer error message.
        if !slices.Contains(responseTypes, "code") {
                return nil, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "response_types must include 'code'",
                }
        }
        for _, rt := range responseTypes {
                if !allowedResponseTypes[rt] {
                        return nil, &DCRError{
                                Error:            DCRErrorInvalidClientMetadata,
                                ErrorDescription: "unsupported response_type: " + rt,
                        }
                }
        }
        return responseTypes, nil
}

// ValidateRedirectURI validates a redirect URI per RFC 8252:
// - HTTPS is allowed for any address (web-based redirects)
// - HTTP is only allowed for loopback addresses (127.0.0.1, [::1], localhost)
// - Private-use URI schemes (e.g., cursor://, vscode://) are allowed for native apps
func ValidateRedirectURI(uri string) *DCRError {
        if err := oauthproto.ValidateRedirectURI(uri, oauthproto.RedirectURIPolicyAllowPrivateSchemes); err != nil {
                return &DCRError{
                        Error:            DCRErrorInvalidRedirectURI,
                        ErrorDescription: err.Error(),
                }
        }
        return nil
}

// ValidateScopes validates a slice of already-parsed scope tokens against
// the server's allowed set per RFC 7591 §2.
//
//   - Empty/nil input falls back to the intersection of DefaultScopes with
//     allowedScopes (RFC 7591 §2 permits the server to register a client
//     with a default set of scopes). Defaults the server does not support
//     are returned in droppedDefaults so callers can surface them; only an
//     empty intersection is rejected, because no sensible default exists
//     and the client must declare scope explicitly.
//   - Each requested scope must appear in allowedScopes; otherwise returns
//     invalid_client_metadata.
//   - Duplicates in the input are tolerated and deduplicated per RFC 6749
//     §3.3 (scope is a set of case-sensitive strings).
//
// droppedDefaults is non-nil only on the default-fallback path; requested
// scopes are validated strictly and never dropped.
func ValidateScopes(requestedScopes, allowedScopes []string) (scopes, droppedDefaults []string, dcrErr *DCRError) {
        // Build allowed scope set for O(1) lookup
        allowed := make(map[string]bool, len(allowedScopes))
        for _, s := range allowedScopes {
                allowed[s] = true
        }

        // Deduplicate while validating each requested scope against the
        // allowed set.
        if len(requestedScopes) > 0 {
                seen := make(map[string]bool)
                for _, s := range requestedScopes {
                        if !allowed[s] {
                                return nil, nil, &DCRError{
                                        Error:            DCRErrorInvalidClientMetadata,
                                        ErrorDescription: "unsupported scope: " + s,
                                }
                        }
                        if !seen[s] {
                                seen[s] = true
                                scopes = append(scopes, s)
                        }
                }
                return scopes, nil, nil
        }

        // No scopes requested: fall back to the intersection of DefaultScopes
        // with the allowed set. Requiring every default to be supported would
        // reject any client that omits scope against a server whose
        // scopes_supported does not carry the full default set, even though a
        // perfectly usable subset exists (see issue #6186).
        //
        // Known residual gap (#6186): this fallback draws from DefaultScopes
        // only, never from the rest of allowedScopes. A scope-omitting client
        // against a server whose scopes_supported carries entries outside the
        // default set (e.g. mcp:tools) registers without them — and because
        // protected-resource metadata advertises scopes_supported and MCP
        // guidance tells clients to request all of it absent a narrower hint,
        // such a client can then fail at /oauth/authorize with invalid_scope.
        // Defaulting to allowedScopes itself would close that, but granting
        // every anonymous registration the full advertised set is a policy
        // decision; the explicit opt-in for it is
        // baseline_client_scopes = scopes_supported.
        for _, s := range DefaultScopes {
                if allowed[s] {
                        scopes = append(scopes, s)
                } else {
                        droppedDefaults = append(droppedDefaults, s)
                }
        }
        if len(scopes) == 0 {
                return nil, nil, &DCRError{
                        Error: DCRErrorInvalidClientMetadata,
                        ErrorDescription: "none of the default scopes (" + strings.Join(DefaultScopes, " ") +
                                ") are supported by this server; the client must declare scope explicitly",
                }
        }
        return scopes, droppedDefaults, nil
}

// UnionScopes returns the union of requested and baseline scopes, preserving
// the order of requested first, then appending any baseline scopes not already
// present. Duplicates are removed. Returns nil when the result is empty.
//
// Both inputs must already be validated by the caller. UnionScopes does not
// filter empty strings or validate scope syntax — it only deduplicates and
// merges in stable order.
func UnionScopes(requested, baseline []string) []string {
        seen := make(map[string]bool, len(requested)+len(baseline))
        out := make([]string, 0, len(requested)+len(baseline))
        for _, s := range requested {
                if !seen[s] {
                        seen[s] = true
                        out = append(out, s)
                }
        }
        for _, s := range baseline {
                if !seen[s] {
                        seen[s] = true
                        out = append(out, s)
                }
        }
        if len(out) == 0 {
                return nil
        }
        return out
}

// ValidatePublicGrantTypes validates the grant_types for a public OAuth client,
// applying the same rules as DCR: authorization_code must be present, and all
// declared values must be in the allowed set. Returns the validated slice (with
// defaults applied when nil/empty) or a *DCRError on violation.
func ValidatePublicGrantTypes(grantTypes []string) ([]string, *DCRError) {
        return validateGrantTypes(grantTypes, oauthproto.TokenEndpointAuthMethodNone)
}

// ValidatePublicResponseTypes validates the response_types for a public OAuth
// client, applying the same rules as DCR: code must be present and all declared
// values must be in the allowed set (unless grantTypes is device_code-only, which
// requires no response_types). Returns the validated slice (with defaults
// applied when nil/empty) or a *DCRError on violation.
func ValidatePublicResponseTypes(responseTypes, grantTypes []string) ([]string, *DCRError) {
        return validateResponseTypes(responseTypes, oauthproto.TokenEndpointAuthMethodNone, grantTypes)
}

// FilterPublicGrantTypes returns the subset of grantTypes this server supports
// for public clients, dropping unsupported entries instead of rejecting the
// whole set the way ValidatePublicGrantTypes does.
//
// This is the right semantics for CIMD: a Client ID Metadata Document
// describes the client's capabilities across every authorization server it
// talks to, and the client cannot tailor it per server — VS Code, for
// example, declares the device_code grant alongside authorization_code, and
// rejecting the document over a grant type this flow never uses breaks the
// client entirely. A DCR request, by contrast, is addressed to this server
// specifically, so ValidatePublicGrantTypes' rejection remains the correct
// feedback on that path.
//
// The surviving set must still include "authorization_code": a client whose
// supported grant types do not intersect with the only redemption flow this
// server offers cannot function against it, and a clear error beats a client
// that registers and then fails every token request. nil/empty input gets the
// same defaults as DCR.
func FilterPublicGrantTypes(grantTypes []string) ([]string, *DCRError) {
        // Clone so callers that store the result (e.g. in a cached fosite client)
        // never alias the package-level default.
        if len(grantTypes) == 0 {
                return slices.Clone(defaultGrantTypes), nil
        }
        filtered := make([]string, 0, len(grantTypes))
        for _, gt := range grantTypes {
                if allowedGrantTypes[gt] {
                        filtered = append(filtered, gt)
                }
        }
        if !slices.Contains(filtered, oauthproto.GrantTypeAuthorizationCode) {
                return nil, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "grant_types must include 'authorization_code'",
                }
        }
        return filtered, nil
}

// FilterPublicResponseTypes returns the subset of responseTypes this server
// supports for public clients, dropping unsupported entries instead of
// rejecting the whole set. Same reasoning as FilterPublicGrantTypes: a CIMD
// document declares capabilities across all servers, so an entry this server
// does not support must not be fatal — but "code" must survive the filter,
// since it is the only response type this server can serve. nil/empty input
// gets the same defaults as DCR.
func FilterPublicResponseTypes(responseTypes []string) ([]string, *DCRError) {
        // Clone for the same non-aliasing reason as FilterPublicGrantTypes.
        if len(responseTypes) == 0 {
                return slices.Clone(defaultResponseTypes), nil
        }
        filtered := make([]string, 0, len(responseTypes))
        for _, rt := range responseTypes {
                if allowedResponseTypes[rt] {
                        filtered = append(filtered, rt)
                }
        }
        if !slices.Contains(filtered, "code") {
                return nil, &DCRError{
                        Error:            DCRErrorInvalidClientMetadata,
                        ErrorDescription: "response_types must include 'code'",
                }
        }
        return filtered, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package registration

import (
        "context"
        "crypto/sha256"
        "crypto/subtle"
        "errors"

        "github.com/ory/fosite"
)

// SHA256Hasher is a fosite.Hasher for client secrets based on plain SHA-256.
//
// Client secrets issued by this server are 32 bytes of crypto/rand output,
// base64url-encoded, and are never client-chosen (the DCR request struct has
// no client_secret field). At 256 bits of CSPRNG entropy there is nothing for
// a password-stretching KDF to protect: brute force is infeasible regardless
// of hash cost. What a slow KDF *would* do is give an unauthenticated caller
// a CPU-amplification lever — one bcrypt verify per wrong-secret attempt
// against /oauth/token, in the same process as the proxy data path — so the
// work factor protects the attacker, not the secret.
var SHA256Hasher fosite.Hasher = sha256Hasher{}

type sha256Hasher struct{}

// Compare compares a SHA-256 hash with the plaintext data in constant time.
func (sha256Hasher) Compare(_ context.Context, hash, data []byte) error {
        dataHash := sha256.Sum256(data)
        if len(hash) != len(dataHash) || subtle.ConstantTimeCompare(hash, dataHash[:]) != 1 {
                return errors.New("client secret mismatch")
        }
        return nil
}

// Hash returns the raw SHA-256 digest of data.
func (sha256Hasher) Hash(_ context.Context, data []byte) ([]byte, error) {
        sum := sha256.Sum256(data)
        return sum[:], nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package registration

import (
        "fmt"
        "slices"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// SPIFFEClient is the immutable OAuth client representation of a configured
// SPIFFE principal association. It is neither public nor secret-bearing.
// JWT-SVID authentication of configured associations is implemented but not
// yet wired to a trust bundle source (see docs/arch/18-spiffe-association-declarations.md);
// X.509-SVID credential validation remains pending.
type SPIFFEClient struct {
        BackChannelOnlyMarker
        id         string
        grantTypes fosite.Arguments
        scopes     fosite.Arguments
        audiences  []string
        resources  []string
}

// NewSPIFFEClient creates an immutable, secretless, non-public client for a
// configured SPIFFE association. GetAudience returns the configured audience
// allowlist. resources is the independent RFC 8707 resource allowlist,
// available via Resources(); it may be empty.
func NewSPIFFEClient(id string, scopes, audiences, resources []string) (*SPIFFEClient, error) {
        if id == "" {
                return nil, fmt.Errorf("SPIFFE client ID is required")
        }
        if len(scopes) == 0 || len(audiences) == 0 {
                return nil, fmt.Errorf("SPIFFE client scopes and audiences are required")
        }

        return &SPIFFEClient{
                id:         id,
                grantTypes: fosite.Arguments{oauthproto.GrantTypeTokenExchange},
                scopes:     slices.Clone(scopes),
                audiences:  slices.Clone(audiences),
                resources:  slices.Clone(resources),
        }, nil
}

// GetID returns the configured association client ID.
func (c *SPIFFEClient) GetID() string { return c.id }

// GetHashedSecret returns nil because no OAuth client secret is assigned.
// JWT-SVID authentication does not use a client secret, and X.509-SVID
// credential validation remains pending.
func (*SPIFFEClient) GetHashedSecret() []byte { return nil }

// GetRedirectURIs returns nil because SPIFFE clients do not use authorization redirects.
func (*SPIFFEClient) GetRedirectURIs() []string { return nil }

// GetGrantTypes returns a copy of the configured grant types.
func (c *SPIFFEClient) GetGrantTypes() fosite.Arguments { return slices.Clone(c.grantTypes) }

// GetResponseTypes returns nil because SPIFFE clients do not use authorization responses.
func (*SPIFFEClient) GetResponseTypes() fosite.Arguments { return nil }

// GetScopes returns a copy of the configured scopes.
func (c *SPIFFEClient) GetScopes() fosite.Arguments { return slices.Clone(c.scopes) }

// Audiences returns a copy of the allowed RFC 8693 audience request values.
func (c *SPIFFEClient) Audiences() []string { return slices.Clone(c.audiences) }

// Resources returns a copy of the allowed RFC 8707 resource request values.
// Resources and audiences are independent request dimensions; the token
// exchange handler checks the "resource" request parameter against this,
// not against GetAudience.
func (c *SPIFFEClient) Resources() []string { return slices.Clone(c.resources) }

// GetAudience returns the allowed RFC 8693 audience request values.
func (c *SPIFFEClient) GetAudience() fosite.Arguments { return slices.Clone(c.audiences) }

// IsPublic returns false so Fosite does not treat unauthenticated requests as
// public-client requests. JWT-SVID authentication remains separate from the
// still-pending X.509-SVID credential validation.
func (*SPIFFEClient) IsPublic() bool { return false }

var _ fosite.Client = (*SPIFFEClient)(nil)

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package session provides OAuth session management for the authorization server.
// Sessions link issued access tokens to upstream identity provider tokens,
// enabling token exchange and refresh operations.
package session

import (
        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"
        "github.com/ory/fosite/token/jwt"
)

// Factory is a function that creates a new session given subject, IDP session ID, and client ID.
// This allows storage implementations to create sessions without importing this package directly.
//
// The factory is primarily used during deserialization where the clientID may be empty
// because the client_id claim is preserved in the JWT claims Extra map from the original
// serialized session data.
type Factory func(subject, idpSessionID, clientID string) fosite.Session

// UserClaims holds optional user profile claims from the upstream IDP.
// Using a struct avoids parameter-ordering mistakes in function signatures
// that accept multiple string parameters.
type UserClaims struct {
        // Name is the user's display name (OIDC "name" claim).
        Name string
        // Email is the user's email address (OIDC "email" claim).
        Email string
}

// UpstreamSession is an interface for sessions that support IDP linking and JWT claims.
// It embeds oauth2.JWTSessionContainer (which includes fosite.Session, GetJWTClaims,
// and GetJWTHeader) and adds IDP session tracking.
// This interface is used by the storage layer for serialization/deserialization.
type UpstreamSession interface {
        oauth2.JWTSessionContainer
        GetIDPSessionID() string
}

// TokenSessionIDClaimKey is the JWT claim key for the token session ID.
// This links JWT access tokens to stored upstream IDP tokens.
// We use "tsid" instead of "sid" to avoid confusion with OIDC session management
// which defines "sid" for different purposes (RFC 7519, OIDC Session Management).
const TokenSessionIDClaimKey = "tsid"

// NoUpstreamSessionClaimKey marks a token this authorization server minted
// without any link to an upstream IdP login, so no upstream credential exists
// for it. Today only the RFC 7523 JWT-bearer grant sets it: its subject is a
// synthetic machine identity (issuer#subject) asserted by a separate trust
// root, with no end user and no upstream session.
//
// It exists so a resource server can distinguish "this token was deliberately
// minted without an upstream session" from "this identity was never enriched
// with upstream credentials" — states that are otherwise indistinguishable, as
// both leave Identity.UpstreamTokens nil. pkg/authz/authorizers/cedar reads it
// to decide whether request-token claims may drive policy. The RFC 8693
// delegation grant needs no marker: its "act" claim already says so.
//
// The name is collision-resistant per RFC 7519 Section 4.3, so a JWT from
// another issuer the deployment trusts cannot assert it by accident. It is a
// private claim; RFC 9068 Section 2.2.3 permits additional claims on an
// access token.
const NoUpstreamSessionClaimKey = "https://toolhive.dev/no_upstream_session"

// ClientIDClaimKey is the JWT claim key for the OAuth client ID.
// This identifies which client was issued the token.
const ClientIDClaimKey = "client_id"

// NameClaimKey is the JWT claim key for the user's display name.
// Per OIDC Core Section 5.1.
const NameClaimKey = "name"

// EmailClaimKey is the JWT claim key for the user's email address.
// Per OIDC Core Section 5.1.
const EmailClaimKey = "email"

const subjectIdentifierClaimKey = "sub_id"

// UpstreamSubjectClaimKey is the collision-resistant claim name (RFC 7519 §4.2) of
// the claim that carries the upstream subject a consumer can key on across upstream
// client registrations. It is emitted together with the RFC 9493 sub_id claim and
// under the same disclosure policy.
const UpstreamSubjectClaimKey = "https://toolhive.dev/upstream_subject"

// UpstreamSubject is the verified identity of the upstream OIDC login that
// established the local user. Issuer and Subject are the exact iss and sub of the
// validated upstream ID token. StableClaim names the claim the upstream is
// configured to key users on (its subject_claim, or "sub"), and StableValue is that
// claim's value from the same ID token. All four are set together or not at all.
//
// upstream.Identity carries the same four values as its Verified* fields, because
// the upstream package cannot import this one. verifiedUpstreamIdentity in
// handlers/callback.go converts between them; a member added here must be added
// there and to upstream.Identity too, or IsComplete never holds and the claims
// are silently dropped.
type UpstreamSubject struct {
        Issuer      string `json:"iss,omitempty"`
        Subject     string `json:"sub,omitempty"`
        StableClaim string `json:"stable_claim,omitempty"`
        StableValue string `json:"stable_value,omitempty"`
}

// IsComplete reports whether every member is set. Only a complete value is
// stored or disclosed.
func (u UpstreamSubject) IsComplete() bool {
        return u.Issuer != "" && u.Subject != "" && u.StableClaim != "" && u.StableValue != ""
}

// Session extends fosite's JWT session with an IDP session reference.
// This allows the authorization server to link issued tokens to
// upstream IDP tokens stored separately.
//
// Most methods are provided by the embedded *oauth2.JWTSession. This type
// adds IDP session tracking and verified upstream identity provenance; Clone()
// preserves those fields along with the embedded JWT session.
//
// Concurrency: Sessions are designed to be request-scoped and are not
// safe for concurrent access from multiple goroutines. This follows
// Fosite's design pattern where sessions are created per-request,
// populated by handlers, and then persisted to storage. The storage
// layer is responsible for thread-safe access to stored sessions.
type Session struct {
        *oauth2.JWTSession

        // UpstreamSessionID links this session to stored upstream IDP tokens.
        // This ID is used to retrieve the IDP tokens from storage.
        UpstreamSessionID string

        // VerifiedUpstream is preserved server-side through code/refresh storage. It is
        // not recorded at new logins when SubjectIdentifierAudiences is empty, and
        // becomes the sub_id and upstream subject claims only for listed audiences.
        VerifiedUpstream UpstreamSubject `json:"verified_upstream,omitzero"`
}

// New creates a new Session with the given subject, IDP session ID, client ID, and
// optional user profile claims.
//
// Parameters:
//   - subject: The OAuth subject (user identifier). May be empty for placeholder sessions.
//   - idpSessionID: Links to upstream IDP tokens in storage. If provided, it will be
//     included in the JWT claims as "tsid" to allow proxy middleware to look up
//     upstream IDP tokens.
//   - clientID: The OAuth client ID. When provided, it will be included in the JWT
//     claims as "client_id" for binding verification per RFC 9068.
//   - claims: Optional user profile claims (name, email) from the upstream IDP.
//     Included in the JWT per OIDC Core Section 5.1.
//
// ClientID handling:
//   - For token issuance (authorize handler): Pass the client ID to ensure RFC 9068
//     compliance. The client_id claim will be embedded in the JWT.
//   - For token exchange (token handler): May be empty - Fosite copies claims from
//     the stored authorize session, preserving the original client_id.
//   - For deserialization: May be empty - the client_id claim is preserved in the
//     JWT claims Extra map from the serialized session data.
//
// Note: The remaining RFC 9068 required claims (iss, aud, exp, iat, jti) are
// populated by Fosite during token generation. This session only initializes
// custom claims that are not part of Fosite's standard JWT handling.
func New(subject, idpSessionID, clientID string, claims UserClaims) *Session {
        // Initialize the Extra map for JWT claims
        claimsExtra := make(map[string]any)

        // Add tsid claim if idpSessionID is provided
        if idpSessionID != "" {
                claimsExtra[TokenSessionIDClaimKey] = idpSessionID
        }

        // Add client_id claim for binding verification per RFC 9068
        // This may be empty for placeholder sessions or deserialization;
        // in those cases the claim is preserved from the original session.
        if clientID != "" {
                claimsExtra[ClientIDClaimKey] = clientID
        }

        // Add user profile claims per OIDC Core Section 5.1
        if claims.Name != "" {
                claimsExtra[NameClaimKey] = claims.Name
        }
        if claims.Email != "" {
                claimsExtra[EmailClaimKey] = claims.Email
        }

        return &Session{
                JWTSession: &oauth2.JWTSession{
                        JWTClaims: &jwt.JWTClaims{
                                Subject: subject,
                                Extra:   claimsExtra,
                        },
                        JWTHeader: &jwt.Headers{
                                Extra: make(map[string]any),
                        },
                        Subject: subject, // Also set on JWTSession for fosite compatibility
                },
                UpstreamSessionID: idpSessionID,
        }
}

// Clone creates a deep copy of the session.
// This overrides the embedded JWTSession.Clone() to also copy UpstreamSessionID
// and the verified upstream identity.
func (s *Session) Clone() fosite.Session {
        if s == nil {
                return nil
        }

        var jwtSession *oauth2.JWTSession
        if s.JWTSession != nil {
                jwtSession = s.JWTSession.Clone().(*oauth2.JWTSession)
        }

        return &Session{
                JWTSession:        jwtSession,
                UpstreamSessionID: s.UpstreamSessionID,
                VerifiedUpstream:  s.VerifiedUpstream,
        }
}

// ClearSubjectIdentifier removes any persisted or source-token sub_id and
// upstream subject claim before the current grant decides whether it is
// eligible to issue them. It mutates the receiver's claims in place; a caller
// holding a session loaded from storage must Clone it first.
func (s *Session) ClearSubjectIdentifier() {
        if s == nil || s.JWTSession == nil || s.JWTClaims == nil {
                return
        }
        delete(s.JWTClaims.Extra, subjectIdentifierClaimKey)
        delete(s.JWTClaims.Extra, UpstreamSubjectClaimKey)
}

// SetSubjectIdentifier sets the RFC 9493 sub_id claim and the upstream subject
// claim from the verified upstream identity. It does nothing unless the identity
// is complete. It mutates the receiver's claims in place; a caller holding a
// session loaded from storage must Clone it first.
func (s *Session) SetSubjectIdentifier(upstream UpstreamSubject) {
        if s == nil || s.JWTSession == nil || s.JWTClaims == nil || !upstream.IsComplete() {
                return
        }
        if s.JWTClaims.Extra == nil {
                s.JWTClaims.Extra = make(map[string]any)
        }
        s.JWTClaims.Extra[subjectIdentifierClaimKey] = map[string]string{
                "format": "iss_sub",
                "iss":    upstream.Issuer,
                "sub":    upstream.Subject,
        }
        s.JWTClaims.Extra[UpstreamSubjectClaimKey] = map[string]string{
                "iss":   upstream.Issuer,
                "claim": upstream.StableClaim,
                "value": upstream.StableValue,
        }
}

// GetIDPSessionID returns the IDP session ID.
func (s *Session) GetIDPSessionID() string {
        if s == nil {
                return ""
        }
        return s.UpstreamSessionID
}

// SetIDPSessionID sets the IDP session ID.
func (s *Session) SetIDPSessionID(id string) {
        if s == nil {
                return
        }
        s.UpstreamSessionID = id
}

// Compile-time interface compliance check.
var _ UpstreamSession = (*Session)(nil)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package server

import (
        "context"
        "net/http"
        "net/url"
        "slices"
        "strings"
        "time"

        "github.com/ory/fosite"
        "github.com/spiffe/go-spiffe/v2/bundle/jwtbundle"
        "github.com/spiffe/go-spiffe/v2/svid/jwtsvid"

        spiffeauth "github.com/stacklok/toolhive/pkg/authserver/spiffe"
)

// maxSPIFFEJWTAssertionRemainingValidity permits the recommended five-minute
// JWT-SVID lifetime plus the one-minute clock skew tolerated by go-spiffe.
const maxSPIFFEJWTAssertionRemainingValidity = 6 * time.Minute

// SPIFFEClientResolver resolves a verified SPIFFE identity to its configured
// OAuth client. It is the seam that lets the client-authentication strategy
// here reach the association registry and storage constructed in package
// authserver, which this package cannot import (authserver imports server).
// clientID is an optional exact selector; an empty value asks the resolver to
// derive the configured client from the verified SPIFFE identity.
// One signature covers both X.509 and JWT credentials, with method as an
// explicit discriminator, so the two arms share a single resolution path
// instead of each inventing its own. spiffeID is passed explicitly rather
// than pulled from ctx, so an identity resolved from the wrong source cannot
// be mistaken for a verified one.
type SPIFFEClientResolver func(
        ctx context.Context, spiffeID, clientID string, method spiffeauth.SPIFFEAuthenticationMethod,
) (fosite.Client, error)

func newSPIFFEClientAuthenticationStrategy(
        defaultStrategy fosite.ClientAuthenticationStrategy,
        issuer string,
        jwtBundleSource jwtbundle.Source,
        resolver SPIFFEClientResolver,
) fosite.ClientAuthenticationStrategy {
        return func(ctx context.Context, r *http.Request, form url.Values) (fosite.Client, error) {
                // Without SPIFFE trust, every request goes to the default strategy
                // untouched — this arm must not change shared-dispatcher behavior for
                // entirely non-SPIFFE requests (e.g. RFC 7523 private-key JWT).
                if resolver == nil {
                        return defaultStrategy(ctx, r, form)
                }
                // An explicit assertion type takes precedence over an ambient mTLS identity.
                // A repeated form key must be checked in full: form.Get would only see the
                // first value, letting a SPIFFE assertion type hidden behind an earlier
                // value slip through to the default strategy. A duplicated
                // client_assertion_type is rejected as ambiguous by
                // authenticateSPIFFEJWTClient itself, once SPIFFE is actually selected.
                if slices.Contains(form["client_assertion_type"], spiffeauth.SPIFFEJWTAssertionType) {
                        return authenticateSPIFFEJWTClient(ctx, r, form, issuer, jwtBundleSource, resolver)
                }
                if _, ok := spiffeauth.SPIFFEIDFromContext(ctx); ok {
                        return nil, fosite.ErrInvalidClient.WithHint("SPIFFE X.509 client authentication is not implemented")
                }
                return defaultStrategy(ctx, r, form)
        }
}

// authenticateSPIFFEJWTClient validates a SPIFFE JWT-SVID client assertion and
// resolves it to its configured OAuth client. It fails closed on any
// malformed request field, verification failure, or resolver rejection, and
// never logs the assertion or its claims.
func authenticateSPIFFEJWTClient(
        ctx context.Context,
        r *http.Request,
        form url.Values,
        issuer string,
        jwtBundleSource jwtbundle.Source,
        resolver SPIFFEClientResolver,
) (fosite.Client, error) {
        assertionType, ok := exactNonEmptyFormValue(form, "client_assertion_type")
        if !ok || assertionType != spiffeauth.SPIFFEJWTAssertionType {
                return nil, fosite.ErrInvalidRequest
        }
        assertion, ok := exactNonEmptyFormValue(form, "client_assertion")
        if !ok {
                return nil, fosite.ErrInvalidRequest
        }
        clientID, ok := optionalExactNonEmptyFormValue(form, "client_id")
        if !ok {
                return nil, fosite.ErrInvalidRequest
        }
        if rejectedSPIFFEJWTRequest(r, form, assertion) {
                return nil, fosite.ErrInvalidClient
        }
        if jwtBundleSource == nil {
                return nil, fosite.ErrInvalidClient
        }

        svid, err := jwtsvid.ParseAndValidate(assertion, jwtBundleSource, []string{issuer})
        if err != nil || !validSPIFFEJWTIdentityClaims(svid, issuer) {
                return nil, fosite.ErrInvalidClient
        }
        if svid.Expiry.After(time.Now().Add(maxSPIFFEJWTAssertionRemainingValidity)) {
                return nil, fosite.ErrInvalidClient
        }
        client, err := resolver(ctx, svid.ID.String(), clientID, spiffeauth.SPIFFEAuthenticationMethodJWT)
        if !validResolvedSPIFFEClient(client, err, clientID) {
                return nil, fosite.ErrInvalidClient
        }
        return client, nil
}

// validSPIFFEJWTIdentityClaims reports whether svid carries exactly the AS as
// its sole audience. It does not require an `iss` claim: RFC 7519 states iss
// is OPTIONAL, and the SPIFFE JWT-SVID spec does not mandate it either — a
// conformant SVID signed by a bundle-trusted key can omit it entirely.
func validSPIFFEJWTIdentityClaims(svid *jwtsvid.SVID, audience string) bool {
        return svid != nil && len(svid.Audience) == 1 && svid.Audience[0] == audience
}

func validResolvedSPIFFEClient(client fosite.Client, err error, requestedClientID string) bool {
        return err == nil && client != nil && (requestedClientID == "" || client.GetID() == requestedClientID)
}

// exactNonEmptyFormValue returns the sole value for key, rejecting an absent,
// duplicated, or whitespace-only field. Duplicated fields are rejected rather
// than taking the first or last value, since either choice could be steered
// by an attacker who controls only one of the duplicates.
func exactNonEmptyFormValue(form url.Values, key string) (string, bool) {
        values, ok := form[key]
        if !ok || len(values) != 1 || strings.TrimSpace(values[0]) == "" {
                return "", false
        }
        return values[0], true
}

// optionalExactNonEmptyFormValue returns an absent field as an empty value.
// A supplied field must contain exactly one non-whitespace value, which is
// returned without normalization.
func optionalExactNonEmptyFormValue(form url.Values, key string) (string, bool) {
        if !form.Has(key) {
                return "", true
        }
        return exactNonEmptyFormValue(form, key)
}

// rejectedSPIFFEJWTRequest reports whether a request mixes SPIFFE JWT
// authentication with another client credential (HTTP Basic or a client
// secret form field), or carries an oversized assertion. SPIFFE JWT
// authentication must be the sole credential in the request; allowing a
// second one to also be present would let an attacker probe which the
// server accepts.
func rejectedSPIFFEJWTRequest(r *http.Request, form url.Values, assertion string) bool {
        return hasBasicAuthorization(r) || form.Has("client_secret") || len(assertion) > 16*1024
}

func hasBasicAuthorization(r *http.Request) bool {
        for _, authorization := range r.Header.Values("Authorization") {
                fields := strings.Fields(authorization)
                if len(fields) > 0 && strings.EqualFold(fields[0], "Basic") {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package tokenexchange

import (
        "context"
        "encoding/json"
        "errors"
        "log/slog"
        "net/http"
        "net/url"
        "slices"

        "github.com/ory/fosite"
        "github.com/ory/x/errorsx"

        "github.com/stacklok/toolhive/pkg/authserver/authorizationdetails"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

const authorizationDetailsParam = "authorization_details"

// Reason codes logged when authorization_details is rejected. They are fixed
// strings: neither the payload nor a hook's error text is ever logged at WARN.
const (
        rejectMalformed    = "malformed"
        rejectRepeated     = "repeated"
        rejectHookRejected = "hook_rejected"
        rejectHookFailed   = "hook_failed"
)

// parseAuthorizationDetails runs the early, cheap shape checks. It returns a
// nil Pending when the parameter is ignored: no registry is configured, or the
// parameter is absent or empty.
func (h *Handler) parseAuthorizationDetails(form url.Values, clientID string) (*authorizationdetails.Pending, error) {
        if h.authorizationDetailsRegistry == nil {
                return nil, nil
        }
        values := form[authorizationDetailsParam]
        if len(values) == 0 || (len(values) == 1 && values[0] == "") {
                return nil, nil
        }
        if len(values) > 1 {
                warnAuthorizationDetailsRejected(clientID, rejectRepeated)
                return nil, &fosite.RFC6749Error{
                        ErrorField:       string(authorizationdetails.CodeInvalidRequest),
                        DescriptionField: "The authorization_details parameter must not be repeated.",
                        CodeField:        http.StatusBadRequest,
                }
        }
        pending, err := h.authorizationDetailsRegistry.Parse(values[0], oauthproto.GrantTypeTokenExchange)
        if err != nil {
                warnAuthorizationDetailsRejected(clientID, rejectMalformed)
                return nil, invalidAuthorizationDetailsError()
        }
        return pending, nil
}

// grantAuthorizationDetails asks the registered hooks what to grant and stores
// the granted JSON in the request form for PopulateTokenEndpointResponse.
func (*Handler) grantAuthorizationDetails(
        ctx context.Context,
        pending *authorizationdetails.Pending,
        requester fosite.AccessRequester,
        params *formParams,
        claims *ValidatedClaims,
        actorSub string,
) error {
        form := requester.GetRequestForm()
        clientID := requester.GetClient().GetID()

        granted, err := pending.Grant(ctx, authorizationdetails.Request{
                ClientID:           clientID,
                Actor:              actorSub,
                ActorTokenPresent:  params.actorToken != "",
                RequestedTokenType: form.Get("requested_token_type"),
                RequestedScopes:    slices.Clone([]string(requester.GetRequestedScopes())),
                GrantedAudience:    slices.Clone([]string(requester.GetGrantedAudience())),
                Subject: authorizationdetails.Subject{
                        Issuer:   claims.Issuer,
                        Subject:  claims.Subject,
                        ClientID: claims.ClientID,
                        JWTID:    claims.JWTID,
                        Audience: slices.Clone(claims.Audience),
                        IssuedAt: claims.IssuedAt,
                        Expiry:   claims.Expiry,
                        External: claims.ExternalIssuer != "",
                        Claims:   deepCopyClaims(claims.Claims),
                },
        })
        if err != nil {
                return authorizationDetailsGrantError(clientID, err)
        }

        // The form is dropped by fosite's Sanitize before the request is stored,
        // so the granted JSON reaches only this request's response.
        form.Set(authorizationDetailsParam, string(granted))
        return nil
}

// echoAuthorizationDetails adds the granted authorization_details to the
// response. The value is the granted JSON that grantAuthorizationDetails
// stored in the form; without a registry the parameter is never echoed.
func (h *Handler) echoAuthorizationDetails(requester fosite.AccessRequester, responder fosite.AccessResponder) {
        if h.authorizationDetailsRegistry == nil {
                return
        }
        if granted := requester.GetRequestForm().Get(authorizationDetailsParam); granted != "" {
                responder.SetExtra(authorizationDetailsParam, json.RawMessage(granted))
        }
}

// authorizationDetailsGrantError maps a Grant failure to an OAuth error and
// logs the rejection once at WARN.
func authorizationDetailsGrantError(clientID string, err error) error {
        var typed authorizationdetails.Error
        switch {
        case errors.As(err, &typed):
                slog.Warn("authorization_details rejected",
                        "client_id", clientID, "reason", rejectHookRejected,
                        "code", string(typed.Code), "hook_reason", typed.Reason)
                return &fosite.RFC6749Error{
                        ErrorField:       string(typed.Code),
                        DescriptionField: typed.Reason,
                        CodeField:        http.StatusBadRequest,
                }
        default:
                warnAuthorizationDetailsRejected(clientID, rejectHookFailed)
                // The error text may hold internal data, so it is never logged or
                // returned to the client.
                return errorsx.WithStack(fosite.ErrServerError.WithDescription("The authorization_details could not be processed."))
        }
}

func warnAuthorizationDetailsRejected(clientID, reason string) {
        slog.Warn("authorization_details rejected", "client_id", clientID, "reason", reason)
}

func invalidAuthorizationDetailsError() *fosite.RFC6749Error {
        return &fosite.RFC6749Error{
                ErrorField:       string(authorizationdetails.CodeInvalidAuthorizationDetails),
                DescriptionField: "The authorization_details parameter is invalid.",
                CodeField:        http.StatusBadRequest,
        }
}

// deepCopyClaims copies a JSON-decoded claim map so a hook cannot modify the
// validated claims. Scalars are immutable; only maps and slices are cloned.
func deepCopyClaims(claims map[string]any) map[string]any {
        if claims == nil {
                return nil
        }
        out := make(map[string]any, len(claims))
        for key, value := range claims {
                out[key] = deepCopyValue(value)
        }
        return out
}

func deepCopyValue(value any) any {
        switch v := value.(type) {
        case map[string]any:
                return deepCopyClaims(v)
        case []any:
                out := make([]any, len(v))
                for i, item := range v {
                        out[i] = deepCopyValue(item)
                }
                return out
        default:
                return value
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package tokenexchange

import (
        "fmt"
        "time"

        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"

        "github.com/stacklok/toolhive/pkg/authserver/authorizationdetails"
        "github.com/stacklok/toolhive/pkg/authserver/server"
)

// NewSharedTrustedIssuerValidator builds the single MultiIssuerTokenValidator
// that backs a server's trusted-issuer handling. buildProvider constructs it
// whenever any trusted issuer is configured — not only when the RFC 7523
// JWT-bearer grant is also enabled — and passes it to
// FactoryWithSharedTrustedIssuerValidator (and, when the JWT-bearer grant is
// enabled, JWTBearerIssuanceFactory), both of which now require it for a
// non-empty trusted-issuer set rather than building their own. One shared
// instance keeps a single JWKS cache/goroutine set per issuer and, crucially,
// gives the server one validator to Close on shutdown so those goroutines are
// released (see MultiIssuerTokenValidator.Close). Returns (nil, nil) when
// trustedIssuers is empty, in which case no external validator is built.
func NewSharedTrustedIssuerValidator(
        config *server.AuthorizationServerConfig, trustedIssuers []TrustedIssuer,
) (*MultiIssuerTokenValidator, error) {
        if len(trustedIssuers) == 0 {
                return nil, nil
        }
        if config == nil {
                return nil, fmt.Errorf("authorization server config is required when trusted issuers are configured")
        }
        selfValidator, err := NewSelfIssuedTokenValidator(config.PublicJWKS(), config.GetAccessTokenIssuer(), config.AllowedAudiences)
        if err != nil {
                return nil, fmt.Errorf("failed to create self validator: %w", err)
        }
        return NewMultiIssuerTokenValidator(selfValidator, config.GetAccessTokenIssuer(), trustedIssuers, config.AllowedAudiences)
}

// FactoryOption configures a token-exchange Handler factory.
type FactoryOption func(*factoryOptions)

type factoryOptions struct {
        authorizationDetailsRegistry *authorizationdetails.Registry
}

// WithAuthorizationDetailsRegistry configures the factory to pass the registry
// through to the token-exchange Handler. With no registry (or a nil one) the
// authorization_details parameter is ignored.
func WithAuthorizationDetailsRegistry(registry *authorizationdetails.Registry) FactoryOption {
        return func(o *factoryOptions) { o.authorizationDetailsRegistry = registry }
}

// Factory returns a server.Factory that creates a token exchange Handler.
// The delegationLifespan parameter sets the maximum lifetime for delegated tokens;
// the actual lifetime is the minimum of this value and the subject token's remaining lifetime.
// Returns an error if delegationLifespan is not in (0, server.MaxAccessTokenLifespan]: a zero
// or negative value would produce delegated tokens with an expiry already in the past, and a
// value above the access token ceiling would only be caught at request time by the per-request cap.
//
// trustedIssuers must be empty here. A MultiIssuerTokenValidator owns
// per-issuer JWKS refresh worker pools that only its Close releases, and the
// bare Factory has no way to hand that instance back to the caller for
// shutdown (it is built at fosite-compose time, from config not available at
// this call). Passing a non-empty set therefore returns an error rather than
// silently building a validator whose workers leak; callers with trusted
// issuers must build it via NewSharedTrustedIssuerValidator, hold it for Close,
// and pass it to FactoryWithSharedTrustedIssuerValidator. With no trusted
// issuers the self-issued validator is used directly, preserving prior
// behavior exactly.
//
// configuredDelegateClients is the operator-configured list of delegate
// client IDs (Config.DelegateClients, projected down to just their
// ClientIDs by the caller). An empty list preserves existing behavior
// exactly. The trust source here is server config, not client storage: the
// set is read once at process construction, so removing a client from
// config revokes its trust on the next restart rather than requiring any
// explicit revocation step against storage.
func Factory(
        delegationLifespan time.Duration, trustedIssuers []TrustedIssuer, configuredDelegateClients []string,
        options ...FactoryOption,
) (server.Factory, error) {
        return FactoryWithSharedTrustedIssuerValidator(
                delegationLifespan, trustedIssuers, configuredDelegateClients, nil, options...)
}

// FactoryWithSharedTrustedIssuerValidator is Factory with a shared
// external-issuer validator. shared is used as the subject-token validator when
// non-nil; it is REQUIRED whenever trustedIssuers is non-empty (an error is
// returned otherwise), because a locally-built MultiIssuerTokenValidator's JWKS
// refresh workers would have no owner to Close them — see the error below.
// Callers build it once with NewSharedTrustedIssuerValidator, hold it for Close,
// and can reuse the same instance across the RFC 8693 token-exchange and RFC
// 7523 JWT-bearer grants.
func FactoryWithSharedTrustedIssuerValidator(
        delegationLifespan time.Duration, trustedIssuers []TrustedIssuer, configuredDelegateClients []string,
        shared *MultiIssuerTokenValidator, options ...FactoryOption,
) (server.Factory, error) {
        if delegationLifespan <= 0 || delegationLifespan > server.MaxAccessTokenLifespan {
                return nil, fmt.Errorf("tokenexchange: delegationLifespan must be between %v and %v, got %v",
                        time.Duration(0), server.MaxAccessTokenLifespan, delegationLifespan)
        }
        for _, id := range configuredDelegateClients {
                if id == "" {
                        return nil, fmt.Errorf("tokenexchange: configuredDelegateClients must not contain an empty client ID")
                }
        }
        // A trusted-issuer validator owns per-issuer JWKS refresh worker pools that
        // only its Close releases, but the returned closure (built at fosite-compose
        // time, from a config not available here) cannot hand that instance back to
        // the caller for shutdown. Requiring the caller to build it up front via
        // NewSharedTrustedIssuerValidator and pass it as shared is the only
        // construction path that stays releasable — fail loudly rather than silently
        // build a leaked one.
        if shared == nil && len(trustedIssuers) > 0 {
                return nil, fmt.Errorf("tokenexchange: trusted issuers require a shared validator built via " +
                        "NewSharedTrustedIssuerValidator so its JWKS refresh workers can be released on shutdown")
        }
        resolved := resolveFactoryOptions(options)
        return func(config *server.AuthorizationServerConfig, storage fosite.Storage, strategy any) (any, error) {
                selfValidator, err := NewSelfIssuedTokenValidator(config.PublicJWKS(), config.GetAccessTokenIssuer(), config.AllowedAudiences)
                if err != nil {
                        return nil, fmt.Errorf("tokenexchange: failed to create subject token validator: %w", err)
                }

                // shared is guaranteed non-nil whenever trustedIssuers is non-empty
                // (checked above), so the trusted-issuer path always uses the
                // caller-owned, closeable validator; this closure never constructs one
                // whose JWKS workers nothing can release. The IIFE keeps validator a
                // single immutable assignment (go-style).
                validator := func() SubjectTokenValidator {
                        if shared != nil {
                                return shared
                        }
                        return selfValidator
                }()

                // Use the embedded *fosite.Config for HandleHelper and handlerConfig
                // because AuthorizationServerConfig shadows GetAccessTokenLifespan() without
                // a context parameter, which doesn't satisfy fosite's provider interfaces.
                atStrategy, ok := strategy.(oauth2.AccessTokenStrategy)
                if !ok {
                        return nil, fmt.Errorf("tokenexchange: strategy does not implement oauth2.AccessTokenStrategy (got %T)", strategy)
                }
                atStorage, ok := storage.(oauth2.AccessTokenStorage)
                if !ok {
                        return nil, fmt.Errorf("tokenexchange: storage does not implement oauth2.AccessTokenStorage (got %T)", storage)
                }
                return &Handler{
                        HandleHelper: &oauth2.HandleHelper{
                                AccessTokenStrategy: atStrategy,
                                AccessTokenStorage:  atStorage,
                                Config:              config.Config,
                        },
                        validator:                    validator,
                        selfValidator:                selfValidator,
                        issuer:                       config.GetAccessTokenIssuer(),
                        delegationLifespan:           delegationLifespan,
                        config:                       config.Config,
                        allowedAudiences:             config.AllowedAudiences,
                        configuredDelegateClients:    configuredDelegateClients,
                        authorizationDetailsRegistry: resolved.authorizationDetailsRegistry,
                }, nil
        }, nil
}

func resolveFactoryOptions(options []FactoryOption) factoryOptions {
        var resolved factoryOptions
        for _, option := range options {
                option(&resolved)
        }
        return resolved
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package tokenexchange

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "net/url"
        "slices"
        "strings"
        "time"

        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"
        "github.com/ory/x/errorsx"

        coreaudit "github.com/stacklok/toolhive-core/audit"
        "github.com/stacklok/toolhive/pkg/authserver/authorizationdetails"
        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// Compile-time check that Handler implements fosite.TokenEndpointHandler.
var _ fosite.TokenEndpointHandler = (*Handler)(nil)

// maxDelegationDepth bounds how many nested "act" entries a subject token's
// delegation chain may carry. Without a cap, a client could repeatedly
// re-exchange a delegated token to grow the chain without limit, bloating the
// issued token and the load it places on downstream resource servers that
// parse it.
//
// maxActClaimSize bounds the serialized normalized act data (iss/sub) after
// chainToAct strips non-identity claims, preventing a shallow prior chain from
// inflating every re-issued token.
const (
        maxDelegationDepth = 10
        maxActClaimSize    = 8 * 1024
)

// anyDelegateClient is the TrustedIssuer.AllowedDelegateClients entry that
// explicitly opts an issuer into "any ToolHive confidential client holding
// the token-exchange grant may act as its delegate" — the behavior that used
// to be the implicit default whenever the field was left empty. See
// checkDelegationConsent's doc comment for why that default was removed:
// permissiveness must now be declared with this wildcard, not left to an
// absent field.
const anyDelegateClient = "*"

// Handler implements RFC 8693 token exchange for user-to-agent delegation.
//
// When an authenticated OAuth client (the acting agent) presents a user's JWT
// as subject_token, the handler validates the token and issues a delegated JWT
// with sub=user and an act claim containing the client's identity, per
// RFC 8693 Section 4.1.
//
// Subject tokens are intentionally reusable within their lifetime: per
// RFC 8693's security considerations, a token exchange does not invalidate
// the subject token, so the same subject token may be exchanged more than
// once. Replay is bounded by the delegated token's lifetime cap
// (min(subject-remaining, delegation)), not by single-use tracking; per-jti
// single-use enforcement is deferred to the broader M2M/sender-constrained-
// token effort.
type Handler struct {
        *oauth2.HandleHelper
        validator          SubjectTokenValidator // for subject tokens (multi-issuer)
        selfValidator      SubjectTokenValidator // for actor tokens (self-issued only)
        issuer             string
        delegationLifespan time.Duration
        config             tokenExchangeConfig
        allowedAudiences   []string
        // configuredDelegateClients holds the operator-configured delegate-client
        // IDs (Config.DelegateClients, ID-only).
        configuredDelegateClients []string
        // authorizationDetailsRegistry validates RFC 9396 details on token exchange.
        authorizationDetailsRegistry *authorizationdetails.Registry
}

// formParams holds the validated RFC 8693 form parameters extracted from a
// token exchange request.
type formParams struct {
        subjectToken string
        actorToken   string // empty if not provided
}

// tokenExchangeConfig defines the configuration interface needed by the handler.
type tokenExchangeConfig interface {
        fosite.ScopeStrategyProvider
        fosite.AudienceStrategyProvider
        fosite.AccessTokenLifespanProvider
}

// CanHandleTokenEndpointRequest returns true if the request's grant_type is
// the RFC 8693 token exchange grant type.
func (*Handler) CanHandleTokenEndpointRequest(_ context.Context, requester fosite.AccessRequester) bool {
        return requester.GetGrantTypes().ExactOne(oauthproto.GrantTypeTokenExchange)
}

// CanSkipClientAuth always returns false: RFC 8693 token exchange requires
// client authentication, and other handlers decide their own requirements.
func (*Handler) CanSkipClientAuth(_ context.Context, _ fosite.AccessRequester) bool {
        return false
}

// HandleTokenEndpointRequest validates the token exchange request parameters,
// verifies the subject token, and constructs a delegated session with the act claim.
//
// The delegated token's lifetime is the minimum of the subject token's remaining
// lifetime and the configured delegation lifespan.
func (h *Handler) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }

        client := requester.GetClient()

        // The client MUST be confidential — only authenticated confidential clients
        // may act on behalf of a user.
        if client.IsPublic() {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                        "The OAuth 2.0 Client is marked as public and is thus not allowed to use authorization grant 'token-exchange'."))
        }

        // Reject clients not registered for this grant type up front, before any
        // subject-token validation work. PopulateTokenEndpointResponse repeats
        // this check as defense-in-depth.
        if !client.GetGrantTypes().Has(oauthproto.GrantTypeTokenExchange) {
                return errorsx.WithStack(fosite.ErrUnauthorizedClient.WithHint(
                        "The OAuth 2.0 Client is not allowed to use authorization grant 'token-exchange'."))
        }

        // The authenticated client is the acting party ("actor"). The client is
        // already authenticated by fosite's client authentication strategy before
        // this handler runs.
        actorID := client.GetID()
        form := requester.GetRequestForm()

        // Validate required RFC 8693 form parameters.
        params, err := validateFormParams(form)
        if err != nil {
                return err
        }

        pendingDetails, err := h.parseAuthorizationDetails(form, actorID)
        if err != nil {
                return err
        }

        // Validate requested_token_type per RFC 8693 Section 2.1: if the client
        // requests a token type the server does not support, the request must fail.
        requestedTokenType := form.Get("requested_token_type")
        if requestedTokenType != "" && requestedTokenType != oauthproto.TokenTypeAccessToken {
                return errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf(
                        "The 'requested_token_type' value %q is not supported. This server only issues %q.",
                        requestedTokenType, oauthproto.TokenTypeAccessToken))
        }

        // Validate the subject token against the configured token validator.
        validatedClaims, err := h.validator.Validate(ctx, params.subjectToken)
        if err != nil {
                slog.Debug("Subject token validation failed",
                        "error", err,
                        "actor", actorID,
                )
                // RFC 8693 §2.2.2 mandates invalid_request here: "if either the
                // subject_token or actor_token are invalid for any reason, or are
                // unacceptable based on policy ... The value of the error parameter
                // MUST be the invalid_request error code." checkDelegationConsent
                // below deliberately returns invalid_grant instead for its own
                // failures — a documented deviation: RFC 6749 §5.2's invalid_grant
                // covers "was issued to another client", §2.2.2 permits other error
                // codes for other failures, and both Keycloak and Hydra follow this
                // same split.
                return errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The subject token is invalid or could not be verified."))
        }

        // Resolve actor identity: explicit actor_token or authenticated client.
        actorSub, err := h.resolveActorIdentity(ctx, params, client)
        if err != nil {
                return err
        }

        if err := h.authorizeDelegation(ctx, requester, client, params, validatedClaims, actorSub, pendingDetails); err != nil {
                return err
        }

        // Build the delegated session with the user's identity and the
        // authenticated client's identity. The third argument becomes the
        // issued token's RFC 9068 client_id, so it must identify the client the
        // token was actually issued to (client.GetID()) — not actorSub, which
        // may name a distinct actor asserted via actor_token and belongs only
        // in the act claim below.
        delegatedSession := session.New(
                delegatedSubject(validatedClaims),
                "", // No IDP session link for delegated tokens.
                client.GetID(),
                session.UserClaims{
                        Name:  validatedClaims.Name,
                        Email: validatedClaims.Email,
                },
        )

        act, externalIssuer, err := buildActClaim(validatedClaims, h.issuer, actorSub)
        if err != nil {
                return err
        }
        delegatedSession.JWTClaims.Extra["act"] = act
        // external_issuer records that this exchange (or, on re-exchange, a
        // prior one in the chain) involved a trusted external issuer,
        // independent of whether a client-namespace actor was also resolved
        // (see buildActClaim's doc comment for why this is not nested inside
        // act, and for how it is carried forward across re-exchanges).
        if externalIssuer != "" {
                delegatedSession.JWTClaims.Extra["external_issuer"] = externalIssuer
        }

        // Compute the delegated token lifetime: the shorter of the subject token's
        // remaining lifetime and the configured delegation lifespan.
        lifetime, err := h.computeLifetime(validatedClaims.Expiry)
        if err != nil {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                        "The subject token has expired."))
        }
        delegatedSession.SetExpiresAt(fosite.AccessToken, time.Now().UTC().Add(lifetime))

        requester.SetSession(delegatedSession)

        slog.Debug("Token exchange request validated",
                "subject", validatedClaims.Subject,
                "actor", actorSub,
                "issuer", validatedClaims.Issuer,
                "subject_token_client", validatedClaims.ExternalActor,
                "subject_token_client_id", validatedClaims.ClientID,
                "lifetime", lifetime.String(),
        )

        return nil
}

// PopulateTokenEndpointResponse issues the delegated access token and sets
// the RFC 8693 issued_token_type in the response. When an authorization_details
// registry is configured and details were granted, it also echoes the granted
// details as the RFC 9396 authorization_details response member.
func (h *Handler) PopulateTokenEndpointResponse(
        ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder,
) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }

        if !requester.GetClient().GetGrantTypes().Has(oauthproto.GrantTypeTokenExchange) {
                return errorsx.WithStack(fosite.ErrUnauthorizedClient.WithHint(
                        "The OAuth 2.0 Client is not allowed to use authorization grant 'token-exchange'."))
        }

        // Use the session's ExpiresAt (set during HandleTokenEndpointRequest) as the
        // authoritative lifetime. This respects the min(subject_remaining, delegation)
        // bound computed earlier. Fall back to the configured access token lifespan
        // only if no expiry was set on the session.
        atLifespan := h.config.GetAccessTokenLifespan(ctx)
        if sessionExpiry := requester.GetSession().GetExpiresAt(fosite.AccessToken); !sessionExpiry.IsZero() {
                remaining := time.Until(sessionExpiry)
                if remaining <= 0 {
                        // The session's bound has already elapsed since HandleTokenEndpointRequest
                        // computed it. Fail closed instead of signing a token whose exp is already
                        // in the past while reporting a negative expires_in — fosite's JWT strategy
                        // sets exp directly from this session's expiry (GetExpiresAt), independent
                        // of atLifespan, so letting this fall through would not simply hand out a
                        // longer-lived token, it would issue an already-expired one.
                        return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                                "The delegated token's bound lifetime has already elapsed."))
                }
                if remaining < atLifespan {
                        atLifespan = remaining
                }
        }

        if _, err := h.IssueAccessToken(ctx, atLifespan, requester, responder); err != nil {
                return err
        }

        // Per RFC 8693 Section 2.2.1, the response MUST include issued_token_type.
        responder.SetExtra("issued_token_type", oauthproto.TokenTypeAccessToken)

        h.echoAuthorizationDetails(requester, responder)

        return nil
}

// authorizeDelegation runs consent, then scopes, then audiences, then the
// authorization_details hook, in that order, so the hook sees only a request
// that passed every other check.
func (h *Handler) authorizeDelegation(
        ctx context.Context,
        requester fosite.AccessRequester,
        client fosite.Client,
        params *formParams,
        validatedClaims *ValidatedClaims,
        actorSub string,
        pendingDetails *authorizationdetails.Pending,
) error {
        configuredDelegate := slices.Contains(h.configuredDelegateClients, client.GetID())
        if err := checkDelegationConsent(validatedClaims, client.GetID(), actorSub, configuredDelegate); err != nil {
                return err
        }

        if err := h.grantScopes(ctx, requester, client, validatedClaims); err != nil {
                return err
        }

        if err := h.grantAndBoundAudiences(ctx, requester, client, validatedClaims); err != nil {
                return err
        }

        if pendingDetails != nil {
                if err := h.grantAuthorizationDetails(ctx, pendingDetails, requester, params, validatedClaims, actorSub); err != nil {
                        return err
                }
        }

        return nil
}

// resolveActorIdentity determines the acting party identity: either the
// authenticated OAuth client ID, or — when actor_token is present — the
// distinct actor identity it asserts.
//
// actor_token lets the authenticated client name a more specific actor than
// its own client_id (e.g. a particular agent instance or delegate persona)
// the same way a normal issued token records its client identity: via a
// "client_id" claim. The actor_token's own "client_id" claim MUST match the
// authenticated client ID — this is the client-ID binding check, proving the
// token was minted for this client — while its "sub" claim is the actor
// identity that is returned here and flows into the delegated token's act.sub
// the same place a normal client's identity would go. "sub" is deliberately
// not compared to client.GetID(): requiring equality there would make
// actor_token unable to ever assert an identity different from
// the OAuth client, collapsing it to a no-op self-check. Note that this does
// not by itself grant the asserted actor any extra privilege: the resulting
// actor identity still has to satisfy checkDelegationConsent (may_act,
// ExternalActor, client_id binding, or the configured-delegate exemption)
// like any other actor identity would.
func (h *Handler) resolveActorIdentity(
        ctx context.Context, params *formParams, client fosite.Client,
) (string, error) {
        if params.actorToken != "" {
                // Validate actor_token against the AS's own JWKS (must be self-issued).
                actorClaims, err := h.selfValidator.Validate(ctx, params.actorToken)
                if err != nil {
                        slog.Debug("Actor token validation failed", "error", fmt.Errorf("actor token: %w", err))
                        return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                                "The actor token is invalid or could not be verified."))
                }
                // Client-ID binding: actor_token's client_id claim MUST match the
                // authenticated client ID. This prevents replay attacks where a leaked
                // actor token is presented by a different client. The client ID is
                // always verified by fosite's client authentication before reaching here.
                // RFC 8693 §2.2.2 requires invalid_request when an actor token is
                // unacceptable based on policy.
                if actorClaims.ClientID != client.GetID() {
                        return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                                "The actor token's client_id claim does not match the authenticated client identity."))
                }
                return actorClaims.Subject, nil
        }

        // No actor_token: the authenticated client is the acting party.
        return client.GetID(), nil
}

// validateFormParams validates the required RFC 8693 form parameters and returns
// the parsed parameters on success.
func validateFormParams(form url.Values) (*formParams, error) {
        subjectToken := form.Get("subject_token")
        if subjectToken == "" {
                return nil, errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The 'subject_token' parameter is required for token exchange."))
        }

        subjectTokenType := form.Get("subject_token_type")
        if subjectTokenType == "" {
                return nil, errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The 'subject_token_type' parameter is required for token exchange."))
        }

        // id_token is intentionally excluded here, the same as for actor_token
        // tokens below: an ID token's claim conventions differ from an access
        // token's (e.g. "aud" names the relying-party client, not a resource) and
        // this validator applies neither a distinct validation profile nor an
        // id_token-specific claim mapping — accepting the type would let a client
        // declare "id_token" while the token is validated exactly like an access
        // token, silently ignoring the declared type's semantics.
        switch subjectTokenType {
        case oauthproto.TokenTypeAccessToken, oauthproto.TokenTypeJWT:
                // Valid subject token types.
        default:
                return nil, errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf(
                        "The 'subject_token_type' value %q is not supported. Use %q or %q.",
                        subjectTokenType, oauthproto.TokenTypeAccessToken, oauthproto.TokenTypeJWT))
        }

        actorToken := form.Get("actor_token")
        actorTokenType := form.Get("actor_token_type")

        // actor_token_type without actor_token is invalid.
        if actorTokenType != "" && actorToken == "" {
                return nil, errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The 'actor_token_type' parameter requires 'actor_token' to be present."))
        }

        // actor_token requires actor_token_type.
        if actorToken != "" && actorTokenType == "" {
                return nil, errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The 'actor_token_type' parameter is required when 'actor_token' is present."))
        }

        // Validate actor_token_type if present.
        // Note: id_token is intentionally excluded for actor tokens. An actor presents
        // a bearer credential (access_token/jwt), not an identity assertion (id_token).
        if actorTokenType != "" {
                switch actorTokenType {
                case oauthproto.TokenTypeAccessToken, oauthproto.TokenTypeJWT:
                        // Valid actor token types.
                default:
                        return nil, errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf(
                                "The 'actor_token_type' value %q is not supported. Use %q or %q.",
                                actorTokenType, oauthproto.TokenTypeAccessToken, oauthproto.TokenTypeJWT))
                }
        }

        return &formParams{
                subjectToken: subjectToken,
                actorToken:   actorToken,
        }, nil
}

// computeLifetime returns the minimum of the subject token's remaining lifetime
// and the configured delegation lifespan. Returns an error if the subject token
// has already expired.
func (h *Handler) computeLifetime(subjectExpiry time.Time) (time.Duration, error) {
        remaining := time.Until(subjectExpiry)
        if remaining <= 0 {
                return 0, fmt.Errorf("subject token expired %v ago", -remaining)
        }

        if remaining < h.delegationLifespan {
                return remaining, nil
        }
        return h.delegationLifespan, nil
}

// delegatedSubject returns the "sub" to embed in the delegated token.
//
// Self-issued tokens pass through unchanged: their subject is already a
// UUID from UserResolver.ResolveUser's (providerID, providerSubject)
// mapping (pkg/authserver/server/handlers/user.go), ToolHive's own native
// subject namespace.
//
// A subject token from a trusted external issuer is qualified as
// "<issuerURL>#<sub>" instead, because the issued token's "sub" is the only
// signal Cedar's extractClientIDFromClaims (pkg/authz/authorizers/cedar/core.go)
// keys on — it reads "sub" alone, not "iss" — and the external issuer
// chooses that value, not ToolHive. Without qualification, an external
// issuer could pick a "sub" equal to a native user's UUID and mint a
// delegated token indistinguishable from that user's, entirely without
// malicious intent on either issuer's part: two trusted issuers merely
// choosing overlapping subject values collide with no attacker involved.
//
// "#" is an unambiguous separator: an issuer URL cannot itself contain a
// fragment (the OIDC/OAuth "iss" value is a URL without a fragment
// component, and this validator's own issuer comparison is an exact string
// match against that value), so strings.Cut(sub, "#") on the qualified
// subject always recovers (issuerURL, externalSub) uniquely. A native UUID
// can never collide with a qualified value, since a UUID contains no "#".
//
// Re-exchanging a previously-issued delegated token does not double-qualify:
// such a token's "iss" is always this server's own issuer (fosite stamps
// "iss" from the server's config at issuance — see session.New's doc
// comment), so Validate routes it to the self-issued path, which never
// calls this branch.
func delegatedSubject(validatedClaims *ValidatedClaims) string {
        if validatedClaims.ExternalIssuer == "" {
                return validatedClaims.Subject
        }
        return validatedClaims.ExternalIssuer + "#" + validatedClaims.Subject
}

// buildActClaim assembles the RFC 8693 Section 4.1 "act" claim for the
// delegated token, and separately reports the external issuer of the
// original external hop (if any) for the caller to record as its own
// top-level claim. This is either the external issuer this exchange itself
// resolved, or — when this exchange has none of its own (a self-issued
// subject token being re-exchanged) — the value already carried in the
// subject token's own "external_issuer" claim, so the fact that an external
// issuer was ever involved survives across re-exchanges (up to whatever
// depth maxDelegationDepth still permits) instead of disappearing after one
// hop.
//
// externalIssuer is deliberately NOT nested inside act: RFC 8693 §4.1
// defines act as identifying a party that acted, and an issuer alone names
// no party — "the combination of the two claims iss and sub might be
// necessary to uniquely identify an actor" (§4.1) treats iss only as a
// qualifier for an already-identified sub, never as a standalone actor.
// Nesting a bare {"iss": ...} entry under act.act would present a
// non-actor as a prior actor to any RFC-8693-aware consumer that walks the
// chain -- including this codebase's own pkg/audit, whose
// extractDelegationChainFromIdentity is documented as extracting "the full
// chain of acting parties." A may_act-bearing or ActorMatcher-only external
// token has no client-namespace actor to report (see ExternalActor's doc
// comment), so for those the act claim stops at the outer, genuinely-acting
// hop, and the external issuer is reported only via the returned string.
//
// The outermost act.sub is always actorID (the ToolHive client) — every
// downstream consumer reads that as "who is acting", and it must not change
// regardless of how the subject token was obtained.
//
// Extracted from HandleTokenEndpointRequest rather than inlined: the external
// provenance nesting, the prior-chain depth gate, and the encoded-size gate
// below add several branches that have nothing to do with the surrounding
// request plumbing, and multiple of them reject the request outright.
func buildActClaim(validatedClaims *ValidatedClaims, issuer, actorID string) (map[string]any, string, error) {
        act := map[string]any{"iss": issuer, "sub": actorID}
        nestUnder := act
        newLevels := 1

        // When the subject token came from a trusted external issuer
        // (ValidatedClaims.ExternalIssuer is set only there — see
        // multi_issuer_validator.go), nest that issuer's allowlisted actor one
        // level in, when one was resolved: RFC 8693 §4.1 anticipates exactly
        // this, "the combination of the two claims 'iss' and 'sub' might be
        // necessary to uniquely identify an actor." A may_act-bearing or
        // ActorMatcher-only external token has no such actor to nest (see
        // ExternalActor's doc comment) — its act claim stays a single hop, and
        // the caller-returned externalIssuer below is this exchange's only
        // record that an external issuer was involved.
        if validatedClaims.ExternalActor != "" {
                external := map[string]any{"iss": validatedClaims.ExternalIssuer, "sub": validatedClaims.ExternalActor}
                act["act"] = external
                nestUnder = external
                newLevels = 2
        }

        // If the subject token itself carries an "act" claim (i.e. it is a
        // previously-delegated token being re-exchanged), nest it so the full
        // delegation chain remains auditable rather than being discarded.
        // newLevels accounts for the external wrapper above, if any, so the
        // resulting chain never exceeds maxDelegationDepth regardless of how
        // many levels this exchange itself adds.
        if priorAct, ok := validatedClaims.Extra["act"]; ok && priorAct != nil {
                // Parse with the shared audit-side parser rather than a bespoke walker: it
                // reports both the chain depth and any RFC 8693 Section 4.1 conformance
                // violation, so a non-object act cannot slip past the depth gate and be
                // re-minted. A JSON-null act is filtered by the guard above: it asserts no
                // delegation and must not read as malformed.
                chain := coreaudit.ParseDelegationChain(priorAct, maxDelegationDepth)
                if chain.Malformed {
                        // MalformedReason is a closed, value-free enum; it is surfaced to the
                        // client and MUST stay that way — never interpolate claim contents here.
                        return nil, "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf(
                                "The subject token's delegation chain is malformed (%s).", chain.MalformedReason))
                }
                // len(Chain) is the prior chain's depth: the parser appends exactly one
                // hop per level and stops at maxDelegationDepth, so it is
                // min(depth, maxDelegationDepth). newLevels is what this exchange adds
                // on top — one for the acting client, two when an external issuer is
                // also nested — so the resulting chain never exceeds the cap.
                if len(chain.Chain)+newLevels > maxDelegationDepth {
                        return nil, "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                                "The subject token's delegation chain is too deep."))
                }
                // Rebuild the nested act from the parsed chain rather than nesting
                // priorAct verbatim: the subject token's prior act claim may itself
                // have originated from an untrusted external issuer (e.g. a
                // re-exchange of a token whose own nested act content we never
                // validated), and nesting it verbatim would re-sign that content
                // under ToolHive's own key, presenting it as vouched-for. Rebuilding
                // from chain.Chain keeps only iss/sub per hop — extraClaims in
                // coreaudit are dropped by the parser into an in-memory-only map and
                // never resurface here — so unknown or untrusted per-hop claims are
                // normalized away while the audit trail of who-acted-for-whom
                // (Section 4.1's stated purpose for act) is preserved.
                //
                // nestUnder is the innermost entry this exchange added: the acting
                // client normally, or the external-issuer entry when one was nested
                // above, so the prior chain hangs below the provenance rather than
                // displacing it.
                normalizedAct, err := chainToAct(chain.Chain)
                if err != nil {
                        return nil, "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                                "The subject token's delegation chain contains an empty actor."))
                }
                nestUnder["act"] = normalizedAct
        }
        encodedAct, err := json.Marshal(act)
        if err != nil {
                return nil, "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The subject token's delegation chain cannot be serialized."))
        }
        if len(encodedAct) > maxActClaimSize {
                return nil, "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The subject token's delegation chain is too large."))
        }

        // Under the current handler structure there is only ever one external
        // hop, so "original" and "most recent" coincide today; if a design ever
        // introduces multiple external-issuer hops, this will need to change to
        // keep tracking the original rather than the latest.
        externalIssuer := validatedClaims.ExternalIssuer
        if externalIssuer == "" {
                if prior, ok := validatedClaims.Extra["external_issuer"].(string); ok {
                        externalIssuer = prior
                }
        }
        return act, externalIssuer, nil
}

// chainToAct rebuilds a nested RFC 8693 act structure from a parsed
// delegation chain, keeping only each hop's iss/sub — never the hop's extra
// claims, which coreaudit.ParseDelegationChain deliberately keeps out of the
// returned struct's serializable fields. This is what strips arbitrary
// content an untrusted prior issuer may have placed in a nested act before
// it gets re-signed under ToolHive's own key. hops is outermost-first (see
// coreaudit.DelegationChain's doc comment); the result nests them in that
// same order. It rejects any hop without an issuer or subject. Returns nil for
// an empty chain.
func chainToAct(hops []coreaudit.DelegatedActor) (map[string]any, error) {
        var nested map[string]any
        for i := len(hops) - 1; i >= 0; i-- {
                if hops[i].Issuer == "" && hops[i].Subject == "" {
                        return nil, fmt.Errorf("delegation hop %d has no issuer or subject", i)
                }

                hop := make(map[string]any, 3)
                if hops[i].Issuer != "" {
                        hop["iss"] = hops[i].Issuer
                }
                if hops[i].Subject != "" {
                        hop["sub"] = hops[i].Subject
                }
                if nested != nil {
                        hop["act"] = nested
                }
                nested = hop
        }
        return nested, nil
}

// delegateClientAllowed reports whether actorID may use an issuer's
// AllowedDelegateClients-bound consent grant (a may_act.sub match or external
// actor authorization): either the list contains the wildcard
// anyDelegateClient, or it contains actorID directly. An empty list denies
// every client — validateTrustedIssuer rejects that configuration outright
// for any TrustedIssuer, so the empty case here is an extra fail-closed
// backstop, not the path an operator is expected to reach. Callers must only
// invoke this for a claim that actually went through that per-issuer
// validation (ExternalIssuer != ""); a self-issued token's
// AllowedDelegateClients is always nil, since that field has no equivalent
// concept on the self-issued path.
func delegateClientAllowed(allowedDelegateClients []string, actorID string) bool {
        return slices.Contains(allowedDelegateClients, anyDelegateClient) ||
                slices.Contains(allowedDelegateClients, actorID)
}

// checkDelegationConsent enforces the RFC 8693 §4.4 delegation consent check.
//
// There are three consent sources, checked in order:
//
//  1. may_act: if present, it is the authoritative consent signal — only the
//     party named in may_act.sub may delegate. The client_id binding is
//     skipped in this case because may_act enables cross-client delegation
//     (the token was issued to client A but authorizes client B to act).
//     ValidatedClaims.AllowedDelegateClients is additionally enforced here,
//     but ONLY when ExternalIssuer is set: that field is a TrustedIssuer
//     concept with no self-issued equivalent, so a self-issued may_act (a
//     ToolHive user token delegating to a ToolHive client directly) is
//     bound by may_act.sub alone, the same as before this check existed.
//     On the external path, may_act bypasses AllowedActors and ActorMatcher,
//     not per-client containment — delegateClientAllowed still applies there.
//
//  2. External actor authorization: if may_act is absent but the multi-issuer
//     validator has already established consent for this token
//     (multi_issuer_validator.go), it did so by matching the subject token's
//     actor claim against that issuer's operator-configured AllowedActors or
//     ActorMatcher. An allowlisted actor claim — even when configured as
//     "client_id" — names a client in the EXTERNAL issuer's namespace, not ToolHive's, so it must never be compared
//     against actorID the way case 3 below compares ValidatedClaims.ClientID.
//     This case MUST be checked before case 3: it is not a fallback for an
//     empty ClientID, it is a distinct, already-verified consent signal that
//     takes priority whenever it is set, even if ClientID also happens to be
//     populated.
//
//     ValidatedClaims.AllowedDelegateClients binds this to specific ToolHive
//     clients: actorID must appear in it (or it must contain the wildcard
//     anyDelegateClient), checked below by delegateClientAllowed. Without
//     this, the allowlist would authorize "this external client's tokens may
//     be exchanged", not "...by this particular ToolHive client" — every
//     ToolHive confidential client holding the token-exchange grant would be
//     delegation-equivalent with respect to an allowlisted external actor, so
//     compromise of the weakest such client would be as good as compromise
//     of all of them (see #5989). validateTrustedIssuer rejects an empty
//     AllowedDelegateClients at construction for exactly this reason: an
//     operator must declare permissiveness with the wildcard, not get it by
//     omission. Either way this remains bounded by: the calling client must
//     already possess a valid subject token (it cannot forge one), and
//     grantScopes/grantAndBoundAudiences still narrow the result to what
//     both the client and the subject token are authorized for.
//
//  3. client_id: if neither of the above applies, fall back to client_id
//     binding — the subject token's client_id must match the authenticated
//     client. This prevents a stolen subject token from being exchanged by a
//     different client.
//
//     Exception: a configured delegate client (Handler.configuredDelegateClients,
//     Config.DelegateClients ID-only) is exempt from this binding when the
//     subject token is self-issued (ExternalIssuer == ""). Such a client may
//     present ANY self-issued subject token regardless of which client it was
//     originally issued to. This is deliberately scoped to the self-issued
//     path only — see selfIssuedDelegate below — because delegate clients are
//     operator-declared at server startup, not obtained through self-service
//     registration (DCR), matching the blanket-trust model already accepted
//     for the AllowedDelegateClients wildcard (anyDelegateClient) above.
//     The remaining bounds still apply: grantScopes/grantAndBoundAudiences
//     narrow the result to what both the client and the subject token are
//     authorized for, and the delegated token's lifetime is capped as usual.
//
// If none of the three consent sources apply, the subject token carries no
// verifiable binding to any client at all — this fails closed (CWE-863)
// rather than allowing an unbound token through.
func checkDelegationConsent(validatedClaims *ValidatedClaims, clientID, actorSub string, configuredDelegate bool) error {
        // selfIssuedDelegate is true only when a configured delegate client is
        // presenting a self-issued token (ExternalIssuer == "" rules out the
        // external-issuer path explicitly, rather than relying on ExternalActor
        // always being set there — see multi_issuer_validator.go). Without this,
        // a future change to that file's actor-resolution logic could let a
        // delegate client bypass AllowedDelegateClients (the #5989 protection)
        // on the external path.
        selfIssuedDelegate := configuredDelegate && validatedClaims.ExternalIssuer == ""

        switch {
        case validatedClaims.MayAct != nil:
                // may_act.sub is compared against actorSub, not clientID: this is
                // the one binding that is meant to key off the asserted actor —
                // may_act's whole purpose is authorizing a specific actor, which
                // actor_token lets the authenticated client name distinctly from
                // itself.
                if validatedClaims.MayAct.Sub != actorSub {
                        return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                                "The subject token does not authorize this client to act on behalf of the subject."))
                }
                // AllowedDelegateClients binds to the authenticated ToolHive client
                // (clientID), not the asserted actor — see the doc comment above.
                if validatedClaims.ExternalIssuer != "" && !delegateClientAllowed(validatedClaims.AllowedDelegateClients, clientID) {
                        return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                                "This client is not authorized to exchange subject tokens from the external actor's issuer."))
                }
        case validatedClaims.ExternalActorAuthorized:
                // Consent was already established by the external-issuer validator:
                // either the subject token's actor claim matched this issuer's
                // operator-configured AllowedActors or its complete verified claims
                // matched ActorMatcher. An allowlisted actor claim lives in the external
                // issuer's client namespace, not ToolHive's, so — even when ClientID is
                // also populated (ActorClaim: "client_id") — it must never be compared
                // against clientID or actorSub. This case must be checked before the
                // client_id cases below, not merged with them.
                //
                // AllowedDelegateClients binds this authorization to a specific set of
                // ToolHive clients — see delegateClientAllowed. This must be the
                // authenticated client (clientID), not the asserted actor (actorSub):
                // the allowlist is "this external actor's tokens may be exchanged by
                // this ToolHive client", independent of whatever actor identity that
                // client asserts via actor_token.
                if !delegateClientAllowed(validatedClaims.AllowedDelegateClients, clientID) {
                        return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                                "This client is not authorized to exchange subject tokens from the external actor's issuer."))
                }
        case validatedClaims.ClientID != "" && validatedClaims.ClientID != clientID && !selfIssuedDelegate:
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                        "The subject token was issued to a different client."))
        case validatedClaims.ClientID == "":
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                        "The subject token carries no verifiable client binding: it has neither a 'may_act' claim nor a 'client_id' claim."))
        }
        return nil
}

// grantScopes validates that each requested scope is allowed for both the
// client and the subject token, granting the intersection. The delegated
// token's scope set is the intersection of the client's registered scopes
// and the subject token's granted scopes, preventing a client from
// escalating privileges beyond what the user authorized. A subject token
// without a scope claim grants no scopes.
func (h *Handler) grantScopes(
        ctx context.Context, requester fosite.AccessRequester, client fosite.Client, validatedClaims *ValidatedClaims,
) error {
        subjectScopes := strings.Fields(validatedClaims.Scopes)
        subjectScopeSet := make(map[string]bool, len(subjectScopes))
        for _, s := range subjectScopes {
                subjectScopeSet[s] = true
        }

        for _, scope := range requester.GetRequestedScopes() {
                if !h.config.GetScopeStrategy(ctx)(client.GetScopes(), scope) {
                        return errorsx.WithStack(fosite.ErrInvalidScope.WithHintf(
                                "The OAuth 2.0 Client is not allowed to request scope '%s'.", scope))
                }
                if !subjectScopeSet[scope] {
                        return errorsx.WithStack(fosite.ErrInvalidScope.WithHintf(
                                "The scope '%s' was not granted by the subject token.", scope))
                }
                requester.GrantScope(scope)
        }
        return nil
}

// grantAndBoundAudiences resolves the delegated token's audience from the
// request (explicit "audience", RFC 8707 "resource", or the configured
// default) and then bounds the result by the subject token.
//
// The delegated token must not target a resource the subject token was not
// itself valid for: every granted audience must be covered by the subject
// token's audience — delegation narrows, never broadens, the resource
// boundary. This mirrors grantScopes, which bounds scopes by the subject
// token; without it a client registered for audiences A and B could exchange a
// user token minted for A into a token for B, an escalation the user never
// consented to.
func (h *Handler) grantAndBoundAudiences(
        ctx context.Context, requester fosite.AccessRequester, client fosite.Client, validatedClaims *ValidatedClaims,
) error {
        if err := h.grantAudiences(ctx, requester, client); err != nil {
                return err
        }
        if err := h.grantResourceAudience(ctx, requester, client); err != nil {
                return err
        }
        if err := h.grantDefaultAudience(ctx, requester, client); err != nil {
                return err
        }
        return ensureAudienceSubsetOfSubject(requester.GetGrantedAudience(), validatedClaims.Audience)
}

// grantAudiences validates that the requested audience is allowed for this
// client and grants it.
func (h *Handler) grantAudiences(ctx context.Context, requester fosite.AccessRequester, client fosite.Client) error {
        if err := h.config.GetAudienceStrategy(ctx)(client.GetAudience(), requester.GetRequestedAudience()); err != nil {
                return errorsx.WithStack(err)
        }
        for _, aud := range requester.GetRequestedAudience() {
                requester.GrantAudience(aud)
        }
        return nil
}

// resourceScopedClient is implemented by fosite.Client types (e.g.
// SPIFFEClient) that maintain an RFC 8707 resource allowlist independent of
// GetAudience's RFC 8693 audience allowlist. Resources and audiences are
// independent request dimensions: permission in one must never imply
// permission in the other. A client that doesn't implement this interface
// falls back to GetAudience for the resource check, preserving prior
// behavior for clients (e.g. DelegateClient) that only ever had one list.
type resourceScopedClient interface {
        Resources() []string
}

// grantResourceAudience validates the RFC 8707 resource parameter against
// the server's allowedAudiences and the client's own registered resource
// allowlist, and grants it as an additional audience claim, binding the
// issued token to a specific resource server (e.g., an MCP server).
//
// Per RFC 8707 §2, a request MAY carry multiple resource parameters; this
// handler deliberately narrows that to at most one, since a delegated MCP
// token is expected to target a single resource.
func (h *Handler) grantResourceAudience(ctx context.Context, requester fosite.AccessRequester, client fosite.Client) error {
        resources := requester.GetRequestForm()["resource"]
        if len(resources) > 1 {
                return errorsx.WithStack(server.ErrInvalidTarget.WithHint(
                        "Multiple resource parameters are not supported."))
        }
        if len(resources) == 1 && resources[0] != "" {
                resource := resources[0]
                if err := server.ValidateAudienceURI(resource); err != nil {
                        return errorsx.WithStack(err)
                }
                if err := server.ValidateAudienceAllowed(resource, h.allowedAudiences); err != nil {
                        return errorsx.WithStack(err)
                }
                // The resource parameter is RFC 8707's mechanism for requesting a
                // resource-bound token, so it must be subject to the client's own
                // registered allowlist — otherwise a client could bypass its
                // registration simply by using "resource" instead of "audience".
                allowedResources := client.GetAudience()
                if rc, ok := client.(resourceScopedClient); ok {
                        allowedResources = rc.Resources()
                }
                if err := h.config.GetAudienceStrategy(ctx)(allowedResources, []string{resource}); err != nil {
                        return errorsx.WithStack(server.ErrInvalidTarget.WithHintf(
                                "The client is not permitted to request a token for resource %q.", resource))
                }
                requester.GrantAudience(resource)
        }
        return nil
}

// grantDefaultAudience ensures the delegated token carries at least one
// audience, since RFC 9068 §4 makes "aud" a required claim. grantAudiences and
// grantResourceAudience only grant an audience when the client explicitly
// requests one; if neither was requested, this defaults to the server's sole
// configured audience when unambiguous (mirroring handlers/token.go's
// authorization_code fallback), or rejects the request when no unambiguous
// default exists rather than silently issuing an audience-less token. The
// default audience is checked against the client's registered audiences
// like any explicitly requested audience, so a client is never granted an
// audience it wasn't registered for.
func (h *Handler) grantDefaultAudience(ctx context.Context, requester fosite.AccessRequester, client fosite.Client) error {
        if len(requester.GetGrantedAudience()) > 0 {
                return nil
        }
        if len(h.allowedAudiences) != 1 {
                return errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "An explicit 'resource' or 'audience' parameter is required because no unambiguous default audience is configured."))
        }
        defaultAudience := h.allowedAudiences[0]
        if err := h.config.GetAudienceStrategy(ctx)(client.GetAudience(), []string{defaultAudience}); err != nil {
                return errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf(
                        "The client is not permitted to request a token for the default audience %q; "+
                                "an explicit 'resource' or 'audience' parameter is required.", defaultAudience))
        }
        slog.Debug("no resource parameter, defaulting to sole allowed audience",
                "audience", defaultAudience,
        )
        requester.GrantAudience(defaultAudience)
        return nil
}

// ensureAudienceSubsetOfSubject verifies that every audience granted to the
// delegated token is covered by the subject token's own audience. A subject
// token always carries at least one audience, so an empty subject audience
// here can only reject — but what guarantees that non-empty audience differs
// by path. On the self-issued path, SelfIssuedTokenValidator rejects tokens
// whose aud does not intersect this server's AllowedAudiences. On the
// external path, validateExternalToken instead requires aud to contain the
// issuer's own configured ExpectedAudience, which has no required
// relationship to AllowedAudiences — see the TrustedIssuers field doc
// comment in pkg/authserver/config.go for the operator-facing consequence.
func ensureAudienceSubsetOfSubject(granted, subjectAud []string) error {
        subj := make(map[string]bool, len(subjectAud))
        for _, a := range subjectAud {
                subj[a] = true
        }
        for _, a := range granted {
                if !subj[a] {
                        return errorsx.WithStack(server.ErrInvalidTarget.WithHintf(
                                "The delegated token audience %q is not covered by the subject token's audience.", a))
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package tokenexchange

import (
        "context"
        "time"

        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"
        "github.com/ory/x/errorsx"

        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// idJAGReplayPurpose namespaces ID-JAG replay keys apart from plain
// JWT-bearer assertions in storage.AssertionJWTConsumer. The two handlers
// split one grant type by assertion typ, so a given assertion can only ever
// be consumed under one purpose; the separate namespace just keeps that true
// by construction rather than by the typ split alone.
const idJAGReplayPurpose = "id-jag"

// IDJAGHandler implements the bound RFC 7523 JWT-bearer grant for Identity
// Assertion Authorization Grant assertions
// (draft-ietf-oauth-identity-assertion-authz-grant, "ID-JAG" — the inbound
// half of Cross App Access; the requesting half is
// pkg/vmcp/auth/strategies/xaa.go). It claims exactly the assertions
// JWTBearerHandler declines: those whose JOSE typ header is
// "oauth-id-jag+jwt".
//
// "Bound" is the trust-model difference from the plain handler. An ID-JAG
// names the OAuth client that may redeem it — the IdP mints the draft
// §4.4.1 client_id claim from its administrator-configured resource
// connection — so this handler never skips client authentication
// (CanSkipClientAuth) and requires the authenticated client to be that
// client. Confidential clients authenticate fully; a public client is
// identified by client_id, which is sufficient binding here because the
// assertion itself is the primary credential: single-use (required jti),
// short-lived (maxAssertionAge), and audience-pinned to this server
// (acceptedAudiences). The plain handler's synthetic client is not needed —
// the redeeming client is a real registered one.
//
// The redeeming client's registered grant_types metadata is deliberately not
// consulted: with open dynamic client registration, self-asserted metadata
// authorizes nothing. What authorizes the grant is the per-issuer policy
// plus the client_id binding the IdP administrator configured.
type IDJAGHandler struct {
        // core supplies the shared assertion/policy/issuance machinery
        // (validator, per-issuer policies, replay consumer, HandleHelper). It is
        // a named field, not an embedded type, so no fosite.TokenEndpointHandler
        // method is ever promoted from it: a promoted method would gate on the
        // core's own CanHandleTokenEndpointRequest (method values in Go bind at
        // the embedded receiver, there is no virtual dispatch) and reject every
        // ID-JAG request as unknown.
        core *JWTBearerHandler
}

// newIDJAGIssuanceHandler constructs the production bound handler over the
// same issuance core as newJWTBearerIssuanceHandler; see that constructor
// for the dependency requirements.
func newIDJAGIssuanceHandler(
        validator JWTBearerAssertionValidator, tokenEndpoint string, consumer storage.AssertionJWTConsumer,
        config *fosite.Config, strategy oauth2.AccessTokenStrategy, tokenStorage oauth2.AccessTokenStorage,
        trustedIssuers []TrustedIssuer,
) (*IDJAGHandler, error) {
        core, err := newJWTBearerIssuanceHandler(
                validator, tokenEndpoint, consumer, config, strategy, tokenStorage, trustedIssuers)
        if err != nil {
                return nil, err
        }
        return &IDJAGHandler{core: core}, nil
}

// CanHandleTokenEndpointRequest claims only recognized ID-JAG assertions —
// the exact complement of JWTBearerHandler's matcher, so for any given
// jwt-bearer request at most one of the two handlers is responsible.
func (*IDJAGHandler) CanHandleTokenEndpointRequest(_ context.Context, requester fosite.AccessRequester) bool {
        return requester.GetGrantTypes().ExactOne(oauthproto.GrantTypeJWTBearer) &&
                assertionType(requester.GetRequestForm().Get("assertion")) == idJAGJWTType
}

// CanSkipClientAuth never permits skipping client authentication: an ID-JAG
// is issued TO a specific client (its client_id claim), so fosite must have
// authenticated — or, for a public client, at least identified via client_id
// — the caller before HandleTokenEndpointRequest can check the binding.
// fosite.NewAccessRequest returns the client-authentication error whenever
// this is false and authentication failed, so the handler below only ever
// runs with a resolved client.
func (*IDJAGHandler) CanSkipClientAuth(context.Context, fosite.AccessRequester) bool {
        return false
}

// HandleTokenEndpointRequest validates the ID-JAG assertion, its client
// binding, and per-issuer policy, then prepares a bounded access-token
// session.
func (h *IDJAGHandler) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }
        assertion, err := validateJWTBearerAssertionForm(requester.GetRequestForm())
        if err != nil {
                return err
        }
        claims, err := h.core.validator.ValidateJWTBearerAssertion(ctx, assertion, h.core.tokenEndpoint)
        if err != nil {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The ID-JAG assertion is invalid or could not be verified."))
        }
        // Same fail-closed nil-map/nil-consumer posture as the plain handler;
        // see the comments in JWTBearerHandler.HandleTokenEndpointRequest.
        policy, ok := h.core.policies[claims.Issuer]
        if !ok {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The ID-JAG assertion issuer is not enabled for this grant."))
        }
        if h.core.consumer == nil {
                return errorsx.WithStack(fosite.ErrServerError.WithHint("The JWT-bearer grant is not fully configured."))
        }
        // Unlike a plain RFC 7523 assertion (where many IdPs omit jti and the
        // replay key falls back to an assertion hash), the ID-JAG draft requires
        // jti, and it is what makes the assertion single-use. Absence means the
        // token is not a conforming ID-JAG; reject rather than fall back.
        if claims.JWTID == "" {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The ID-JAG assertion must carry a 'jti' claim."))
        }
        if err := h.checkClientBinding(requester, claims); err != nil {
                return err
        }
        resource, err := validateJWTBearerPolicy(requester.GetRequestForm(), claims, policy)
        if err != nil {
                return err
        }
        // Consume before issuing — fails closed like the plain handler: if
        // issuance fails past this point the assertion stays consumed rather
        // than becoming replayable.
        if err := h.core.consumeAssertion(ctx, idJAGReplayPurpose, claims.Issuer, claims.JWTID, claims.Expiry); err != nil {
                return err
        }
        lifetime, err := h.core.assertionBoundedLifetime(ctx, claims.Expiry)
        if err != nil {
                return err
        }
        // Subject is namespaced by issuer exactly like the plain handler's, so a
        // resource server sees one consistent subject form for every
        // assertion-derived token regardless of which handler minted it. The
        // client, unlike there, is the real authenticated one — fosite resolved
        // it (CanSkipClientAuth is always false), so requester already carries
        // it and no synthetic replacement happens.
        issuedSession := session.New(claims.Issuer+"#"+claims.Subject, "", requester.GetClient().GetID(), session.UserClaims{})
        // An ID-JAG links to no upstream IdP login at THIS server, same as a
        // plain assertion; see session.NoUpstreamSessionClaimKey.
        issuedSession.JWTClaims.Extra[session.NoUpstreamSessionClaimKey] = true
        // Carry the assertion's actor chain (RFC 8693 §4.1 "act" — for an
        // Okta-minted ID-JAG, the requesting agent's identity) into the issued
        // token, so a resource server or auditor can still tell WHICH agent
        // acted for the subject after redemption. Copied verbatim: act is
        // defined as a nested-JSON claim and buildValidatedClaims left it
        // untouched in Extra.
        if act, ok := claims.Extra["act"]; ok {
                issuedSession.JWTClaims.Extra["act"] = act
        }
        issuedSession.SetExpiresAt(fosite.AccessToken, time.Now().UTC().Add(lifetime))
        requester.GrantAudience(resource)
        requester.SetSession(issuedSession)
        return nil
}

// PopulateTokenEndpointResponse issues only an access token, sharing the
// plain handler's session-bounded issuance.
func (h *IDJAGHandler) PopulateTokenEndpointResponse(
        ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder,
) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }
        return h.core.populateAccessTokenResponse(ctx, requester, responder)
}

// checkClientBinding enforces the draft's §4.4.1 client_id continuity: the
// assertion's client_id claim (minted by the IdP from its resource
// connection) must name the client fosite resolved for this request. Both
// sides are required to be non-empty first so the check can never
// degenerate into an empty-equals-empty pass — fosite initializes an
// unauthenticated request's client to an ID-less placeholder, and this
// handler must not depend on that staying unreachable.
func (*IDJAGHandler) checkClientBinding(requester fosite.AccessRequester, claims *ValidatedClaims) error {
        if claims.ClientID == "" {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The ID-JAG assertion must carry a 'client_id' claim."))
        }
        client := requester.GetClient()
        if client == nil || client.GetID() == "" {
                return errorsx.WithStack(fosite.ErrInvalidClient.WithHint("The ID-JAG grant requires an authenticated or identified client."))
        }
        if claims.ClientID != client.GetID() {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The ID-JAG assertion was not issued to the requesting client."))
        }
        return nil
}

// IDJAGIssuanceFactory builds the production bound ID-JAG handler. It is
// registered alongside JWTBearerIssuanceFactory whenever any trusted issuer
// enables the JWT-bearer grant: the two handlers share the per-issuer policy
// and split one grant type by assertion typ, so enabling the grant enables
// both assertion forms. See JWTBearerIssuanceFactory for the shared
// validator's ownership contract.
func IDJAGIssuanceFactory(trustedIssuers []TrustedIssuer, shared *MultiIssuerTokenValidator) (server.Factory, error) {
        return jwtBearerGrantFactory(trustedIssuers, shared,
                func(
                        validator JWTBearerAssertionValidator, tokenEndpoint string, consumer storage.AssertionJWTConsumer,
                        config *fosite.Config, strategy oauth2.AccessTokenStrategy, tokenStorage oauth2.AccessTokenStorage,
                        resolvedIssuers []TrustedIssuer,
                ) (any, error) {
                        return newIDJAGIssuanceHandler(
                                validator, tokenEndpoint, consumer, config, strategy, tokenStorage, resolvedIssuers)
                })
}

var _ fosite.TokenEndpointHandler = (*IDJAGHandler)(nil)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package tokenexchange

import (
        "context"
        "crypto/sha256"
        "encoding/base64"
        "encoding/binary"
        "errors"
        "fmt"
        "net/http"
        "net/url"
        "slices"
        "strings"
        "time"

        "github.com/go-jose/go-jose/v4"
        "github.com/go-jose/go-jose/v4/jwt"
        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"
        "github.com/ory/x/errorsx"

        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

const (
        idJAGJWTType                = "oauth-id-jag+jwt"
        jwtBearerReplayPurpose      = "jwt-bearer"
        jwtBearerNoJTIReplayPurpose = "jwt-bearer-no-jti" //nolint:gosec // G101: replay-store purpose label, not a credential
)

// JWTBearerAssertionValidator verifies cryptographic and registered claims of a
// plain RFC 7523 JWT-bearer assertion.
type JWTBearerAssertionValidator interface {
        ValidateJWTBearerAssertion(ctx context.Context, rawToken, tokenEndpoint string) (*ValidatedClaims, error)
}

// JWTBearerHandler implements the unbound RFC 7523 JWT-bearer grant.
type JWTBearerHandler struct {
        *oauth2.HandleHelper
        validator     JWTBearerAssertionValidator
        tokenEndpoint string
        consumer      storage.AssertionJWTConsumer
        config        tokenExchangeConfig
        policies      map[string]*JWTBearerGrantPolicy
}

// newJWTBearerHandler constructs the shared plain-assertion-matching core of
// JWTBearerHandler: CanHandleTokenEndpointRequest, CanSkipClientAuth, and
// request-form/assertion-type validation. It is not, by itself, safe to wire
// into a real token endpoint — it has no policy, replay-tracking, or
// issuance dependencies. Kept unexported because of that: the only
// production path is newJWTBearerIssuanceHandler, which always supplies
// them. In-package tests use it directly only to exercise this seam in
// isolation (see jwt_bearer_handler_test.go); HandleTokenEndpointRequest
// fails closed rather than succeeding when consumer/policies are unset (see
// below), so an accidental use of this constructor alone can never look like
// a working grant.
func newJWTBearerHandler(validator JWTBearerAssertionValidator, tokenEndpoint string) (*JWTBearerHandler, error) {
        if validator == nil {
                return nil, errors.New("JWT-bearer assertion validator must not be nil")
        }
        if tokenEndpoint == "" {
                return nil, errors.New("JWT-bearer token endpoint must not be empty")
        }
        endpoint, err := url.ParseRequestURI(tokenEndpoint)
        if err != nil || endpoint.Scheme == "" || endpoint.Host == "" {
                return nil, fmt.Errorf("JWT-bearer token endpoint is invalid: %q", tokenEndpoint)
        }
        return &JWTBearerHandler{validator: validator, tokenEndpoint: tokenEndpoint}, nil
}

func newJWTBearerIssuanceHandler(
        validator JWTBearerAssertionValidator, tokenEndpoint string, consumer storage.AssertionJWTConsumer,
        config *fosite.Config, strategy oauth2.AccessTokenStrategy, tokenStorage oauth2.AccessTokenStorage,
        trustedIssuers []TrustedIssuer,
) (*JWTBearerHandler, error) {
        handler, err := newJWTBearerHandler(validator, tokenEndpoint)
        if err != nil {
                return nil, err
        }
        if consumer == nil {
                return nil, errors.New("JWT-bearer storage must implement storage.AssertionJWTConsumer")
        }
        if config == nil || strategy == nil || tokenStorage == nil {
                return nil, errors.New("JWT-bearer issuance dependencies must not be nil")
        }
        policies := make(map[string]*JWTBearerGrantPolicy)
        for _, issuer := range trustedIssuers {
                if issuer.JWTBearerGrant != nil {
                        policies[issuer.IssuerURL] = issuer.JWTBearerGrant
                }
        }
        if len(policies) == 0 {
                return nil, errors.New("JWT-bearer issuance requires at least one enabled trusted issuer")
        }
        handler.HandleHelper = &oauth2.HandleHelper{
                AccessTokenStrategy: strategy,
                AccessTokenStorage:  tokenStorage,
                Config:              config,
        }
        handler.consumer = consumer
        handler.config = config
        handler.policies = policies
        return handler, nil
}

// CanHandleTokenEndpointRequest only claims plain assertions. A recognized ID-JAG
// assertion belongs to the bound handler (IDJAGHandler); malformed and
// unsupported typ values remain this handler's responsibility to reject.
func (*JWTBearerHandler) CanHandleTokenEndpointRequest(_ context.Context, requester fosite.AccessRequester) bool {
        return requester.GetGrantTypes().ExactOne(oauthproto.GrantTypeJWTBearer) &&
                assertionType(requester.GetRequestForm().Get("assertion")) != idJAGJWTType
}

// CanSkipClientAuth permits a credential-free plain JWT-bearer assertion.
// fosite.NewAccessRequest always calls AuthenticateClient first and only
// discards its error when this returns true (access_request_handler.go),
// and AuthenticateClient falls back to reading HTTP Basic credentials off
// the raw request when no form-based client_assertion is present
// (client_authentication.go). If this only inspected the form, a request
// carrying a valid assertion plus a wrong HTTP Basic password would be
// silently accepted, since the resulting clientErr would be discarded
// without ever being looked at. So the raw *http.Request that fosite stashes
// under RequestContextKey must be checked for Basic credentials too; treat a
// missing/malformed request as "credentials might be present" so this fails
// closed rather than open.
func (*JWTBearerHandler) CanSkipClientAuth(ctx context.Context, requester fosite.AccessRequester) bool {
        form := requester.GetRequestForm()
        if assertionType(form.Get("assertion")) == idJAGJWTType ||
                form.Get("client_id") != "" ||
                form.Get("client_secret") != "" ||
                form.Get("client_assertion") != "" ||
                form.Get("client_assertion_type") != "" {
                return false
        }
        req, ok := ctx.Value(fosite.RequestContextKey).(*http.Request)
        if !ok || req == nil {
                return false
        }
        if _, _, ok := req.BasicAuth(); ok {
                return false
        }
        return true
}

// HandleTokenEndpointRequest validates policy and prepares a bounded access-token session.
func (h *JWTBearerHandler) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }
        assertion, err := validateJWTBearerAssertionForm(requester.GetRequestForm())
        if err != nil {
                return err
        }
        segments, err := canonicalAssertionSegments(assertion)
        if err != nil {
                return err
        }
        if err := validateAssertionType(assertion); err != nil {
                return err
        }
        claims, err := h.validator.ValidateJWTBearerAssertion(ctx, assertion, h.tokenEndpoint)
        if err != nil {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The JWT bearer assertion is invalid or could not be verified."))
        }
        // h.policies is nil when this handler was built by newJWTBearerHandler
        // directly rather than newJWTBearerIssuanceHandler (the seam tests use it
        // this way) — a nil map lookup below always misses, so this request is
        // rejected rather than treated as authorized. Only newJWTBearerIssuanceHandler
        // is wired into a real token endpoint; see its doc comment.
        policy, ok := h.policies[claims.Issuer]
        if !ok {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The JWT bearer assertion issuer is not enabled for this grant."))
        }
        if h.consumer == nil {
                return errorsx.WithStack(fosite.ErrServerError.WithHint("The JWT-bearer grant is not fully configured."))
        }
        resource, err := validateJWTBearerPolicy(requester.GetRequestForm(), claims, policy)
        if err != nil {
                return err
        }
        // Consume before issuing. This intentionally fails closed: if issuance fails
        // after this point, the assertion remains consumed rather than becoming
        // replayable.
        replayKey := assertionReplayKey(segments, claims.JWTID)
        purpose := jwtBearerReplayPurpose
        if claims.JWTID == "" {
                purpose = jwtBearerNoJTIReplayPurpose
        }
        if err := h.consumeAssertion(ctx, purpose, claims.Issuer, replayKey, claims.Expiry); err != nil {
                return err
        }
        lifetime, err := h.assertionBoundedLifetime(ctx, claims.Expiry)
        if err != nil {
                return err
        }
        clientID := jwtBearerClientID(claims.Issuer, claims.Subject)
        issuedSession := session.New(claims.Issuer+"#"+claims.Subject, "", clientID, session.UserClaims{})
        // This grant links to no upstream IdP login, so the issued token can never
        // carry an upstream credential. Say so in the token rather than leaving a
        // resource server to infer it from the absence of a tsid claim, which is
        // also how an unenriched or foreign-issued identity looks. See
        // session.NoUpstreamSessionClaimKey.
        issuedSession.JWTClaims.Extra[session.NoUpstreamSessionClaimKey] = true
        issuedSession.SetExpiresAt(fosite.AccessToken, time.Now().UTC().Add(lifetime))
        // This grant skips client authentication (CanSkipClientAuth), so fosite
        // never populates requester's client — attach a synthetic one so every
        // storage backend's marshal path (which calls GetClient unconditionally)
        // has a real, non-nil client to serialize.
        accessRequest, ok := requester.(*fosite.AccessRequest)
        if !ok {
                return errorsx.WithStack(fosite.ErrServerError.WithHint("The JWT-bearer grant requires a *fosite.AccessRequest."))
        }
        // Always replace any resolved caller client. Fosite accepts a supplied
        // public client_id without a secret; retaining it could renew a DCR client
        // registration TTL after this clientless grant succeeds.
        accessRequest.Client = storage.NewSyntheticClient(clientID)
        requester.GrantAudience(resource)
        requester.SetSession(issuedSession)
        return nil
}

func validateJWTBearerPolicy(form url.Values, claims *ValidatedClaims, policy *JWTBearerGrantPolicy) (string, error) {
        if claims.IssuedAt.IsZero() || claims.Expiry.IsZero() || claims.Expiry.Sub(claims.IssuedAt) > policy.maxAssertionAge {
                return "", errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The JWT bearer assertion exceeds the configured maximum age."))
        }
        resources, present := form["resource"]
        if !present || len(resources) != 1 || resources[0] == "" {
                return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "Exactly one resource parameter is required for the JWT-bearer grant."))
        }
        if err := server.ValidateAudienceURI(resources[0]); err != nil {
                return "", errorsx.WithStack(err)
        }
        for _, binding := range policy.SubjectBindings {
                if binding.Subject == claims.Subject {
                        for _, resource := range binding.AllowedResources {
                                if resource == resources[0] {
                                        return resource, nil
                                }
                        }
                        return "", errorsx.WithStack(&fosite.RFC6749Error{
                                ErrorField: "invalid_target",
                                CodeField:  http.StatusBadRequest,
                        })
                }
        }
        return "", errorsx.WithStack(fosite.ErrInvalidGrant.WithHint(
                "The JWT bearer assertion subject is not configured for this grant."))
}

// PopulateTokenEndpointResponse issues only an access token.
func (h *JWTBearerHandler) PopulateTokenEndpointResponse(
        ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder,
) error {
        if !h.CanHandleTokenEndpointRequest(ctx, requester) {
                return errorsx.WithStack(fosite.ErrUnknownRequest)
        }
        return h.populateAccessTokenResponse(ctx, requester, responder)
}

// populateAccessTokenResponse issues an access token bounded by the session
// expiry HandleTokenEndpointRequest set. Shared by both JWT-bearer handlers
// (plain and ID-JAG-bound), each of which gates on its own
// CanHandleTokenEndpointRequest before calling this.
func (h *JWTBearerHandler) populateAccessTokenResponse(
        ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder,
) error {
        if h.HandleHelper == nil {
                return errorsx.WithStack(fosite.ErrInvalidRequest.WithHint("JWT-bearer token issuance is not configured."))
        }
        lifetime := h.config.GetAccessTokenLifespan(ctx)
        if expiry := requester.GetSession().GetExpiresAt(fosite.AccessToken); !expiry.IsZero() {
                remaining := time.Until(expiry)
                if remaining <= 0 {
                        return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The JWT bearer assertion has expired."))
                }
                if remaining < lifetime {
                        lifetime = remaining
                }
        }
        _, err := h.IssueAccessToken(ctx, lifetime, requester, responder)
        return err
}

// consumeAssertion marks an assertion's replay key used, mapping consumer
// errors onto the grant's wire errors. Shared by both JWT-bearer handlers;
// purpose keeps their replay namespaces separate.
func (h *JWTBearerHandler) consumeAssertion(
        ctx context.Context, purpose, issuer, replayKey string, expiry time.Time,
) error {
        if err := h.consumer.ConsumeAssertionJWT(ctx, purpose, issuer, replayKey, expiry); err != nil {
                if errors.Is(err, fosite.ErrJTIKnown) {
                        return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The JWT bearer assertion has already been used."))
                }
                // A non-ErrJTIKnown failure here is a storage/operational problem (e.g. a
                // replay-store outage), not evidence the caller sent a bad assertion —
                // it must not be indistinguishable from a genuine invalid_grant.
                return errorsx.WithStack(fosite.ErrServerError.WithHint("The JWT bearer assertion could not be consumed.").
                        WithWrap(err).WithDebug(err.Error()))
        }
        return nil
}

// assertionBoundedLifetime returns the issued token's lifetime: the
// configured access-token lifespan, capped by the assertion's remaining
// validity so a token derived from an assertion can never outlive it.
func (h *JWTBearerHandler) assertionBoundedLifetime(ctx context.Context, expiry time.Time) (time.Duration, error) {
        remaining := time.Until(expiry)
        if remaining <= 0 {
                return 0, errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The JWT bearer assertion has expired."))
        }
        lifetime := h.config.GetAccessTokenLifespan(ctx)
        if remaining < lifetime {
                lifetime = remaining
        }
        return lifetime, nil
}

func jwtBearerClientID(issuer, subject string) string {
        digest := sha256.Sum256([]byte(issuer + "\x00" + subject))
        return storage.SyntheticClientIDPrefix + "jwt-bearer-" + base64.RawURLEncoding.EncodeToString(digest[:])
}

// canonicalAssertionSegments rejects alternate encodings of a compact JWT before
// it reaches the validator. The decoded protected header and payload are reused
// for replay keying; the signature is checked for canonical form but not keyed.
func canonicalAssertionSegments(assertion string) ([2][]byte, error) {
        var replaySegments [2][]byte
        parts := strings.Split(assertion, ".")
        if len(parts) != 3 {
                return replaySegments, errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The assertion is not a valid signed JWT."))
        }
        for i, part := range parts {
                decoded, err := base64.RawURLEncoding.DecodeString(part)
                if err != nil || base64.RawURLEncoding.EncodeToString(decoded) != part {
                        return replaySegments, errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The assertion is not a valid signed JWT."))
                }
                if i < len(replaySegments) {
                        replaySegments[i] = decoded
                }
        }
        return replaySegments, nil
}

// assertionReplayKey uses jti when present; otherwise it hashes the decoded
// protected header and payload with a length-delimited header. Excluding the
// signature makes re-signings of the same assertion share a replay slot. The
// no-jti purpose keeps these keys separate from issuer-supplied jti values.
func assertionReplayKey(segments [2][]byte, jti string) string {
        if jti != "" {
                return jti
        }
        framed := binary.BigEndian.AppendUint64(nil, uint64(len(segments[0])))
        framed = append(framed, segments[0]...)
        framed = append(framed, segments[1]...)
        digest := sha256.Sum256(framed)
        return "jwt-bearer-noJTI-" + base64.RawURLEncoding.EncodeToString(digest[:])
}

func validateJWTBearerAssertionForm(form url.Values) (string, error) {
        assertions, ok := form["assertion"]
        if !ok || len(assertions) != 1 || assertions[0] == "" {
                return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint(
                        "The 'assertion' parameter is required exactly once for the JWT-bearer grant."))
        }
        return assertions[0], nil
}

func validateAssertionType(assertion string) error {
        parsed, err := jwt.ParseSigned(assertion, allowedSignatureAlgorithms)
        if err != nil {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The assertion is not a valid signed JWT."))
        }
        if len(parsed.Headers) != 1 {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The assertion must contain exactly one JOSE signature."))
        }
        typ, ok := parsed.Headers[0].ExtraHeaders[jose.HeaderType]
        if !ok {
                return nil
        }
        value, ok := typ.(string)
        if !ok {
                return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The assertion JOSE typ header must be a string."))
        }
        if value == "" || value == "JWT" || value == idJAGJWTType {
                return nil
        }
        return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("The assertion JOSE typ header is not supported."))
}

func assertionType(assertion string) string {
        parsed, err := jwt.ParseSigned(assertion, allowedSignatureAlgorithms)
        if err != nil || len(parsed.Headers) != 1 {
                return ""
        }
        typ, _ := parsed.Headers[0].ExtraHeaders[jose.HeaderType].(string)
        return typ
}

// validateJWTBearerAssertionClaims validates the assertion's registered
// claims. acceptedAudiences is the set of "this AS" identity strings the
// assertion's "aud" must intersect — see JWTBearerGrantPolicy.AcceptedAudiences.
// "jti" is not required: many real-world IdPs (e.g. Microsoft Entra ID
// client_credentials tokens) never emit one. When absent, assertionReplayKey
// hashes the decoded protected header and payload for the replay check.
func validateJWTBearerAssertionClaims(claims jwt.Claims, issuer string, acceptedAudiences []string) error {
        if claims.Expiry == nil {
                return errors.New("assertion is missing required 'exp' claim")
        }
        if claims.IssuedAt == nil {
                return errors.New("assertion is missing required 'iat' claim")
        }
        if claims.Subject == "" {
                return errors.New("assertion is missing required 'sub' claim")
        }
        if !audienceIntersects(claims.Audience, acceptedAudiences) {
                return fmt.Errorf("assertion audience must include one of the accepted authorization server identities %q", acceptedAudiences)
        }
        if err := claims.ValidateWithLeeway(jwt.Expected{Issuer: issuer}, externalClockSkewLeeway); err != nil {
                return fmt.Errorf("assertion claims validation failed: %w", err)
        }
        if claims.Expiry.Time().Before(time.Now()) {
                return errors.New("assertion has expired")
        }
        return nil
}

// audienceIntersects reports whether any value in audience is a member of
// accepted. RFC 7523 §3 permits a multi-valued assertion "aud"; membership
// (not equality) is deliberate so the assertion can additionally name other
// audiences (e.g. an IdP-specific application ID) without losing its match
// against one of this AS's own accepted identities.
func audienceIntersects(audience jwt.Audience, accepted []string) bool {
        for _, value := range audience {
                if slices.Contains(accepted, value) {
                        return true
                }
        }
        return false
}

// JWTBearerIssuanceFactory builds the production RFC 7523 handler. It is only
// registered by composition when a trusted issuer opts into the grant.
//
// shared is used as the JWTBearerAssertionValidator and is REQUIRED whenever
// trustedIssuers is non-empty (an error is returned otherwise). The RFC 8693
// token-exchange Factory and this one are usually enabled for the same trusted
// issuers, and each MultiIssuerTokenValidator registers its own per-issuer
// jwkfetch.Cache and background refresh goroutines; sharing one instance avoids
// doubling that cost, and — since the validator's JWKS workers are released
// only by its Close — keeps them owned by the caller rather than built and
// abandoned inside this compose-time closure. Build it once with
// NewSharedTrustedIssuerValidator and hold it for Close.
func JWTBearerIssuanceFactory(trustedIssuers []TrustedIssuer, shared *MultiIssuerTokenValidator) (server.Factory, error) {
        return jwtBearerGrantFactory(trustedIssuers, shared,
                func(
                        validator JWTBearerAssertionValidator, tokenEndpoint string, consumer storage.AssertionJWTConsumer,
                        config *fosite.Config, strategy oauth2.AccessTokenStrategy, tokenStorage oauth2.AccessTokenStorage,
                        resolvedIssuers []TrustedIssuer,
                ) (any, error) {
                        return newJWTBearerIssuanceHandler(
                                validator, tokenEndpoint, consumer, config, strategy, tokenStorage, resolvedIssuers)
                })
}

// jwtBearerGrantFactory is the shared core of JWTBearerIssuanceFactory and
// IDJAGIssuanceFactory: policy resolution, the shared-validator ownership
// requirement, and the compose-time dependency and audience checks are
// identical for both handlers — only the handler constructed at the end
// differs, supplied as build. Both factories are registered from the same
// issuer policies, so every check here runs once per factory; that is
// deliberate redundancy, not waste — each factory must stay independently
// safe to register alone.
func jwtBearerGrantFactory(
        trustedIssuers []TrustedIssuer,
        shared *MultiIssuerTokenValidator,
        build func(
                validator JWTBearerAssertionValidator, tokenEndpoint string, consumer storage.AssertionJWTConsumer,
                config *fosite.Config, strategy oauth2.AccessTokenStrategy, tokenStorage oauth2.AccessTokenStorage,
                resolvedIssuers []TrustedIssuer,
        ) (any, error),
) (server.Factory, error) {
        resolvedIssuers, err := ResolveJWTBearerGrantPolicies(trustedIssuers)
        if err != nil {
                return nil, fmt.Errorf("JWT-bearer trusted issuers: %w", err)
        }
        // Unconditionally required (unlike the token-exchange Factory, which validly
        // supports self-issued-only): JWT-bearer issuance is meaningful only against
        // trusted issuers, so there is no no-issuer case that would legitimately
        // leave shared nil. Requiring it also keeps the validator's JWKS workers
        // owned by the caller for Close rather than built and abandoned here.
        if shared == nil {
                return nil, fmt.Errorf("JWT-bearer: a shared validator built via NewSharedTrustedIssuerValidator " +
                        "is required so its JWKS refresh workers can be released on shutdown")
        }
        return func(config *server.AuthorizationServerConfig, rawStorage fosite.Storage, strategy any) (any, error) {
                consumer, err := assertionJWTConsumer(rawStorage)
                if err != nil {
                        return nil, err
                }
                atStrategy, ok := strategy.(oauth2.AccessTokenStrategy)
                if !ok {
                        return nil, fmt.Errorf("JWT-bearer strategy does not implement oauth2.AccessTokenStrategy (got %T)", strategy)
                }
                atStorage, ok := rawStorage.(oauth2.AccessTokenStorage)
                if !ok {
                        return nil, fmt.Errorf("JWT-bearer storage does not implement oauth2.AccessTokenStorage (got %T)", rawStorage)
                }
                for _, issuer := range resolvedIssuers {
                        if issuer.JWTBearerGrant == nil {
                                continue
                        }
                        for _, binding := range issuer.JWTBearerGrant.SubjectBindings {
                                for _, resource := range binding.AllowedResources {
                                        if !slices.Contains(config.AllowedAudiences, resource) {
                                                return nil, fmt.Errorf("JWT-bearer resource %q is not an allowed audience", resource)
                                        }
                                }
                        }
                        // accepted_audiences identifies this authorization server, not a
                        // resource; checked here (this is the runtime choke point every
                        // JWTBearerIssuanceFactory call goes through) as well as inside
                        // NewSharedTrustedIssuerValidator, so a caller that builds shared
                        // separately is still covered.
                        for _, audience := range issuer.JWTBearerGrant.AcceptedAudiences {
                                if slices.Contains(config.AllowedAudiences, audience) {
                                        return nil, fmt.Errorf(
                                                "JWT-bearer accepted_audiences value %q must not also be a configured "+
                                                        "resource audience (allowed_audiences)", audience)
                                }
                        }
                }
                // shared is guaranteed non-nil (checked when this factory was built), so
                // the validator is always the caller-owned, closeable one — this closure
                // never builds a MultiIssuerTokenValidator whose JWKS workers leak.
                var validator JWTBearerAssertionValidator = shared
                return build(validator, config.TokenURL, consumer, config.Config, atStrategy, atStorage, resolvedIssuers)
        }, nil
}

// assertionJWTConsumer resolves rawStorage's AssertionJWTConsumer capability.
// Every decorator in the actual composition chain must itself implement (and
// forward, one level down) AssertionJWTConsumer for this to succeed -- there
// is no automatic bypass via storage.Unwrap. That is deliberate: a decorator
// that sits between rawStorage and the innermost backend gets an explicit,
// visible opportunity to intercept or audit assertion-JWT consumption, rather
// than being silently skipped past. See CIMDStorageDecorator.ConsumeAssertionJWT
// (and the shape guard that NewCIMDShapeGuardStorage returns when CIMD is
// disabled), SPIFFEStorageDecorator.ConsumeAssertionJWT and
// ClientAdmissionStorage.ConsumeAssertionJWT for the production decorators
// that forward it today. A future decorator that omits this method
// breaks visibly, right here, with an error naming the offending type -- not
// silently, by having its ConsumeAssertionJWT logic (if any) never run.
func assertionJWTConsumer(rawStorage fosite.Storage) (storage.AssertionJWTConsumer, error) {
        consumer, ok := rawStorage.(storage.AssertionJWTConsumer)
        if !ok {
                return nil, fmt.Errorf("JWT-bearer storage %T does not implement storage.AssertionJWTConsumer", rawStorage)
        }
        return consumer, nil
}

var _ fosite.TokenEndpointHandler = (*JWTBearerHandler)(nil)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package tokenexchange

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

        celgo "cel.dev/cel-go/cel"
        "github.com/go-jose/go-jose/v4"
        "github.com/go-jose/go-jose/v4/jwt"
        "github.com/jwx-go/jwkfetch/v4"
        "github.com/lestrrat-go/httprc/v3"
        "github.com/lestrrat-go/jwx/v4/jwk"

        "github.com/stacklok/toolhive-core/cel"
        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

const (
        // httpTimeout is the timeout for HTTP requests to external OIDC endpoints.
        httpTimeout = 10 * time.Second

        // maxResponseBodySize is the maximum size of HTTP response bodies read
        // from external OIDC discovery documents AND JWKS fetches (1 MiB). This
        // prevents resource exhaustion from unexpectedly large responses. The
        // discovery read enforces it directly via io.LimitReader; jwkfetch enforces
        // it for every initial and background JWKS fetch.
        maxResponseBodySize = 1 << 20

        // maxJWKSKeys caps the number of keys accepted from an external JWKS to
        // prevent CPU amplification from a hostile endpoint serving many keys.
        maxJWKSKeys = 100

        // minKidRefreshInterval bounds how often refreshOnUnknownKid forces an
        // issuer's jwkfetch.Cache to fetch its JWKS ahead of jwx's own background
        // refresh schedule. verifySignature's kidMatched check runs before the
        // subject token's signature is trusted, so without this floor a client
        // presenting a syntactically valid JWT that merely names a made-up kid
        // could force a fresh fetch to the external IdP on every attempt.
        minKidRefreshInterval = 30 * time.Second

        // jwksFetchFailureBackoff bounds how often ensureRegistered retries a
        // JWKS fetch for an issuer that has never yet succeeded. Key resolution
        // runs before the subject token's signature is checked, so without this
        // an authenticated client holding the token-exchange grant could force a
        // real outbound fetch to a broken external IdP on every single request.
        jwksFetchFailureBackoff = 30 * time.Second

        // jwksRefreshInterval is the fixed interval at which each issuer's
        // jwkfetch.Cache re-fetches its JWKS in the background. It is passed to
        // Register via jwkfetch.WithConstantInterval, which makes the resource
        // ignore the response's Cache-Control/Expires headers entirely rather
        // than merely bounding them with WithMaxInterval — deliberately:
        // absent this override, httprc derives the interval from those
        // headers, clamped to [DefaultMinInterval, DefaultMaxInterval] =
        // [15m, 30 days], so a hostile or misconfigured issuer could otherwise
        // extend our own key-retention window up to a month simply by setting
        // a long max-age.
        jwksRefreshInterval = 5 * time.Minute

        // defaultActorClaim is the claim read to identify the client that
        // requested an external subject token, when a TrustedIssuer does not
        // configure ActorClaim. This matches Microsoft Entra v2 and many other
        // OIDC providers' convention for the authorized-party claim.
        //
        // "client_id" is the normative claim: RFC 8693 §4.3 names it as the
        // party a token was issued to, and RFC 9068 §2.2 makes it a REQUIRED
        // access-token claim. "azp" is the default here only because Entra,
        // Okta, and other major providers omit "client_id" from their access
        // tokens in practice — "azp" itself appears nowhere in RFC 9068; it is
        // an OIDC Core claim defined for ID tokens, not access tokens.
        defaultActorClaim = "azp"

        // externalClockSkewLeeway widens "nbf"/"iat"/"exp" acceptance for external
        // subject tokens to tolerate clock skew between the external IdP and
        // ToolHive. validateExternalToken independently rejects an expired
        // subject token by comparing its exp against time.Now() with zero
        // tolerance, so this leeway only protects against spurious
        // not-yet-valid/issued-in-the-future rejections — it never causes an
        // expired token to be accepted.
        externalClockSkewLeeway = 60 * time.Second
)

// actorClaimsNotInExtra are claims assignClaim drops or reroutes, so they
// never reach ValidatedClaims.Extra. Only "client_id" has an explicit
// fallback in resolveAllowedActor; configuring ActorClaim to any of these
// would make every external token look like it's missing the claim.
var actorClaimsNotInExtra = []string{
        "sub", "iss", "aud", "exp", "iat", "nbf", "jti", "name", "email", "scope", "scp", "may_act",
}

// Compile-time check that MultiIssuerTokenValidator implements SubjectTokenValidator.
var _ SubjectTokenValidator = (*MultiIssuerTokenValidator)(nil)

// Compile-time check that MultiIssuerTokenValidator implements JWTBearerAssertionValidator.
var _ JWTBearerAssertionValidator = (*MultiIssuerTokenValidator)(nil)

// JWTBearerSubjectBinding restricts a JWT-bearer assertion subject to the
// resources for which it may obtain an access token.
type JWTBearerSubjectBinding struct {
        Subject          string   `json:"subject" yaml:"subject"`
        AllowedResources []string `json:"allowed_resources" yaml:"allowed_resources"`
}

// JWTBearerGrantPolicy enables RFC 7523 JWT-bearer assertions for an issuer.
// MaxAssertionAge uses Go's duration syntax in serialized RunConfig.
type JWTBearerGrantPolicy struct {
        MaxAssertionAge string                    `json:"max_assertion_age" yaml:"max_assertion_age"`
        SubjectBindings []JWTBearerSubjectBinding `json:"subject_bindings" yaml:"subject_bindings"`
        // AcceptedAudiences is the set of "this AS" identity strings an
        // assertion's "aud" claim must intersect — e.g. to support migrating
        // this server's issuer/token-endpoint URL, or exposing it under more
        // than one valid name. Each value uniquely identifies this
        // authorization server for this grant; it is NOT a resource/API
        // identifier — a bare resource audience is deliberately not accepted
        // here, that would let any RFC 8707 resource-scoped token satisfy the
        // grant instead of only tokens minted for this AS. Defaults to
        // [tokenEndpoint] when empty, preserving prior exact-match behavior.
        AcceptedAudiences []string `json:"accepted_audiences,omitempty" yaml:"accepted_audiences,omitempty"`

        maxAssertionAge time.Duration
}

// TrustedIssuer configures an external OIDC issuer whose tokens are accepted
// as RFC 8693 subject tokens or RFC 7523 JWT-bearer assertions.
//
// This type is reused verbatim as the wire schema for
// authserver.RunConfig.TrustedIssuers (deliberately, to avoid a parallel
// type that drifts — see the go-style rule against that). Its JSON/YAML
// tags are therefore part of the serialized RunConfig, which is reflected
// into docs/server/swagger.*; adding, renaming, or retagging a field here is
// a schema change, not a purely internal one.
type TrustedIssuer struct {
        // Name optionally identifies this trust declaration for canonical issuer_ref references.
        Name string `json:"name,omitempty" yaml:"name,omitempty"`
        // IssuerURL is the expected "iss" claim value (exact match).
        IssuerURL string `json:"issuer_url" yaml:"issuer_url"`
        // ExpectedAudience is the expected "aud" claim value that must appear
        // in an RFC 8693 subject token's audience list (a resource/API identifier,
        // not a client ID — required for delegation unless JWTBearerGrant is
        // configured; see looksLikeResourceIdentifier). RFC 7523 assertions use
        // the token endpoint as their audience instead.
        //
        // This legacy field is deprecated; configure RFC 8693 policy under
        // inbound_grants.token_exchange.issuer_policies.
        // See docs/arch/17-token-exchange-delegation.md ("ID/access-token
        // discrimination") for why and its limits.
        ExpectedAudience string `json:"expected_audience" yaml:"expected_audience"`
        // JWKSURL is the URL to fetch the issuer's JSON Web Key Set from.
        // If empty, it is resolved via OIDC discovery at {IssuerURL}/.well-known/openid-configuration.
        JWKSURL string `json:"jwks_url,omitempty" yaml:"jwks_url,omitempty"`
        // InsecureAllowHTTP permits plain-HTTP OIDC discovery and JWKS fetches
        // for THIS issuer only. Development and testing only — never set in
        // production. Does not relax the private-IP guard; see AllowPrivateIPs.
        // Deliberately per-issuer: this server's own InsecureAllowHTTP must not
        // silently permit plaintext discovery for every trusted external issuer
        // too — a network attacker who can intercept that traffic could
        // substitute a JWKS and forge subject tokens for that issuer's
        // namespace.
        InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"`
        // AllowPrivateIPs permits OIDC discovery and JWKS fetches for THIS
        // issuer to resolve to a private or loopback address. Use only when the
        // issuer is hosted inside the same cluster and has no public endpoint.
        AllowPrivateIPs bool `json:"allow_private_ips,omitempty" yaml:"allow_private_ips,omitempty"`
        // CAFilePath is the path to a PEM CA bundle added to the system roots when
        // fetching this issuer's OIDC discovery document and JWKS. Trust is additive
        // and scoped to this issuer: the public roots still apply, and no other
        // issuer's client is affected.
        CAFilePath string `json:"ca_file_path,omitempty" yaml:"ca_file_path,omitempty"`
        // ActorClaim names the claim identifying the client that requested the
        // subject token from THIS EXTERNAL ISSUER (used by AllowedActors below).
        // Values are in the external issuer's namespace, NOT ToolHive client
        // IDs. Defaults to "azp"; use "appid" for Microsoft Entra v1, "cid" for
        // Okta. The special value "client_id" reads ValidatedClaims.ClientID
        // instead of Extra (assignClaim routes it to that field) — it is still
        // the external token's client_id claim, not a ToolHive one.
        ActorClaim string `json:"actor_claim,omitempty" yaml:"actor_claim,omitempty"`
        // AllowedActors is the allowlist of ActorClaim values authorized to
        // exchange a subject token from this issuer when it carries no
        // "may_act" claim. ActorMatcher can additionally authorize a token by
        // matching its complete verified claims map; either signal is sufficient.
        // When both are empty, only may_act-bearing tokens are accepted, and only
        // if AllowMayAct is also true for this issuer. By itself names no
        // ToolHive client — see AllowedDelegateClients and
        // docs/arch/17-token-exchange-delegation.md ("Accepted limitations" #1).
        AllowedActors []string `json:"allowed_actors,omitempty" yaml:"allowed_actors,omitempty"`
        // ActorMatcher is an admin-authored CEL expression evaluated against the
        // complete signature-verified JWT claims map as "claims". A true result
        // authorizes delegation alongside AllowedActors; a syntax or type error
        // fails configuration validation. An expression that compiles but does
        // not return bool is NOT caught at that point, though — it compiles
        // successfully and is only rejected the first time it is evaluated
        // against a real token, denying that token (and every one after it, since
        // the expression will never return bool). Any other runtime evaluation
        // error denies the token the same way.
        ActorMatcher string `json:"actor_matcher,omitempty" yaml:"actor_matcher,omitempty"`
        // AllowedDelegateClients restricts which ToolHive client IDs may
        // exchange a subject token from this issuer, for BOTH consent paths.
        // Required (validateTrustedIssuer rejects empty/absent); "*" permits
        // any confidential client holding the grant. See
        // docs/arch/17-token-exchange-delegation.md ("Accepted limitations" #1).
        //nolint:lll // field tags require full JSON+YAML names
        AllowedDelegateClients []string `json:"allowed_delegate_clients,omitempty" yaml:"allowed_delegate_clients,omitempty"`
        // AllowMayAct permits this external issuer's may_act claim to authorize
        // delegation. It defaults to false; external issuers must be opted in
        // explicitly because may_act bypasses AllowedActors and ActorMatcher. It
        // does not affect self-issued subject tokens. When enabled,
        // AllowedDelegateClients must name specific ToolHive clients rather than
        // use the wildcard.
        AllowMayAct bool `json:"allow_may_act,omitempty" yaml:"allow_may_act,omitempty"`
        // JWTBearerGrant optionally enables the plain RFC 7523 JWT-bearer grant.
        // It accepts assertions from this issuer without client authentication and
        // limits their maximum age, subjects, and RFC 8707 resources. It is
        // independent from RFC 8693 delegation policy.
        //
        // This legacy field is deprecated; configure RFC 7523 policy under
        // inbound_grants.jwt_bearer.issuer_policies.
        JWTBearerGrant *JWTBearerGrantPolicy `json:"jwt_bearer_grant,omitempty" yaml:"jwt_bearer_grant,omitempty"`
}

// MultiIssuerTokenValidator validates subject tokens from the authorization
// server itself or from configured external OIDC issuers, delegating
// self-issued tokens to SelfIssuedTokenValidator and resolving external
// issuers' JWKS (via OIDC discovery if needed) to verify signature and
// claims.
//
// A valid signature and audience alone would authorize ToolHive as a
// resource, not any particular client, as a delegate — a confused-deputy
// risk (CWE-863). validateExternalToken therefore requires a "may_act"
// claim or authorization by AllowedActors or ActorMatcher before returning
// successfully. See docs/arch/17-token-exchange-delegation.md for the full
// consent-signal precedence and trust model.
type MultiIssuerTokenValidator struct {
        selfIssuer    string
        selfValidator *SelfIssuedTokenValidator
        issuers       map[string]*externalIssuerConfig

        // cancel tears down the validator-scoped context that every per-issuer
        // jwksCache's background worker pool is rooted in. Close calls it after
        // shutting each cache down individually; see Close. It is always set for a
        // validator returned by NewMultiIssuerTokenValidator, even one with no
        // external issuers.
        cancel context.CancelFunc
}

// externalIssuerConfig holds the configuration and cached state for an external
// OIDC issuer. The embedded TrustedIssuer is treated as immutable after
// construction: resolveActorAuthorization reads TrustedIssuer.AllowedActors
// and actorMatcher (and discoverJWKSURL/fetchJWKS read httpClient) on every
// validation without holding mu, since mu protects JWKS state only. Mutating
// TrustedIssuer fields in place after NewMultiIssuerTokenValidator returns is
// a data race.
type externalIssuerConfig struct {
        TrustedIssuer

        // actorMatcher is compiled once from TrustedIssuer.ActorMatcher during
        // construction and is immutable thereafter.
        actorMatcher *cel.CompiledExpression

        // httpClient is dedicated to this issuer, built once at construction
        // time from its own InsecureAllowHTTP/AllowPrivateIPs. A single
        // validator-wide client couldn't enforce per-issuer SSRF/transport
        // policy: http.Client.CheckRedirect and Transport.DialContext have no
        // way to know which issuer's fetch they are guarding.
        httpClient *http.Client

        // jwksCache is this issuer's own jwkfetch.Cache, constructed with
        // httpClient above. A cache per issuer, rather than one shared across every
        // configured issuer, is what makes two issuers resolving to the same
        // jwks_url (e.g. two Microsoft Entra v1 tenants, which share one
        // tenant-independent JWKS endpoint) a non-event: httprc keys a cached
        // resource by URL alone, so a shared cache would make the second issuer
        // silently inherit the first one's *http.Client — defeating
        // InsecureAllowHTTP/AllowPrivateIPs's per-issuer guarantee for it. Splitting
        // the cache per issuer makes that collision unrepresentable instead of
        // guarding against it.
        jwksCache *jwkfetch.Cache

        mu sync.Mutex
        // jwksURL is resolved from OIDC discovery, or copied from
        // TrustedIssuer.JWKSURL when hand-configured. Once set, it is never
        // cleared: unlike the JWKS document itself (which jwksCache keeps fresh
        // on its own schedule — see ensureRegistered), a change to
        // the *endpoint URL* an issuer serves its keys from requires a process
        // restart to pick up. Neither Microsoft Entra nor Okta documents rotating
        // that URL, only the keys served at it, and every other TrustedIssuer
        // field already requires a restart to take effect, so this is
        // consistent with the rest of this config's lifecycle.
        jwksURL string
        // fetched is true once at least one JWKS fetch has succeeded for this
        // issuer since process start. Until then, ensureRegistered forces a
        // synchronous fetch on every call, since there is no cached value yet to
        // fall back on and jwx's own background schedule offers no way to wait
        // for its result.
        fetched bool
        // lastKidRefresh is the last time refreshOnUnknownKid forced a fetch;
        // see minKidRefreshInterval.
        lastKidRefresh time.Time
        // fetchFailedAt and fetchErr gate retries once registered but never fetched:
        // key resolution runs before signature verification, so without this an
        // authenticated client can otherwise drive one real outbound fetch per
        // token-exchange request just by naming an issuer whose endpoint is
        // down. fetchErr is served directly while the gate is closed, so the
        // caller still gets a specific error instead of a generic "try later".
        // Not consulted once fetched is true: a healthy issuer, or one serving
        // stale-but-valid keys through a later background-refresh failure, must
        // never be gated.
        fetchFailedAt time.Time
        fetchErr      error
}

// newActorMatcherEngine creates a CEL engine for admin-authored actor matcher
// expressions. The only available variable is "claims", a map[string]any.
func newActorMatcherEngine() *cel.Engine {
        return cel.NewEngine(
                celgo.Variable("claims", celgo.MapType(celgo.StringType, celgo.DynType)),
        )
}

// compileActorMatcher compiles an optional admin-authored actor matcher.
//
// It does not check the expression's output type: toolhive-core's
// cel.Engine has no accessor for the checked AST's static type, only the
// compiled *cel.CompiledExpression. An expression that type-checks but
// does not return bool (e.g. "claims.foo" where foo is a string) therefore
// compiles successfully here — it is only rejected later, at evaluation
// time, by EvaluateBool's runtime type assertion (see
// resolveActorAuthorization), which denies the token on any evaluation
// error including a wrong result type.
func compileActorMatcher(actorMatcher string) (*cel.CompiledExpression, error) {
        if actorMatcher == "" {
                return nil, nil
        }

        expr, err := newActorMatcherEngine().Compile(actorMatcher)
        if err != nil {
                return nil, fmt.Errorf("invalid actor_matcher: %w", err)
        }

        return expr, nil
}

// ResolveJWTBearerGrantPolicies clones grant policies and parses their duration
// fields once at the RunConfig-to-runtime boundary.
func ResolveJWTBearerGrantPolicies(issuers []TrustedIssuer) ([]TrustedIssuer, error) {
        resolved := slices.Clone(issuers)
        for i := range resolved {
                if resolved[i].JWTBearerGrant == nil {
                        continue
                }
                policy := cloneJWTBearerGrantPolicy(resolved[i].JWTBearerGrant)
                age, err := time.ParseDuration(policy.MaxAssertionAge)
                if err != nil || age <= 0 {
                        return nil, fmt.Errorf("trusted_issuers[%d].jwt_bearer_grant.max_assertion_age must be a positive duration", i)
                }
                policy.maxAssertionAge = age
                resolved[i].JWTBearerGrant = policy
        }
        return resolved, nil
}

// NewMultiIssuerTokenValidator creates a validator that accepts tokens from
// the authorization server itself and from the provided trusted external
// issuers. Returns an error if selfValidator is nil, selfIssuer is empty,
// any TrustedIssuer is invalid (see validateTrustedIssuer), or an issuer's
// dedicated HTTP client cannot be built (see newExternalIssuerConfig for
// why each issuer gets its own).
func NewMultiIssuerTokenValidator(
        selfValidator *SelfIssuedTokenValidator,
        selfIssuer string,
        trustedIssuers []TrustedIssuer,
        allowedAudiences []string,
) (_ *MultiIssuerTokenValidator, retErr error) {
        if selfValidator == nil {
                return nil, errors.New("selfValidator must not be nil")
        }
        if selfIssuer == "" {
                return nil, errors.New("selfIssuer must not be empty")
        }

        // Root every per-issuer JWKS worker pool in one validator-scoped context so
        // Close can tear them all down; see (*MultiIssuerTokenValidator).Close and
        // newExternalIssuerConfig.
        ctx, cancel := context.WithCancel(context.Background())
        issuers := make(map[string]*externalIssuerConfig, len(trustedIssuers))
        // If construction fails after some issuers' caches have already started,
        // release them rather than leaking their workers for the life of the
        // process (buildProvider runs after buildUpstreams, so a later failure that
        // abandons a half-built validator is reachable). Same cancel-first, single
        // shared-deadline shape as Close.
        defer func() {
                if retErr != nil {
                        cancel()
                        shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), httpTimeout)
                        defer cancelShutdown()
                        for issuerURL, issuerConfig := range issuers {
                                // Log rather than drop: retErr is the dominant signal, but a
                                // pool that fails to drain during this cleanup would otherwise
                                // leave no diagnostic.
                                if err := issuerConfig.shutdownJWKSCache(shutdownCtx); err != nil {
                                        slog.Warn("failed to shut down JWKS cache during validator construction cleanup",
                                                "issuer", issuerURL, "error", err)
                                }
                        }
                }
        }()

        for _, ti := range trustedIssuers {
                if err := validateTrustedIssuer(ti, selfIssuer, issuers, allowedAudiences); err != nil {
                        return nil, err
                }
                if len(ti.AllowedActors) == 0 && ti.ActorMatcher == "" {
                        message := func() string {
                                if ti.AllowMayAct {
                                        return "Trusted issuer has no allowed actors or actor matcher configured; " +
                                                "only may_act-bearing subject tokens from it will be accepted"
                                }
                                return "Trusted issuer has no allowed actors or actor matcher configured; " +
                                        "no subject tokens from it will be accepted"
                        }()
                        slog.Warn(message, "issuer", ti.IssuerURL)
                }
                if !looksLikeResourceIdentifier(ti.ExpectedAudience) {
                        slog.Warn("Trusted issuer's expected_audience does not look like a resource identifier "+
                                "(no URI scheme, e.g. \"https://\" or \"api://\"); it may actually be a client ID. "+
                                "rejectIDTokenClaims only catches an ID token that carries at_hash/c_hash, which OIDC "+
                                "does not require, so this issuer's audience check may not distinguish an ID token "+
                                "from an access token",
                                "issuer", ti.IssuerURL, "expected_audience", ti.ExpectedAudience,
                        )
                }

                issuerConfig, err := newExternalIssuerConfig(ctx, ti)
                if err != nil {
                        return nil, err
                }
                issuers[ti.IssuerURL] = issuerConfig
        }

        return &MultiIssuerTokenValidator{
                selfIssuer:    selfIssuer,
                selfValidator: selfValidator,
                issuers:       issuers,
                cancel:        cancel,
        }, nil
}

// Close shuts down every per-issuer jwkfetch.Cache, stopping the background JWKS
// refresh worker pool (and its ~3 goroutines) each one runs. It cancels the
// validator-scoped context those pools share — signalling them all to stop at
// once — then waits, under a single shared httpTimeout budget, for each cache
// to drain. Cancelling up front (rather than after the loop) means the pools
// unwind in parallel, and the one shared deadline bounds the total wait by
// httpTimeout rather than httpTimeout×N even if a pool ignores cancellation.
// This is the same order the construction-failure path uses. It is safe to call
// more than once and on a validator with no external issuers; the validator
// must not be used after Close.
//
// A server holds its MultiIssuerTokenValidator and calls this from Close and
// its construction error path (see pkg/authserver), so neither a normal
// shutdown nor a failed reconstruction leaks these workers.
// authserver.CloseIdleConnections deliberately does NOT reach it: that path
// must stay safe to call on a still-serving server, which needs the workers to
// keep refreshing external issuers' keys.
func (v *MultiIssuerTokenValidator) Close() error {
        if v == nil {
                return nil
        }
        if v.cancel != nil {
                v.cancel()
        }
        ctx, cancel := context.WithTimeout(context.Background(), httpTimeout)
        defer cancel()
        var errs []error
        for issuerURL, issuerConfig := range v.issuers {
                if err := issuerConfig.shutdownJWKSCache(ctx); err != nil {
                        errs = append(errs, fmt.Errorf("issuer %s: %w", issuerURL, err))
                }
        }
        return errors.Join(errs...)
}

func cloneJWTBearerGrantPolicy(policy *JWTBearerGrantPolicy) *JWTBearerGrantPolicy {
        if policy == nil {
                return nil
        }
        clone := *policy
        clone.AcceptedAudiences = slices.Clone(policy.AcceptedAudiences)
        clone.SubjectBindings = make([]JWTBearerSubjectBinding, len(policy.SubjectBindings))
        for i, binding := range policy.SubjectBindings {
                clone.SubjectBindings[i] = JWTBearerSubjectBinding{
                        Subject:          binding.Subject,
                        AllowedResources: slices.Clone(binding.AllowedResources),
                }
        }
        if clone.maxAssertionAge == 0 {
                clone.maxAssertionAge, _ = time.ParseDuration(clone.MaxAssertionAge)
        }
        return &clone
}

// newExternalIssuerConfig builds the *externalIssuerConfig for a single
// already-validated TrustedIssuer: a dedicated HTTP client (scoped to that
// issuer's own InsecureAllowHTTP/AllowPrivateIPs) and its own body-size-capped
// jwkfetch.Cache rooted in ctx. Called once per issuer from
// NewMultiIssuerTokenValidator's constructor loop, after validateTrustedIssuer
// and the startup warnings have already run for ti. ctx is the validator-scoped
// context (see NewMultiIssuerTokenValidator); cancelling it, or the validator's
// Close, tears down the cache's worker pool.
func newExternalIssuerConfig(ctx context.Context, ti TrustedIssuer) (*externalIssuerConfig, error) {
        // Clone AllowedActors and AllowedDelegateClients so a caller mutating
        // their original slices in place (e.g. a future config reload) cannot
        // race with the unsynchronized reads in resolveActorAuthorization and
        // validateExternalToken, which run on every validation without
        // holding externalIssuerConfig.mu (that mutex guards JWKS state only).
        ti.AllowedActors = slices.Clone(ti.AllowedActors)
        ti.AllowedDelegateClients = slices.Clone(ti.AllowedDelegateClients)
        ti.JWTBearerGrant = cloneJWTBearerGrantPolicy(ti.JWTBearerGrant)

        actorMatcher, err := compileActorMatcher(ti.ActorMatcher)
        if err != nil {
                return nil, fmt.Errorf("issuer_url %q: %w", ti.IssuerURL, err)
        }

        // Deliberately networking.NewHttpClientBuilder(), not
        // NewHostScopedClientBuilder: that helper ORs
        // INSECURE_DISABLE_URL_VALIDATION and an auto-localhost exemption into
        // BOTH the HTTP-scheme and private-IP gates, so an unrelated env var —
        // or a trusted issuer that merely happens to be on localhost — would
        // silently widen AllowPrivateIPs regardless of what the operator set,
        // defeating the point of splitting the two flags per issuer.
        //
        // Keep-alive connections are disabled: this client dials jwks_uri, a host taken
        // from an untrusted discovery document, only on the fixed
        // jwksRefreshInterval schedule plus occasional on-demand refreshes —
        // no hot path here to trade the per-dial SSRF check away for.
        builder := networking.NewHttpClientBuilder().
                WithInsecureAllowHTTP(ti.InsecureAllowHTTP).
                WithPrivateIPs(ti.AllowPrivateIPs).
                WithTimeout(httpTimeout).
                WithDisableKeepAlives(true)
        if ti.CAFilePath != "" {
                builder = builder.WithSystemRootsPlusCABundle(ti.CAFilePath)
        }
        httpClient, err := builder.Build()
        if err != nil {
                return nil, fmt.Errorf("issuer_url %q: failed to build HTTP client: %w", ti.IssuerURL, err)
        }
        // Guard against a discovery/JWKS redirect hop landing on a
        // different, unvetted host — the same policy the transparent proxy
        // data path applies to a response derived from an untrusted remote
        // server (see SameHostRedirectPolicy's doc comment).
        httpClient.CheckRedirect = networking.SameHostRedirectPolicy()

        // jwkfetch applies this cap to every initial and background JWKS fetch.
        // Unlike a transport-level truncation, it reads at most one byte beyond the
        // limit and returns a specific body-too-large error before parsing.

        // One jwkfetch.Cache per issuer (see externalIssuerConfig.jwksCache's doc
        // comment for why), each running its own background worker pool
        // (jwkfetch.NewCache -> httprc.Client.Start). WithWorkers(1) caps that pool to
        // one worker per issuer instead of httprc's default five — budget roughly
        // three goroutines per issuer including its controller loop and wait-group
        // waiter.
        //
        // The pool is rooted in ctx, the validator-scoped context, so it outlives
        // any single call but is torn down when the validator's Close cancels ctx
        // and shuts each cache down (see shutdownJWKSCache). Before this wiring the
        // pools ran for the life of the process, one leaked set per issuer on every
        // server reconstruction; pkg/auth/token.go's TokenValidator still has that
        // shape but roots its cache in a caller-supplied context.
        jwksCache, err := jwkfetch.NewCache(
                ctx,
                httprc.NewClient(httprc.WithWorkers(1)),
                jwkfetch.WithHTTPClient(httpClient),
                jwkfetch.WithMaxBodySize(maxResponseBodySize),
        )
        if err != nil {
                return nil, fmt.Errorf("issuer_url %q: failed to create JWKS cache: %w", ti.IssuerURL, err)
        }

        return &externalIssuerConfig{
                TrustedIssuer: ti,
                actorMatcher:  actorMatcher,
                jwksURL:       ti.JWKSURL,
                httpClient:    httpClient,
                jwksCache:     jwksCache,
        }, nil
}

// shutdownJWKSCache stops this issuer's JWKS refresh worker pool, waiting for
// its goroutines to drain until ctx expires. ctx must be fresh — not the
// validator-scoped one the cache was started with, which Close cancels — or
// Shutdown returns immediately without waiting; see httprc's
// Controller.ShutdownContext. Callers pass one shared deadline context so a set
// of N issuers drains within a single timeout budget, not N of them. A nil
// cache (an externalIssuerConfig built by ValidateTrustedIssuers, which never
// starts one) is a no-op, as is a double shutdown.
func (c *externalIssuerConfig) shutdownJWKSCache(ctx context.Context) error {
        if c.jwksCache == nil {
                return nil
        }
        return c.jwksCache.Shutdown(ctx)
}

// ValidateTrustedIssuers runs every structural check
// NewMultiIssuerTokenValidator performs on trustedIssuers — required fields,
// self-issuer collision, duplicate issuers, ActorClaim reachability, and
// ActorMatcher compilation — without constructing a validator or any
// per-issuer HTTP client. Config validation calls this to fail before the live upstream DCR registration
// and storage creation that run between RunConfig.Validate and server
// construction; NewMultiIssuerTokenValidator repeats the same checks at
// server startup as defence in depth. Both route through validateTrustedIssuer,
// so the two can't drift out of sync.
func ValidateTrustedIssuers(trustedIssuers []TrustedIssuer, selfIssuer string, allowedAudiences []string) error {
        issuers := make(map[string]*externalIssuerConfig, len(trustedIssuers))
        for _, ti := range trustedIssuers {
                if err := validateTrustedIssuer(ti, selfIssuer, issuers, allowedAudiences); err != nil {
                        return err
                }
                issuers[ti.IssuerURL] = &externalIssuerConfig{TrustedIssuer: ti}
        }
        return nil
}

// Validate parses the raw JWT to extract the issuer claim, then routes validation
// to either the self-issued validator or the appropriate external issuer validator.
// Returns an error if the issuer is not trusted.
func (v *MultiIssuerTokenValidator) Validate(ctx context.Context, rawToken string) (*ValidatedClaims, error) {
        // Parse the JWT without verification to peek at the issuer claim.
        // This is safe because we verify the signature in a subsequent step.
        issuer, err := peekIssuer(rawToken)
        if err != nil {
                return nil, fmt.Errorf("failed to determine token issuer: %w", err)
        }

        if issuer == v.selfIssuer {
                return v.selfValidator.Validate(ctx, rawToken)
        }

        issuerConfig, ok := v.issuers[issuer]
        if !ok {
                return nil, fmt.Errorf("untrusted issuer: %q", issuer)
        }

        return v.validateExternalToken(ctx, rawToken, issuerConfig)
}

// ValidateJWTBearerAssertion verifies a plain RFC 7523 assertion issued by an
// external issuer explicitly enabled for the JWT-bearer grant. It deliberately does not apply RFC 8693
// delegation consent (AllowedActors, ActorMatcher, or may_act): a JWT-bearer
// assertion's authorization policy is separate from delegation policy.
//
// The assertion audience must intersect the issuer's configured
// JWTBearerGrant.AcceptedAudiences, or just tokenEndpoint when that is unset.
// The caller is responsible for the phase-specific subject binding, maximum
// age, replay, and issuance policy after this cryptographic verification
// succeeds.
func (v *MultiIssuerTokenValidator) ValidateJWTBearerAssertion(
        ctx context.Context, rawToken, tokenEndpoint string,
) (*ValidatedClaims, error) {
        if tokenEndpoint == "" {
                return nil, errors.New("token endpoint must not be empty")
        }

        issuer, err := peekIssuer(rawToken)
        if err != nil {
                return nil, fmt.Errorf("failed to determine assertion issuer: %w", err)
        }
        issuerConfig, ok := v.issuers[issuer]
        if !ok {
                return nil, fmt.Errorf("untrusted assertion issuer: %q", issuer)
        }
        if issuerConfig.JWTBearerGrant == nil {
                return nil, fmt.Errorf("assertion issuer is not enabled for the JWT-bearer grant: %q", issuer)
        }

        standardClaims, extraClaims, err := v.verifyExternalSignature(ctx, rawToken, issuerConfig)
        if err != nil {
                return nil, err
        }

        if err := rejectIDTokenClaims(extraClaims); err != nil {
                return nil, err
        }

        acceptedAudiences := []string{tokenEndpoint}
        if issuerConfig.JWTBearerGrant != nil && len(issuerConfig.JWTBearerGrant.AcceptedAudiences) > 0 {
                acceptedAudiences = issuerConfig.JWTBearerGrant.AcceptedAudiences
        }
        if err := validateJWTBearerAssertionClaims(standardClaims, issuerConfig.IssuerURL, acceptedAudiences); err != nil {
                return nil, err
        }
        return buildValidatedClaims(standardClaims, extraClaims), nil
}

// verifyExternalSignature parses and verifies a JWT using the configured
// issuer's cached JWKS. Both RFC 8693 subject tokens and RFC 7523 assertions
// use this path so unknown-kid refreshes and key-use/algorithm restrictions
// cannot diverge between grant types.
func (v *MultiIssuerTokenValidator) verifyExternalSignature(
        ctx context.Context, rawToken string, issuerConfig *externalIssuerConfig,
) (jwt.Claims, map[string]any, error) {
        parsedToken, err := jwt.ParseSigned(rawToken, allowedSignatureAlgorithms)
        if err != nil {
                return jwt.Claims{}, nil, fmt.Errorf("subject token is not a valid JWT: %w", err)
        }

        jwks, err := v.lookupJWKS(ctx, issuerConfig)
        if err != nil {
                return jwt.Claims{}, nil, fmt.Errorf("failed to fetch JWKS for issuer %s: %w", issuerConfig.IssuerURL, err)
        }

        standardClaims, extraClaims, kidMatched, err := verifySignature(parsedToken, jwks)
        if err != nil && !kidMatched {
                // The token's kid isn't among the keys we have cached — possibly a
                // legitimate rotation the issuer's five-minute refresh hasn't caught
                // up with yet.
                // Force an immediate re-fetch and retry once before giving up; a
                // spoofed-kid attempt (kidMatched true) never reaches this branch,
                // since a refresh can't change what signature the token was made
                // with.
                v.refreshOnUnknownKid(ctx, issuerConfig)
                if refreshedJWKS, lookupErr := v.lookupJWKS(ctx, issuerConfig); lookupErr == nil {
                        standardClaims, extraClaims, _, err = verifySignature(parsedToken, refreshedJWKS)
                }
        }
        if err != nil {
                return jwt.Claims{}, nil, err
        }
        return standardClaims, extraClaims, nil
}

// validateExternalToken verifies a JWT from a trusted external issuer by
// fetching the issuer's JWKS (with caching) and validating the signature and claims.
func (v *MultiIssuerTokenValidator) validateExternalToken(
        ctx context.Context,
        rawToken string,
        issuerConfig *externalIssuerConfig,
) (*ValidatedClaims, error) {
        standardClaims, extraClaims, err := v.verifyExternalSignature(ctx, rawToken, issuerConfig)
        if err != nil {
                return nil, err
        }

        // Validate issuer and audience, tolerating externalClockSkewLeeway on
        // nbf/iat/exp. Expiry is enforced strictly (no leeway) below.
        expected := jwt.Expected{
                Issuer:      issuerConfig.IssuerURL,
                AnyAudience: jwt.Audience{issuerConfig.ExpectedAudience},
        }
        if err := standardClaims.ValidateWithLeeway(expected, externalClockSkewLeeway); err != nil {
                return nil, fmt.Errorf("subject token claims validation failed: %w", err)
        }

        // Subject is required for delegation.
        if standardClaims.Subject == "" {
                return nil, fmt.Errorf("subject token is missing required 'sub' claim")
        }

        // Expiry is required so the delegated token can be bounded by the subject
        // token's remaining lifetime.
        if standardClaims.Expiry == nil {
                return nil, errors.New("subject token is missing required 'exp' claim")
        }

        // An expired subject token is never acceptable, despite the leeway
        // above: it cannot bound the delegated token's lifetime.
        if standardClaims.Expiry.Time().Before(time.Now()) {
                return nil, errors.New("subject token has expired")
        }

        // Reject a subject token that is actually an ID token — see
        // rejectIDTokenClaims's doc comment.
        if err := rejectIDTokenClaims(extraClaims); err != nil {
                return nil, err
        }

        if err := checkMayActAllowed(extraClaims, v.selfIssuer, issuerConfig); err != nil {
                return nil, err
        }

        claims := buildValidatedClaims(standardClaims, extraClaims)

        // Recorded for every external token, even a may_act-bearing one that
        // leaves ExternalActor unset — without this, a may_act-authorized
        // external token would be indistinguishable from a self-issued exchange
        // downstream, which is backwards: this is the path that most needs an
        // audit trail. Set from issuerConfig, already matched against the
        // validated "iss" claim above — never from token content.
        claims.ExternalIssuer = issuerConfig.IssuerURL

        // Set unconditionally, including for a may_act-bearing token: an issuer
        // that can emit may_act but has no way to populate AllowedActors (e.g.
        // Entra, Okta) would otherwise have no per-client containment on that
        // path (see #5989). checkDelegationConsent enforces it against the
        // authenticated client; this validator never sees that client ID.
        claims.AllowedDelegateClients = issuerConfig.AllowedDelegateClients

        // A may_act claim is authoritative and enforced by checkDelegationConsent
        // against the authenticated client — validateMayActShape above has
        // already confirmed its own "iss" (if any) names this server, so
        // may_act.sub is in ToolHive's own client namespace. AllowedActors and
        // ActorMatcher below are skipped entirely when MayAct is set. Otherwise,
        // either the resolved external actor claim or the matcher evaluated against
        // the complete signature-verified JWT claims authorizes delegation. Neither
        // is ever compared against the authenticated ToolHive client directly.
        if claims.MayAct == nil {
                actor, authorized, err := resolveActorAuthorization(issuerConfig, claims, extraClaims)
                if err != nil {
                        return nil, err
                }
                claims.ExternalActor = actor
                claims.ExternalActorAuthorized = authorized
        }

        return claims, nil
}

// checkMayActAllowed rejects a may_act claim from an issuer that hasn't
// opted in (issuerConfig.AllowMayAct false) before validating the claim's
// shape: a malformed may_act on a disabled-by-default issuer should surface
// this clearer, more-actionable message rather than "malformed 'may_act'
// claim" — both paths reject the token either way, but this one names the
// actual reason for the common case.
//
// selfIssuer (not issuerConfig.IssuerURL) is the correct comparison for
// may_act's optional "iss" once shape validation runs: that member
// identifies the namespace of may_act.sub, always compared against a
// ToolHive client ID regardless of which issuer validated the surrounding
// token. requireIss is true here — unlike the self-issued path, the external
// path cannot leave "iss" optional; see validateMayActShape's doc comment.
func checkMayActAllowed(extraClaims map[string]any, selfIssuer string, issuerConfig *externalIssuerConfig) error {
        rawMayAct, ok := extraClaims["may_act"]
        if ok && rawMayAct != nil && !issuerConfig.AllowMayAct {
                return fmt.Errorf(
                        "subject token from issuer %q carries a may_act claim, but allow_may_act is disabled",
                        issuerConfig.IssuerURL)
        }
        return validateMayActShape(extraClaims, selfIssuer, true)
}

// ensureRegistered resolves issuerConfig.jwksURL — via OIDC discovery on
// first use — and registers it with issuerConfig's own jwkfetch.Cache (see
// registerOrRefresh for the Register-vs-Refresh decision). Discovery
// happens at most once per issuer for the life of the process; see
// externalIssuerConfig's jwksURL doc comment.
//
// The JWKS fetch is retried until one succeeds (fetched stays false
// otherwise), gated by jwksFetchFailureBackoff: key resolution runs before
// the subject token's signature is checked, so without this gate an
// authenticated client could force a real outbound attempt on every
// request to an issuer whose endpoint is down. The gate only applies
// before the first successful fetch — once fetched is true, a healthy
// issuer or one serving stale-but-valid keys through a later refresh
// failure is unaffected. While closed, the last error is replayed
// directly rather than retried.
//
// Every background refresh, not just the first request-triggered one,
// is capped by the cache's maxResponseBodySize limit despite having no
// request or client authentication of its own — without that cap a
// compromised issuer could force an oversized allocation on every autonomous
// refresh, indefinitely.
func (v *MultiIssuerTokenValidator) ensureRegistered(ctx context.Context, issuerConfig *externalIssuerConfig) error {
        issuerConfig.mu.Lock()
        defer issuerConfig.mu.Unlock()

        if issuerConfig.fetched {
                return nil
        }
        if time.Since(issuerConfig.fetchFailedAt) < jwksFetchFailureBackoff {
                return issuerConfig.fetchErr
        }

        if err := v.registerOrRefresh(ctx, issuerConfig); err != nil {
                issuerConfig.fetchErr = err
                issuerConfig.fetchFailedAt = time.Now()
                return err
        }
        issuerConfig.fetched = true
        issuerConfig.fetchErr = nil
        return nil
}

// registerOrRefresh performs the actual discovery/registration/fetch
// attempt for issuerConfig; ensureRegistered holds issuerConfig.mu across
// this call, single-flighting it per issuer so concurrent validations
// don't pile up N redundant fetches.
//
// Whether jwksURL is already registered is asked of
// issuerConfig.jwksCache.IsRegistered directly, never remembered in a
// field: Register's own registration step is a channel send to the
// cache's backend goroutine and can fail after enqueue but before receipt
// (context deadline), in which case nothing was actually registered. A
// locally remembered "we called Register" flag can't distinguish that
// from "registered, only the fetch failed", and would wrongly keep
// retrying via Refresh — which errors on a URL the cache never heard of —
// forever after. Asking the cache directly is authoritative either way.
//
// IsRegistered makes no network call but isn't free: it's a round-trip
// over that same channel, so it blocks if the backend is busy and returns
// false (not an error) on context expiry. Its own timeout below matters
// for that reason — a false from an expired context just routes to
// Register, whose "already registered" error is transient and absorbed by
// ensureRegistered's backoff gate.
func (v *MultiIssuerTokenValidator) registerOrRefresh(ctx context.Context, issuerConfig *externalIssuerConfig) error {
        // Detach from the caller's request context throughout this function: net/http
        // cancels ctx when the client disconnects, and this runs before the subject
        // token's signature is even checked, so an aborted connection must not cut off
        // work other in-flight validations of this issuer are waiting on (mu, held by
        // the caller), nor let repeating the abort drive unbounded outbound requests
        // to the external IdP.
        detached := context.WithoutCancel(ctx)

        if issuerConfig.jwksURL == "" {
                discoverCtx, cancel := context.WithTimeout(detached, httpTimeout)
                jwksURL, err := v.discoverJWKSURL(discoverCtx, issuerConfig)
                cancel()
                if err != nil {
                        return fmt.Errorf("OIDC discovery failed for %s: %w", issuerConfig.IssuerURL, err)
                }
                issuerConfig.jwksURL = jwksURL
        }

        // Validate the JWKS URL here, in the single choke point every
        // registration passes through — whether it was hand-configured on
        // TrustedIssuer or just discovered above. A configured JWKSURL never
        // reaches discoverJWKSURL, so checking only there would leave
        // hand-configured URLs unvalidated.
        if err := ValidateJWKSURL(issuerConfig.jwksURL, issuerConfig.InsecureAllowHTTP, issuerConfig.AllowPrivateIPs); err != nil {
                return fmt.Errorf("jwks_url for issuer %s is invalid: %w", issuerConfig.IssuerURL, err)
        }

        registeredCtx, cancel := context.WithTimeout(detached, httpTimeout)
        registered := issuerConfig.jwksCache.IsRegistered(registeredCtx, issuerConfig.jwksURL)
        cancel()

        if registered {
                // Already registered with this issuer's own cache, on a prior call
                // whose own fetch never completed successfully — the only way this
                // can be true, now that each issuer has its own cache. Register
                // would error on an already-tracked URL, so retry via Refresh
                // instead.
                fetchCtx, cancel := context.WithTimeout(detached, httpTimeout)
                defer cancel()
                if _, err := issuerConfig.jwksCache.Refresh(fetchCtx, issuerConfig.jwksURL); err != nil {
                        return fmt.Errorf("failed to fetch JWKS for issuer %s: %w", issuerConfig.IssuerURL, err)
                }
                return nil
        }

        // A newly created httprc.Resource is always scheduled to fetch
        // immediately, so Register's own default WithWaitReady(true) blocks on
        // that single automatic fetch. An explicit Refresh call right here would
        // race it and issue a genuine second outbound request — that's why the
        // registered branch above, not this one, is where Refresh is used.
        fetchCtx, cancel := context.WithTimeout(detached, httpTimeout)
        defer cancel()
        if err := issuerConfig.jwksCache.Register(
                fetchCtx,
                issuerConfig.jwksURL,
                jwkfetch.WithConstantInterval(jwksRefreshInterval),
        ); err != nil {
                return fmt.Errorf("failed to register JWKS for issuer %s: %w", issuerConfig.IssuerURL, err)
        }
        return nil
}

// lookupJWKS returns issuerConfig's current JWKS, registering and fetching it
// first if this is the first use (see ensureRegistered). jwkfetch.Cache serves the
// last successfully fetched Set even while a later background refresh is
// failing (httprc only stores a value after a successful fetch), so a
// transient outage at an issuer that has already been reached once no longer
// surfaces as a validation failure.
// Each call converts the cached Set into a fresh *jose.JSONWebKeySet —
// dropping UnsupportedKey placeholders first (see bridgeJWKSet) — the two
// libraries' key representations aren't shared, so nothing here is visible
// to, or mutable by, any other concurrent caller.
func (v *MultiIssuerTokenValidator) lookupJWKS(
        ctx context.Context,
        issuerConfig *externalIssuerConfig,
) (*jose.JSONWebKeySet, error) {
        if err := v.ensureRegistered(ctx, issuerConfig); err != nil {
                return nil, err
        }

        set, err := issuerConfig.jwksCache.Lookup(ctx, issuerConfig.jwksURL)
        if err != nil {
                return nil, fmt.Errorf("failed to lookup JWKS: %w", err)
        }

        jwks, err := bridgeJWKSet(set)
        if err != nil {
                return nil, err
        }

        if len(jwks.Keys) == 0 {
                return nil, errors.New("JWKS contains no keys")
        }
        if len(jwks.Keys) > maxJWKSKeys {
                return nil, fmt.Errorf("JWKS contains too many keys: %d (max %d)", len(jwks.Keys), maxJWKSKeys)
        }

        return jwks, nil
}

// bridgeJWKSet converts a jwx jwk.Set into the go-jose jose.JSONWebKeySet
// shape verifySignature works with, by round-tripping through the standard
// JWKS JSON both types serialize to and from. This is the simplest correct
// conversion between the two libraries' key representations — it doesn't
// require hand-mapping every key type's fields (RSA, EC, OKP, ...) between
// jwk.Key and jose.JSONWebKey.
//
// UnsupportedKey placeholders are dropped first. jwx v4 retains unparsable
// JWKS entries (unknown kty, RSA below the 2048-bit floor, malformed
// parameters) as placeholders whose original JSON round-trips losslessly.
// Feeding those to go-jose either fails the whole set or rehydrates a key
// jwx already refused; skipping them keeps the usable keys and applies the
// same parse-time floor on this path as TokenValidator's jwk.Export.
func bridgeJWKSet(set jwk.Set) (*jose.JSONWebKeySet, error) {
        usable := jwk.NewSet()
        for _, key := range set.All() {
                if jwk.IsUnsupportedKey(key) {
                        continue
                }
                if err := usable.AddKey(key); err != nil {
                        return nil, fmt.Errorf("failed to collect JWKS key: %w", err)
                }
        }

        raw, err := json.Marshal(usable)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal JWKS: %w", err)
        }
        var jwks jose.JSONWebKeySet
        if err := json.Unmarshal(raw, &jwks); err != nil {
                return nil, fmt.Errorf("failed to parse JWKS: %w", err)
        }
        return &jwks, nil
}

// refreshOnUnknownKid forces issuerConfig's own jwkfetch.Cache to re-fetch its
// JWKS immediately, ahead of jwx's own background refresh schedule, when a
// subject token names a kid the last cached JWKS doesn't have — the
// situation a legitimate key rotation produces. Gated by minKidRefreshInterval
// and single-flighted via issuerConfig.mu, the same mutex ensureRegistered
// uses for its own initial fetch.
//
// Errors are logged, not returned: the caller (validateExternalToken) has
// already failed signature verification once and will simply fail again if
// the refresh didn't produce a usable key, which is the correct outcome for
// a genuinely invalid token.
func (*MultiIssuerTokenValidator) refreshOnUnknownKid(ctx context.Context, issuerConfig *externalIssuerConfig) {
        issuerConfig.mu.Lock()
        defer issuerConfig.mu.Unlock()

        if time.Since(issuerConfig.lastKidRefresh) < minKidRefreshInterval {
                return
        }
        issuerConfig.lastKidRefresh = time.Now()

        fetchCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*httpTimeout)
        defer cancel()
        if _, err := issuerConfig.jwksCache.Refresh(fetchCtx, issuerConfig.jwksURL); err != nil {
                slog.Debug("JWKS refresh on unknown kid failed", "issuer", issuerConfig.IssuerURL, "error", err)
        }
}

// peekIssuer parses a JWT without signature verification to extract the "iss" claim.
func peekIssuer(rawToken string) (string, error) {
        token, err := jwt.ParseSigned(rawToken, allowedSignatureAlgorithms)
        if err != nil {
                return "", fmt.Errorf("subject token is not a valid JWT: %w", err)
        }

        var claims jwt.Claims
        if err := token.UnsafeClaimsWithoutVerification(&claims); err != nil {
                return "", fmt.Errorf("failed to extract claims from subject token: %w", err)
        }

        if claims.Issuer == "" {
                return "", fmt.Errorf("subject token is missing 'iss' claim")
        }

        return claims.Issuer, nil
}

// discoverJWKSURL performs OIDC discovery to resolve the JWKS URL for an issuer.
// It fetches the OpenID Connect discovery document at {issuerURL}/.well-known/openid-configuration
// and extracts the jwks_uri field, using issuerConfig's own dedicated HTTP client.
//
// IssuerURL itself may legitimately carry a trailing slash (see
// validateTrustedIssuerURL in pkg/authserver/config.go — Microsoft Entra ID
// v1 is one real-world example), so it is trimmed before the well-known path
// is appended here, per OIDC Discovery §4.1. The comparison against the
// discovery document's own "issuer" below intentionally uses the untrimmed
// IssuerURL: per §4.3 that comparison must be exact, since it stands in for
// the eventual match against the token's "iss".
func (*MultiIssuerTokenValidator) discoverJWKSURL(ctx context.Context, issuerConfig *externalIssuerConfig) (string, error) {
        discoveryURL := strings.TrimSuffix(issuerConfig.IssuerURL, "/") + "/.well-known/openid-configuration"

        req, err := http.NewRequestWithContext(ctx, http.MethodGet, discoveryURL, nil)
        if err != nil {
                return "", fmt.Errorf("failed to create discovery request: %w", err)
        }

        resp, err := issuerConfig.httpClient.Do(req)
        if err != nil {
                return "", fmt.Errorf("discovery request failed: %w", err)
        }
        defer func() { _ = resp.Body.Close() }()

        if resp.StatusCode != http.StatusOK {
                _, _ = io.Copy(io.Discard, resp.Body)
                return "", fmt.Errorf("discovery endpoint returned status %d", resp.StatusCode)
        }

        body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBodySize))
        if err != nil {
                return "", fmt.Errorf("failed to read discovery response: %w", err)
        }

        var doc oauthproto.OIDCDiscoveryDocument
        if err := json.Unmarshal(body, &doc); err != nil {
                return "", fmt.Errorf("failed to parse discovery document: %w", err)
        }

        if doc.Issuer != issuerConfig.IssuerURL {
                return "", fmt.Errorf("discovery document issuer does not match configured issuer")
        }

        if doc.JWKSURI == "" {
                return "", fmt.Errorf("discovery document missing 'jwks_uri'")
        }

        // The returned URL is validated by the caller (ensureRegistered), which is
        // the single choke point covering both discovered and pre-configured
        // JWKS URLs.
        return doc.JWKSURI, nil
}

// ValidateTrustedIssuerURL checks that issuerURL is a valid trusted external
// OIDC issuer identifier. Trusted issuers require HTTPS unless their own
// insecureAllowHTTP opt-in is set; unlike this server's issuer, localhost is
// not exempt. Query, fragment, and userinfo are forbidden, while a trailing
// slash is permitted for providers such as Microsoft Entra ID v1.
func ValidateTrustedIssuerURL(issuerURL string, insecureAllowHTTP bool) error {
        if issuerURL == "" {
                return errors.New("issuer_url is required")
        }
        parsed, err := url.Parse(issuerURL)
        if err != nil {
                return errors.New("invalid URL")
        }
        if parsed.Scheme == "" {
                return errors.New("scheme is required")
        }
        if parsed.Hostname() == "" {
                return errors.New("host is required")
        }
        if parsed.RawQuery != "" {
                return errors.New("must not contain query component")
        }
        if parsed.Fragment != "" {
                return errors.New("must not contain fragment component")
        }
        if parsed.User != nil {
                return errors.New("must not contain userinfo (credentials in the URL)")
        }
        if parsed.Scheme != "https" && (parsed.Scheme != "http" || !insecureAllowHTTP) {
                return errors.New("scheme must be https (or http with insecure_allow_http)")
        }
        return nil
}

// ValidateJWKSURL checks that jwksURL parses, has a host, uses HTTPS unless
// insecureAllowHTTP permits plain HTTP — and only exactly the "http" scheme,
// not any other non-https scheme such as "file" or "ftp" — and, when the
// host is an IP literal, is not a private or loopback address unless
// allowPrivateIPs permits that. Both flags come from the specific
// TrustedIssuer being fetched (see ensureRegistered), never from a validator-wide
// or self-issuer setting. This prevents SSRF attacks where a compromised
// discovery document — or a hand-configured jwks_url — points to internal
// services.
//
// This is shared by endpoint-only config-time validation and the runtime fetch
// choke point. Both use the issuer's allowPrivateIPs policy for literal IP
// hosts; runtime additionally protects DNS resolution on every outbound fetch.
func ValidateJWKSURL(jwksURL string, insecureAllowHTTP, allowPrivateIPs bool) error {
        u, err := url.Parse(jwksURL)
        if err != nil {
                return errors.New("invalid URL")
        }

        if u.Hostname() == "" {
                return errors.New("host is required")
        }

        if u.Fragment != "" {
                return errors.New("must not contain fragment component")
        }

        // Unlike issuer_url, a jwks_url carrying userinfo would actually work —
        // net/http turns it into a Basic auth header on every JWKS fetch — which
        // is precisely why it is rejected rather than tolerated: it would put a
        // live credential in the RunConfig or in any log that quotes the URL. A
        // JWKS endpoint is public by definition (it serves verification keys), so
        // there is no legitimate reason to authenticate to one.
        if u.User != nil {
                return errors.New("must not contain userinfo (credentials in the URL)")
        }

        if u.Scheme != "https" && (u.Scheme != "http" || !insecureAllowHTTP) {
                return fmt.Errorf("must use HTTPS, got %q", u.Scheme)
        }

        host := u.Hostname()
        ip := net.ParseIP(host)
        if ip != nil && !allowPrivateIPs && networking.IsPrivateIP(ip) {
                return errors.New("must not point to a private or loopback address")
        }

        return nil
}

func validateTrustedIssuerEndpoints(ti TrustedIssuer) error {
        if err := ValidateTrustedIssuerURL(ti.IssuerURL, ti.InsecureAllowHTTP); err != nil {
                return err
        }
        if ti.JWKSURL != "" {
                if err := ValidateJWKSURL(ti.JWKSURL, ti.InsecureAllowHTTP, ti.AllowPrivateIPs); err != nil {
                        return fmt.Errorf("jwks_url: %w", err)
                }
        }
        if ti.AllowPrivateIPs && ti.JWKSURL == "" {
                return errors.New("allow_private_ips requires jwks_url to be set explicitly")
        }
        return nil
}

// validateTrustedIssuer checks a single TrustedIssuer for structural validity
// before it is admitted into issuers: required fields, no collision with
// selfIssuer or an already-registered issuer, an ActorClaim that
// resolveActorAuthorization can actually read from Extra (or ClientID), a
// compilable ActorMatcher, and a well-formed AllowedDelegateClients (non-empty,
// no empty entries, and the wildcard anyDelegateClient never combined with
// specific client IDs or AllowMayAct).
//
// Error messages name TrustedIssuer's wire keys (issuer_url,
// expected_audience, actor_claim), not its Go field names: TrustedIssuer is
// the serialized RunConfig.TrustedIssuers schema (see its doc comment), so an
// operator's YAML uses these keys and should see them echoed back, not Go
// identifiers they never wrote.
func validateTrustedIssuer(
        ti TrustedIssuer, selfIssuer string, issuers map[string]*externalIssuerConfig, allowedAudiences []string,
) error {
        if err := validateTrustedIssuerEndpoints(ti); err != nil {
                return err
        }
        if ti.ExpectedAudience == "" && ti.JWTBearerGrant == nil {
                return fmt.Errorf("issuer_url %q: expected_audience is required when JWT-bearer grant is disabled", ti.IssuerURL)
        }
        if ti.IssuerURL == selfIssuer {
                return fmt.Errorf("issuer_url %q: must not equal the authorization server's own issuer; "+
                        "self-issued tokens are already handled separately", ti.IssuerURL)
        }
        if _, dup := issuers[ti.IssuerURL]; dup {
                return errors.New("issuer_url configured more than once")
        }
        if ti.ActorClaim != "" && slices.Contains(actorClaimsNotInExtra, ti.ActorClaim) {
                return fmt.Errorf(
                        "issuer_url %q: actor_claim %q is not supported "+
                                `(use "client_id" or a non-registered claim such as "azp", "appid", "cid")`,
                        ti.IssuerURL, ti.ActorClaim)
        }
        if _, err := compileActorMatcher(ti.ActorMatcher); err != nil {
                return fmt.Errorf("issuer_url %q: %w", ti.IssuerURL, err)
        }
        if err := validateDelegationPolicy(ti); err != nil {
                return err
        }
        if err := validateJWTBearerGrantPolicy(ti, allowedAudiences); err != nil {
                return err
        }
        if ti.AllowMayAct && slices.Contains(ti.AllowedDelegateClients, anyDelegateClient) {
                return fmt.Errorf(
                        "issuer_url %q: allow_may_act must not be enabled when allowed_delegate_clients contains the wildcard %q",
                        ti.IssuerURL, anyDelegateClient)
        }
        return nil
}

// validateDelegationPolicy validates fields used exclusively by RFC 8693.
func validateDelegationPolicy(ti TrustedIssuer) error {
        hasDelegationSignal := ti.ActorClaim != "" || ti.ActorMatcher != "" || ti.AllowMayAct || len(ti.AllowedActors) > 0
        if ti.ExpectedAudience == "" && !hasDelegationSignal {
                return nil
        }
        if ti.ExpectedAudience == "" && hasDelegationSignal {
                return fmt.Errorf(
                        "issuer_url %q: expected_audience is required when actor_claim, actor_matcher, allow_may_act, or allowed_actors is set",
                        ti.IssuerURL)
        }
        return validateAllowedDelegateClients(ti)
}

func validateJWTBearerGrantPolicy(ti TrustedIssuer, allowedAudiences []string) error {
        if ti.JWTBearerGrant == nil {
                return nil
        }
        policy := ti.JWTBearerGrant
        age, err := time.ParseDuration(policy.MaxAssertionAge)
        if err != nil || age <= 0 {
                return fmt.Errorf("issuer_url %q: jwt_bearer_grant.max_assertion_age must be a positive duration", ti.IssuerURL)
        }
        if err := validateJWTBearerSubjectBindings(ti, policy.SubjectBindings); err != nil {
                return err
        }
        return validateJWTBearerAcceptedAudiences(ti, policy.AcceptedAudiences, allowedAudiences)
}

func validateJWTBearerSubjectBindings(ti TrustedIssuer, bindings []JWTBearerSubjectBinding) error {
        if len(bindings) == 0 {
                return fmt.Errorf("issuer_url %q: jwt_bearer_grant.subject_bindings is required", ti.IssuerURL)
        }
        seenSubjects := make(map[string]struct{}, len(bindings))
        for _, binding := range bindings {
                if binding.Subject == "" || binding.Subject == anyDelegateClient {
                        return fmt.Errorf("issuer_url %q: jwt_bearer_grant.subject_bindings contains an empty or wildcard subject", ti.IssuerURL)
                }
                if _, duplicate := seenSubjects[binding.Subject]; duplicate {
                        return fmt.Errorf(
                                "issuer_url %q: jwt_bearer_grant.subject_bindings contains duplicate subject %q",
                                ti.IssuerURL, binding.Subject)
                }
                seenSubjects[binding.Subject] = struct{}{}
                if len(binding.AllowedResources) == 0 {
                        return fmt.Errorf("issuer_url %q: jwt_bearer_grant subject %q has no allowed_resources", ti.IssuerURL, binding.Subject)
                }
                for _, resource := range binding.AllowedResources {
                        if resource == "" || resource == anyDelegateClient {
                                return fmt.Errorf(
                                        "issuer_url %q: jwt_bearer_grant subject %q has an empty or wildcard resource",
                                        ti.IssuerURL, binding.Subject)
                        }
                        if err := server.ValidateAudienceURI(resource); err != nil {
                                return fmt.Errorf(
                                        "issuer_url %q: jwt_bearer_grant subject %q resource %q is invalid: %w",
                                        ti.IssuerURL, binding.Subject, resource, err)
                        }
                }
        }
        return nil
}

// validateJWTBearerAcceptedAudiences also rejects any accepted_audiences
// entry that is a configured resource audience (allowedAudiences).
// AcceptedAudiences identifies this authorization server itself (see its
// doc comment); an entry shared with a resource audience would let an
// upstream access token minted for that resource — not for this AS — pass
// ValidateJWTBearerAssertion's audience check and mint a ToolHive token for
// whatever resource the subject binding allows, bypassing the
// client-authenticated RFC 8693 delegation route entirely.
func validateJWTBearerAcceptedAudiences(ti TrustedIssuer, acceptedAudiences, allowedAudiences []string) error {
        for _, audience := range acceptedAudiences {
                if audience == "" {
                        return fmt.Errorf("issuer_url %q: jwt_bearer_grant.accepted_audiences must not contain an empty value", ti.IssuerURL)
                }
                if err := server.ValidateAudienceURI(audience); err != nil {
                        return fmt.Errorf(
                                "issuer_url %q: jwt_bearer_grant.accepted_audiences value %q is invalid: %w",
                                ti.IssuerURL, audience, err)
                }
                if slices.Contains(allowedAudiences, audience) {
                        return fmt.Errorf(
                                "issuer_url %q: jwt_bearer_grant.accepted_audiences value %q must not also be a "+
                                        "configured resource audience (allowed_audiences); accepted_audiences identifies "+
                                        "this authorization server, not a resource, or a resource-scoped access token "+
                                        "could be replayed as a JWT-bearer assertion", ti.IssuerURL, audience)
                }
        }
        return nil
}

// validateAllowedDelegateClients checks the structural and narrowing
// invariants for TrustedIssuer.AllowedDelegateClients, split out of
// validateTrustedIssuer to keep that function's cyclomatic complexity down.
func validateAllowedDelegateClients(ti TrustedIssuer) error {
        if len(ti.AllowedDelegateClients) == 0 {
                return fmt.Errorf(
                        "issuer_url %q: allowed_delegate_clients is required; set it to [%q] to permit any ToolHive "+
                                "confidential client holding the token-exchange grant, or list specific client IDs to bind "+
                                "delegation to them — an issuer with no entries can never authorize a delegated exchange",
                        ti.IssuerURL, anyDelegateClient)
        }
        if slices.Contains(ti.AllowedDelegateClients, "") {
                return fmt.Errorf(
                        "issuer_url %q: allowed_delegate_clients must not contain an empty client ID", ti.IssuerURL)
        }
        if slices.Contains(ti.AllowedDelegateClients, anyDelegateClient) && len(ti.AllowedDelegateClients) > 1 {
                return fmt.Errorf(
                        "issuer_url %q: allowed_delegate_clients must not combine the wildcard %q with specific client IDs",
                        ti.IssuerURL, anyDelegateClient)
        }
        if ti.AllowMayAct && slices.Contains(ti.AllowedDelegateClients, anyDelegateClient) {
                return fmt.Errorf(
                        "issuer_url %q: allow_may_act must not be enabled when allowed_delegate_clients contains the wildcard %q",
                        ti.IssuerURL, anyDelegateClient)
        }
        return nil
}

// looksLikeResourceIdentifier reports whether aud is shaped like a resource/API
// identifier (an absolute URI, e.g. "https://api.example.com" or
// "api://<app-id>") rather than a bare opaque string or GUID, which is more
// likely a client ID. This is a heuristic warning signal only — Entra v1
// legitimately uses a bare app-ID GUID as an access token's "aud" (see
// TrustedIssuer.ExpectedAudience's doc comment), so a bare value is not
// rejected, only flagged for the operator to double-check.
func looksLikeResourceIdentifier(aud string) bool {
        return strings.Contains(aud, "://")
}

// resolveActorAuthorization authorizes an external token when either its
// configured actor claim matches AllowedActors or ActorMatcher evaluates true
// against the complete signature-verified JWT claims map. It returns the
// resolved actor, when one exists, separately from authorization because a
// matcher need not depend on an actor claim. Called only when the subject token
// carries no may_act claim — see validateExternalToken.
func resolveActorAuthorization(
        issuerConfig *externalIssuerConfig,
        claims *ValidatedClaims,
        verifiedClaims map[string]any,
) (string, bool, error) {
        actor, actorErr := resolveAllowedActor(issuerConfig, claims)
        if actorErr == nil {
                return actor, true, nil
        }

        if issuerConfig.actorMatcher != nil {
                matched, err := issuerConfig.actorMatcher.EvaluateBool(map[string]any{"claims": verifiedClaims})
                if err != nil {
                        // The raw CEL error is never logged or returned: it can quote
                        // runtime values from the claims map (e.g. timestamp(claims.foo)
                        // on a malformed string embeds that string in the error),
                        // which would otherwise leak signed claim data into logs and,
                        // via HandleTokenEndpointRequest's own error logging, further
                        // downstream.
                        slog.Debug("Trusted issuer actor matcher evaluation failed; denying subject token",
                                "issuer", issuerConfig.IssuerURL)
                        return "", false, fmt.Errorf("actor matcher evaluation failed for issuer %q", issuerConfig.IssuerURL)
                } else if matched {
                        return actor, true, nil
                }
        }

        if len(issuerConfig.AllowedActors) == 0 && issuerConfig.actorMatcher != nil {
                return "", false, fmt.Errorf("subject token from issuer %q was not authorized by actor matcher", issuerConfig.IssuerURL)
        }

        return "", false, actorErr
}

// resolveAllowedActor resolves the issuer's configured actor claim from
// claims and checks it against issuerConfig.AllowedActors, returning the
// matched value on success.
func resolveAllowedActor(issuerConfig *externalIssuerConfig, claims *ValidatedClaims) (string, error) {
        claimName := issuerConfig.ActorClaim
        if claimName == "" {
                claimName = defaultActorClaim
        }

        // "client_id" is routed to ValidatedClaims.ClientID by assignClaim, so it
        // never appears in Extra; without this fallback a plausible operator
        // config (ActorClaim: "client_id") would silently reject all traffic.
        var raw any
        if claimName == "client_id" {
                raw = claims.ClientID
        } else {
                raw = claims.Extra[claimName]
        }

        actor, ok := raw.(string)
        if !ok || actor == "" {
                return "", fmt.Errorf(
                        "subject token from issuer %q is missing or has an invalid %q claim required for delegation consent",
                        issuerConfig.IssuerURL, claimName)
        }

        if !slices.Contains(issuerConfig.AllowedActors, actor) {
                return "", fmt.Errorf(
                        "subject token from issuer %q names actor %q in claim %q, which is not in the allowed actors list",
                        issuerConfig.IssuerURL, actor, claimName)
        }

        return actor, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package tokenexchange implements RFC 8693 token exchange and RFC 7523
// JWT-bearer grants for the authorization server. It validates subject tokens
// issued by the same authorization server and trusted external JWT assertions,
// enabling delegated and assertion-based access-token issuance.
package tokenexchange

import (
        "context"
        "errors"
        "fmt"
        "slices"
        "strings"
        "time"

        "github.com/go-jose/go-jose/v4"
        "github.com/go-jose/go-jose/v4/jwt"
)

// allowedSignatureAlgorithms lists the signing algorithms accepted during JWT verification.
// This restricts which algorithms the validator will accept, preventing algorithm
// confusion attacks where an attacker might try to use a weaker algorithm.
var allowedSignatureAlgorithms = []jose.SignatureAlgorithm{
        jose.ES256,
        jose.ES384,
        jose.RS256,
        jose.RS384,
        jose.RS512,
        jose.EdDSA,
}

// ValidatedClaims holds the verified claims extracted from a subject token.
// All fields are populated from a successfully validated JWT that was issued
// by this authorization server.
type ValidatedClaims struct {
        // Subject is the user identity from the "sub" claim (required for delegation).
        Subject string
        // Issuer is the token issuer from the "iss" claim.
        Issuer string
        // Audience is the list of intended recipients from the "aud" claim.
        Audience []string
        // Expiry is the token expiration time from the "exp" claim.
        Expiry time.Time
        // IssuedAt is the token issuance time from the "iat" claim.
        IssuedAt time.Time
        // JWTID is the unique token identifier from the "jti" claim.
        JWTID string
        // Name is the user's display name from the custom "name" claim.
        Name string
        // Email is the user's email address from the custom "email" claim.
        Email string
        // ClientID is the OAuth client ID from the custom "client_id" claim.
        ClientID string
        // Scopes is the space-delimited scope string assembled from the "scope"
        // or "scp" claim. RFC 9068 §2.2.1 spells it "scope" as a JSON string, but
        // fosite's default JWT claims strategy (token/jwt/claims_jwt.go) writes
        // scopes as a JSON array under "scp" instead, unless ScopeField is
        // explicitly set to String or Both — this server does not set it, so a
        // genuine ToolHive-issued access token used as a subject token carries
        // "scp", not "scope". When both are present, "scope" wins. Empty if the
        // subject token carries neither claim.
        Scopes string
        // MayAct holds the authorized actor from the "may_act" claim (RFC 8693 §4.4).
        // Nil when the subject token does not carry a may_act claim.
        MayAct *MayActClaim
        // ExternalActor is the external actor claim value when the allowlist
        // authorization path matched. It is empty when ActorMatcher alone
        // authorized the token, because a matcher need not identify an actor.
        // It is never populated from token claims by buildValidatedClaims or
        // assignClaim, and is empty for self-issued and may_act-bearing tokens.
        ExternalActor string
        // ExternalActorAuthorized reports that the external issuer authorized
        // delegation through AllowedActors or ActorMatcher. It is separate from
        // ExternalActor because ActorMatcher can authorize without an actor claim.
        // It is false for self-issued and may_act-bearing tokens.
        ExternalActorAuthorized bool
        // ExternalIssuer is set by the external-issuer validation path
        // (validateExternalToken in multi_issuer_validator.go) to that issuer's
        // IssuerURL, for EVERY external token it validates — unlike
        // ExternalActor, this is set regardless of whether the token carries a
        // may_act claim. It exists so the handler can record provenance (which
        // issuer a delegation actually originated from) even for a
        // may_act-bearing external token, which leaves ExternalActor unset.
        // Empty for self-issued tokens. Like ExternalActor, it is never
        // populated from token claims by buildValidatedClaims or assignClaim —
        // it comes from the already-validated issuer config the token matched,
        // not from anything token-supplied, so it cannot be spoofed via claims.
        ExternalIssuer string
        // AllowedDelegateClients is set for EVERY external token — like
        // ExternalIssuer and unlike ExternalActor, it does not depend on whether
        // the token carries a may_act claim (see validateExternalToken). That
        // matters: the may_act path bypasses AllowedActors and ActorMatcher
        // entirely, so it is the path that most needs this restriction to still
        // apply. It is set to that issuer's configured
        // TrustedIssuer.AllowedDelegateClients, and is never populated
        // from token claims by buildValidatedClaims or assignClaim, and the
        // validator never compares it against anything: the validator does not
        // know the authenticated ToolHive client, so checkDelegationConsent
        // (handler.go) is the one that checks actorID against this list. Nil
        // means the issuer did not configure AllowedDelegateClients, which is
        // permissive — any ToolHive client may use the allowlisted external
        // actor, same as before this field existed.
        AllowedDelegateClients []string
        // Extra contains all non-standard claims not captured by other fields.
        Extra map[string]any
        // Claims holds every verified claim as decoded, registered ones included.
        // It is shared with the validator's decode result: treat it as read-only
        // and deep-copy before handing it to code that may modify it.
        Claims map[string]any
}

// MayActClaim represents the RFC 8693 §4.4 may_act claim from a subject token.
// It identifies the party authorized to act on behalf of the subject.
type MayActClaim struct {
        Sub string `json:"sub"`
        // Iss, when present, qualifies Sub's namespace. RFC 8693 §4.4: "the
        // combination of the two claims 'iss' and 'sub' are sometimes necessary
        // to uniquely identify an authorized actor." Without it, Sub alone could
        // name a party in some other issuer's namespace while still being
        // compared against a ToolHive client ID by checkDelegationConsent — a
        // namespace-confusion bug, not just a missing feature. validateMayActShape
        // requires Iss, when present, to equal this authorization server's own
        // issuer, so by the time checkDelegationConsent reads Sub it is
        // guaranteed to be in ToolHive's own client namespace.
        Iss string `json:"iss,omitempty"`
}

// Compile-time check that SelfIssuedTokenValidator implements SubjectTokenValidator.
var _ SubjectTokenValidator = (*SelfIssuedTokenValidator)(nil)

// SelfIssuedTokenValidator validates subject tokens presented during RFC 8693 token exchange.
// It verifies that the token was issued by this authorization server by checking
// the signature against the server's own JWKS, and validates standard JWT claims.
type SelfIssuedTokenValidator struct {
        publicJWKS       *jose.JSONWebKeySet
        issuer           string
        allowedAudiences []string
}

// NewSelfIssuedTokenValidator creates a new validator for subject tokens.
// The jwks parameter must be non-nil and contain only the authorization server's
// public signing keys (e.g. AuthorizationServerConfig.PublicJWKS) — the validator
// only ever verifies signatures, so it must not be handed private key material.
// The issuer parameter is the expected "iss" claim value. allowedAudiences is the
// set of audiences this server accepts in a subject token's "aud" claim; per the
// same secure default as AuthorizationServerConfig.AllowedAudiences, an empty
// allowedAudiences rejects every subject token rather than skipping the check.
func NewSelfIssuedTokenValidator(
        jwks *jose.JSONWebKeySet, issuer string, allowedAudiences []string,
) (*SelfIssuedTokenValidator, error) {
        if jwks == nil {
                return nil, fmt.Errorf("JWKS must not be nil")
        }
        if issuer == "" {
                return nil, fmt.Errorf("issuer must not be empty")
        }
        return &SelfIssuedTokenValidator{
                publicJWKS:       jwks,
                issuer:           issuer,
                allowedAudiences: allowedAudiences,
        }, nil
}

// Validate parses and verifies a raw JWT subject token.
// It checks the signature against the server's JWKS, validates issuer and audience,
// ensures the token is not expired, and requires a subject claim for delegation.
//
// The subject token's "aud" claim is checked against allowedAudiences, but this
// validator deliberately does not require that the authorization server itself
// (i.e. the token endpoint) be among those audiences. RFC 8693 leaves subject-token
// validation criteria out of scope, and ToolHive's vMCP flow legitimately exchanges
// tokens addressed to a downstream/upstream resource rather than the AS. The residual
// cross-resource risk is mitigated elsewhere: the token is still pinned to the
// server-wide allowedAudiences, and the handler enforces the requested resource
// against the client's registered audiences.
//
// Returns the validated claims on success, or a descriptive error on failure.
func (v *SelfIssuedTokenValidator) Validate(_ context.Context, rawToken string) (*ValidatedClaims, error) {
        parsedToken, err := jwt.ParseSigned(rawToken, allowedSignatureAlgorithms)
        if err != nil {
                return nil, fmt.Errorf("subject token is not a valid JWT: %w", err)
        }

        // kidMatched is discarded here: it exists to let the external-issuer path
        // (validateExternalToken in multi_issuer_validator.go) decide whether an
        // unknown kid is worth an on-demand JWKS refresh. This validator's JWKS
        // never comes over HTTP — it's this authorization server's own
        // PublicJWKS() — so there is nothing to refresh.
        standardClaims, extraClaims, _, err := verifySignature(parsedToken, v.publicJWKS)
        if err != nil {
                return nil, err
        }

        // Validate issuer and expiry. Audience is intentionally excluded from
        // jwt.Expected here — go-jose's AnyAudience check is skipped entirely when
        // empty, which is the opposite of this server's secure default (empty
        // AllowedAudiences means no audience is permitted). Audience is checked
        // explicitly below so an empty allowlist fails closed.
        expected := jwt.Expected{
                Issuer: v.issuer,
        }
        if err := standardClaims.ValidateWithLeeway(expected, 0); err != nil {
                return nil, fmt.Errorf("subject token claims validation failed: %w", err)
        }

        if err := validateAudience(standardClaims.Audience, v.allowedAudiences); err != nil {
                return nil, fmt.Errorf("subject token claims validation failed: %w", err)
        }

        // Expiry is required for delegation — a subject token without an expiry
        // cannot be safely bounded by the delegation lifetime.
        if standardClaims.Expiry == nil {
                return nil, fmt.Errorf("subject token is missing required 'exp' claim")
        }

        // Subject is required for delegation — the token must identify the user.
        if standardClaims.Subject == "" {
                return nil, fmt.Errorf("subject token is missing required 'sub' claim")
        }

        // Reject a subject token that is actually an ID token — see
        // rejectIDTokenClaims's doc comment.
        if err := rejectIDTokenClaims(extraClaims); err != nil {
                return nil, err
        }

        // If may_act is present, it must be well-formed. A present-but-malformed
        // may_act is treated as an invalid token, not an absent one — fail closed
        // to prevent silent downgrade to client_id. requireIss is false here: see
        // validateMayActShape's doc comment for why the self-issued path doesn't
        // need it.
        if err := validateMayActShape(extraClaims, v.issuer, false); err != nil {
                return nil, err
        }

        return buildValidatedClaims(standardClaims, extraClaims), nil
}

// verifySignature attempts to verify the JWT signature using the provided public keys.
// If the JWT header names a key ID (kid) that matches a key in the JWKS, only that
// key is tried: a token naming a kid but signed by a different key is a spoofed-kid
// attempt and must fail, not fall back to verifying against the rest of the keyset.
// When there is no kid, or the kid doesn't match any key in the JWKS, every key whose
// declared "use"/"alg" (RFC 7517 §4.2/§4.4) is compatible with this signature is
// tried in turn — a key marked for a different use (e.g. "enc") or a different
// algorithm is skipped rather than tried, per RFC 8725 §3.12's recommendation to
// scope keys to a single purpose as a substitution defense. Returns on the first
// successful verification. On failure, the last verification error is wrapped.
//
// The returned kidMatched is true only when the token's kid header was found among
// publicJWKS's keys — regardless of whether the signature itself then verified. The
// external-issuer path (multi_issuer_validator.go) uses this to tell "this key
// genuinely isn't in our cached JWKS yet, possibly a rotation" (kidMatched false)
// apart from "this kid exists here but the signature is wrong" (kidMatched true,
// which a refresh can never fix and must not retry).
func verifySignature(
        token *jwt.JSONWebToken,
        publicJWKS *jose.JSONWebKeySet,
) (jwt.Claims, map[string]any, bool, error) {
        candidates := publicJWKS.Keys
        var kidMatched bool
        if len(token.Headers) > 0 && token.Headers[0].KeyID != "" {
                if matched := publicJWKS.Key(token.Headers[0].KeyID); len(matched) > 0 {
                        candidates = matched
                        kidMatched = true
                }
        }

        var headerAlg string
        if len(token.Headers) > 0 {
                headerAlg = token.Headers[0].Algorithm
        }

        var lastErr error
        var filtered int
        for _, key := range candidates {
                if key.Use != "" && key.Use != "sig" {
                        filtered++
                        continue
                }
                if key.Algorithm != "" && key.Algorithm != headerAlg {
                        filtered++
                        continue
                }
                var claims jwt.Claims
                extra := map[string]any{}
                err := token.Claims(key, &claims, &extra)
                if err == nil {
                        return claims, extra, kidMatched, nil
                }
                lastErr = err
        }
        if lastErr != nil {
                return jwt.Claims{}, nil, kidMatched, fmt.Errorf("subject token signature verification failed: %w", lastErr)
        }
        if filtered > 0 {
                return jwt.Claims{}, nil, kidMatched, fmt.Errorf(
                        "subject token signature verification failed: no compatible keys in JWKS "+
                                "(%d key(s) present but all skipped by use/alg filter)",
                        len(candidates))
        }
        return jwt.Claims{}, nil, kidMatched, fmt.Errorf("subject token signature verification failed: no keys in JWKS")
}

// validateAudience checks that tokenAudience intersects allowedAudiences.
// Per this server's secure default (see AuthorizationServerConfig.AllowedAudiences),
// an empty allowedAudiences means no audience is permitted, so the check fails
// closed rather than skipping validation.
func validateAudience(tokenAudience jwt.Audience, allowedAudiences []string) error {
        if len(allowedAudiences) == 0 {
                return fmt.Errorf("no audiences are configured on this server")
        }
        for _, aud := range tokenAudience {
                if slices.Contains(allowedAudiences, aud) {
                        return nil
                }
        }
        return fmt.Errorf("token audience %v does not match any allowed audience", []string(tokenAudience))
}

// buildValidatedClaims constructs a ValidatedClaims from standard and extra claims.
func buildValidatedClaims(
        standard jwt.Claims,
        extra map[string]any,
) *ValidatedClaims {
        vc := &ValidatedClaims{
                Subject:  standard.Subject,
                Issuer:   standard.Issuer,
                Audience: []string(standard.Audience),
                JWTID:    standard.ID,
                Extra:    make(map[string]any),
                Claims:   extra,
        }

        if standard.Expiry != nil {
                vc.Expiry = standard.Expiry.Time()
        }
        if standard.IssuedAt != nil {
                vc.IssuedAt = standard.IssuedAt.Time()
        }

        // Extract well-known custom claims and collect the rest into Extra.
        // Standard JWT registered claims are filtered out since they are already
        // captured in the structured fields above.
        for k, val := range extra {
                assignClaim(vc, k, val)
        }

        // Fall back to "scp" only if "scope" (handled above) didn't already set
        // Scopes. This must run after the loop, not inside assignClaim's per-key
        // switch, because map iteration order is random — assignClaim can't tell
        // whether a not-yet-seen "scope" claim will still show up.
        //
        // Without this fallback, a genuine self-issued subject token (which
        // carries "scp", not "scope") reads as scope-less and any scoped
        // exchange fails with invalid_scope — scoped delegation against a
        // self-issued token was effectively unusable before this fallback.
        if vc.Scopes == "" {
                if scp, ok := extra["scp"]; ok {
                        vc.Scopes = scpToScopeString(scp)
                }
        }

        return vc
}

// scpToScopeString normalizes an "scp" claim value into a space-delimited
// scope string. fosite's default JWT claims strategy writes "scp" as a JSON
// array (see the Scopes field doc comment), which json.Unmarshal decodes as
// []any — but a plain space-delimited string is accepted too, in case an
// issuer writes it that way instead. Non-string array elements are skipped
// rather than rejected: a malformed entry degrades to a smaller scope set
// instead of failing subject-token validation outright.
func scpToScopeString(val any) string {
        switch v := val.(type) {
        case string:
                return v
        case []any:
                scopes := make([]string, 0, len(v))
                for _, item := range v {
                        if s, ok := item.(string); ok {
                                scopes = append(scopes, s)
                        }
                }
                return strings.Join(scopes, " ")
        default:
                return ""
        }
}

// assignClaim routes one non-standard JWT claim onto its structured
// ValidatedClaims field, or into Extra if it isn't a well-known claim.
// Registered JWT claims (sub, iss, aud, exp, iat, nbf, jti) are dropped —
// they're already captured in buildValidatedClaims's structured fields.
// "scp" is likewise dropped here, though its value is read separately (see
// buildValidatedClaims's post-loop fallback) since its precedence relative
// to "scope" can't be decided from a single, order-independent key/value
// pair. Keep actorClaimsNotInExtra (multi_issuer_validator.go) in sync with
// the cases here.
func assignClaim(vc *ValidatedClaims, key string, val any) {
        switch key {
        case "name":
                if s, ok := val.(string); ok {
                        vc.Name = s
                }
        case "email":
                if s, ok := val.(string); ok {
                        vc.Email = s
                }
        case "client_id":
                if s, ok := val.(string); ok {
                        vc.ClientID = s
                }
        case "scope":
                if s, ok := val.(string); ok {
                        vc.Scopes = s
                }
        case "scp":
                // Handled by buildValidatedClaims after its call to assignClaim for
                // every key completes, so "scope" (if present) wins regardless of
                // map iteration order. Case exists here only to keep "scp" out of
                // Extra, like every other well-known claim below.
        case "may_act":
                if m, ok := val.(map[string]any); ok {
                        if s, ok := m["sub"].(string); ok {
                                mayAct := &MayActClaim{Sub: s}
                                if iss, ok := m["iss"].(string); ok {
                                        mayAct.Iss = iss
                                }
                                vc.MayAct = mayAct
                        }
                }
        case "sub", "iss", "aud", "exp", "iat", "nbf", "jti":
                // Skip registered JWT claims — already in structured fields.
        default:
                vc.Extra[key] = val
        }
}

// validateMayActShape rejects a present-but-malformed may_act claim. Both
// validation paths (self-issued here, external in multi_issuer_validator.go)
// call this before buildValidatedClaims: assignClaim silently drops a
// malformed may_act rather than surfacing it, and each path's fallback
// consent signal when may_act is absent (client_id here, the actor allowlist
// externally) is weaker than what a well-formed may_act would have granted.
// Letting a malformed claim fall through would silently widen consent
// instead of rejecting the token.
//
// selfIssuer is this authorization server's own issuer identifier. When
// may_act carries an "iss" member, RFC 8693 §4.4 uses it together with "sub"
// to identify the actor's namespace; this server only ever grants delegation
// to its own clients (checkDelegationConsent compares may_act.sub against a
// ToolHive client ID), so an "iss" naming any other issuer would mean sub is
// being read out of the wrong namespace. Requiring iss == selfIssuer when
// present — same fail-closed treatment as a malformed sub — prevents that
// namespace confusion.
//
// requireIss makes "iss" mandatory rather than merely constrained when
// present. The self-issued path passes false: the subject token's own "iss"
// already equals selfIssuer (Validate routed it here on exactly that basis),
// so may_act.iss would be redundant to require. The external path
// (multi_issuer_validator.go) passes true: there, an external issuer that
// emits no "iss" on may_act could otherwise authorize ANY ToolHive client by
// naming its ID in a bare may_act.sub, entirely bypassing this issuer's
// AllowedActors/AllowedDelegateClients allowlists — the one consent path
// that does. No major IdP is known to emit may_act at all (Microsoft Entra,
// Okta, and Google do not; Keycloak emits "act", a different claim), so
// reaching this path already requires an operator configuring a custom
// claim mapper for a nonstandard claim — an operator capable of adding "sub"
// to that mapper is assumed capable of adding "iss" too. That assumption is
// not verified against any real IdP's behavior in this repo; it is the
// rationale for requireIss, not a tested fact.
func validateMayActShape(extra map[string]any, selfIssuer string, requireIss bool) error {
        raw, ok := extra["may_act"]
        if !ok || raw == nil {
                return nil
        }
        m, ok := raw.(map[string]any)
        if !ok {
                return errors.New("subject token has malformed 'may_act' claim: expected a JSON object")
        }
        if sub, ok := m["sub"].(string); !ok || sub == "" {
                return errors.New("subject token has malformed 'may_act' claim: missing or invalid 'sub'")
        }
        rawIss, present := m["iss"]
        if !present {
                if requireIss {
                        return errors.New("subject token has malformed 'may_act' claim: missing required 'iss'")
                }
                return nil
        }
        iss, ok := rawIss.(string)
        if !ok || iss == "" {
                return errors.New("subject token has malformed 'may_act' claim: invalid 'iss'")
        }
        if iss != selfIssuer {
                return fmt.Errorf(
                        "subject token has malformed 'may_act' claim: 'iss' %q does not match this authorization server's issuer",
                        iss)
        }
        return nil
}

// rejectIDTokenClaims rejects a subject token that carries "at_hash" or
// "c_hash" — OIDC Core §3.3.2.11 / §3.3.2.10 define both for ID tokens only,
// never present on an access token. This is a sound POSITIVE detector (their
// presence proves an ID token) but is NOT EXHAUSTIVE: both claims are
// OPTIONAL even on a valid authorization-code-flow ID token, so an ID token
// that omits them passes this check undetected. A "typ: at+jwt" header check
// is not viable as a replacement: Entra v1/v2, Okta, and Google all emit bare
// "JWT" for access tokens too. Deliberately does NOT check "nonce": Microsoft
// Graph access tokens carry one, so rejecting on its presence would reject
// legitimate Entra access tokens.
//
// The other, complementary layer is the audience-shape check on the external
// path: TrustedIssuer.ExpectedAudience (multi_issuer_validator.go) is
// documented as MUST be a resource identifier, not a client ID, since an ID
// token's "aud" is the requesting client's own ID. That check is
// operator-dependent, not self-enforcing — see ExpectedAudience's doc
// comment and looksLikeResourceIdentifier's startup warning. Between the two,
// ID/access-token discrimination for an external issuer is partial: an
// operator who mis-sets ExpectedAudience to a client ID, presented with a
// real IdP whose ID tokens omit at_hash/c_hash (both legal per OIDC Core),
// has no remaining defense here. The blast radius is bounded elsewhere: an
// admitted ID token carries no "scope"/"scp" claim, so grantScopes
// (handler.go) either rejects the exchange with invalid_scope or grants a
// zero-scope delegated token.
//
// A third discriminator was considered and rejected: an ID token's "aud"
// MUST equal its own "azp"/authorized-party claim (OIDC Core §2, §3.1.3.7
// step 9), so rejecting a token whose "aud" equals the issuer's configured
// actor claim looked like a candidate. It is not shipped, because it also
// fires on a legitimate access token: an app whose API and client share one
// registration — e.g. an Entra v1 app that calls its own exposed API, or any
// setup where ExpectedAudience is misconfigured to the app's bare
// client/resource GUID, which is exactly the ambiguous case this file's
// audience-shape check already warns about — issues access tokens whose
// "aud" is that shared ID and whose azp/appid is the very same ID. Rejecting
// on aud == azp would then reject that provider's genuine access tokens, not
// just its ID tokens. A false-positive risk of that shape is worse than the
// gap it would close.
func rejectIDTokenClaims(extra map[string]any) error {
        if _, ok := extra["at_hash"]; ok {
                return errors.New("subject token carries an 'at_hash' claim: ID tokens are not accepted as subject tokens")
        }
        if _, ok := extra["c_hash"]; ok {
                return errors.New("subject token carries a 'c_hash' claim: ID tokens are not accepted as subject tokens")
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "context"
        "crypto/x509"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "reflect"
        "time"

        josev3 "github.com/go-jose/go-jose/v3"
        "github.com/ory/fosite"
        "github.com/ory/fosite/compose"
        "github.com/spiffe/go-spiffe/v2/bundle/jwtbundle"
        "github.com/spiffe/go-spiffe/v2/bundle/x509bundle"
        "golang.org/x/time/rate"

        "github.com/stacklok/toolhive/pkg/authserver/internal/consentpolicy"
        oauthserver "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/deviceflow"
        "github.com/stacklok/toolhive/pkg/authserver/server/handlers"
        "github.com/stacklok/toolhive/pkg/authserver/server/keys"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange"
        spiffeauth "github.com/stacklok/toolhive/pkg/authserver/spiffe"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/authserver/upstream"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// server is the internal implementation of the Server interface.
type server struct {
        handler http.Handler
        // oauthHandler is the handler behind handler, kept so Close can release
        // the metric callback it registered.
        oauthHandler *handlers.Handler
        // storage is the active storage backend, potentially wrapped by decorators
        // such as CIMDStorageDecorator. Code that needs the concrete type must walk
        // the Unwrap() chain rather than asserting directly.
        storage storage.Storage
        // dcrStore is the same storage.Storage value asserted to
        // storage.DCRCredentialStore. The assertion runs once at construction
        // (newServer) so DCRStore() is a field read rather than re-asserting on
        // every call, and a backend that does not implement DCRCredentialStore
        // is rejected at boot rather than at first DCR resolve.
        dcrStore storage.DCRCredentialStore
        // upstreamRefresher is the single shared refresher instance constructed in
        // newServer. Storing the interface type (not *upstreamTokenRefresher) keeps
        // the nil-return contract: newUpstreamTokenRefresher returns a true nil
        // interface when there are no upstreams, so callers can check == nil safely.
        upstreamRefresher storage.UpstreamTokenRefresher
        upstreams         []handlers.NamedUpstream
        // trustedIssuerValidator is the single MultiIssuerTokenValidator shared by
        // the token-exchange and JWT-bearer handlers, built in buildProvider when
        // TrustedIssuers are configured (nil otherwise). Held here so Close can shut
        // down its per-issuer JWKS refresh worker pools; nothing else releases them.
        trustedIssuerValidator *tokenexchange.MultiIssuerTokenValidator
        // spiffeBundleSource supplies both X.509 and JWT bundles and owns the
        // refresh workers for the server lifetime.
        spiffeBundleSource *spiffeMultiDomainBundleSource
}

// DefaultUpstreamFactory creates the production upstream provider based on type.
// For OIDC providers, it creates an OIDCProviderImpl with discovery and ID token
// validation. For OAuth2 providers, it creates a BaseOAuth2Provider.
//
// It is exported so a Config.UpstreamFactory implementation can delegate to the
// built-in behavior for the upstreams it does not want to handle itself.
func DefaultUpstreamFactory(ctx context.Context, cfg *UpstreamConfig) (upstream.OAuth2Provider, error) {
        switch cfg.Type {
        case UpstreamProviderTypeOIDC:
                return upstream.NewOIDCProvider(ctx, cfg.OIDCConfig)
        case UpstreamProviderTypeOAuth2:
                return upstream.NewOAuth2Provider(cfg.OAuth2Config)
        default:
                return nil, fmt.Errorf("unsupported upstream type: %s", cfg.Type)
        }
}

// buildUpstreams constructs the ordered upstream provider list using
// cfg.UpstreamFactory, or DefaultUpstreamFactory when the caller supplied none.
// A provider that fails to construct aborts the whole list — but the providers
// already built have live connection pools, so they are drained before the
// error is returned.
func buildUpstreams(ctx context.Context, cfg Config) ([]handlers.NamedUpstream, error) {
        factory := cfg.UpstreamFactory
        if factory == nil {
                factory = DefaultUpstreamFactory
        }
        upstreams := make([]handlers.NamedUpstream, 0, len(cfg.Upstreams))
        for i := range cfg.Upstreams {
                upCfg := &cfg.Upstreams[i]
                slog.Debug("creating upstream IDP provider", "type", upCfg.Type, "name", upCfg.Name)
                provider, err := factory(ctx, upCfg)
                if err != nil {
                        closeUpstreamIdleConnections(upstreams)
                        return nil, fmt.Errorf("failed to create upstream provider %q: %w", upCfg.Name, err)
                }
                // A nil provider would reach the authorization chain and panic on the
                // first /oauth/authorize rather than failing at boot, so a custom
                // factory cannot use one to drop an upstream.
                if isNilProvider(provider) {
                        closeUpstreamIdleConnections(upstreams)
                        return nil, fmt.Errorf("upstream factory returned a nil provider for upstream %q", upCfg.Name)
                }
                upstreams = append(upstreams, handlers.NamedUpstream{
                        Name: upCfg.Name, DisplayName: upCfg.DisplayName, Provider: provider,
                })
                slog.Debug("upstream IDP provider configured", "type", upCfg.Type, "name", upCfg.Name)
        }
        return upstreams, nil
}

// isNilProvider reports whether an OAuth2Provider is unusable — either a nil
// interface value or an interface holding a nil pointer.
func isNilProvider(provider upstream.OAuth2Provider) bool {
        if provider == nil {
                return true
        }
        v := reflect.ValueOf(provider)
        return v.Kind() == reflect.Pointer && v.IsNil()
}

// releaseOnConstructionError runs newServer's error-path cleanup: it drains the
// upstream idle connections and shuts down the trusted-issuer validator's JWKS
// worker pools (when one was built). Errors are logged, not returned — retErr is
// what the caller acts on, but a pool that fails to drain here would otherwise
// leave its goroutines running with no diagnostic.
func releaseOnConstructionError(upstreams []handlers.NamedUpstream, validator *tokenexchange.MultiIssuerTokenValidator) {
        closeUpstreamIdleConnections(upstreams)
        if validator != nil {
                if err := validator.Close(); err != nil {
                        slog.Warn("failed to shut down trusted-issuer validator during server construction cleanup", "error", err)
                }
        }
}

// closeUpstreamIdleConnections drains the pooled idle connections of every
// upstream that implements the optional upstream.IdleConnectionCloser
// capability; see that interface for why it is optional and what it exempts.
func closeUpstreamIdleConnections(upstreams []handlers.NamedUpstream) {
        for _, u := range upstreams {
                if closer, ok := u.Provider.(upstream.IdleConnectionCloser); ok {
                        slog.Debug("closing upstream idle connections", "name", u.Name)
                        closer.CloseIdleConnections()
                }
        }
}

func closeSPIFFEBundleSourceOnConstructionError(source *spiffeMultiDomainBundleSource, retErr *error) {
        if *retErr == nil || source == nil {
                return
        }
        if err := source.Close(); err != nil {
                slog.Warn("failed to shut down SPIFFE bundle source during server construction cleanup", "error", err)
        }
}

func spiffeX509BundleSource(source *spiffeMultiDomainBundleSource) x509bundle.Source {
        if source == nil {
                return nil
        }
        return source
}

func spiffeJWTBundleSource(source *spiffeMultiDomainBundleSource) jwtbundle.Source {
        if source == nil {
                return nil
        }
        return source
}

// loadKeyMaterial returns the signing key and the public keys for the JWKS.
// Every key the provider reports (signing key plus any fallbacks configured
// for rotation) must reach the published JWKS, or promoting a fallback key to
// primary invalidates every outstanding token instead of opening the
// documented rotation overlap window (#6451).
func loadKeyMaterial(ctx context.Context, provider keys.KeyProvider) (*keys.SigningKeyData, []*keys.PublicKeyData, error) {
        signingKey, err := provider.SigningKey(ctx)
        if err != nil {
                return nil, nil, fmt.Errorf("failed to get signing key: %w", err)
        }
        publicKeys, err := provider.PublicKeys(ctx)
        if err != nil {
                return nil, nil, fmt.Errorf("failed to get public keys: %w", err)
        }
        return signingKey, publicKeys, nil
}

// authorizationServerParams maps cfg and the loaded key material onto the
// fosite provider parameters. newServer and the test provider builder both use
// it so handler selection cannot drift between them. spiffeRegistry and
// bundleSource may be nil.
func authorizationServerParams(
        cfg Config, signingKey *keys.SigningKeyData, additionalPublicKeys []*keys.PublicKeyData,
        spiffeRegistry *SPIFFEAssociationRegistry, bundleSource *spiffeMultiDomainBundleSource, stor storage.Storage,
) *oauthserver.AuthorizationServerParams {
        return &oauthserver.AuthorizationServerParams{
                Issuer:                                    cfg.Issuer,
                AccessTokenLifespan:                       cfg.AccessTokenLifespan,
                RefreshTokenLifespan:                      cfg.RefreshTokenLifespan,
                AuthCodeLifespan:                          cfg.AuthCodeLifespan,
                HMACSecrets:                               cfg.HMACSecrets,
                SigningKeyID:                              signingKey.KeyID,
                SigningKeyAlgorithm:                       signingKey.Algorithm,
                SigningKey:                                signingKey.Key,
                AdditionalPublicKeys:                      additionalPublicKeys,
                ScopesSupported:                           cfg.ScopesSupported,
                BaselineClientScopes:                      cfg.BaselineClientScopes,
                AllowedAudiences:                          cfg.AllowedAudiences,
                SubjectIdentifierAudiences:                cfg.SubjectIdentifierAudiences,
                AccessTokenType:                           cfg.AccessTokenType,
                AuthorizationEndpointBaseURL:              cfg.AuthorizationEndpointBaseURL,
                CIMDEnabled:                               cfg.CIMDEnabled,
                AllowConfidentialClientRegistration:       cfg.AllowConfidentialClientRegistration,
                AllowPrivateKeyJWTRegistration:            cfg.AllowPrivateKeyJWTRegistration,
                HasStaticDelegateClients:                  len(cfg.DelegateClients) > 0,
                InsecureAllowHTTP:                         cfg.InsecureAllowHTTP,
                InsecureAllowConfidentialOverLoopbackHTTP: cfg.InsecureAllowConfidentialOverLoopbackHTTP,
                ForceConfidentialRedirectURIs:             cfg.ForceConfidentialRedirectURIs,
                DisableTokenExchange:                      cfg.DisableTokenExchange,
                JWTBearerGrantEnabled:                     JWTBearerGrantEnabled(cfg.TrustedIssuers),
                DeviceFlowEnabled:                         cfg.DeviceFlowEnabled,
                DeviceCodeInterval:                        cfg.DeviceCodeInterval,
                SPIFFEClientResolver:                      newSPIFFEClientResolver(spiffeRegistry, stor),
                SPIFFEX509BundleSource:                    spiffeX509BundleSource(bundleSource),
                SPIFFEJWTBundleSource:                     spiffeJWTBundleSource(bundleSource),
        }
}

// newServer creates a new OAuth authorization server.
// The opts parameter allows injecting dependencies for testing.
func newServer(ctx context.Context, cfg Config, stor storage.Storage) (_ *server, retErr error) {
        slog.Debug("initializing OAuth authorization server")

        // Apply defaults to config
        if err := cfg.applyDefaults(); err != nil {
                return nil, fmt.Errorf("failed to apply config defaults: %w", err)
        }

        // Validate config
        if err := cfg.Validate(); err != nil {
                return nil, fmt.Errorf("invalid config: %w", err)
        }

        logConfidentialClientStartup(cfg.AllowConfidentialClientRegistration)

        stor, dcrStore, bundleSource, spiffeRegistry, err := prepareServerStorageAndRegisterClients(ctx, cfg, stor)
        if err != nil {
                return nil, err
        }
        defer closeSPIFFEBundleSourceOnConstructionError(bundleSource, &retErr)

        slog.Debug("creating OAuth2 configuration")

        signingKey, additionalPublicKeys, err := loadKeyMaterial(ctx, cfg.KeyProvider)
        if err != nil {
                return nil, err
        }

        oauthParams := authorizationServerParams(cfg, signingKey, additionalPublicKeys, spiffeRegistry, bundleSource, stor)
        authServerConfig, err := oauthserver.NewAuthorizationServerConfig(oauthParams)
        if err != nil {
                return nil, fmt.Errorf("failed to create OAuth2 config: %w", err)
        }

        slog.Debug("oauth2 configuration created",
                "access_token_lifespan", cfg.AccessTokenLifespan,
                "refresh_token_lifespan", cfg.RefreshTokenLifespan,
                "auth_code_lifespan", cfg.AuthCodeLifespan,
        )

        // Build ordered upstream provider list from all configured upstreams.
        upstreams, err := buildUpstreams(ctx, cfg)
        if err != nil {
                return nil, err
        }
        // trustedIssuerValidator is assigned from buildProvider below; the deferred
        // cleanup captures it by reference so a failure after it is built releases
        // its per-issuer JWKS worker pools rather than leaking them.
        var trustedIssuerValidator *tokenexchange.MultiIssuerTokenValidator
        // Defense in depth: the failure returns below would otherwise abandon
        // upstream providers holding a live pool, or the trusted-issuer validator's
        // JWKS worker pools. The upstream drain is not reachable today (Validate
        // covers every precondition they check) — it is here so a future step that
        // fails after touching the network, or a new error return added above it, is
        // covered by default; the validator shutdown covers the reachable failure
        // between buildProvider and the return below.
        defer func() {
                if retErr != nil {
                        releaseOnConstructionError(upstreams, trustedIssuerValidator)
                }
        }()

        // Run one-shot bulk migration of legacy data before handler construction.
        // TODO(migration): Remove once all deployments have upgraded past this version.
        if err := runLegacyMigration(ctx, stor, cfg.Upstreams); err != nil {
                return nil, err
        }

        // Give the handler a refresher so the authorization chain can transparently
        // refresh an expired upstream leg during login instead of skipping it and
        // failing later at MCP-request token-swap time. The same instance is stored
        // on the server so UpstreamTokenRefresher() returns the identical object,
        // and handed to the upstream refresh guard, ensuring every path shares one
        // process-local singleflight.Group keyed by opaque storage-row identity.
        refresher := newUpstreamTokenRefresher(upstreams, stor, cfg.RefreshTokenLifespan)

        // Create fosite provider with the configured storage decorators.
        slog.Debug("creating fosite OAuth2 provider")
        fositeProvider, trustedIssuerValidator, err := buildProvider(cfg, authServerConfig, stor, refresher, upstreams)
        if err != nil {
                return nil, fmt.Errorf("failed to create fosite OAuth2 provider: %w", err)
        }

        handlerInstance, err := handlers.NewHandler(fositeProvider, authServerConfig, stor, upstreams,
                buildHandlerOptions(refresher, cfg.UpstreamFilter, cfg.RegistrationRateLimit, cfg.ConsentExemptClientIDs)...)
        if err != nil {
                return nil, fmt.Errorf("failed to create handler: %w", err)
        }

        // Create HTTP handler serving all endpoints
        router := handlerInstance.Routes()

        slog.Debug("oauth authorization server initialized",
                "issuer", cfg.Issuer,
                "consent_exemption_enabled", len(cfg.ConsentExemptClientIDs) > 0,
                "consent_exempt_client_count", len(cfg.ConsentExemptClientIDs),
        )

        return &server{
                handler:                router,
                oauthHandler:           handlerInstance,
                storage:                stor,
                dcrStore:               dcrStore,
                upstreams:              upstreams,
                upstreamRefresher:      refresher,
                trustedIssuerValidator: trustedIssuerValidator,
                spiffeBundleSource:     bundleSource,
        }, nil
}

// prepareServerStorageAndRegisterClients decorates storage, validates configured
// consent-exempt clients against the prepared registration view, and registers
// delegate clients before returning the storage handles used by server construction.
func prepareServerStorageAndRegisterClients(
        ctx context.Context, cfg Config, stor storage.Storage,
) (prepared storage.Storage, dcrStore storage.DCRCredentialStore,
        bundleSource *spiffeMultiDomainBundleSource, spiffeRegistry *SPIFFEAssociationRegistry, retErr error,
) {
        if stor == nil {
                return nil, nil, nil, nil, fmt.Errorf("storage is required")
        }

        // Storage no longer embeds DCRCredentialStore (the embed widened secret
        // reach to every Storage consumer); obtain the DCR-capable handle via an
        // explicit assertion at the boundary. The per-backend
        // `var _ DCRCredentialStore = (*MemoryStorage)(nil)` /
        // `var _ DCRCredentialStore = (*RedisStorage)(nil)` checks make this
        // provably safe for the production backends; surfacing a bad backend as
        // a constructor error keeps misconfiguration fail-loud at boot rather
        // than at first DCR resolve.
        baseStore := storage.Unwrap(stor)
        var ok bool
        dcrStore, ok = baseStore.(storage.DCRCredentialStore)
        if !ok {
                return nil, nil, nil, nil, fmt.Errorf("storage backend %T does not implement storage.DCRCredentialStore", baseStore)
        }

        bundleSource, retErr = newSPIFFEMultiDomainBundleSource(ctx, cfg.SPIFFETrust)
        if retErr != nil {
                return nil, nil, nil, nil, fmt.Errorf("create SPIFFE bundle source: %w", retErr)
        }
        defer closeSPIFFEBundleSourceOnConstructionError(bundleSource, &retErr)

        admitted, err := newClientAdmissionStorage(cfg, stor)
        if err != nil {
                return nil, nil, nil, nil, err
        }
        prepared, spiffeRegistry, retErr = decorateStorageForSPIFFE(ctx, cfg, admitted)
        if retErr != nil {
                return nil, nil, nil, nil, retErr
        }
        if retErr = validateConsentExemptClients(ctx, prepared, cfg.ConsentExemptClientIDs); retErr != nil {
                return nil, nil, nil, nil, retErr
        }
        if retErr = registerDelegateClients(ctx, prepared, cfg.DelegateClients); retErr != nil {
                return nil, nil, nil, nil, retErr
        }
        return prepared, dcrStore, bundleSource, spiffeRegistry, nil
}

// delegateClientIDs returns the client IDs of the configured delegate clients.
func delegateClientIDs(delegateClients []DelegateClient) []string {
        ids := make([]string, len(delegateClients))
        for i, c := range delegateClients {
                ids[i] = c.ClientID
        }
        return ids
}

// newClientAdmissionStorage wraps stor so a configured client is served only
// while cfg declares it. Consent-exempt IDs deliberately do not vouch. It must
// sit below the CIMD decorator: callbackClient's fallback (ClientStoreBelowCIMD)
// and CIMD's persisted-row fallback read below that decorator and must still
// pass admission.
func newClientAdmissionStorage(cfg Config, stor storage.Storage) (storage.Storage, error) {
        admission, err := storage.NewClientAdmissionStorage(stor, delegateClientIDs(cfg.DelegateClients))
        if err != nil {
                return nil, fmt.Errorf("create client admission storage: %w", err)
        }
        return admission, nil
}

func validateConsentExemptClients(ctx context.Context, stor storage.Storage, clientIDs []string) error {
        for _, clientID := range clientIDs {
                client, err := stor.GetClient(ctx, clientID)
                if err != nil {
                        return fmt.Errorf("failed to resolve consent-exempt client %q: %w", clientID, err)
                }
                if !consentpolicy.IsEligibleClient(client) {
                        return fmt.Errorf("consent-exempt client %q is missing or ineligible", clientID)
                }
        }
        return nil
}

func registerDelegateClients(ctx context.Context, stor storage.Storage, delegateClients []DelegateClient) error {
        for _, delegateClient := range delegateClients {
                client, err := registration.NewStaticDelegateClient(registration.Config{
                        ID:         delegateClient.ClientID,
                        Secret:     delegateClient.ClientSecret,
                        GrantTypes: []string{oauthproto.GrantTypeTokenExchange},
                        Scopes:     delegateClient.Scopes,
                        Audience:   delegateClient.Audiences,
                })
                if err != nil {
                        return fmt.Errorf("failed to create delegate client %q: %w", delegateClient.ClientID, err)
                }
                if err := stor.ReconcileConfiguredClient(ctx, client); err != nil {
                        return fmt.Errorf("failed to register delegate client %q: %w", delegateClient.ClientID, err)
                }
                slog.Warn("delegate client has blanket self-issued token exchange rights: "+
                        "it may exchange any user's self-issued ToolHive token for a delegated token, "+
                        "with no per-subject binding to the client that originally obtained that token",
                        "client_id", delegateClient.ClientID, "scopes", delegateClient.Scopes, "audiences", delegateClient.Audiences)
        }
        return nil
}

// decorateStorageForSPIFFE resolves the immutable association registry from the
// validated SPIFFE trust model and installs the static overlay outside CIMD.
// A nil cfg.SPIFFETrust means no SPIFFE associations are configured, which
// yields a nil registry and leaves the storage chain unchanged. The returned
// registry is also wired onto the fosite client-authentication strategy by the
// caller (see newSPIFFEClientResolver); it is nil when no SPIFFE trust is
// configured.
func decorateStorageForSPIFFE(
        ctx context.Context, cfg Config, stor storage.Storage,
) (storage.Storage, *SPIFFEAssociationRegistry, error) {
        registry, err := NewSPIFFEAssociationRegistry(cfg.SPIFFETrust)
        if err != nil {
                return nil, nil, fmt.Errorf("create SPIFFE association registry: %w", err)
        }

        // Install dynamic CIMD lookup before the static SPIFFE overlay so configured
        // clients always take precedence over remotely resolved HTTPS client IDs.
        stor, err = decorateStorageForCIMD(cfg, stor)
        if err != nil {
                return nil, nil, err
        }
        clients, err := registry.staticClients()
        if err != nil {
                return nil, nil, fmt.Errorf("build SPIFFE static clients: %w", err)
        }
        stor, err = storage.NewSPIFFEStorageDecorator(ctx, stor, clients)
        if err != nil {
                return nil, nil, fmt.Errorf("initialize SPIFFE client overlay: %w", err)
        }
        return stor, registry, nil
}

// newSPIFFEClientResolver returns the resolver the SPIFFE client-authentication
// strategy uses to turn a verified SPIFFE identity into its configured OAuth
// client. Both credential types resolve through this one function so that an
// X.509-SVID and a JWT-SVID for the same association cannot reach different
// authorization outcomes.
//
// It returns a nil resolver when no SPIFFE trust is configured. Returning a
// non-nil func closing over a nil registry would make the strategy's
// resolver != nil check pass when nothing is actually configured.
//
// Method uses the shared leaf-package type because package server cannot
// import authserver (authserver imports server). Package authserver aliases
// that type, so both sides share one set of method values.
func newSPIFFEClientResolver(
        registry *SPIFFEAssociationRegistry, stor storage.Storage,
) oauthserver.SPIFFEClientResolver {
        if registry == nil {
                return nil
        }
        return func(
                ctx context.Context, spiffeID, clientID string, method spiffeauth.SPIFFEAuthenticationMethod,
        ) (fosite.Client, error) {
                principal, err := registry.Resolve(spiffeID, clientID, method)
                if err != nil {
                        return nil, err
                }
                return stor.GetClient(ctx, principal.ClientID())
        }
}

// decorateStorageForCIMD wraps stor with the CIMD decorator when CIMD is
// enabled, so GetClient calls for HTTPS client_id values are intercepted at
// the fosite level (not just the handler level).
//
// When CIMD is disabled, stor is instead wrapped with
// storage.NewCIMDShapeGuardStorage so that a URL-shaped client_id can never
// resolve from a stale row a prior CIMD-enabled period may have
// write-through persisted (see cimdShapeGuardStorage's doc comment for why
// disabling CIMD alone does not evict or invalidate such a row).
func decorateStorageForCIMD(cfg Config, stor storage.Storage) (storage.Storage, error) {
        if !cfg.CIMDEnabled {
                return storage.NewCIMDShapeGuardStorage(stor), nil
        }
        if len(cfg.BaselineClientScopes) > 0 {
                slog.Warn("CIMD is enabled with baseline_client_scopes configured; "+
                        "any third-party client resolved via CIMD will also receive these scopes — "+
                        "ensure they are scopes you would grant by default to any unknown client",
                        "baseline_client_scopes", cfg.BaselineClientScopes)
        }
        decorated, err := storage.NewCIMDStorageDecorator(stor, storage.CIMDDecoratorConfig{
                Enabled:      true,
                CacheMaxSize: cfg.CIMDCacheMaxSize,
                FallbackTTL:  cfg.CIMDCacheFallbackTTL,
                // A refresh arrives about one access-token lifetime after the client's
                // last use, so any smaller bound never covers a refresh: on defaults the
                // cache entry is ~55m past its 5m expiry when the refresh arrives.
                MaxStale:             cfg.AccessTokenLifespan,
                ScopesSupported:      cfg.ScopesSupported,
                BaselineClientScopes: cfg.BaselineClientScopes,
                AllowedAudiences:     cfg.AllowedAudiences,
        })
        if err != nil {
                return nil, fmt.Errorf("failed to initialize CIMD storage decorator: %w", err)
        }
        return decorated, nil
}

// JWTBearerGrantEnabled reports whether any trusted issuer has the RFC 7523
// JWT-bearer grant configured.
func JWTBearerGrantEnabled(trustedIssuers []tokenexchange.TrustedIssuer) bool {
        for _, issuer := range trustedIssuers {
                if issuer.JWTBearerGrant != nil {
                        return true
                }
        }
        return false
}

// buildProvider assembles the fosite OAuth2 provider, registering the RFC 8693
// token-exchange handler and/or the RFC 7523 JWT-bearer handler as extension
// grants alongside the standard grants -- whichever of the two are enabled
// for cfg (token exchange can be disabled via canonical inbound grants
// configuration; JWT-bearer is enabled per JWTBearerGrantEnabled).
//
// refresher is the server's shared upstream token refresher (nil without
// upstreams) and upstreams are the resolved upstreams it refreshes, both used
// by the built-in upstream refresh guard.
//
// It returns the shared MultiIssuerTokenValidator (nil when no TrustedIssuers
// are configured, or when neither enabled grant would use it) so newServer
// can hold it and release its per-issuer JWKS worker pools on shutdown. On
// its own error paths it shuts that validator down before returning, since
// the caller never receives it.
func buildProvider(
        cfg Config, authServerConfig *oauthserver.AuthorizationServerConfig, stor storage.Storage,
        refresher storage.UpstreamTokenRefresher, upstreams []handlers.NamedUpstream,
) (_ fosite.OAuth2Provider, _ *tokenexchange.MultiIssuerTokenValidator, retErr error) {
        delegateIDs := delegateClientIDs(cfg.DelegateClients)
        jwtBearerEnabled := JWTBearerGrantEnabled(cfg.TrustedIssuers)

        // Built once, up front, and handed to whichever factories below actually
        // need it: otherwise each factory closure would build its own
        // MultiIssuerTokenValidator over the same trusted issuers at
        // fosite-compose time, doubling every issuer's JWKS cache and background
        // refresh goroutines — and, buried in a handler, leaving them unreachable
        // for shutdown. authServerConfig is the exact *AuthorizationServerConfig
        // each factory closure would otherwise receive at call time (see
        // createProvider/NewAuthorizationServer), so building it here first is
        // equivalent. Skipped entirely (shared stays nil) when neither grant that
        // consumes it is enabled -- trusted issuers configured with token
        // exchange disabled and JWT-bearer not enabled is a reachable, if
        // pointless, configuration, and building the validator anyway would
        // start live per-issuer JWKS refresh goroutines with no consumer for the
        // life of the server. NewSharedTrustedIssuerValidator itself also
        // returns nil when there are no trusted issuers.
        var shared *tokenexchange.MultiIssuerTokenValidator
        if !cfg.DisableTokenExchange || jwtBearerEnabled {
                var err error
                shared, err = tokenexchange.NewSharedTrustedIssuerValidator(authServerConfig, cfg.TrustedIssuers)
                if err != nil {
                        return nil, nil, fmt.Errorf("failed to create shared trusted-issuer validator: %w", err)
                }
        }
        // Release the validator's JWKS worker pools if we fail before returning it
        // to newServer, which otherwise owns its shutdown.
        defer func() {
                if retErr != nil {
                        if err := shared.Close(); err != nil {
                                slog.Warn("failed to shut down trusted-issuer validator during provider build cleanup", "error", err)
                        }
                }
        }()

        factories := make([]oauthserver.Factory, 0, 4)
        if !cfg.DisableTokenExchange {
                tokenExchangeFactory, err := tokenexchange.FactoryWithSharedTrustedIssuerValidator(
                        cfg.DelegationTokenLifespan, cfg.TrustedIssuers, delegateIDs, shared,
                        tokenexchange.WithAuthorizationDetailsRegistry(cfg.AuthorizationDetailsRegistry),
                )
                if err != nil {
                        return nil, nil, fmt.Errorf("failed to create token exchange factory: %w", err)
                }
                factories = append(factories, tokenExchangeFactory)
        }
        if jwtBearerEnabled {
                jwtBearerFactories, err := buildJWTBearerFactories(cfg.TrustedIssuers, shared)
                if err != nil {
                        return nil, nil, err
                }
                factories = append(factories, jwtBearerFactories...)
        }
        if cfg.DeviceFlowEnabled {
                deviceFlowFactory, err := buildDeviceFlowFactory(cfg, stor)
                if err != nil {
                        return nil, nil, err
                }
                factories = append(factories, deviceFlowFactory)
        }
        // Refresh-grant validators go last so each runs after the built-in refresh
        // handler has validated the grant.
        validators, err := refreshGrantValidatorChain(cfg, stor, refresher, upstreams)
        if err != nil {
                return nil, nil, err
        }
        factories = append(factories, refreshGrantHandlerFactories(validators)...)
        provider, err := createProvider(authServerConfig, stor, factories...)
        if err != nil {
                return nil, nil, err
        }
        return provider, shared, nil
}

// refreshGrantValidatorChain returns the refresh-grant validators in run
// order: the built-in upstream refresh guard (unless disabled or there is no
// refresher, because no upstreams are configured), then the embedder's
// Config.RefreshGrantHandlers. The guard's set of configured providers comes
// from upstreams, the same resolved list the refresher was built from.
func refreshGrantValidatorChain(
        cfg Config, stor storage.Storage, refresher storage.UpstreamTokenRefresher, upstreams []handlers.NamedUpstream,
) ([]fosite.TokenEndpointHandler, error) {
        if cfg.DisableUpstreamRefreshGuard || refresher == nil {
                return cfg.RefreshGrantHandlers, nil
        }
        upstreamNames := make([]string, len(upstreams))
        for i, u := range upstreams {
                upstreamNames[i] = u.Name
        }
        guard, err := newUpstreamRefreshGuard(stor, refresher, upstreamNames)
        if err != nil {
                return nil, fmt.Errorf("failed to create upstream refresh guard: %w", err)
        }
        return append([]fosite.TokenEndpointHandler{guard}, cfg.RefreshGrantHandlers...), nil
}

// buildJWTBearerFactories builds the RFC 7523 JWT-bearer factory together
// with the bound ID-JAG factory. The bound ID-JAG handler rides the same
// per-issuer JWT-bearer policy: enabling the grant enables both assertion
// forms, split by JOSE typ (plain assertions to JWTBearerHandler,
// oauth-id-jag+jwt to IDJAGHandler).
func buildJWTBearerFactories(
        trustedIssuers []tokenexchange.TrustedIssuer, shared *tokenexchange.MultiIssuerTokenValidator,
) ([]oauthserver.Factory, error) {
        jwtBearerFactory, err := tokenexchange.JWTBearerIssuanceFactory(trustedIssuers, shared)
        if err != nil {
                return nil, fmt.Errorf("failed to create JWT-bearer factory: %w", err)
        }
        idJAGFactory, err := tokenexchange.IDJAGIssuanceFactory(trustedIssuers, shared)
        if err != nil {
                return nil, fmt.Errorf("failed to create ID-JAG factory: %w", err)
        }
        return []oauthserver.Factory{jwtBearerFactory, idJAGFactory}, nil
}

// buildDeviceFlowFactory builds the RFC 8628 device authorization grant
// factory, validating that stor supports the device-code storage the grant
// requires.
func buildDeviceFlowFactory(cfg Config, stor storage.Storage) (oauthserver.Factory, error) {
        deviceStore, ok := storage.Unwrap(stor).(storage.DeviceCodeStorage)
        if !ok {
                return nil, fmt.Errorf("device flow enabled but storage backend %T does not implement storage.DeviceCodeStorage", stor)
        }
        interval := cfg.DeviceCodeInterval
        if interval <= 0 {
                interval = oauthserver.DefaultDeviceCodeInterval
        }
        return deviceflow.Factory(deviceStore, interval), nil
}

// buildHandlerOptions assembles the handlers.Option list for NewHandler: the
// refresher is always wired, the filter is added only when configured so a nil
// filter preserves walking every upstream, the register rate limit only when
// configured so a nil one keeps the handler default, and consent-exempt IDs
// only when the caller supplied a non-empty allowlist.
func buildHandlerOptions(
        refresher storage.UpstreamTokenRefresher,
        filter handlers.UpstreamFilter,
        registrationRateLimit *RegistrationRateLimitRunConfig,
        consentExemptClientIDs []string,
) []handlers.Option {
        opts := []handlers.Option{handlers.WithUpstreamRefresher(refresher)}
        if filter != nil {
                opts = append(opts, handlers.WithUpstreamFilter(filter))
        }
        if registrationRateLimit != nil {
                opts = append(opts, handlers.WithRegisterRateLimit(
                        rate.Limit(registrationRateLimit.RequestsPerSecond), registrationRateLimit.Burst))
        }
        if len(consentExemptClientIDs) > 0 {
                opts = append(opts, handlers.WithConsentExemptClientIDs(consentExemptClientIDs))
        }
        return opts
}

// Handler returns the HTTP handler that serves all OAuth/OIDC endpoints.
func (s *server) Handler() http.Handler {
        return s.handler
}

// IDPTokenStorage returns the IDP token storage interface.
func (s *server) IDPTokenStorage() storage.UpstreamTokenStorage {
        return s.storage
}

// DCRStore returns the persistent DCR credential store the server is wired
// against. See the Server interface doc for SECURITY and lifecycle notes.
func (s *server) DCRStore() storage.DCRCredentialStore {
        return s.dcrStore
}

// UpstreamTokenRefresher returns the single shared refresher constructed in
// newServer. Both the handler's chain-walk path and the runtime token-swap
// path use this instance, ensuring concurrent refreshes for the same logical
// upstream-token row are deduplicated by one process-local singleflight.Group.
func (s *server) UpstreamTokenRefresher() storage.UpstreamTokenRefresher {
        return s.upstreamRefresher
}

// newUpstreamTokenRefresher builds a refresher over the given upstreams, or nil
// when there are none. Called once during newServer so the returned instance
// can be shared between the handler and the UpstreamTokenRefresher() accessor.
func newUpstreamTokenRefresher(
        upstreams []handlers.NamedUpstream,
        stor storage.UpstreamTokenStorage,
        refreshTokenLifespan time.Duration,
) storage.UpstreamTokenRefresher {
        if len(upstreams) == 0 {
                return nil
        }
        providers := make(map[string]upstream.OAuth2Provider, len(upstreams))
        for _, u := range upstreams {
                providers[u.Name] = u.Provider
        }
        return &upstreamTokenRefresher{
                providers:            providers,
                storage:              stor,
                refreshTokenLifespan: refreshTokenLifespan,
        }
}

// Compile-time check that the concrete server satisfies the capability the
// CloseIdleConnections function detects, so that call can never degrade to a
// no-op for a server built by New.
var _ idleConnectionCloser = (*server)(nil)
var _ spiffeX509AuthorityLister = (*server)(nil)

// SPIFFEX509Authorities returns the current X.509 authorities configured for
// SPIFFE client authentication.
func (s *server) SPIFFEX509Authorities() []*x509.Certificate {
        if s.spiffeBundleSource == nil {
                return nil
        }
        return s.spiffeBundleSource.SPIFFEX509Authorities()
}

// CloseIdleConnections drains the upstream providers' pooled connections; see
// the package-level CloseIdleConnections function for the contract.
func (s *server) CloseIdleConnections() {
        closeUpstreamIdleConnections(s.upstreams)
}

// Close releases resources held by the server: it drains upstream idle
// connections, unregisters the handler's metric callback, shuts down its
// SPIFFE bundle and trusted-issuer refresh workers, and closes storage. Shutdown errors are joined so neither hides
// the other.
func (s *server) Close() error {
        slog.Debug("closing OAuth authorization server")
        s.CloseIdleConnections()
        var errs []error
        if s.oauthHandler != nil {
                if err := s.oauthHandler.Close(); err != nil {
                        errs = append(errs, err)
                }
        }
        if s.spiffeBundleSource != nil {
                if err := s.spiffeBundleSource.Close(); err != nil {
                        errs = append(errs, fmt.Errorf("failed to shut down SPIFFE bundle source: %w", err))
                }
        }
        if s.trustedIssuerValidator != nil {
                if err := s.trustedIssuerValidator.Close(); err != nil {
                        errs = append(errs, fmt.Errorf("failed to shut down trusted-issuer validator: %w", err))
                }
        }
        if err := s.storage.Close(); err != nil {
                errs = append(errs, err)
        }
        return errors.Join(errs...)
}

// createProvider creates a fosite OAuth2Provider configured for the authorization code flow.
//
// Fosite is an OAuth 2.0 framework that implements the protocol details. We use
// server.NewAuthorizationServer which accepts server.Factory functions to register
// grant type handlers. The standard compose factories are wrapped via wrapComposeFactory
// and any extra factories (e.g., token exchange) are appended.
//
// The provider is configured with:
//   - JWT strategy for access tokens (asymmetric signing, distributed validation via JWKS)
//   - HMAC strategy for authorization codes and refresh tokens (symmetric, internal only)
//   - Authorization code grant (RFC 6749 Section 4.1)
//   - Refresh token grant (RFC 6749 Section 6)
//   - PKCE (RFC 7636) for public client security
//   - Any extra factories passed in (e.g., RFC 8693 token exchange, or the
//     embedder's Config.RefreshGrantHandlers), registered after the built-ins
func createProvider(
        authServerConfig *oauthserver.AuthorizationServerConfig,
        stor storage.Storage,
        extraFactories ...oauthserver.Factory,
) (fosite.OAuth2Provider, error) {
        slog.Debug("configuring fosite OAuth2 provider",
                "key_id", authServerConfig.SigningKey.KeyID,
                "algorithm", authServerConfig.SigningKey.Algorithm,
        )

        // Convert go-jose/v4 JWK to go-jose/v3 JWK for fosite compatibility.
        // Fosite v0.49.0 depends on go-jose/v3, while we use v4 internally.
        // This ensures the "kid" (key ID) is included in JWT headers so resource
        // servers can look up the correct public key from our JWKS endpoint.
        signingKeyV4 := authServerConfig.SigningKey
        signingKeyV3 := &josev3.JSONWebKey{
                Key:       signingKeyV4.Key,
                KeyID:     signingKeyV4.KeyID,
                Algorithm: signingKeyV4.Algorithm,
                Use:       signingKeyV4.Use,
        }

        // Create a composed token strategy:
        // - JWT strategy (outer): signs access tokens as JWTs using the asymmetric signing key
        // - HMAC strategy (inner): signs authorization codes and refresh tokens using HMACSecret
        //
        // Access tokens are JWTs so resource servers can validate them without calling us.
        // Auth codes and refresh tokens are opaque HMAC tokens since only we validate them.
        jwtStrategy := compose.NewOAuth2JWTStrategy(
                func(_ context.Context) (interface{}, error) { return signingKeyV3, nil },
                compose.NewOAuth2HMACStrategy(authServerConfig.Config),
                authServerConfig.Config,
        )
        // This signer signs only access tokens: auth codes and refresh tokens use
        // the HMAC strategy, and no ID-token strategy is registered. The type is
        // set at signing time, not in the session, because storage persists the
        // session's JWT header and a per-session typ would survive a configuration
        // change across refreshes. Tokens signed through this signer get typ: at+jwt
        // unless they set their own typ, so any future JWT that is not an access
        // token, such as an ID-JAG (#6898), must set its own typ or sign with its
        // own signer.
        if authServerConfig.AccessTokenType == oauthserver.AccessTokenTypeATJWT {
                jwtStrategy.Signer = oauthserver.NewATJWTAccessTokenSigner(jwtStrategy.Signer)
        }

        commonStrategy := &compose.CommonStrategy{CoreStrategy: jwtStrategy}

        // Wrap fosite's compose factories to match server.Factory signature.
        factories := []oauthserver.Factory{
                wrapComposeFactory(compose.OAuth2AuthorizeExplicitFactory), // Authorization code grant
                wrapComposeFactory(compose.OAuth2RefreshTokenGrantFactory), // Refresh token grant
                wrapComposeFactory(compose.OAuth2PKCEFactory),              // PKCE for public clients
        }
        factories = append(factories, extraFactories...)

        return oauthserver.NewAuthorizationServer(
                authServerConfig,
                stor,
                commonStrategy,
                factories...,
        )
}

// runLegacyMigration runs one-shot Redis data migrations before handlers are
// constructed. It is a no-op for non-Redis backends and passes through any
// decorator wrapping so the concrete type can be reached. The DCR client index
// backfill runs before /oauth/register is served, so the cap counts clients
// written before the index existed.
//
// A backfill failure is logged and startup continues, unlike a legacy upstream
// migration failure. Skipping the backfill only leaves the cap undercounting
// rows written before the index existed, which is how the server behaved
// before the cap, and the next start fills in whatever was missed (ZADD NX
// makes a rerun safe). Skipping an upstream migration would break existing
// data.
func runLegacyMigration(ctx context.Context, stor storage.Storage, upstreams []UpstreamConfig) error {
        base := storage.Unwrap(stor)
        rs, ok := base.(*storage.RedisStorage)
        if !ok {
                return nil
        }
        if err := rs.BackfillDCRClientIndex(ctx); err != nil {
                slog.Warn("DCR client index backfill failed; the DCR client cap may undercount existing clients until the next restart",
                        "error", err)
        }
        for i := range upstreams {
                upCfg := &upstreams[i]
                if err := rs.MigrateLegacyUpstreamData(ctx, upCfg.Name, string(upCfg.Type)); err != nil {
                        return fmt.Errorf("legacy data migration failed for upstream %q: %w", upCfg.Name, err)
                }
        }
        return nil
}

// logConfidentialClientStartup emits an Info line naming the consequence of
// enabling confidential-client DCR. Combining it with cleartext HTTP is
// rejected by Config.Validate (see ValidateConfidentialClientTransport), so
// that combination can no longer reach this function.
func logConfidentialClientStartup(allowConfidentialClientRegistration bool) {
        if !allowConfidentialClientRegistration {
                return
        }
        slog.Info("confidential-client dynamic registration is enabled: " +
                "this server issues client secrets over unauthenticated dynamic registration")
}

// wrapComposeFactory adapts a compose.Factory to a server.Factory.
// Compose factories take (fosite.Configurator, interface{}, interface{}) while
// server factories take (*AuthorizationServerConfig, fosite.Storage, any).
// The embedded *fosite.Config satisfies fosite.Configurator.
func wrapComposeFactory(cf compose.Factory) oauthserver.Factory {
        return func(config *oauthserver.AuthorizationServerConfig, storage fosite.Storage, strategy any) (any, error) {
                return cf(config.Config, storage, strategy), nil
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package spiffeauth provides shared SPIFFE authentication primitives for the authorization server.
package spiffeauth

import (
        "context"

        "github.com/spiffe/go-spiffe/v2/spiffeid"
)

// SPIFFEJWTAssertionType identifies SPIFFE JWT-SVID client assertions.
const SPIFFEJWTAssertionType = "urn:ietf:params:oauth:client-assertion-type:jwt-spiffe"

type contextKey struct{}

// ContextWithSPIFFEID returns a new context with the SPIFFE ID stored.
func ContextWithSPIFFEID(ctx context.Context, id spiffeid.ID) context.Context {
        return context.WithValue(ctx, contextKey{}, id)
}

// SPIFFEIDFromContext retrieves the SPIFFE ID from the context.
// It returns the zero value and false if no SPIFFE ID is present.
func SPIFFEIDFromContext(ctx context.Context) (spiffeid.ID, bool) {
        id, ok := ctx.Value(contextKey{}).(spiffeid.ID)
        return id, ok
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package spiffeauth

import (
        "crypto/x509"
        "fmt"
        "log/slog"
        "net/http"
        "path"

        "github.com/spiffe/go-spiffe/v2/spiffeid"
        "github.com/spiffe/go-spiffe/v2/svid/x509svid"
)

// Middleware extracts a claimed SPIFFE ID from token-endpoint client certificates that contain a SPIFFE URI SAN.
// Certificates without a SPIFFE URI SAN pass through unchanged; credential validation remains the responsibility of
// the client-authentication strategy.
func Middleware(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                if r.URL.Path != "/oauth/token" || r.TLS == nil || len(r.TLS.PeerCertificates) == 0 {
                        next.ServeHTTP(w, r)
                        return
                }

                cert := r.TLS.PeerCertificates[0]
                if !hasSPIFFEURI(cert) {
                        next.ServeHTTP(w, r)
                        return
                }

                id, err := SPIFFEIDFromCertificate(cert)
                if err != nil {
                        w.Header().Set("Content-Type", "application/json")
                        w.Header().Set("Cache-Control", "no-store")
                        w.WriteHeader(http.StatusUnauthorized)
                        if _, err := w.Write([]byte(`{"error":"invalid_client"}`)); err != nil {
                                slog.Debug("failed to write OAuth error response", "error", err)
                        }
                        return
                }
                next.ServeHTTP(w, r.WithContext(ContextWithSPIFFEID(r.Context(), id)))
        })
}

// SPIFFEIDFromCertificate returns the single non-root SPIFFE URI SAN in cert.
func SPIFFEIDFromCertificate(cert *x509.Certificate) (spiffeid.ID, error) {
        id, err := x509svid.IDFromCert(cert)
        if err != nil {
                return spiffeid.ID{}, err
        }
        if id.Path() == "" || path.Clean(id.Path()) != id.Path() {
                return spiffeid.ID{}, fmt.Errorf("SPIFFE ID path is invalid")
        }
        return id, nil
}

func hasSPIFFEURI(cert *x509.Certificate) bool {
        for _, uri := range cert.URIs {
                if uri != nil && uri.Scheme == "spiffe" {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "crypto/sha256"
        "encoding/hex"
        "fmt"
        "slices"
        "strings"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
)

// SPIFFEAssociationRegistry is the immutable runtime index of validated SPIFFE
// associations. It selects policy only; it neither accepts nor authenticates a
// SPIFFE credential.
type SPIFFEAssociationRegistry struct {
        byPattern  map[string]SPIFFEClientAuthConfig
        byClientID map[string]SPIFFEClientAuthConfig
}

// NewSPIFFEAssociationRegistry creates an immutable lookup registry from a
// trust configuration. A nil trust configuration represents an absent SPIFFE
// configuration and returns nil without enabling SPIFFE clients.
//
// It rejects any two associations whose patterns could both match the same
// concrete SPIFFE ID (including an exact duplicate, which is the identical
// pattern matching itself) — associationForSPIFFEID resolves wildcard
// patterns by an unordered map scan, so an unrejected overlap would make
// resolution depend on random iteration order.
func NewSPIFFEAssociationRegistry(trust *SPIFFETrustConfig) (*SPIFFEAssociationRegistry, error) {
        if trust == nil {
                return nil, nil
        }

        associations := trust.Associations()
        registry := &SPIFFEAssociationRegistry{
                byPattern:  make(map[string]SPIFFEClientAuthConfig, len(associations)),
                byClientID: make(map[string]SPIFFEClientAuthConfig, len(associations)),
        }
        for _, association := range associations {
                if _, exists := registry.byClientID[association.ClientID()]; exists {
                        return nil, fmt.Errorf("duplicate SPIFFE association client ID %q", association.ClientID())
                }
                for existingPattern := range registry.byPattern {
                        if spiffePatternsOverlap(association.Principal(), existingPattern) {
                                return nil, fmt.Errorf(
                                        "SPIFFE association pattern %q overlaps with existing pattern %q: resolution would be ambiguous",
                                        association.Principal(), existingPattern,
                                )
                        }
                }
                registry.byPattern[association.Principal()] = association.clone()
                registry.byClientID[association.ClientID()] = association.clone()
        }
        return registry, nil
}

// NormalizedSPIFFEPrincipal is an immutable, validated identity selected from a
// SPIFFE association. It represents policy only; it does not authenticate a
// credential or enable an authentication method.
type NormalizedSPIFFEPrincipal struct {
        clientID      string
        spiffeID      string
        trustDomain   string
        authMethod    SPIFFEAuthenticationMethod
        authorization SPIFFEAuthorizationPolicy
}

// ClientID returns the configured OAuth client ID.
func (p NormalizedSPIFFEPrincipal) ClientID() string { return p.clientID }

// SPIFFEID returns the canonical concrete SPIFFE ID.
func (p NormalizedSPIFFEPrincipal) SPIFFEID() string { return p.spiffeID }

// TrustDomain returns the canonical SPIFFE trust domain.
func (p NormalizedSPIFFEPrincipal) TrustDomain() string { return p.trustDomain }

// AuthenticationMethod returns the selected credential-method discriminator.
func (p NormalizedSPIFFEPrincipal) AuthenticationMethod() SPIFFEAuthenticationMethod {
        return p.authMethod
}

// AuthorizationPolicy returns a defensive copy of the selected policy.
func (p NormalizedSPIFFEPrincipal) AuthorizationPolicy() SPIFFEAuthorizationPolicy {
        return SPIFFEAuthorizationPolicy{
                grantTypes: slices.Clone(p.authorization.grantTypes),
                scopes:     slices.Clone(p.authorization.scopes),
                resources:  slices.Clone(p.authorization.resources),
                audiences:  slices.Clone(p.authorization.audiences),
        }
}

// Resolve selects a configured association for a canonical concrete SPIFFE ID,
// requested OAuth client ID, and explicitly selected authentication method.
// When clientID is empty, the client ID is derived entirely from the
// resolved SPIFFE association (the "optional client ID" case); otherwise
// clientID must match the association's client ID. It fails closed for
// unknown IDs, mismatched client ownership, and disabled methods. This
// method does not validate a credential.
func (r *SPIFFEAssociationRegistry) Resolve(
        spiffeID, clientID string, method SPIFFEAuthenticationMethod,
) (NormalizedSPIFFEPrincipal, error) {
        if r == nil {
                return NormalizedSPIFFEPrincipal{}, fmt.Errorf("no SPIFFE associations are configured")
        }
        canonicalID, err := normalizeSPIFFEPrincipal(spiffeID, false)
        if err != nil {
                return NormalizedSPIFFEPrincipal{}, fmt.Errorf("invalid SPIFFE ID: %w", err)
        }
        parsedID, err := parseSPIFFEID(canonicalID)
        if err != nil {
                return NormalizedSPIFFEPrincipal{}, fmt.Errorf("invalid SPIFFE ID: %w", err)
        }
        association, ok := r.associationForSPIFFEID(canonicalID)
        if !ok {
                return NormalizedSPIFFEPrincipal{}, fmt.Errorf("no SPIFFE association for ID %q", canonicalID)
        }
        if clientID != "" && association.ClientID() != clientID {
                return NormalizedSPIFFEPrincipal{}, fmt.Errorf("SPIFFE ID is not associated with client ID %q", clientID)
        }
        if !containsSPIFFEAuthenticationMethod(association.Methods(), method) {
                return NormalizedSPIFFEPrincipal{}, fmt.Errorf(
                        "SPIFFE authentication method %q is not enabled for client ID %q", method, association.ClientID(),
                )
        }

        return NormalizedSPIFFEPrincipal{
                clientID:      association.ClientID(),
                spiffeID:      canonicalID,
                trustDomain:   parsedID.TrustDomain().String(),
                authMethod:    method,
                authorization: association.AuthorizationPolicy(),
        }, nil
}

func (r *SPIFFEAssociationRegistry) associationForSPIFFEID(spiffeID string) (SPIFFEClientAuthConfig, bool) {
        if association, ok := r.byPattern[spiffeID]; ok {
                return association.clone(), true
        }
        for pattern, association := range r.byPattern {
                if strings.HasSuffix(pattern, "/*") && matchSPIFFEPrincipalPattern(pattern, spiffeID) {
                        return association.clone(), true
                }
        }
        return SPIFFEClientAuthConfig{}, false
}

func containsSPIFFEAuthenticationMethod(methods []SPIFFEAuthenticationMethod, wanted SPIFFEAuthenticationMethod) bool {
        for _, method := range methods {
                if method == wanted {
                        return true
                }
        }
        return false
}

// spiffeStaticClient embeds the concrete *registration.SPIFFEClient (not the
// fosite.Client interface) so that embedding still promotes Resources() and
// the BackChannelOnlyMarker -- an interface embedding would silently drop
// both, since neither is part of the fosite.Client method set.
type spiffeStaticClient struct {
        *registration.SPIFFEClient
        identityFingerprint string
}

func (c spiffeStaticClient) IdentityFingerprint() string { return c.identityFingerprint }

// staticClients builds immutable OAuth clients and their association identity for
// durable placeholder reconciliation.
func (r *SPIFFEAssociationRegistry) staticClients() (map[string]fosite.Client, error) {
        if r == nil {
                return nil, nil
        }
        clients := make(map[string]fosite.Client, len(r.byClientID))
        for clientID, association := range r.byClientID {
                policy := association.AuthorizationPolicy()
                client, err := registration.NewSPIFFEClient(
                        association.ClientID(), policy.Scopes(), policy.Audiences(), policy.Resources(),
                )
                if err != nil {
                        return nil, fmt.Errorf("SPIFFE client %q: %w", clientID, err)
                }
                clients[clientID] = spiffeStaticClient{
                        SPIFFEClient:        client,
                        identityFingerprint: fingerprintSPIFFEAssociation(association),
                }
        }
        return clients, nil
}

// fingerprintSPIFFEAssociation hashes the SPIFFE association identity (trust
// domain reference, principal pattern, and accepted methods) so two
// associations with a differing identity durably reconcile as different
// clients rather than being silently accepted as the same one. Each value is
// length-prefixed with an unambiguous decimal-and-colon marker so
// concatenation can't alias two different value sequences to the same hash,
// and methods are sorted first so their original order never affects the
// result.
//
// Note: TrustDomainRef() hashes the trust-domain *reference name* (e.g.
// "production"), not the trust domain value itself. In practice this doesn't
// weaken the fingerprint because a principal pattern always embeds its trust
// domain, so a genuine trust-domain difference still changes the hash via the
// principal -- but repointing a trust domain ref's bundle source to a
// different CA while keeping the same ref name and trust domain would not be
// detected by this fingerprint alone.
func fingerprintSPIFFEAssociation(association SPIFFEClientAuthConfig) string {
        hash := sha256.New()
        write := func(value string) {
                _, _ = fmt.Fprintf(hash, "%d:", len(value))
                _, _ = hash.Write([]byte(value))
        }
        write(association.TrustDomainRef())
        write(association.Principal())
        methods := association.Methods()
        slices.Sort(methods)
        for _, method := range methods {
                write(string(method))
        }
        return hex.EncodeToString(hash.Sum(nil))
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "context"
        "crypto/x509"
        "errors"
        "fmt"
        "log/slog"
        "os"
        "sort"
        "sync"
        "sync/atomic"
        "time"

        "github.com/spiffe/go-spiffe/v2/bundle/jwtbundle"
        "github.com/spiffe/go-spiffe/v2/bundle/spiffebundle"
        "github.com/spiffe/go-spiffe/v2/bundle/x509bundle"
        "github.com/spiffe/go-spiffe/v2/spiffeid"
)

const (
        spiffeBundleMaxStaleAge      = 24 * time.Hour
        spiffeBundleFilePollInterval = 30 * time.Second
)

// spiffeMultiDomainBundleSource provides live trust bundles for the configured
// SPIFFE trust domains. It is both an x509bundle.Source and jwtbundle.Source.
type spiffeMultiDomainBundleSource struct {
        byTrustDomain map[spiffeid.TrustDomain]*spiffeLiveBundleSource
        cancel        context.CancelFunc
        workers       sync.WaitGroup
        closers       []func() error
        closeOnce     sync.Once
        closeErr      error
}

type spiffeLiveBundleSource struct {
        x509    x509bundle.Source
        jwt     jwtbundle.Source
        methods map[SPIFFEAuthenticationMethod]struct{}
}

type bundleHolder struct {
        current atomic.Pointer[bundleSnapshot]
}

type bundleSnapshot struct {
        bundle    *spiffebundle.Bundle
        fetchedAt time.Time
}

// store publishes bundle as a complete valid snapshot. Once a sequence number
// has been observed, it rejects a missing, lower, or conflicting equal sequence
// so a replayed reload cannot refresh the last-known-good age.
func (h *bundleHolder) store(bundle *spiffebundle.Bundle) error {
        for {
                current := h.current.Load()
                if err := validateBundleSnapshotUpdate(current, bundle); err != nil {
                        return err
                }
                next := &bundleSnapshot{bundle: bundle, fetchedAt: time.Now()}
                if h.current.CompareAndSwap(current, next) {
                        return nil
                }
        }
}

func validateBundleSnapshotUpdate(current *bundleSnapshot, bundle *spiffebundle.Bundle) error {
        if current == nil || current.bundle == nil {
                return nil
        }
        currentSequence, currentHasSequence := current.bundle.SequenceNumber()
        if !currentHasSequence {
                return nil
        }
        sequence, hasSequence := bundle.SequenceNumber()
        if !hasSequence {
                return fmt.Errorf("SPIFFE bundle omits sequence number after sequence %d was observed", currentSequence)
        }
        if sequence < currentSequence {
                return fmt.Errorf("SPIFFE bundle sequence %d is lower than current sequence %d", sequence, currentSequence)
        }
        if sequence == currentSequence && !bundle.Equal(current.bundle) {
                return fmt.Errorf("SPIFFE bundle sequence %d conflicts with the current bundle", sequence)
        }
        return nil
}

// newSPIFFEMultiDomainBundleSource constructs all configured bundle sources and
// completes their initial load before returning. The supplied context controls
// refresh workers. Initial regular-file reads are synchronous and cannot be
// canceled.
//
//nolint:unused // wired into server construction in the next PR of the stack
func newSPIFFEMultiDomainBundleSource(
        ctx context.Context, cfg *SPIFFETrustConfig,
) (*spiffeMultiDomainBundleSource, error) {
        if cfg == nil || len(cfg.trustDomains) == 0 {
                return nil, nil
        }

        runtimeCtx, cancel := context.WithCancel(ctx)
        source := &spiffeMultiDomainBundleSource{
                byTrustDomain: make(map[spiffeid.TrustDomain]*spiffeLiveBundleSource, len(cfg.trustDomains)),
                cancel:        cancel,
        }

        trustDomainNames := make([]string, 0, len(cfg.trustDomains))
        for name := range cfg.trustDomains {
                trustDomainNames = append(trustDomainNames, name)
        }
        sort.Strings(trustDomainNames)

        for _, name := range trustDomainNames {
                domain := cfg.trustDomains[name]
                trustDomain, err := spiffeid.TrustDomainFromString(domain.TrustDomain())
                if err != nil {
                        return nil, closeSPIFFEBundleSourceAfterError(source, fmt.Errorf("SPIFFE trust domain %q: %w", name, err))
                }
                if _, exists := source.byTrustDomain[trustDomain]; exists {
                        return nil, closeSPIFFEBundleSourceAfterError(source, fmt.Errorf("duplicate SPIFFE trust domain %q", trustDomain))
                }

                live, err := newSPIFFELiveBundleSource(runtimeCtx, trustDomain, domain, source)
                if err != nil {
                        return nil, closeSPIFFEBundleSourceAfterError(source, fmt.Errorf("load SPIFFE trust bundle for %q: %w", trustDomain, err))
                }
                source.byTrustDomain[trustDomain] = live
        }

        return source, nil
}

// closeSPIFFEBundleSourceAfterError preserves the construction failure while
// reporting a secondary resource-cleanup failure.
//
//nolint:unused // wired into server construction in the next PR of the stack
func closeSPIFFEBundleSourceAfterError(source *spiffeMultiDomainBundleSource, primary error) error {
        if err := source.Close(); err != nil {
                slog.Warn("failed to clean up SPIFFE bundle source after construction failure", "error", err)
        }
        return primary
}

// SPIFFEX509Authorities returns the current X.509 authorities for every
// trust domain that enables X.509 authentication. A bundle that cannot be
// read during rotation is skipped so one stale domain does not hide others.
func (s *spiffeMultiDomainBundleSource) SPIFFEX509Authorities() []*x509.Certificate {
        if s == nil {
                return nil
        }
        authorities := make([]*x509.Certificate, 0)
        domains := make([]spiffeid.TrustDomain, 0, len(s.byTrustDomain))
        for trustDomain := range s.byTrustDomain {
                domains = append(domains, trustDomain)
        }
        sort.Slice(domains, func(i, j int) bool { return domains[i].String() < domains[j].String() })
        for _, trustDomain := range domains {
                live := s.byTrustDomain[trustDomain]
                if live == nil {
                        continue
                }
                if _, enabled := live.methods[SPIFFEAuthenticationMethodX509]; !enabled {
                        continue
                }
                bundle, err := s.GetX509BundleForTrustDomain(trustDomain)
                if err != nil {
                        slog.Debug("failed to load SPIFFE X.509 bundle authorities", "trust_domain", trustDomain, "error", err)
                        continue
                }
                if bundle != nil {
                        authorities = append(authorities, bundle.X509Authorities()...)
                }
        }
        return authorities
}

func (s *spiffeMultiDomainBundleSource) GetX509BundleForTrustDomain(
        trustDomain spiffeid.TrustDomain,
) (*x509bundle.Bundle, error) {
        live, err := s.sourceFor(trustDomain, SPIFFEAuthenticationMethodX509)
        if err != nil {
                return nil, err
        }
        return live.x509.GetX509BundleForTrustDomain(trustDomain)
}

func (s *spiffeMultiDomainBundleSource) GetJWTBundleForTrustDomain(
        trustDomain spiffeid.TrustDomain,
) (*jwtbundle.Bundle, error) {
        live, err := s.sourceFor(trustDomain, SPIFFEAuthenticationMethodJWT)
        if err != nil {
                return nil, err
        }
        return live.jwt.GetJWTBundleForTrustDomain(trustDomain)
}

// Close stops all bundle refresh workers and releases their clients. It is safe
// to call more than once.
func (s *spiffeMultiDomainBundleSource) Close() error {
        if s == nil {
                return nil
        }
        s.closeOnce.Do(func() {
                s.cancel()
                s.workers.Wait()
                errs := make([]error, 0, len(s.closers))
                for index := len(s.closers) - 1; index >= 0; index-- {
                        if err := s.closers[index](); err != nil {
                                errs = append(errs, err)
                        }
                }
                s.closeErr = errors.Join(errs...)
        })
        return s.closeErr
}

func newSPIFFELiveBundleSource(
        runtimeCtx context.Context,
        trustDomain spiffeid.TrustDomain,
        domain SPIFFETrustDomain,
        multi *spiffeMultiDomainBundleSource,
) (*spiffeLiveBundleSource, error) {
        methods := make(map[SPIFFEAuthenticationMethod]struct{}, len(domain.Methods()))
        for _, method := range domain.Methods() {
                methods[method] = struct{}{}
        }
        live := &spiffeLiveBundleSource{methods: methods}

        switch domain.BundleSource().Type() {
        case SPIFFEBundleSourceTypeFile:
                path := domain.BundleSource().Path()
                bundle, err := loadSPIFFEBundleFile(trustDomain, path)
                if err != nil {
                        return nil, fmt.Errorf("initial bundle load: %w", err)
                }
                if err := validateSPIFFEBundleAuthorities(bundle, trustDomain, methods); err != nil {
                        return nil, fmt.Errorf("initial bundle validation: %w", err)
                }
                holder := &bundleHolder{}
                if err := holder.store(bundle); err != nil {
                        return nil, fmt.Errorf("store initial SPIFFE bundle: %w", err)
                }
                live.x509 = holder
                live.jwt = holder
                multi.workers.Add(1)
                go pollFileBundle(runtimeCtx, trustDomain, path, holder, live.methods, &multi.workers)
        case SPIFFEBundleSourceTypeEndpoint, SPIFFEBundleSourceTypeWorkloadAPI:
                return nil, fmt.Errorf("unsupported SPIFFE bundle source type %q: no loader in this build", domain.BundleSource().Type())
        default:
                return nil, fmt.Errorf("unsupported SPIFFE bundle source type %q", domain.BundleSource().Type())
        }
        return live, nil
}

func loadSPIFFEBundleFile(trustDomain spiffeid.TrustDomain, path string) (*spiffebundle.Bundle, error) {
        info, err := os.Stat(path)
        if err != nil {
                return nil, fmt.Errorf("stat SPIFFE bundle file %q: %w", path, err)
        }
        if !info.Mode().IsRegular() {
                return nil, fmt.Errorf("SPIFFE bundle path %q is not a regular file", path)
        }
        return spiffebundle.Load(trustDomain, path)
}

func validateSPIFFEBundleAuthorities(
        bundle *spiffebundle.Bundle,
        trustDomain spiffeid.TrustDomain,
        methods map[SPIFFEAuthenticationMethod]struct{},
) error {
        if _, enabled := methods[SPIFFEAuthenticationMethodX509]; enabled {
                x509Bundle, err := bundle.GetX509BundleForTrustDomain(trustDomain)
                if err != nil {
                        return fmt.Errorf("X.509 bundle: %w", err)
                }
                if len(x509Bundle.X509Authorities()) == 0 {
                        return fmt.Errorf("X.509 bundle has no authorities")
                }
        }
        if _, enabled := methods[SPIFFEAuthenticationMethodJWT]; enabled {
                jwtBundle, err := bundle.GetJWTBundleForTrustDomain(trustDomain)
                if err != nil {
                        return fmt.Errorf("JWT bundle: %w", err)
                }
                if len(jwtBundle.JWTAuthorities()) == 0 {
                        return fmt.Errorf("JWT bundle has no authorities")
                }
        }
        return nil
}

func reloadSPIFFEBundleFile(
        trustDomain spiffeid.TrustDomain,
        path string,
        holder *bundleHolder,
        methods map[SPIFFEAuthenticationMethod]struct{},
) error {
        bundle, err := loadSPIFFEBundleFile(trustDomain, path)
        if err != nil {
                return err
        }
        if err := validateSPIFFEBundleAuthorities(bundle, trustDomain, methods); err != nil {
                return err
        }
        return holder.store(bundle)
}

func (s *spiffeMultiDomainBundleSource) sourceFor(
        trustDomain spiffeid.TrustDomain, method SPIFFEAuthenticationMethod,
) (*spiffeLiveBundleSource, error) {
        live, ok := s.byTrustDomain[trustDomain]
        if !ok {
                return nil, fmt.Errorf("no SPIFFE bundle source configured for trust domain %q", trustDomain)
        }
        if _, ok := live.methods[method]; !ok {
                return nil, fmt.Errorf("SPIFFE authentication method %q is not enabled for trust domain %q", method, trustDomain)
        }
        return live, nil
}

func (h *bundleHolder) GetX509BundleForTrustDomain(trustDomain spiffeid.TrustDomain) (*x509bundle.Bundle, error) {
        bundle, err := h.bundle(trustDomain)
        if err != nil {
                return nil, err
        }
        return bundle.GetX509BundleForTrustDomain(trustDomain)
}

func (h *bundleHolder) GetJWTBundleForTrustDomain(trustDomain spiffeid.TrustDomain) (*jwtbundle.Bundle, error) {
        bundle, err := h.bundle(trustDomain)
        if err != nil {
                return nil, err
        }
        return bundle.GetJWTBundleForTrustDomain(trustDomain)
}

func (h *bundleHolder) bundle(trustDomain spiffeid.TrustDomain) (*spiffebundle.Bundle, error) {
        snapshot := h.current.Load()
        if snapshot == nil || snapshot.bundle == nil {
                return nil, fmt.Errorf("no SPIFFE bundle loaded for trust domain %q", trustDomain)
        }
        if time.Since(snapshot.fetchedAt) > spiffeBundleMaxStaleAge {
                return nil, fmt.Errorf("SPIFFE bundle for trust domain %q is older than %s", trustDomain, spiffeBundleMaxStaleAge)
        }
        return snapshot.bundle, nil
}

// pollFileBundle reloads a locally mounted SPIFFE bundle file on a fixed
// interval. A ConfigMap-mounted file is updated via a symlink swap the
// kubelet performs on its own sync period, which polling observes correctly
// and inotify on the mounted path frequently misses. A failed reload is
// logged and the last known good bundle is retained.
func pollFileBundle(
        ctx context.Context,
        trustDomain spiffeid.TrustDomain,
        path string,
        holder *bundleHolder,
        methods map[SPIFFEAuthenticationMethod]struct{},
        worker *sync.WaitGroup,
) {
        defer worker.Done()
        failed := false
        ticker := time.NewTicker(spiffeBundleFilePollInterval)
        defer ticker.Stop()
        for {
                select {
                case <-ctx.Done():
                        return
                case <-ticker.C:
                }

                err := reloadSPIFFEBundleFile(trustDomain, path, holder, methods)
                if err != nil {
                        if !failed {
                                slog.Warn("SPIFFE bundle file reload failed; retaining last known good bundle",
                                        "trust_domain", trustDomain, "path", path, "error", err)
                                failed = true
                        }
                        continue
                }
                if failed {
                        slog.Debug("SPIFFE bundle file reload recovered", "trust_domain", trustDomain)
                        failed = false
                }
        }
}

var _ x509bundle.Source = (*spiffeMultiDomainBundleSource)(nil)
var _ jwtbundle.Source = (*spiffeMultiDomainBundleSource)(nil)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "context"
        "fmt"

        "github.com/stacklok/toolhive/pkg/authserver/storage"
)

// PreflightSPIFFEStaticClientCollisions rejects static client IDs that already
// exist in durable storage. It runs before upstream DCR registration so a
// collision cannot leave an orphaned upstream registration.
func PreflightSPIFFEStaticClientCollisions(ctx context.Context, base storage.Storage, trust *SPIFFETrustConfig) error {
        registry, err := NewSPIFFEAssociationRegistry(trust)
        if err != nil {
                return fmt.Errorf("create SPIFFE association registry: %w", err)
        }
        clients, err := registry.staticClients()
        if err != nil {
                return err
        }
        return storage.PreflightSPIFFEStaticClientCollisions(ctx, base, clients)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
        "fmt"
        "net"
        "net/url"
        "path/filepath"
        "slices"
        "strings"

        "github.com/spiffe/go-spiffe/v2/spiffeid"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        spiffeauth "github.com/stacklok/toolhive/pkg/authserver/spiffe"
        "github.com/stacklok/toolhive/pkg/authserver/storage"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

const (
        // SPIFFEAuthenticationMethodX509 authenticates a workload with an X.509-SVID.
        SPIFFEAuthenticationMethodX509 = spiffeauth.SPIFFEAuthenticationMethodX509
        // SPIFFEAuthenticationMethodJWT authenticates a workload with a JWT-SVID.
        SPIFFEAuthenticationMethodJWT = spiffeauth.SPIFFEAuthenticationMethodJWT

        // SPIFFEGrantTypeTokenExchange is the only SPIFFE client grant supported by
        // this configuration surface.
        SPIFFEGrantTypeTokenExchange = oauthproto.GrantTypeTokenExchange

        // SPIFFEBundleSourceTypeEndpoint selects a HTTPS SPIFFE Bundle Endpoint.
        SPIFFEBundleSourceTypeEndpoint SPIFFEBundleSourceType = "bundle_endpoint"
        // SPIFFEBundleSourceTypeFile selects a locally mounted SPIFFE trust bundle.
        SPIFFEBundleSourceTypeFile SPIFFEBundleSourceType = "file"
        // SPIFFEBundleSourceTypeWorkloadAPI selects the local SPIFFE Workload API.
        SPIFFEBundleSourceTypeWorkloadAPI SPIFFEBundleSourceType = "workload_api"

        // SPIFFEBundleEndpointProfileHTTPSWeb authenticates the bundle endpoint's
        // TLS connection with a Web PKI certificate (the SPIFFE Bundle Endpoint
        // "https_web" profile).
        SPIFFEBundleEndpointProfileHTTPSWeb SPIFFEBundleEndpointProfile = "https_web"
        // SPIFFEBundleEndpointProfileHTTPSSPIFFE authenticates the bundle
        // endpoint's TLS connection with an X.509-SVID trusted by a separately
        // distributed root (the SPIFFE Bundle Endpoint "https_spiffe" profile).
        SPIFFEBundleEndpointProfileHTTPSSPIFFE SPIFFEBundleEndpointProfile = "https_spiffe"
)

// SPIFFEAuthenticationMethod identifies the credential type permitted for a
// SPIFFE workload.
type SPIFFEAuthenticationMethod = spiffeauth.SPIFFEAuthenticationMethod

// SPIFFEBundleSourceType identifies the selected trust-bundle source.
type SPIFFEBundleSourceType string

// SPIFFEBundleEndpointProfile identifies how a SPIFFE Bundle Endpoint's TLS
// connection is authenticated, per the SPIFFE Federation specification.
type SPIFFEBundleEndpointProfile string

// SPIFFETrustDomainRunConfig declares one SPIFFE trust domain. Credential
// authentication is deliberately outside this configuration step.
type SPIFFETrustDomainRunConfig struct {
        // Name uniquely identifies this declaration and is referenced by
        // InboundGrants.SPIFFEClientAuth entries.
        Name string `json:"name" yaml:"name"`

        // TrustDomain is the SPIFFE trust domain accepted by this declaration.
        TrustDomain string `json:"trust_domain" yaml:"trust_domain"`

        // Methods explicitly enables the supported credential types for this trust
        // domain. No authentication method is enabled when the list is empty.
        Methods []SPIFFEAuthenticationMethod `json:"methods" yaml:"methods"`

        // BundleSource declares exactly one trust-bundle source and is validated for
        // shape here. The standalone file loader exists, but source initialization
        // is not wired into SPIFFE authentication.
        BundleSource SPIFFEBundleSourceRunConfig `json:"bundle_source" yaml:"bundle_source"`
}

// SPIFFEBundleSourceRunConfig is a discriminated bundle-source declaration.
// Type determines which, and only which, source payload may be set.
type SPIFFEBundleSourceRunConfig struct {
        Type SPIFFEBundleSourceType `json:"type" yaml:"type"`

        Endpoint    *SPIFFEBundleEndpointSourceRunConfig    `json:"endpoint,omitempty" yaml:"endpoint,omitempty"`
        File        *SPIFFEFileBundleSourceRunConfig        `json:"file,omitempty" yaml:"file,omitempty"`
        WorkloadAPI *SPIFFEWorkloadAPIBundleSourceRunConfig `json:"workload_api,omitempty" yaml:"workload_api,omitempty"`
}

// SPIFFEBundleEndpointSourceRunConfig declares a HTTPS SPIFFE Bundle Endpoint.
type SPIFFEBundleEndpointSourceRunConfig struct {
        URL string `json:"url" yaml:"url"`

        // Profile selects how the endpoint's TLS connection is authenticated:
        // SPIFFEBundleEndpointProfileHTTPSWeb (Web PKI) or
        // SPIFFEBundleEndpointProfileHTTPSSPIFFE (a separately distributed
        // X.509-SVID root). Required, since an endpoint loader needs this profile
        // to choose the trust anchor for its initial connection.
        Profile SPIFFEBundleEndpointProfile `json:"profile" yaml:"profile"`
}

// SPIFFEFileBundleSourceRunConfig declares a trust bundle read from a local
// file, typically a mounted ConfigMap or Secret containing a SPIFFE JWKS
// trust-bundle document.
type SPIFFEFileBundleSourceRunConfig struct {
        Path string `json:"path" yaml:"path"`
}

// SPIFFEWorkloadAPIBundleSourceRunConfig selects the local SPIFFE Workload API.
// It deliberately has no payload; loading and deployment details are deferred
// to the bundle-loading implementation.
type SPIFFEWorkloadAPIBundleSourceRunConfig struct{}

// InboundGrantsRunConfig declares canonical inbound grant configuration for
// separately declared trust roots. SPIFFE client authentication entries live
// here, alongside other inbound grant purposes, so SPIFFE is not a parallel
// trust path, and so client authentication and grant-family enablement
// (token exchange, JWT-bearer) remain independently configurable.
type InboundGrantsRunConfig struct {
        // SPIFFEClientAuth associates SPIFFE principal patterns with explicit OAuth
        // client identities and permissions. See SPIFFEClientAuthRunConfig.
        SPIFFEClientAuth []SPIFFEClientAuthRunConfig `json:"spiffe_client_auth,omitempty" yaml:"spiffe_client_auth,omitempty"`

        // TokenExchange configures RFC 8693 inbound clients and issuer policies.
        TokenExchange *TokenExchangeInboundGrantRunConfig `json:"token_exchange,omitempty" yaml:"token_exchange,omitempty"`

        // JWTBearer configures RFC 7523 issuer policies.
        JWTBearer *JWTBearerInboundGrantRunConfig `json:"jwt_bearer,omitempty" yaml:"jwt_bearer,omitempty"`
}

// SPIFFEClientAuthRunConfig associates one SPIFFE principal pattern from a
// declared trust domain with an explicit OAuth client identity and permissions.
type SPIFFEClientAuthRunConfig struct {
        // TrustDomainRef identifies the SPIFFE trust-domain declaration governing
        // this association policy.
        TrustDomainRef string `json:"trust_domain_ref" yaml:"trust_domain_ref"`

        // PrincipalPattern is a concrete SPIFFE ID or a terminal /* pattern within
        // the declared trust domain.
        PrincipalPattern string `json:"principal_pattern" yaml:"principal_pattern"`

        // ClientID is the explicit OAuth client_id. It is never derived from a
        // SPIFFE ID.
        ClientID string `json:"client_id" yaml:"client_id"`

        Methods []SPIFFEAuthenticationMethod `json:"methods" yaml:"methods"`

        // Resources are RFC 8707 resource indicators this association may
        // request. Must be a subset of the server's allowed_audiences allowlist
        // (RunConfig.AllowedAudiences) — the same RFC 8707 resource-URI list
        // DelegateClientRunConfig.Audiences is validated against. Distinct from
        // Audiences: a resource permission does not imply the same value is also
        // a permitted token audience, or vice versa.
        Resources []string `json:"resources,omitempty" yaml:"resources,omitempty"`

        // Audiences are RFC 8693 token audiences this association may request.
        // This is an independent request dimension from Resources: it is not
        // bounded by allowed_audiences (which is an RFC 8707 resource-URI list)
        // and may contain non-URI logical audience identifiers.
        Audiences []string `json:"audiences" yaml:"audiences"`

        // Scopes are OAuth scopes granted to this association. They must be a
        // subset of the server's effective supported scopes.
        Scopes []string `json:"scopes" yaml:"scopes"`

        // GrantTypes are the OAuth grant types this association may use. Client
        // authentication does not by itself confer any grant.
        GrantTypes []string `json:"grant_types" yaml:"grant_types"`
}

// SPIFFEAuthorizationPolicy is the immutable authorization policy selected by a
// validated SPIFFE association.
type SPIFFEAuthorizationPolicy struct {
        grantTypes []string
        scopes     []string
        resources  []string
        audiences  []string
}

// GrantTypes returns a copy of the permitted OAuth grant types.
func (p SPIFFEAuthorizationPolicy) GrantTypes() []string { return slices.Clone(p.grantTypes) }

// Scopes returns a copy of the permitted OAuth scopes.
func (p SPIFFEAuthorizationPolicy) Scopes() []string { return slices.Clone(p.scopes) }

// Resources returns a copy of the permitted RFC 8707 resource indicators.
func (p SPIFFEAuthorizationPolicy) Resources() []string { return slices.Clone(p.resources) }

// Audiences returns a copy of the permitted RFC 8693 token audiences.
func (p SPIFFEAuthorizationPolicy) Audiences() []string { return slices.Clone(p.audiences) }

// matchSPIFFEPrincipalPattern reports whether principal matches pattern. A
// terminal /* matches descendants only at a path-segment boundary: /agent/*
// matches /agent/one but not /agent or /agent-two.
func matchSPIFFEPrincipalPattern(pattern, principal string) bool {
        normalizedPattern, err := normalizeSPIFFEPrincipal(pattern, true)
        if err != nil {
                return false
        }
        normalizedPrincipal, err := normalizeSPIFFEPrincipal(principal, false)
        if err != nil {
                return false
        }
        if !strings.HasSuffix(normalizedPattern, "/*") {
                return normalizedPattern == normalizedPrincipal
        }
        return strings.HasPrefix(normalizedPrincipal, strings.TrimSuffix(normalizedPattern, "*"))
}

// NewSPIFFETrustConfig validates and normalizes SPIFFE trust declarations into
// an immutable runtime model. It does not load trust bundles or authenticate
// credentials. The zero value SPIFFETrustConfig{} is also valid and denotes no
// SPIFFE trust domains or associations configured — the same value this
// constructor returns for empty input — so callers do not need to avoid
// constructing it directly.
func NewSPIFFETrustConfig(
        trustDomains []SPIFFETrustDomainRunConfig,
        inboundGrants *InboundGrantsRunConfig,
        scopesSupported []string,
        allowedAudiences []string,
) (*SPIFFETrustConfig, error) {
        trustDomainByName, err := validateSPIFFETrust(trustDomains, inboundGrants, scopesSupported, allowedAudiences)
        if err != nil {
                return nil, err
        }
        associations := inboundAssociations(inboundGrants)
        config := &SPIFFETrustConfig{
                associations: make([]SPIFFEClientAuthConfig, 0, len(associations)),
                trustDomains: make(map[string]SPIFFETrustDomain, len(trustDomainByName)),
        }
        for _, domain := range trustDomains {
                validated := trustDomainByName[domain.Name]
                config.trustDomains[domain.Name] = SPIFFETrustDomain{
                        trustDomain:  validated.trustDomain.String(),
                        methods:      sortedSPIFFEMethods(validated.methods),
                        bundleSource: normalizeSPIFFEBundleSource(domain.BundleSource),
                }
        }
        for _, association := range associations {
                principal, err := normalizeSPIFFEPrincipal(association.PrincipalPattern, true)
                if err != nil {
                        return nil, err
                }
                config.associations = append(config.associations, SPIFFEClientAuthConfig{
                        trustDomainRef: association.TrustDomainRef,
                        principal:      principal,
                        clientID:       association.ClientID,
                        methods:        slices.Clone(association.Methods),
                        authorization: SPIFFEAuthorizationPolicy{
                                grantTypes: slices.Clone(association.GrantTypes),
                                scopes:     slices.Clone(association.Scopes),
                                resources:  slices.Clone(association.Resources),
                                audiences:  slices.Clone(association.Audiences),
                        },
                })
        }
        return config, nil
}

// ValidateSPIFFETrust validates SPIFFE trust declarations without fetching
// bundles or validating credentials.
func ValidateSPIFFETrust(
        trustDomains []SPIFFETrustDomainRunConfig,
        inboundGrants *InboundGrantsRunConfig,
        scopesSupported []string,
        allowedAudiences []string,
) error {
        _, err := validateSPIFFETrust(trustDomains, inboundGrants, scopesSupported, allowedAudiences)
        return err
}

// SPIFFETrustDomain is the immutable normalized form of a declared SPIFFE
// trust domain: its canonical trust domain string, enabled credential
// methods, and bundle-source declaration. Future X.509/JWT-SVID validators
// look this up by declaration name instead of re-parsing the raw RunConfig,
// so the canonical trust domain and enabled methods have exactly one
// authoritative source.
type SPIFFETrustDomain struct {
        trustDomain  string
        methods      []SPIFFEAuthenticationMethod
        bundleSource SPIFFEBundleSourceConfig
}

// TrustDomain returns the canonical SPIFFE trust domain string.
func (d SPIFFETrustDomain) TrustDomain() string { return d.trustDomain }

// Methods returns a copy of the credential methods this trust domain enables.
func (d SPIFFETrustDomain) Methods() []SPIFFEAuthenticationMethod { return slices.Clone(d.methods) }

// BundleSource returns the selected, immutable bundle-source declaration.
func (d SPIFFETrustDomain) BundleSource() SPIFFEBundleSourceConfig { return d.bundleSource }

// SPIFFEBundleSourceConfig is the normalized bundle-source discriminator.
type SPIFFEBundleSourceConfig struct {
        sourceType SPIFFEBundleSourceType
        endpoint   string
        profile    SPIFFEBundleEndpointProfile
        path       string
}

// Type returns the selected bundle-source type.
func (c SPIFFEBundleSourceConfig) Type() SPIFFEBundleSourceType { return c.sourceType }

// Endpoint returns the configured Bundle Endpoint URL, or an empty string for
// a Workload API or file source.
func (c SPIFFEBundleSourceConfig) Endpoint() string { return c.endpoint }

// Profile returns the configured Bundle Endpoint authentication profile, or
// an empty string for a Workload API or file source.
func (c SPIFFEBundleSourceConfig) Profile() SPIFFEBundleEndpointProfile { return c.profile }

// Path returns the configured local bundle file path, or an empty string for
// a Bundle Endpoint or Workload API source.
func (c SPIFFEBundleSourceConfig) Path() string { return c.path }

func normalizeSPIFFEBundleSource(source SPIFFEBundleSourceRunConfig) SPIFFEBundleSourceConfig {
        switch {
        case source.Endpoint != nil:
                return SPIFFEBundleSourceConfig{
                        sourceType: source.Type,
                        endpoint:   source.Endpoint.URL,
                        profile:    source.Endpoint.Profile,
                }
        case source.File != nil:
                return SPIFFEBundleSourceConfig{sourceType: source.Type, path: source.File.Path}
        default:
                return SPIFFEBundleSourceConfig{sourceType: source.Type}
        }
}

// SPIFFETrustConfig is the immutable normalized SPIFFE trust model used at
// runtime. Its zero value is valid and denotes no SPIFFE trust domains or
// associations configured — external packages may safely construct
// SPIFFETrustConfig{} directly.
type SPIFFETrustConfig struct {
        associations []SPIFFEClientAuthConfig
        trustDomains map[string]SPIFFETrustDomain
}

// TrustDomain returns the normalized trust-domain record declared under name,
// and whether a declaration by that name exists.
func (c *SPIFFETrustConfig) TrustDomain(name string) (SPIFFETrustDomain, bool) {
        if c == nil {
                return SPIFFETrustDomain{}, false
        }
        domain, ok := c.trustDomains[name]
        return domain, ok
}

// Associations returns a defensive copy of the normalized association policies.
func (c *SPIFFETrustConfig) Associations() []SPIFFEClientAuthConfig {
        if c == nil {
                return nil
        }
        associations := make([]SPIFFEClientAuthConfig, len(c.associations))
        for i, association := range c.associations {
                associations[i] = association.clone()
        }
        return associations
}

// SPIFFEClientAuthConfig is an immutable normalized association and policy.
type SPIFFEClientAuthConfig struct {
        trustDomainRef string
        principal      string
        clientID       string
        methods        []SPIFFEAuthenticationMethod
        authorization  SPIFFEAuthorizationPolicy
}

// TrustDomainRef returns the configured trust-domain declaration name.
func (c SPIFFEClientAuthConfig) TrustDomainRef() string { return c.trustDomainRef }

// Principal returns the canonical SPIFFE ID or terminal wildcard policy pattern.
func (c SPIFFEClientAuthConfig) Principal() string { return c.principal }

// ClientID returns the configured OAuth client ID.
func (c SPIFFEClientAuthConfig) ClientID() string { return c.clientID }

// Methods returns a copy of permitted authentication methods.
func (c SPIFFEClientAuthConfig) Methods() []SPIFFEAuthenticationMethod {
        return slices.Clone(c.methods)
}

// AuthorizationPolicy returns a defensive copy of the association policy.
func (c SPIFFEClientAuthConfig) AuthorizationPolicy() SPIFFEAuthorizationPolicy {
        return SPIFFEAuthorizationPolicy{
                grantTypes: slices.Clone(c.authorization.grantTypes),
                scopes:     slices.Clone(c.authorization.scopes),
                resources:  slices.Clone(c.authorization.resources),
                audiences:  slices.Clone(c.authorization.audiences),
        }
}

func (c SPIFFEClientAuthConfig) clone() SPIFFEClientAuthConfig {
        c.methods = slices.Clone(c.methods)
        c.authorization = c.AuthorizationPolicy()
        return c
}

// inboundAssociations returns the configured SPIFFE client-auth entries, or
// nil when inboundGrants is nil.
func inboundAssociations(inboundGrants *InboundGrantsRunConfig) []SPIFFEClientAuthRunConfig {
        if inboundGrants == nil {
                return nil
        }
        return inboundGrants.SPIFFEClientAuth
}

// validateSPIFFETrust validates SPIFFE trust declarations and their inbound
// associations without fetching bundles or validating credentials. It fails
// closed on ambiguity because later authentication must not depend on
// configuration order. It returns the validated trust-domain records so
// NewSPIFFETrustConfig can retain them without re-parsing.
func validateSPIFFETrust(
        trustDomains []SPIFFETrustDomainRunConfig,
        inboundGrants *InboundGrantsRunConfig,
        scopesSupported []string,
        allowedResources []string,
) (map[string]validatedSPIFFETrustDomain, error) {
        trustDomainByName, err := validateSPIFFETrustDomains(trustDomains)
        if err != nil {
                return nil, err
        }
        associations := inboundAssociations(inboundGrants)
        if len(associations) == 0 {
                if len(trustDomains) != 0 {
                        return nil, fmt.Errorf("inbound_grants.spiffe_client_auth is required when spiffe_trust_domains is configured")
                }
                return trustDomainByName, nil
        }
        if err := validateSPIFFEClientAuth(
                associations, trustDomainByName, scopesSupported, allowedResources,
        ); err != nil {
                return nil, err
        }
        return trustDomainByName, nil
}

// sortedSPIFFEMethods returns the enabled methods in a stable order, since
// methods is stored as a set during validation.
func sortedSPIFFEMethods(methods map[SPIFFEAuthenticationMethod]struct{}) []SPIFFEAuthenticationMethod {
        sorted := make([]SPIFFEAuthenticationMethod, 0, len(methods))
        for _, method := range []SPIFFEAuthenticationMethod{SPIFFEAuthenticationMethodX509, SPIFFEAuthenticationMethodJWT} {
                if _, ok := methods[method]; ok {
                        sorted = append(sorted, method)
                }
        }
        return sorted
}

type validatedSPIFFETrustDomain struct {
        trustDomain spiffeid.TrustDomain
        methods     map[SPIFFEAuthenticationMethod]struct{}
}

func validateSPIFFETrustDomains(domains []SPIFFETrustDomainRunConfig) (map[string]validatedSPIFFETrustDomain, error) {
        byName := make(map[string]validatedSPIFFETrustDomain, len(domains))
        byTrustDomain := make(map[string]string, len(domains))
        for i, domain := range domains {
                if domain.Name == "" {
                        return nil, fmt.Errorf("spiffe_trust_domains[%d]: name is required", i)
                }
                if _, exists := byName[domain.Name]; exists {
                        return nil, fmt.Errorf("spiffe_trust_domains[%d]: duplicate name %q", i, domain.Name)
                }
                trustDomain, err := parseTrustDomain(domain.TrustDomain)
                if err != nil {
                        return nil, fmt.Errorf("spiffe_trust_domains[%d]: %w", i, err)
                }
                canonicalTrustDomain := trustDomain.String()
                if existingName, exists := byTrustDomain[canonicalTrustDomain]; exists {
                        return nil, fmt.Errorf(
                                "spiffe_trust_domains[%d]: duplicate trust domain %q already declared as %q",
                                i, canonicalTrustDomain, existingName,
                        )
                }
                methods, err := validateMethods(domain.Methods, fmt.Sprintf("spiffe_trust_domains[%d].methods", i))
                if err != nil {
                        return nil, err
                }
                if err := validateSPIFFEBundleSource(domain.BundleSource, i); err != nil {
                        return nil, err
                }
                byName[domain.Name] = validatedSPIFFETrustDomain{trustDomain: trustDomain, methods: methods}
                byTrustDomain[canonicalTrustDomain] = domain.Name
        }
        return byName, nil
}

func validateSPIFFEBundleSource(source SPIFFEBundleSourceRunConfig, index int) error {
        switch source.Type {
        case SPIFFEBundleSourceTypeEndpoint:
                if source.Endpoint == nil || source.File != nil || source.WorkloadAPI != nil {
                        return fmt.Errorf("spiffe_trust_domains[%d].bundle_source: type %q requires endpoint only", index, source.Type)
                }
                return validateSPIFFEBundleEndpoint(*source.Endpoint, index)
        case SPIFFEBundleSourceTypeFile:
                if source.File == nil || source.Endpoint != nil || source.WorkloadAPI != nil {
                        return fmt.Errorf("spiffe_trust_domains[%d].bundle_source: type %q requires file only", index, source.Type)
                }
                return validateSPIFFEBundleFilePath(source.File.Path, index)
        case SPIFFEBundleSourceTypeWorkloadAPI:
                if source.WorkloadAPI == nil || source.Endpoint != nil || source.File != nil {
                        return fmt.Errorf("spiffe_trust_domains[%d].bundle_source: type %q requires workload_api only", index, source.Type)
                }
                return nil
        default:
                return fmt.Errorf("spiffe_trust_domains[%d].bundle_source.type: unknown source type %q", index, source.Type)
        }
}

// validateSPIFFEBundleFilePath validates the context-independent shape of a
// local trust-bundle file path: it must be a non-empty absolute path with no
// traversal segments. The isolated file loader checks that it resolves to a
// regular file and can parse a bundle when it is constructed.
func validateSPIFFEBundleFilePath(path string, index int) error {
        if path == "" || !filepath.IsAbs(path) {
                return fmt.Errorf("spiffe_trust_domains[%d].bundle_source.file.path must be a non-empty absolute path", index)
        }
        for _, segment := range strings.Split(filepath.ToSlash(path), "/") {
                if segment == ".." {
                        return fmt.Errorf("spiffe_trust_domains[%d].bundle_source.file.path must not contain traversal segments", index)
                }
        }
        return nil
}

func validateSPIFFEBundleEndpoint(endpoint SPIFFEBundleEndpointSourceRunConfig, index int) error {
        if err := ValidateSPIFFEBundleEndpoint(endpoint); err != nil {
                return fmt.Errorf("spiffe_trust_domains[%d].bundle_source.endpoint: %w", index, err)
        }
        return nil
}

// ValidateSPIFFEBundleEndpoint validates a SPIFFE Bundle Endpoint URL and its
// TLS-authentication profile. Exported so the operator CRD admission path
// can reject the same structurally invalid endpoints at admission time
// instead of only at reconcile time; fetching or loading a bundle from the
// endpoint remains a separate, later step.
func ValidateSPIFFEBundleEndpoint(endpoint SPIFFEBundleEndpointSourceRunConfig) error {
        endpointURL := endpoint.URL
        u, err := url.ParseRequestURI(endpointURL)
        if err != nil || u.Scheme != "https" || u.Host == "" || u.Hostname() == "" {
                return fmt.Errorf("url must be an absolute HTTPS URL with a valid authority")
        }
        if u.User != nil || u.RawQuery != "" || u.Fragment != "" ||
                strings.Contains(endpointURL, "?") || strings.Contains(endpointURL, "#") ||
                net.ParseIP(u.Hostname()) != nil || networking.IsLoopbackHost(u.Hostname()) {
                return fmt.Errorf(
                        "url must not contain credentials, query, fragment, an IP-literal host, or a loopback host",
                )
        }
        switch endpoint.Profile {
        case SPIFFEBundleEndpointProfileHTTPSWeb, SPIFFEBundleEndpointProfileHTTPSSPIFFE:
                return nil
        default:
                return fmt.Errorf(
                        "profile must be %q or %q", SPIFFEBundleEndpointProfileHTTPSWeb, SPIFFEBundleEndpointProfileHTTPSSPIFFE,
                )
        }
}

func validateSPIFFEClientAuth(
        entries []SPIFFEClientAuthRunConfig,
        trustDomainByName map[string]validatedSPIFFETrustDomain,
        scopesSupported []string,
        allowedResources []string,
) error {
        effectiveScopes := func() []string {
                if len(scopesSupported) == 0 {
                        return registration.DefaultScopes
                }
                return scopesSupported
        }()
        seenPrincipals := make([]string, 0, len(entries))
        clientIDs := make(map[string]struct{}, len(entries))
        referencedTrustDomains := make(map[string]struct{}, len(entries))
        for i, entry := range entries {
                trustDomain, ok := trustDomainByName[entry.TrustDomainRef]
                if entry.TrustDomainRef == "" {
                        return fmt.Errorf("inbound_grants.spiffe_client_auth[%d]: trust_domain_ref is required", i)
                }
                if !ok {
                        return fmt.Errorf("inbound_grants.spiffe_client_auth[%d]: unknown trust_domain_ref %q", i, entry.TrustDomainRef)
                }
                principal, err := validateSPIFFEClientAssociation(entry, i, trustDomain, effectiveScopes, allowedResources)
                if err != nil {
                        return err
                }
                if _, exists := clientIDs[entry.ClientID]; exists {
                        return fmt.Errorf("inbound_grants.spiffe_client_auth[%d]: duplicate client_id %q", i, entry.ClientID)
                }
                if overlapsSPIFFEPatterns(seenPrincipals, principal) {
                        return fmt.Errorf(
                                "inbound_grants.spiffe_client_auth[%d]: "+
                                        "principal_pattern %q overlaps an existing policy",
                                i,
                                entry.PrincipalPattern,
                        )
                }
                seenPrincipals = append(seenPrincipals, principal)
                clientIDs[entry.ClientID] = struct{}{}
                referencedTrustDomains[entry.TrustDomainRef] = struct{}{}
        }
        for name := range trustDomainByName {
                if _, ok := referencedTrustDomains[name]; !ok {
                        return fmt.Errorf("spiffe_trust_domains: trust domain %q is not referenced by spiffe_client_auth", name)
                }
        }
        return nil
}

func validateSPIFFEClientAssociation(
        entry SPIFFEClientAuthRunConfig,
        index int,
        trustDomain validatedSPIFFETrustDomain,
        effectiveScopes []string,
        allowedResources []string,
) (string, error) {
        principal, err := normalizeSPIFFEPrincipal(entry.PrincipalPattern, true)
        if err != nil {
                return "", fmt.Errorf("inbound_grants.spiffe_client_auth[%d].principal_pattern: %w", index, err)
        }
        principalID, err := parseSPIFFEID(strings.TrimSuffix(principal, "/*"))
        if err != nil {
                return "", fmt.Errorf("inbound_grants.spiffe_client_auth[%d].principal_pattern: %w", index, err)
        }
        if !principalID.MemberOf(trustDomain.trustDomain) {
                return "", fmt.Errorf(
                        "inbound_grants.spiffe_client_auth[%d]: principal trust domain does not match %q",
                        index,
                        trustDomain.trustDomain,
                )
        }
        if entry.ClientID == "" {
                return "", fmt.Errorf("inbound_grants.spiffe_client_auth[%d]: client_id is required", index)
        }
        if err := storage.ValidateRegisterableClientID(entry.ClientID); err != nil {
                return "", fmt.Errorf("inbound_grants.spiffe_client_auth[%d]: client_id: %w", index, err)
        }
        if oauthproto.IsClientIDMetadataDocumentURL(entry.ClientID) {
                return "", fmt.Errorf(
                        "inbound_grants.spiffe_client_auth[%d]: client_id must not be a client metadata document URL "+
                                "(reserved for CIMD-resolved clients): %q",
                        index, entry.ClientID,
                )
        }
        fieldPrefix := fmt.Sprintf("inbound_grants.spiffe_client_auth[%d]", index)
        methods, err := validateMethods(entry.Methods, fieldPrefix+".methods")
        if err != nil {
                return "", err
        }
        for method := range methods {
                if _, enabled := trustDomain.methods[method]; !enabled {
                        return "", fmt.Errorf("%s.methods: method %q is not enabled by trust domain %q", fieldPrefix, method, entry.TrustDomainRef)
                }
        }
        if err := validateSPIFFEClientAssociationPermissions(entry, index, fieldPrefix, effectiveScopes, allowedResources); err != nil {
                return "", err
        }
        return principal, nil
}

// validateSPIFFEClientAssociationPermissions validates the resources,
// audiences, scopes, and grant types an association is permitted to request.
// Resources (RFC 8707) and audiences (RFC 8693) are independent request
// dimensions: permission in one must never imply permission in the other, so
// only resources are bounded by the server's allowed_audiences allowlist.
func validateSPIFFEClientAssociationPermissions(
        entry SPIFFEClientAuthRunConfig,
        index int,
        fieldPrefix string,
        effectiveScopes []string,
        allowedResources []string,
) error {
        if err := validateSPIFFEResources(entry.Resources, fieldPrefix+".resources", allowedResources); err != nil {
                return err
        }
        if err := validateDistinctNonEmpty(entry.Audiences, fieldPrefix+".audiences"); err != nil {
                return err
        }
        if err := validateDistinctNonEmpty(entry.Scopes, fieldPrefix+".scopes"); err != nil {
                return err
        }
        if err := registration.ValidateScopeSubset(entry.Scopes, effectiveScopes, fieldPrefix+".scopes"); err != nil {
                return err
        }
        return validateSPIFFEGrants(entry.GrantTypes, index)
}

func parseTrustDomain(trustDomain string) (spiffeid.TrustDomain, error) {
        if trustDomain == "" {
                return spiffeid.TrustDomain{}, fmt.Errorf("trust_domain is required")
        }
        parsed, err := spiffeid.TrustDomainFromString(trustDomain)
        if err != nil {
                return spiffeid.TrustDomain{}, fmt.Errorf("must be a valid SPIFFE trust domain: %w", err)
        }
        return parsed, nil
}

func normalizeSPIFFEPrincipal(principal string, allowPattern bool) (string, error) {
        if principal == "" {
                return "", fmt.Errorf("is required")
        }
        wildcard := strings.HasSuffix(principal, "/*")
        if strings.Contains(principal, "*") && (!allowPattern || !wildcard || strings.Count(principal, "*") != 1) {
                return "", fmt.Errorf("must use only a terminal /* wildcard")
        }
        base := strings.TrimSuffix(principal, "/*")
        id, err := parseSPIFFEID(base)
        if err != nil {
                return "", err
        }
        if id.Path() == "" && !wildcard {
                return "", fmt.Errorf("must include a path or terminal /* wildcard")
        }
        if wildcard {
                return id.String() + "/*", nil
        }
        return id.String(), nil
}

func parseSPIFFEID(principal string) (spiffeid.ID, error) {
        id, err := spiffeid.FromString(principal)
        if err != nil {
                return spiffeid.ID{}, fmt.Errorf("must be a SPIFFE ID: %w", err)
        }
        return id, nil
}

func validateMethods(methods []SPIFFEAuthenticationMethod, field string) (map[SPIFFEAuthenticationMethod]struct{}, error) {
        if len(methods) == 0 {
                return nil, fmt.Errorf("%s is required", field)
        }
        seen := make(map[SPIFFEAuthenticationMethod]struct{}, len(methods))
        for _, method := range methods {
                if method != SPIFFEAuthenticationMethodX509 && method != SPIFFEAuthenticationMethodJWT {
                        return nil, fmt.Errorf("%s: unknown method %q", field, method)
                }
                if _, exists := seen[method]; exists {
                        return nil, fmt.Errorf("%s: duplicate method %q", field, method)
                }
                seen[method] = struct{}{}
        }
        return seen, nil
}

// validateSPIFFEResources validates the RFC 8707 resource indicators an
// association may request: each must be a syntactically valid absolute
// HTTP(S) URI and a member of the server's allowed_audiences allowlist, the
// same RFC 8707 resource-URI list DelegateClientRunConfig.Audiences is
// validated against (see validateDelegateClients in config.go).
func validateSPIFFEResources(values []string, field string, allowedResources []string) error {
        if len(values) == 0 {
                return nil
        }
        if err := ValidateResourceIndicators(values, field); err != nil {
                return err
        }
        for _, value := range values {
                if !slices.Contains(allowedResources, value) {
                        return fmt.Errorf("%s: resource %q is not allowed by allowed_audiences", field, value)
                }
        }
        return nil
}

// ValidateResourceIndicators validates the context-independent shape of RFC
// 8707 resource indicators: each must be a syntactically valid absolute
// HTTP(S) URI with a non-empty host, no userinfo, and no fragment. It does
// not check allowlist membership, which needs the reconcile-time-derived
// allowed_audiences value (see validateSPIFFEResources and
// MCPExternalAuthConfig's own admission-time call for this shape-only half).
func ValidateResourceIndicators(values []string, field string) error {
        if len(values) == 0 {
                return nil
        }
        if err := validateDistinctNonEmpty(values, field); err != nil {
                return err
        }
        for _, value := range values {
                u, err := url.ParseRequestURI(value)
                if err != nil || (u.Scheme != "https" && u.Scheme != "http") || u.Hostname() == "" || u.User != nil ||
                        u.Fragment != "" || strings.Contains(value, "#") {
                        return fmt.Errorf("%s: resource indicator %q must be an absolute HTTP(S) URI without a fragment", field, value)
                }
        }
        return nil
}

func validateDistinctNonEmpty(values []string, field string) error {
        if len(values) == 0 {
                return fmt.Errorf("%s is required", field)
        }
        seen := make(map[string]struct{}, len(values))
        for _, value := range values {
                if value == "" {
                        return fmt.Errorf("%s must not contain an empty value", field)
                }
                if _, exists := seen[value]; exists {
                        return fmt.Errorf("%s: duplicate value %q", field, value)
                }
                seen[value] = struct{}{}
        }
        return nil
}

func validateSPIFFEGrants(grants []string, index int) error {
        if len(grants) != 1 || grants[0] != SPIFFEGrantTypeTokenExchange {
                return fmt.Errorf(
                        "inbound_grants.spiffe_client_auth[%d].grant_types must be exactly [%q]",
                        index, SPIFFEGrantTypeTokenExchange,
                )
        }
        return nil
}

func overlapsSPIFFEPatterns(patterns []string, principal string) bool {
        for _, pattern := range patterns {
                if spiffePatternsOverlap(pattern, principal) {
                        return true
                }
        }
        return false
}

// ValidateSPIFFEPrincipalPattern validates a single principal pattern — a
// concrete SPIFFE ID or a terminal /* wildcard — using the same
// normalization as the runtime association validator, without comparing it
// to any other pattern. Exported so an admission-time caller can validate
// each entry in a set up front and attribute a normalization failure to the
// correct entry, before running SPIFFEPatternsOverlap pairwise across the set.
func ValidateSPIFFEPrincipalPattern(pattern string) error {
        _, err := normalizeSPIFFEPrincipal(pattern, true)
        return err
}

// SPIFFEPatternsOverlap reports whether two principal patterns — each a
// concrete SPIFFE ID or a terminal /* wildcard — overlap, using the same
// normalization (via the go-spiffe parser) and prefix semantics as the
// runtime association validator. Exported so the operator CRD admission
// path can reject overlapping principal patterns (e.g. "/agent/*" and
// "/agent/one") at admission time instead of only at reconcile time.
func SPIFFEPatternsOverlap(first, second string) (bool, error) {
        normalizedFirst, err := normalizeSPIFFEPrincipal(first, true)
        if err != nil {
                return false, err
        }
        normalizedSecond, err := normalizeSPIFFEPrincipal(second, true)
        if err != nil {
                return false, err
        }
        return spiffePatternsOverlap(normalizedFirst, normalizedSecond), nil
}

func spiffePatternsOverlap(first, second string) bool {
        firstWildcard := strings.HasSuffix(first, "/*")
        secondWildcard := strings.HasSuffix(second, "/*")
        if !firstWildcard && !secondWildcard {
                return first == second
        }
        if firstWildcard && secondWildcard {
                return strings.HasPrefix(strings.TrimSuffix(first, "*"), strings.TrimSuffix(second, "*")) ||
                        strings.HasPrefix(strings.TrimSuffix(second, "*"), strings.TrimSuffix(first, "*"))
        }
        if firstWildcard {
                return matchSPIFFEPrincipalPattern(first, second)
        }
        return matchSPIFFEPrincipalPattern(second, first)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "slices"
        "strings"
        "time"

        lru "github.com/hashicorp/golang-lru/v2"
        "github.com/ory/fosite"
        "golang.org/x/sync/singleflight"
        "golang.org/x/time/rate"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/oauthproto"
        "github.com/stacklok/toolhive/pkg/oauthproto/cimd"
)

// ErrCIMDRateLimited is wrapped by the fosite.ErrNotFound error that GetClient
// returns when resolving a CIMD client_id would exceed the decorator's outbound
// fetch budget and no cached entry within MaxStale of its expiry can be served
// instead.
// On a request whose context carries WithPersistedClientFallback (the token
// endpoint), it is also not returned when a persisted client row that still
// matches the configuration exists; that row is served instead.
// Handlers map it to HTTP 429 instead of invalid_client.
var ErrCIMDRateLimited = errors.New("CIMD document fetch rate limit exceeded")

// CIMDStorageDecorator wraps storage.Storage and intercepts GetClient calls
// for HTTPS client_id values, fetching and caching the corresponding Client
// ID Metadata Document instead of requiring prior DCR registration.
//
// All other Storage methods delegate to the underlying storage unchanged,
// except RenewClientTTL (CIMD ids are skipped), and ConsumeAssertionJWT and
// Unwrap, which forward to the wrapped backend. GetClient is the main
// override. DCR clients (opaque IDs) continue to work exactly as before.
//
// Every route that resolves a CIMD client_id (authorize, token, device
// authorization) reaches the outbound fetch and write-through persistence
// through GetClient, so fetchLimiter bounds both for all of them. When the
// limiter refuses a fetch, an expired cache entry is served for up to maxStale
// past its expiry so a known client is not locked out by traffic that drains
// the shared budget. On the token endpoint
// (WithPersistedClientFallback), a refused lookup with no servable cache entry
// is answered from the persisted client row when one exists.
type CIMDStorageDecorator struct {
        Storage                                 // embed full interface — all methods delegate
        sf                   singleflight.Group // deduplicates concurrent fetches for the same URL
        cache                *lru.Cache[string, *cimdCacheEntry]
        ttl                  time.Duration
        scopesSupported      []string // AS-configured scopes; nil means accept any
        baselineClientScopes []string // unioned into every client's scope set, same as DCR
        allowedAudiences     []string // AS-configured audiences granted to every CIMD client
        // maxStale is how long past expiry a cached entry may be served when
        // fetchLimiter refuses a fetch; zero disables stale serving.
        maxStale time.Duration
        // fetchLimiter bounds outbound document fetches (and therefore persisted
        // rows) triggered by unauthenticated callers. Per-process, so N replicas
        // allow N times the rate, like the handlers' registerLimiter. A refusal
        // returns ErrCIMDRateLimited unless an entry within maxStale is cached or,
        // for a WithPersistedClientFallback context, a matching persisted row exists.
        fetchLimiter *rate.Limiter
}

type cimdCacheEntry struct {
        client  fosite.Client
        expires time.Time
}

// CIMDDecoratorConfig holds the configuration for NewCIMDStorageDecorator.
// Using a struct prevents silent swaps of the three adjacent []string fields.
type CIMDDecoratorConfig struct {
        // Enabled returns base unchanged when false, avoiding an allocation.
        Enabled bool
        // CacheMaxSize is the maximum number of documents in the LRU cache (must be >= 1).
        CacheMaxSize int
        // FallbackTTL is the fixed TTL applied to every cache entry.
        FallbackTTL time.Duration
        // MaxStale is how long past its expiry a cached entry may still be served
        // when the fetch rate limiter refuses a re-fetch. Zero disables stale
        // serving; negative values are rejected.
        MaxStale time.Duration
        // ScopesSupported is the AS scope allowlist; see pkg/authserver/config.go
        // applyDefaults for production guarantees. Pass nil in tests only.
        ScopesSupported []string
        // BaselineClientScopes is unioned into every CIMD client's scope set,
        // matching DCR handler behaviour.
        BaselineClientScopes []string
        // AllowedAudiences is the server's own allowed-audience list, granted to
        // every CIMD client exactly as registration.New() grants it to DCR clients
        // (see #3796). CIMD documents don't declare audience, so without this a
        // CIMD client's fosite.Client.Audience stays empty and every refresh_token
        // grant is rejected by fosite's DefaultAudienceMatchingStrategy, which
        // matches the granted audience against the client's own audience list.
        AllowedAudiences []string
}

// NewCIMDStorageDecorator wraps base with CIMD client lookup.
// When cfg.Enabled is false it returns base unchanged (no allocation).
func NewCIMDStorageDecorator(base Storage, cfg CIMDDecoratorConfig) (Storage, error) {
        if !cfg.Enabled {
                return base, nil
        }

        if cfg.CacheMaxSize < 1 {
                return nil, fmt.Errorf("CIMD storage decorator cacheMaxSize must be >= 1, got %d", cfg.CacheMaxSize)
        }

        if cfg.MaxStale < 0 {
                return nil, fmt.Errorf("CIMD storage decorator maxStale must be >= 0, got %s", cfg.MaxStale)
        }

        c, err := lru.New[string, *cimdCacheEntry](cfg.CacheMaxSize)
        if err != nil {
                return nil, fmt.Errorf("failed to create CIMD LRU cache: %w", err)
        }

        return &CIMDStorageDecorator{
                Storage:              base,
                cache:                c,
                ttl:                  cfg.FallbackTTL,
                maxStale:             cfg.MaxStale,
                scopesSupported:      slices.Clone(cfg.ScopesSupported),
                baselineClientScopes: slices.Clone(cfg.BaselineClientScopes),
                allowedAudiences:     slices.Clone(cfg.AllowedAudiences),
                // Same rate as the handlers' registerLimiter: both gate the same kind
                // of unauthenticated persisted-state minting.
                fetchLimiter: rate.NewLimiter(rate.Limit(1), 5),
        }, nil
}

type persistedClientFallbackKey struct{}

// WithPersistedClientFallback marks ctx as belonging to the token endpoint. It
// must not be used anywhere else.
//
// fosite calls Storage.GetClient itself inside NewAccessRequest, so a context
// value is the only seam the handler has. The mark permits
// CIMDStorageDecorator.GetClient, when the fetch limiter refused and no stale
// cache entry is servable, to answer from the persisted client row instead of
// returning ErrCIMDRateLimited. It applies to any URL-shaped GetClient made
// with the returned context. The fallback is read-only and serves a snapshot
// that was validated when it was last fetched; it is not re-validated against
// the document.
func WithPersistedClientFallback(ctx context.Context) context.Context {
        return context.WithValue(ctx, persistedClientFallbackKey{}, true)
}

// PersistedClientFallbackAllowed reports whether ctx carries the mark set by
// WithPersistedClientFallback. The decorator reads it, and tests use it to
// check which endpoints set it.
func PersistedClientFallbackAllowed(ctx context.Context) bool {
        allowed, _ := ctx.Value(persistedClientFallbackKey{}).(bool)
        return allowed
}

// GetClient intercepts HTTPS client_id values to resolve them via CIMD.
// Opaque DCR-issued IDs are delegated to the underlying storage unchanged.
func (d *CIMDStorageDecorator) GetClient(ctx context.Context, id string) (fosite.Client, error) {
        if !oauthproto.IsClientIDMetadataDocumentURL(id) {
                return d.Storage.GetClient(ctx, id)
        }
        client, err := d.fetchOrCached(ctx, id)
        // Outside singleflight and with the caller's ctx on purpose: the mark is
        // per request and must not leak to other callers sharing a flight. The
        // row is deliberately neither cached nor renewed (a cached row would be
        // served to /oauth/authorize as a fresh hit), and every failure returns
        // the original rate-limit error, unlike handlers' callbackClient which
        // returns the base lookup's error.
        if errors.Is(err, ErrCIMDRateLimited) && PersistedClientFallbackAllowed(ctx) {
                row, rowErr := d.Storage.GetClient(ctx, id)
                if rowErr == nil && row != nil && d.persistedRowMatchesConfig(id, row) {
                        slog.DebugContext(ctx, "CIMD document fetch rate limited; serving persisted client row", "client_id", id)
                        return row, nil
                }
        }
        return client, err
}

// persistedRowMatchesConfig reports whether the persisted row for id is a
// snapshot this decorator wrote and still agrees with the current AS
// configuration, failing closed. Only the CIMD fetch writes rows at URL-shaped
// ids, so the row must carry that id, be a public client and carry the
// DCR-issued marker; anything else at such an id was not validated by us.
// Every grant type on the row must survive registration.FilterPublicGrantTypes
// unchanged, i.e. be one the fetch path can still produce (a row may hold a
// subset). That filter is static: the AS configuration cannot narrow it.
//
// A restart is the usual trigger for the persisted-row fallback and the
// operator may have changed AllowedAudiences or ScopesSupported before
// restarting; the row would carry the old values. The audience must equal
// allowedAudiences (order-insensitive, which is all buildFositeClient puts
// there) and, when scopesSupported is configured, every row scope must be in
// scopesSupported or baselineClientScopes. It never mutates the row.
func (d *CIMDStorageDecorator) persistedRowMatchesConfig(id string, client fosite.Client) bool {
        if client.GetID() != id || !client.IsPublic() || !registration.DCRIssued(client) {
                return false
        }
        rowGrants := client.GetGrantTypes()
        allowed, dcrErr := registration.FilterPublicGrantTypes(rowGrants)
        if dcrErr != nil || len(allowed) != len(rowGrants) {
                return false
        }
        rowAud := slices.Sorted(slices.Values(client.GetAudience()))
        wantAud := slices.Sorted(slices.Values(d.allowedAudiences))
        if !slices.Equal(rowAud, wantAud) {
                return false
        }
        if len(d.scopesSupported) == 0 {
                return true
        }
        for _, scope := range client.GetScopes() {
                if !slices.Contains(d.scopesSupported, scope) && !slices.Contains(d.baselineClientScopes, scope) {
                        return false
                }
        }
        return true
}

// RenewClientTTL skips URL-shaped (CIMD) client ids and delegates every other
// id to the wrapped storage. Renewal is redundant for CIMD rows because every
// successful fetch upserts the row and resets its TTL. Not renewing keeps the
// row TTL equal to the time since the last successful fetch, which is the age
// bound for the persisted-row fallback on the token endpoint on both backends.
// Opaque DCR and static ids are renewed as before.
func (d *CIMDStorageDecorator) RenewClientTTL(ctx context.Context, client fosite.Client) error {
        if client != nil && oauthproto.IsClientIDMetadataDocumentURL(client.GetID()) {
                return nil
        }
        return d.Storage.RenewClientTTL(ctx, client)
}

// cimdShapeGuardStorage wraps storage.Storage and rejects URL-shaped
// client_id values at GetClient without ever reading the underlying row.
//
// It exists for the case where CIMD is (or was previously) enabled and then
// disabled: the write-through persistence in CIMDStorageDecorator.fetch may
// have already left a resolved CIMD client's row in the underlying storage,
// and that row's TTL — up to DefaultDCRClientTTL — can outlive the config
// change. Without this guard, a bare Storage's GetClient does not check the
// shape of the id it is asked for, so it would happily resolve the stale row
// from a snapshot the document may have since changed or that the operator
// intended to revoke by turning CIMD off, with no re-fetch and therefore no
// re-validation of scopes, redirect URIs, or auth method. Wrapping the
// storage when CIMD is disabled closes that gap for both backends without
// changing GetClient's behavior for any opaque (non-URL) client_id.
type cimdShapeGuardStorage struct {
        Storage // embed full interface — all methods delegate
}

// GetClient rejects URL-shaped ids with fosite.ErrNotFound before consulting
// the wrapped storage, so a stale CIMD-origin row can never be resolved while
// CIMD is disabled. Opaque DCR-issued and pre-provisioned ids are delegated
// to the underlying storage unchanged.
func (g cimdShapeGuardStorage) GetClient(ctx context.Context, id string) (fosite.Client, error) {
        if oauthproto.IsClientIDMetadataDocumentURL(id) {
                return nil, fosite.ErrNotFound.WithHint("CIMD is disabled on this server")
        }
        return g.Storage.GetClient(ctx, id)
}

// ConsumeAssertionJWT delegates assertion replay consumption to the wrapped
// backend, matching CIMDStorageDecorator's identical method: Storage embeds
// the interface, not the concrete backend, so this narrow capability is not
// promoted automatically and must be forwarded explicitly for callers (e.g.
// the RFC 7523 JWT-bearer grant factory) that need it through this wrapper.
func (g cimdShapeGuardStorage) ConsumeAssertionJWT(
        ctx context.Context, purpose, issuer, jti string, exp time.Time,
) error {
        consumer, ok := g.Storage.(AssertionJWTConsumer)
        if !ok {
                return fmt.Errorf("wrapped storage does not support assertion JWT replay consumption")
        }
        return consumer.ConsumeAssertionJWT(ctx, purpose, issuer, jti, exp)
}

// Unwrap returns the underlying storage so that type assertions (e.g. for
// storage.DCRCredentialStore in server_impl.go) can reach the concrete type.
func (g cimdShapeGuardStorage) Unwrap() Storage {
        return g.Storage
}

// NewCIMDShapeGuardStorage wraps base so that GetClient refuses any
// URL-shaped client_id, regardless of whether a matching row exists in base.
// Callers apply this when CIMD is disabled — see cimdShapeGuardStorage's doc
// comment for why a disabled decorator alone does not close this gap.
func NewCIMDShapeGuardStorage(base Storage) Storage {
        return cimdShapeGuardStorage{Storage: base}
}

// ConsumeAssertionJWT delegates assertion replay consumption to the wrapped
// backend. Storage intentionally does not include this narrow capability, so a
// decorator fails closed when its backend does not provide it.
func (d *CIMDStorageDecorator) ConsumeAssertionJWT(
        ctx context.Context, purpose, issuer, jti string, exp time.Time,
) error {
        consumer, ok := d.Storage.(AssertionJWTConsumer)
        if !ok {
                return fmt.Errorf("wrapped storage does not support assertion JWT replay consumption")
        }
        return consumer.ConsumeAssertionJWT(ctx, purpose, issuer, jti, exp)
}

// Unwrap returns the underlying storage so that type assertions (e.g. for
// storage.DCRCredentialStore in server_impl.go) can reach the concrete type.
func (d *CIMDStorageDecorator) Unwrap() Storage {
        return d.Storage
}

func (d *CIMDStorageDecorator) fetchOrCached(ctx context.Context, id string) (fosite.Client, error) {
        // Check cache first (outside singleflight to avoid holding the group lock for cache hits)
        if entry, ok := d.cache.Get(id); ok && time.Now().Before(entry.expires) {
                return entry.client, nil
        }

        // Deduplicate concurrent fetches for the same URL. The shared fetch uses a
        // context detached from the caller so that one caller cancelling does not
        // abort the in-flight request for other waiters. The HTTP client inside
        // FetchClientMetadataDocument enforces its own 5-second timeout. The detached
        // context still carries context values such as the persisted-fallback mark;
        // nothing inside fetch reads it, and fetch must not call this decorator's
        // GetClient.
        fetchCtx := context.WithoutCancel(ctx)
        result, err, _ := d.sf.Do(id, func() (interface{}, error) {
                // Re-check cache inside singleflight (another goroutine may have populated
                // it). The entry is reused below for the stale fallback.
                entry, ok := d.cache.Get(id)
                if ok && time.Now().Before(entry.expires) {
                        return entry.client, nil
                }
                // Allow, not Wait: fetchCtx is detached from the caller, so a waiting
                // call could never be cancelled. Cache hits and callers sharing this
                // singleflight call cost nothing; a failed fetch costs a token.
                // The stale check sits after the limiter so that with budget left an
                // expired entry is still re-fetched and re-validated; only a refusal
                // falls back to the expired entry, and only within maxStale. A fetch
                // that is attempted and fails returns its error but does not evict the
                // entry, so while the limiter stays empty the pre-failure entry can be
                // served again until expires+maxStale.
                if !d.fetchLimiter.Allow() {
                        // maxStale > 0 makes the zero value explicit: stale serving disabled.
                        if ok && d.maxStale > 0 && time.Now().Before(entry.expires.Add(d.maxStale)) {
                                // DEBUG: callers can trigger this at request rate.
                                slog.DebugContext(ctx, "CIMD document fetch rate limited; serving stale cache entry",
                                        "client_id", id, "age_past_expiry", time.Since(entry.expires))
                                return entry.client, nil
                        }
                        slog.DebugContext(ctx, "CIMD document fetch rate limited", "client_id", id)
                        return nil, fosite.ErrNotFound.WithWrap(ErrCIMDRateLimited)
                }
                return d.fetch(fetchCtx, id)
        })
        if err != nil {
                return nil, err
        }
        client, ok := result.(fosite.Client)
        if !ok {
                return nil, fmt.Errorf("CIMD singleflight returned unexpected type %T", result)
        }
        return client, nil
}

func (d *CIMDStorageDecorator) fetch(ctx context.Context, id string) (fosite.Client, error) {
        doc, err := cimd.FetchClientMetadataDocument(ctx, id)
        if err != nil {
                // Log every rejection in fetch() at Warn: the errors returned here
                // surface to the client as a generic invalid_client/invalid_request
                // with the hint dropped by fosite's production error rendering, so
                // without a server-side log line these failures are undiagnosable
                // (see issue #6186).
                slog.WarnContext(ctx, "CIMD document fetch failed",
                        "client_id", id, "error", err)
                return nil, fosite.ErrNotFound.WithHint("CIMD fetch failed").WithWrap(err)
        }

        // Negotiate the effective token_endpoint_auth_method rather than rejecting
        // outright on a declared-but-unsupported singular value. ErrInvalidClient:
        // the document was fetched successfully but its declared metadata violates
        // AS policy (distinct from ErrNotFound which means the document could not
        // be fetched at all).
        authMethod, ok := negotiateTokenEndpointAuthMethod(doc)
        if !ok {
                slog.WarnContext(ctx, "CIMD client rejected: unsupported token_endpoint_auth_method",
                        "client_id", id, "token_endpoint_auth_method", doc.TokenEndpointAuthMethod)
                return nil, fosite.ErrInvalidClient.WithHintf("CIMD document at %s claims token_endpoint_auth_method %q "+
                        "but this server only supports %q (token_endpoint_auth_methods_supported: %v)",
                        id, doc.TokenEndpointAuthMethod, defaultCIMDTokenEndpointAuthMethod,
                        doc.TokenEndpointAuthMethodsSupported)
        }
        if doc.TokenEndpointAuthMethod != "" && authMethod != doc.TokenEndpointAuthMethod {
                slog.Debug("CIMD: negotiated token_endpoint_auth_method from the client's supported list",
                        "client_id", id, "declared", doc.TokenEndpointAuthMethod, "effective", authMethod)
        }

        // Filter — not reject — grant_types and response_types this AS does not
        // support. A CIMD document describes the client's capabilities across
        // every AS it talks to and cannot be tailored per server (VS Code
        // declares device_code alongside authorization_code, see #6290), so an
        // unsupported entry is ignored rather than failing the whole document.
        // The filters still reject when the intersection lacks the one flow this
        // server offers (authorization_code / code): such a client could never
        // complete a token exchange here, and a clear error at resolution beats
        // a client that resolves and then fails every token request.
        grantTypes, dcrErr := registration.FilterPublicGrantTypes(doc.GrantTypes)
        if dcrErr != nil {
                slog.WarnContext(ctx, "CIMD client rejected: invalid grant_types",
                        "client_id", id, "error", dcrErr.ErrorDescription)
                return nil, fosite.ErrInvalidClient.WithHintf("CIMD document at %s: %s", id, dcrErr.ErrorDescription)
        }
        if len(grantTypes) < len(doc.GrantTypes) {
                slog.Debug("CIMD: ignoring grant_types this server does not support",
                        "client_id", id, "declared", doc.GrantTypes, "effective", grantTypes)
        }
        responseTypes, dcrErr := registration.FilterPublicResponseTypes(doc.ResponseTypes)
        if dcrErr != nil {
                slog.WarnContext(ctx, "CIMD client rejected: invalid response_types",
                        "client_id", id, "error", dcrErr.ErrorDescription)
                return nil, fosite.ErrInvalidClient.WithHintf("CIMD document at %s: %s", id, dcrErr.ErrorDescription)
        }
        if len(responseTypes) < len(doc.ResponseTypes) {
                slog.Debug("CIMD: ignoring response_types this server does not support",
                        "client_id", id, "declared", doc.ResponseTypes, "effective", responseTypes)
        }

        resolvedScopes, err := d.resolveScopes(ctx, id, doc)
        if err != nil {
                return nil, err
        }

        client := registration.MarkDCRIssued(buildFositeClient(fositeClientParams{
                doc:                     doc,
                resolvedScopes:          resolvedScopes,
                grantTypes:              grantTypes,
                responseTypes:           responseTypes,
                tokenEndpointAuthMethod: authMethod,
                allowedAudiences:        d.allowedAudiences,
        }))

        // Best-effort write-through: persist the resolved client in the underlying
        // storage so backends whose session rehydration resolves the client
        // through their own row lookup — Redis's unmarshalRequester, which never
        // consults this decorator — find CIMD clients at the token endpoint (see
        // issue #6187). The client carries the DCR-issued marker (above) so the
        // row gets the same anti-bloat TTL as DCR registrations: unauthenticated
        // traffic can mint these rows (fetchLimiter bounds the rate), so they must
        // never be permanent. UpsertDCRIssuedClient (not RegisterClient) is used because
        // RegisterClient is strictly create-only and would fail every fetch after
        // the first for the same client_id -- this method re-persists on every
        // fresh fetch, keeping the stored snapshot current with the document,
        // while still refusing to clobber a configured/SPIFFE-reconciled client
        // at the same ID. A persistence failure only degrades token-path
        // rehydration, so it must not fail the resolution itself.
        // A backend at its CIMD row cap refuses new rows; unauthenticated
        // callers set that volume, so it stays at Debug.
        if err := d.UpsertDCRIssuedClient(ctx, client); errors.Is(err, ErrClientCapacity) {
                slog.DebugContext(ctx, "CIMD client not persisted: storage at capacity",
                        "client_id", id, "error", err)
        } else if err != nil {
                slog.WarnContext(ctx, "failed to persist resolved CIMD client",
                        "client_id", id, "error", err)
        }

        d.cache.Add(id, &cimdCacheEntry{
                client:  client,
                expires: time.Now().Add(d.ttl),
        })

        return client, nil
}

// resolveScopes computes and validates the client scope list consistent
// with DCR. When d.scopesSupported is configured:
//   - Declared scopes are validated via registration.ValidateScopes (same
//     function as the DCR handler); an unsupported declared scope rejects.
//   - Omitted scope uses ValidateScopes(nil, scopesSupported), which grants
//     the intersection of DefaultScopes with ScopesSupported (matching DCR)
//     and reports the defaults it dropped; only an empty intersection is
//     rejected, in which case the document must declare scope explicitly.
//
// When d.scopesSupported is not configured: no AS-level validation; declared
// scopes are used directly, or nil to let buildFositeClient apply
// DefaultScopes. In both cases d.baselineClientScopes is unioned in after
// validation, matching the DCR handler's behaviour.
func (d *CIMDStorageDecorator) resolveScopes(
        ctx context.Context, id string, doc *cimd.ClientMetadataDocument,
) ([]string, error) {
        var resolvedScopes []string
        var droppedDefaults []string
        if len(d.scopesSupported) > 0 {
                if doc.Scope != "" {
                        computed, _, dcrErr := registration.ValidateScopes(strings.Fields(doc.Scope), d.scopesSupported)
                        if dcrErr != nil {
                                slog.WarnContext(ctx, "CIMD client rejected: invalid scope",
                                        "client_id", id, "error", dcrErr.ErrorDescription)
                                return nil, fosite.ErrInvalidClient.WithHintf("CIMD document at %s: %s", id, dcrErr.ErrorDescription)
                        }
                        resolvedScopes = computed
                } else {
                        // Omitted scope: match DCR — grant the intersection of DefaultScopes
                        // with scopes_supported; reject only when the intersection is empty.
                        computed, dropped, dcrErr := registration.ValidateScopes(nil, d.scopesSupported)
                        if dcrErr != nil {
                                slog.WarnContext(ctx, "CIMD client rejected: no usable default scopes",
                                        "client_id", id, "error", dcrErr.ErrorDescription)
                                return nil, fosite.ErrInvalidClient.WithHint("scope field required")
                        }
                        resolvedScopes = computed
                        droppedDefaults = dropped
                }
        } else if doc.Scope != "" {
                resolvedScopes = strings.Fields(doc.Scope)
        }
        if len(d.baselineClientScopes) > 0 {
                resolvedScopes = registration.UnionScopes(resolvedScopes, d.baselineClientScopes)
        }
        // The document omitted scope and scopes_supported does not carry the full
        // default set: resolution proceeded with the intersection (see issue
        // #6186). Recorded after the baseline union so "scopes" is the set the
        // client actually receives; Debug because the drop is fully determined by
        // startup config — the one-time operator-facing signal lives in
        // Config.applyDefaults.
        if len(droppedDefaults) > 0 {
                slog.DebugContext(ctx, "CIMD document omits scope; granted the intersection of default scopes with scopes_supported",
                        "client_id", id, "scopes", resolvedScopes, "dropped_defaults", droppedDefaults)
        }
        return resolvedScopes, nil
}

// defaultCIMDTokenEndpointAuthMethod is the token endpoint authentication
// method applied when the CIMD document omits token_endpoint_auth_method, and
// the only method this server ever selects when negotiating a fallback for a
// declared-but-unsupported value. Documents whose declared method is neither
// this value nor negotiable to it via negotiateTokenEndpointAuthMethod are
// rejected by fetch() before buildFositeClient is called.
const defaultCIMDTokenEndpointAuthMethod = "none"

// negotiateTokenEndpointAuthMethod resolves the effective token endpoint auth
// method for a CIMD document. The declared singular TokenEndpointAuthMethod is
// used when it is empty or already equal to defaultCIMDTokenEndpointAuthMethod.
// When it names anything else, the plural TokenEndpointAuthMethodsSupported
// (OpenID Connect RP Metadata Choices 1.0 — not part of the CIMD draft) is
// consulted: if it contains defaultCIMDTokenEndpointAuthMethod, the document
// is accepted with that as the negotiated method, since the client itself
// declared willingness to use it. Otherwise negotiation fails and ok is false.
//
// The supported set this function negotiates against is deliberately the
// defaultCIMDTokenEndpointAuthMethod constant alone, never the AS discovery
// document's own token_endpoint_auth_methods_supported: discovery legitimately
// advertises client_secret_basic/client_secret_post when confidential DCR or
// static delegate clients are enabled, and the CIMD draft (§4.1) forbids those
// symmetric methods in CIMD documents outright. Deriving the negotiated set
// from discovery would let a CIMD document negotiate its way into a forbidden
// method, so no such config plumbing is introduced.
func negotiateTokenEndpointAuthMethod(doc *cimd.ClientMetadataDocument) (string, bool) {
        declared := doc.TokenEndpointAuthMethod
        if declared == "" || declared == defaultCIMDTokenEndpointAuthMethod {
                return defaultCIMDTokenEndpointAuthMethod, true
        }
        if slices.Contains(doc.TokenEndpointAuthMethodsSupported, defaultCIMDTokenEndpointAuthMethod) {
                return defaultCIMDTokenEndpointAuthMethod, true
        }
        return "", false
}

// fositeClientParams carries the inputs buildFositeClient turns into a
// fosite.Client. Every field is computed by fetch() beforehand, so the
// builder applies no validation, filtering, or negotiation of its own. A
// struct keeps the four []string values from being swapped silently at the
// call site, for the same reason CIMDDecoratorConfig is one.
type fositeClientParams struct {
        doc *cimd.ClientMetadataDocument
        // resolvedScopes is the already-validated scope list computed by fetch()
        // via registration.ValidateScopes; when empty, DefaultScopes is used —
        // this occurs when the decorator has no ScopesSupported restriction
        // (unconstrained AS).
        resolvedScopes []string
        // grantTypes and responseTypes are the already-filtered lists computed by
        // fetch() via registration.FilterPublicGrantTypes/FilterPublicResponseTypes —
        // the document's declared values with unsupported entries dropped, never
        // empty (the filters apply defaults and reject empty intersections). The
        // stored client therefore carries only grant/response types this server
        // can actually serve, not the document's full declaration.
        grantTypes    []string
        responseTypes []string
        // tokenEndpointAuthMethod is the already-negotiated value computed by
        // fetch() via negotiateTokenEndpointAuthMethod; the builder applies no
        // empty-to-default fallback of its own, so the resolver is the single
        // authority over which method a CIMD-derived client ends up with.
        tokenEndpointAuthMethod string
        // allowedAudiences is the decorator's configured AllowedAudiences, granted
        // to the client verbatim (mirroring registration.New() for DCR clients,
        // #3796): CIMD documents never declare audience, so the AS applies its own
        // audience policy by granting the audiences it would validate a "resource"
        // parameter against, rather than leaving the client's audience list empty.
        // An empty list reproduces the prior nil behavior.
        allowedAudiences []string
}

// buildFositeClient converts a ClientMetadataDocument into a fosite.Client.
// RFC 8252 §7.3 loopback dynamic-port matching for a "http://localhost" redirect
// URI is provided generically by registration.RegisteredLoopbackRedirectURI
// (keyed on IsPublic() + GetRedirectURIs()), so no wrapper type is needed here.
func buildFositeClient(p fositeClientParams) fosite.Client {
        // Scopes were computed and validated by fetch() via registration.ValidateScopes,
        // consistent with the DCR handler. Fall back to DefaultScopes only when the
        // decorator has no ScopesSupported restriction (unconstrained AS).
        scopes := p.resolvedScopes
        if len(scopes) == 0 {
                scopes = slices.Clone(registration.DefaultScopes)
        }

        defaultClient := &fosite.DefaultClient{
                ID:            p.doc.ClientID,
                RedirectURIs:  p.doc.RedirectURIs,
                GrantTypes:    p.grantTypes,
                ResponseTypes: p.responseTypes,
                Scopes:        scopes,
                // CIMD clients don't pre-declare audience; the AS applies its own
                // audience policy by granting the audiences it would validate a
                // "resource" parameter against (see the allowedAudiences field), rather
                // than rejecting all values with an empty list.
                Audience: slices.Clone(p.allowedAudiences),
                Public:   true,
        }

        return &fosite.DefaultOpenIDConnectClient{
                DefaultClient:           defaultClient,
                TokenEndpointAuthMethod: p.tokenEndpointAuthMethod,
        }
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "slices"
        "sync"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// ClientAdmissionStorage hides a stored client that ToolHive wrote from
// configuration but whose configuration no longer declares it, so a configured
// client removed from configuration stops authenticating and stops using its
// grants even though its row is still in the backend.
//
// Registered clients (DCR and CIMD) are always served, and so are unmarked rows
// an embedder wrote itself with RegisterClient, except a confidential row whose
// grant types are exactly [token-exchange] (the legacy delegate shape), which is
// treated as a configured client. See clientClass for how a row is classified.
//
// Among client methods it overrides only GetClient, and no write method
// (ConsumeAssertionJWT and Unwrap only forward): rows are never deleted here,
// which keeps rolling updates safe, since each replica enforces its own
// configuration against the shared store.
//
// It does not revoke access JWTs that were already issued; those stay valid
// until they expire. It does not hide the row from code that reads the backend
// directly, and it does not protect binaries that lack it.
type ClientAdmissionStorage struct {
        Storage
        configured map[string]struct{}
        warned     sync.Map
}

// NewClientAdmissionStorage wraps base. configuredIDs are the IDs this
// process's configuration declares. It fails on a nil base or an empty ID, and
// ignores repeated IDs.
func NewClientAdmissionStorage(base Storage, configuredIDs []string) (*ClientAdmissionStorage, error) {
        if base == nil {
                return nil, errors.New("client admission: base storage is required")
        }
        set := make(map[string]struct{}, len(configuredIDs))
        for _, id := range configuredIDs {
                if id == "" {
                        return nil, errors.New("client admission: configured client ID is required")
                }
                set[id] = struct{}{}
        }
        return &ClientAdmissionStorage{Storage: base, configured: set}, nil
}

// GetClient returns the stored client only when it is admitted. Backend errors
// pass through unchanged.
func (s *ClientAdmissionStorage) GetClient(ctx context.Context, id string) (fosite.Client, error) {
        client, err := s.Storage.GetClient(ctx, id)
        if err != nil {
                return client, err
        }
        if client == nil {
                return nil, notFoundRFC6749Error("Client not found")
        }
        class := classifyClient(client)
        clientID := client.GetID()
        if s.admits(clientID, class) {
                return client, nil
        }
        if _, seen := s.warned.LoadOrStore(clientID, struct{}{}); !seen {
                slog.Warn("stored client is not declared by this server's configuration; treating as unknown",
                        "client_id", clientID, "class", class.String())
        }
        return nil, notFoundRFC6749Error("Client not found")
}

// ConsumeAssertionJWT forwards to the wrapped backend. Storage does not
// include this capability, so it is not promoted automatically.
func (s *ClientAdmissionStorage) ConsumeAssertionJWT(
        ctx context.Context, purpose, issuer, jti string, exp time.Time,
) error {
        consumer, ok := s.Storage.(AssertionJWTConsumer)
        if !ok {
                return fmt.Errorf("wrapped storage %T does not support assertion JWT replay consumption", s.Storage)
        }
        return consumer.ConsumeAssertionJWT(ctx, purpose, issuer, jti, exp)
}

// Unwrap returns the wrapped storage.
func (s *ClientAdmissionStorage) Unwrap() Storage { return s.Storage }

// ClientStoreBelowCIMD returns the storage to use for a client lookup that must
// bypass the CIMD decorator but still pass client admission: the first
// *ClientAdmissionStorage found walking Unwrap, or the innermost store when the
// chain has none (handler unit tests, or an embedder calling handlers.NewHandler
// directly without admission).
func ClientStoreBelowCIMD(s Storage) Storage {
        for {
                if _, ok := s.(*ClientAdmissionStorage); ok {
                        return s
                }
                unwrapper, ok := s.(interface{ Unwrap() Storage })
                if !ok {
                        return s
                }
                s = unwrapper.Unwrap()
        }
}

// admits reports whether a stored client with this ID and class is served.
func (s *ClientAdmissionStorage) admits(id string, c clientClass) bool {
        switch c {
        case classRegistered, classUnmarked:
                return true
        case classConfigured:
                _, ok := s.configured[id]
                return ok
        case classPlaceholder:
                return false
        default:
                return false
        }
}

// clientClass records how a stored client row came to exist, as far as the row
// itself can tell.
type clientClass uint8

const (
        // classPlaceholder marks the inert durable placeholder that reserves a
        // static SPIFFE client ID. It is never an authenticatable client.
        classPlaceholder clientClass = iota + 1

        // classConfigured marks a client ToolHive wrote from configuration. It is
        // served only while this process's configuration declares it.
        classConfigured

        // classRegistered marks a client that carries the registration marker
        // (registration.DCRIssued): dynamic client registration and CIMD documents.
        classRegistered

        // classUnmarked marks every other stored client, such as one an embedder
        // writes into storage itself with RegisterClient.
        classUnmarked
)

// String returns a log-friendly name for the class.
func (c clientClass) String() string {
        switch c {
        case classPlaceholder:
                return "placeholder"
        case classConfigured:
                return "configured"
        case classRegistered:
                return "registered"
        case classUnmarked:
                return "unmarked"
        default:
                return "invalid"
        }
}

// classifyClient reports the class of a stored client. It is total and checks
// Placeholder, Configured, Registered, then Unmarked. The classes are disjoint:
// stamped rows come back as the MarkConfigured wrapper, reconcile refuses
// DCR-marked input, and the legacy shape needs a bare DefaultClient.
func classifyClient(client fosite.Client) clientClass {
        switch {
        case isReservedPlaceholder(client):
                return classPlaceholder
        case registration.IsConfigured(client), isLegacyDelegateShape(client):
                return classConfigured
        case registration.DCRIssued(client):
                return classRegistered
        default:
                return classUnmarked
        }
}

// isLegacyDelegateShape matches the delegate row every binary wrote from
// v0.43.0 (#6320) before the configured stamp existed. On Redis a confidential
// token-exchange-only OIDC client with no auth method reads back as a bare
// DefaultClient and so matches too; memory serves it.
func isLegacyDelegateShape(client fosite.Client) bool {
        c, ok := client.(*fosite.DefaultClient)
        return ok && !c.Public && slices.Equal(c.GetGrantTypes(), []string{oauthproto.GrantTypeTokenExchange})
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import "time"

// Type defines the type of storage backend.
type Type string

const (
        // TypeMemory uses in-memory storage (default).
        TypeMemory Type = "memory"

        // TypeRedis uses Redis-backed storage for distributed deployments.
        TypeRedis Type = "redis"

        // AuthTypeACLUser is the Redis ACL user authentication type.
        // This is currently the only supported auth type for Redis storage.
        AuthTypeACLUser = "aclUser"

        // DefaultCleanupInterval is how often the background cleanup runs.
        DefaultCleanupInterval = 5 * time.Minute

        // DefaultAccessTokenTTL is the default TTL for access tokens when not extractable from session.
        DefaultAccessTokenTTL = 1 * time.Hour

        // DefaultRefreshTokenTTL is the default TTL for refresh tokens when not extractable from session.
        DefaultRefreshTokenTTL = 30 * 24 * time.Hour // 30 days

        // DefaultAuthCodeTTL is the default TTL for authorization codes (RFC 6749 recommendation).
        DefaultAuthCodeTTL = 10 * time.Minute

        // DefaultInvalidatedCodeTTL is how long invalidated codes are kept for replay detection.
        DefaultInvalidatedCodeTTL = 30 * time.Minute

        // DefaultPKCETTL is the default TTL for PKCE requests (same as auth codes).
        DefaultPKCETTL = 10 * time.Minute

        // DefaultDCRClientTTL is the TTL for dynamically registered (DCR-issued)
        // clients, public and confidential alike, once they have been used. This
        // prevents unbounded growth from DCR. Only clients carrying the
        // registration.DCRIssued marker expire; pre-provisioned confidential
        // clients are stored without a TTL. RenewClientTTL sets it on every
        // successful token exchange.
        // It is also the age bound of the CIMD persisted-row fallback on
        // /oauth/token: CIMD rows are not renewed on token success
        // (CIMDStorageDecorator.RenewClientTTL), so changing it changes that bound.
        DefaultDCRClientTTL = 30 * 24 * time.Hour // 30 days

        // DefaultUnusedDCRClientTTL is the TTL a client gets when /oauth/register
        // creates it. /oauth/register is unauthenticated, so a client that is
        // never used for a token exchange expires after this window instead of
        // DefaultDCRClientTTL. The first successful token exchange raises the TTL
        // to DefaultDCRClientTTL through RenewClientTTL.
        DefaultUnusedDCRClientTTL = 24 * time.Hour

        // DefaultMaxDCRClients is the default cap on live clients created by
        // /oauth/register, on both backends. Only clients counted by
        // countsTowardDCRCap are counted: operator-configured and CIMD clients
        // never are. Clients register once and reuse their client_id, so the cap
        // sits far above any legitimate count. Zero means unlimited.
        DefaultMaxDCRClients = 10000
)

// Config configures the storage backend.
type Config struct {
        // Type specifies the storage backend type. Defaults to memory.
        Type Type
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() *Config {
        return &Config{
                Type: TypeMemory,
        }
}

// RunConfig is the serializable storage configuration for RunConfig.
// This is used when the config needs to be passed across process boundaries
// (e.g., in Kubernetes operator).
type RunConfig struct {
        // Type specifies the storage backend type. Defaults to "memory".
        Type string `json:"type,omitempty" yaml:"type,omitempty"`

        // RedisConfig is the Redis-specific configuration when Type is "redis".
        RedisConfig *RedisRunConfig `json:"redis_config,omitempty" yaml:"redis_config,omitempty"`

        // MaxDCRClients caps the number of live clients created by
        // /oauth/register. At the cap, registration fails with HTTP 503 until a
        // client expires; no existing client is evicted. If omitted, defaults to
        // 10000; 0 means unlimited, and a negative value is rejected. On Redis the
        // count is shared by every replica.
        //
        // Sizing: a client never used for a token exchange lives 24 hours, so a
        // caller registering at the full rate keeps up to rate x replicas x 86400
        // unused clients alive. At the default rate of 1/s that is about 86400 per
        // replica: one replica fills the default cap in about 2.8 hours, and new
        // registrations then get 503 for up to 24 hours after the traffic stops.
        // A client promoted by a token exchange holds its slot for 30 days. The
        // operator sets it from the embedded auth server's storage.maxDCRClients.
        MaxDCRClients *int `json:"max_dcr_clients,omitempty" yaml:"max_dcr_clients,omitempty"`
}

// RedisRunConfig is the serializable Redis configuration for RunConfig.
// Exactly one of Addr (standalone/cluster) or SentinelConfig must be set.
// Set ClusterMode to true when Addr points to a Redis Cluster discovery endpoint.
type RedisRunConfig struct {
        // Addr is the Redis server address (host:port). Required for standalone and cluster modes.
        // Mutually exclusive with SentinelConfig.
        Addr string `json:"addr,omitempty" yaml:"addr,omitempty"`

        // ClusterMode enables the Redis Cluster protocol. Requires Addr to be set.
        ClusterMode bool `json:"cluster_mode,omitempty" yaml:"cluster_mode,omitempty"`

        // SentinelConfig contains Sentinel-specific configuration.
        // Mutually exclusive with Addr.
        SentinelConfig *SentinelRunConfig `json:"sentinel_config,omitempty" yaml:"sentinel_config,omitempty"`

        // AuthType selects the Redis authentication mode. "aclUser" is the only
        // authenticated mode. Leave it empty, with a nil ACLUserConfig, for a
        // no-auth connection to a Redis/Valkey instance that has no authentication
        // configured. Setting AuthType to "aclUser" declares authenticated intent:
        // the conversion rejects that pairing with a nil ACLUserConfig rather than
        // downgrading to no-auth. Otherwise presence of ACLUserConfig is what
        // enables authentication.
        AuthType string `json:"auth_type" yaml:"auth_type"`

        // ACLUserConfig contains ACL user authentication configuration.
        // A nil value is a valid no-auth configuration: the store connects without
        // credentials. A populated block whose password resolves to empty is a
        // misconfiguration (mis-keyed or unsynced secret) and is rejected rather
        // than silently downgraded to an unauthenticated connection.
        ACLUserConfig *ACLUserRunConfig `json:"acl_user_config,omitempty" yaml:"acl_user_config,omitempty"`

        // KeyPrefix for multi-tenancy, typically "thv:auth:{ns}:{name}:".
        KeyPrefix string `json:"key_prefix" yaml:"key_prefix"`

        // DialTimeout is the timeout for establishing connections (e.g., "5s").
        DialTimeout string `json:"dial_timeout,omitempty" yaml:"dial_timeout,omitempty"`

        // ReadTimeout is the timeout for read operations (e.g., "3s").
        ReadTimeout string `json:"read_timeout,omitempty" yaml:"read_timeout,omitempty"`

        // WriteTimeout is the timeout for write operations (e.g., "3s").
        WriteTimeout string `json:"write_timeout,omitempty" yaml:"write_timeout,omitempty"`

        // TLS configures TLS for Redis/Valkey master or cluster node connections.
        TLS *RedisTLSRunConfig `json:"tls,omitempty" yaml:"tls,omitempty"`

        // SentinelTLS configures TLS for Sentinel connections. Only applies when SentinelConfig is set.
        SentinelTLS *RedisTLSRunConfig `json:"sentinel_tls,omitempty" yaml:"sentinel_tls,omitempty"`
}

// SentinelRunConfig contains Redis Sentinel configuration.
type SentinelRunConfig struct {
        // MasterName is the name of the Redis Sentinel master.
        MasterName string `json:"master_name" yaml:"master_name"`

        // SentinelAddrs is the list of Sentinel addresses (host:port).
        SentinelAddrs []string `json:"sentinel_addrs" yaml:"sentinel_addrs"`

        // DB is the Redis database number (default: 0).
        DB int `json:"db,omitempty" yaml:"db,omitempty"`
}

// RedisTLSRunConfig holds TLS configuration for Redis connections.
// Presence of this struct enables TLS for the connection type.
type RedisTLSRunConfig struct {
        // InsecureSkipVerify skips certificate verification.
        InsecureSkipVerify bool `json:"insecure_skip_verify,omitempty" yaml:"insecure_skip_verify,omitempty"`

        // CACertFile is the path to a PEM-encoded CA certificate file.
        CACertFile string `json:"ca_cert_file,omitempty" yaml:"ca_cert_file,omitempty"`
}

// ACLUserRunConfig contains Redis ACL user authentication configuration.
// Credentials are read from environment variables for security.
type ACLUserRunConfig struct {
        // UsernameEnvVar is the environment variable containing the Redis username.
        UsernameEnvVar string `json:"username_env_var" yaml:"username_env_var"`

        // PasswordEnvVar is the environment variable containing the Redis password.
        PasswordEnvVar string `json:"password_env_var" yaml:"password_env_var"`
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "fmt"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// ClientCapacityError is the error RegisterClient returns when the DCR client
// cap is reached. It matches ErrClientCapacity under errors.Is.
type ClientCapacityError struct {
        // Cap is the client cap that was reached.
        Cap int
        // RetryAfter is how long until the oldest counted client expires, which
        // is the earliest moment a slot can free up. Zero when unknown.
        RetryAfter time.Duration
}

// Error implements error.
func (e *ClientCapacityError) Error() string {
        return fmt.Sprintf("%s: retry after %s", ErrClientCapacity, e.RetryAfter)
}

// Is makes errors.Is(err, ErrClientCapacity) report true.
func (*ClientCapacityError) Is(target error) bool {
        return target == ErrClientCapacity
}

// newClientCapacityError builds a ClientCapacityError for the reached
// capacity from the oldest counted client's expiry. A past expiry yields a
// zero RetryAfter.
func newClientCapacityError(capacity int, oldestExpiry, now time.Time) *ClientCapacityError {
        return &ClientCapacityError{Cap: capacity, RetryAfter: max(oldestExpiry.Sub(now), 0)}
}

// DCRClientCounter is implemented by storage backends that can report how many
// live clients count against the DCR cap. It backs the live-client gauge and
// is read at metric collection time, never on a request path. Reach it through
// Unwrap, since the CIMD and SPIFFE decorators do not forward it.
type DCRClientCounter interface {
        CountDCRClients(ctx context.Context) (int64, error)
}

// countsTowardDCRCap reports whether client counts against the DCR client
// cap. That is every client carrying the registration.DCRIssued marker except
// a CIMD client: the CIMD decorator persists its resolved clients with the
// same marker so they expire, but they are not minted by /oauth/register and
// are not counted. Operator-configured clients carry no marker.
func countsTowardDCRCap(client fosite.Client) bool {
        return registration.DCRIssued(client) && !oauthproto.IsClientIDMetadataDocumentURL(client.GetID())
}

// ResolvedMaxDCRClients returns the DCR client cap c configures: unset (or a
// nil c) means DefaultMaxDCRClients, zero means unlimited, and a negative
// value is an error.
func (c *RunConfig) ResolvedMaxDCRClients() (int, error) {
        if c == nil || c.MaxDCRClients == nil {
                return DefaultMaxDCRClients, nil
        }
        if *c.MaxDCRClients < 0 {
                return 0, fmt.Errorf("max_dcr_clients must not be negative, got %d", *c.MaxDCRClients)
        }
        return *c.MaxDCRClients, nil
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package storage

import (
        "context"
        "fmt"
        "log/slog"
        "net/url"
        "slices"
        "sync"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
)

// defaultMaxMemoryCIMDClients caps the live CIMD rows MemoryStorage keeps.
// Each distinct client metadata document URL is one row, so real deployments
// sit far below it.
const defaultMaxMemoryCIMDClients = 10000

// timedEntry wraps a value with its creation time for TTL tracking.
type timedEntry[T any] struct {
        value     T
        createdAt time.Time
        expiresAt time.Time
}

// upstreamKey is the composite key for the flat upstream token map.
type upstreamKey struct {
        sessionID    string
        providerName string
}

// assertionJWTKey identifies a consumed assertion JWT without sharing JTIs
// across assertion purposes or authorization-server issuers.
type assertionJWTKey struct {
        purpose string
        issuer  string
        jti     string
}

// MemoryStorage implements the Storage interface with in-memory maps.
// This implementation is thread-safe and suitable for development and testing.
// For production use, consider implementing a persistent storage backend.
//
// # Fosite Storage Design
//
// Token maps store fosite.Requester (not just token strings) because fosite needs
// the full authorization context for validation and introspection. The Requester
// contains the Client, granted scopes, Session (with expiration times), and more.
//
// Maps are keyed by "signature" (cryptographic token identifier) for O(1) token
// lookup. Revocation by "request ID" requires O(n) scan; production implementations
// should maintain a reverse index for efficiency.
type MemoryStorage struct {
        mu sync.RWMutex

        // clients maps client_id -> Client for client lookup (fosite.ClientManager).
        // A SPIFFE static-client durable placeholder (see staticClientPlaceholder
        // in spiffe_decorator.go) is stored here as the live inertPlaceholderClient
        // Go value, so its overridden GetGrantTypes/GetResponseTypes are called
        // directly on every read — unlike RedisStorage, which serializes the
        // client's fields and must re-derive the same guarantee on read-back via
        // storedClient.Reserved (see clientFromStored in redis.go). No equivalent
        // marker is needed here.
        clients map[string]fosite.Client

        // dcrClientExpiry holds the expiry of every client in clients counted by
        // countsTowardDCRCap, and cimdClientExpiry that of every other client
        // carrying the registration.DCRIssued marker (the CIMD rows), mirroring
        // the key TTL RedisStorage gives the same rows. They are split so each
        // cap walks only the clients it bounds. A client in neither never
        // expires. Every read treats an expired client as absent, and
        // cleanupExpired removes it.
        dcrClientExpiry  map[string]time.Time
        cimdClientExpiry map[string]time.Time

        // configured holds the ID of every client in clients that was last
        // written by ReconcileConfiguredClient. It is a side map, not a wrapper
        // around the stored value: wrapping a SPIFFE placeholder would break
        // isReservedPlaceholder and make a second reconcile of it fail. GetClient
        // wraps non-reserved stamped clients on the way out. Every other write,
        // delete and expiry of an ID clears its entry.
        configured map[string]struct{}

        // maxDCRClients caps the live clients counted by countsTowardDCRCap; zero
        // means unlimited. /oauth/register is unauthenticated and every call
        // mints an entry, so an uncapped map is a memory-exhaustion DoS in the
        // same process as the proxy. At the cap RegisterClient rejects with a
        // *ClientCapacityError and evicts nothing: evicting would let a caller
        // registering at the full rate push out every idle real client. A full
        // map drains as entries expire.
        maxDCRClients int

        // maxCIMDClients caps the live CIMD rows UpsertDCRIssuedClient writes.
        // Unauthenticated /oauth/authorize traffic mints them, bounded otherwise
        // only by the per-process CIMD rate limit and their 30-day TTL. At the cap
        // a new row is refused with a *ClientCapacityError. The CIMD decorator's
        // write-through is best-effort, and this backend keeps sessions as Go
        // values rather than rehydrating their client from the row, so a refused
        // row costs at most a document refetch. RedisStorage has no such cap.
        maxCIMDClients int

        // now returns the current time for client expiry.
        now func() time.Time

        // authCodes maps authorization code -> Requester. Codes are one-time-use;
        // invalidatedCodes tracks used codes to return ErrInvalidatedAuthorizeCode.
        authCodes map[string]*timedEntry[fosite.Requester]

        // accessTokens maps token signature -> Requester. The signature is derived
        // from the token value, enabling O(1) lookup when validating bearer tokens.
        accessTokens map[string]*timedEntry[fosite.Requester]

        // refreshTokens maps token signature -> Requester. Linked to access tokens
        // via request ID for token rotation per RFC 6749.
        refreshTokens map[string]*timedEntry[fosite.Requester]

        // pkceRequests maps code signature -> Requester containing the PKCE challenge.
        // Validated during token exchange per RFC 7636.
        pkceRequests map[string]*timedEntry[fosite.Requester]

        // upstreamTokens maps (sessionID, providerName) -> timedEntry for multi-provider support.
        upstreamTokens map[upstreamKey]*timedEntry[*UpstreamTokens]

        // pendingAuthorizations tracks authorization requests awaiting upstream IDP callback
        pendingAuthorizations map[string]*timedEntry[*PendingAuthorization]
        consentSessions       map[string]*timedEntry[*ConsentSession]
        clientApprovals       map[[2]string]*timedEntry[*ClientApproval]

        // deviceRequests maps device_code -> timedEntry[*DeviceRequest]. The
        // canonical store; TTL-bounded like pendingAuthorizations.
        deviceRequests map[string]*timedEntry[*DeviceRequest]

        // deviceRequestsByUserCode is a secondary index, user_code -> device_code,
        // so the verification page can look up a request without scanning
        // deviceRequests. Kept in lockstep with deviceRequests: every store/delete
        // touches both maps under the same lock.
        deviceRequestsByUserCode map[string]string

        // pendingDeviceLogins tracks device-flow verification-page logins
        // awaiting the upstream IDP callback, keyed by state -- the same role
        // pendingAuthorizations plays for the client authorization_code flow.
        pendingDeviceLogins map[string]*timedEntry[*PendingDeviceLogin]

        // pendingDeviceConfirmations tracks resolved device-flow logins awaiting
        // an explicit Approve/Deny decision at the verification page, keyed by
        // an opaque confirmation token.
        pendingDeviceConfirmations map[string]*timedEntry[*PendingDeviceConfirmation]

        // invalidatedCodes tracks auth codes that have been used/invalidated.
        // Kept separate from authCodes to return the Requester with ErrInvalidatedAuthorizeCode.
        invalidatedCodes map[string]*timedEntry[bool]

        // clientAssertionJWTs tracks JTIs to prevent JWT replay attacks per RFC 7523.
        clientAssertionJWTs map[string]time.Time

        // assertionJWTs tracks consumed assertion JWTs by purpose, issuer, and JTI.
        assertionJWTs map[assertionJWTKey]time.Time

        // users maps user ID -> User for user account lookup.
        // Users are not subject to TTL-based cleanup as they represent persistent accounts.
        users map[string]*User

        // providerIdentities maps "providerID:providerSubject" -> ProviderIdentity for identity lookup.
        // This enables O(1) lookup during authentication callbacks.
        providerIdentities map[string]*ProviderIdentity

        // dcrCredentials maps DCRKey -> DCRCredentials for RFC 7591 Dynamic Client
        // Registration credentials. These entries come from OUTBOUND DCR: ToolHive
        // acting as a DCR client to register itself against a configured upstream
        // IdP (see pkg/auth/dcr/store.go, the only caller of StoreDCRCredentialsIfAbsent).
        // This is not reachable from the inbound, unauthenticated /oauth/register
        // handler, so unlike clients (bounded by maxDCRClients and the client expiry maps because
        // every inbound registration call mints an entry), growth here is bounded
        // by the number of configured upstreams, which the operator controls, not
        // by request volume — hence no eviction or cap on this map.
        //
        // Entries are intentionally excluded from the periodic cleanupExpired
        // loop: DCR registrations are long-lived and the authoritative expiry
        // signal is RFC 7591 client_secret_expires_at, which is honored at read
        // time by callers (and by the future Redis backend's SetEX TTL). Growth
        // is bounded by upstream count × distinct scope sets ever registered for
        // each upstream during the process lifetime; for a stable configuration
        // this collapses to the upstream count, but rotating scope sets
        // (operator-driven scope changes, or upstream scopes_supported rotations
        // re-derived by the resolver) accumulate stale entries that survive until
        // process restart. The Redis backend's SetEX TTL mitigates this in
        // production deployments.
        dcrCredentials map[DCRKey]*DCRCredentials

        // cleanupInterval is how often the background cleanup runs
        cleanupInterval time.Duration

        // stopCleanup is used to signal the cleanup goroutine to stop
        stopCleanup chan struct{}

        // cleanupDone is closed when the cleanup goroutine has fully stopped
        cleanupDone chan struct{}
}

// MemoryStorageOption configures a MemoryStorage instance.
type MemoryStorageOption func(*MemoryStorage)

// WithCleanupInterval sets a custom cleanup interval.
func WithCleanupInterval(interval time.Duration) MemoryStorageOption {
        return func(s *MemoryStorage) {
                s.cleanupInterval = interval
        }
}

// WithMaxDCRClients sets the cap on live clients counted by
// countsTowardDCRCap; zero means unlimited. At the cap RegisterClient returns a
// *ClientCapacityError and evicts nothing. Operator-configured and CIMD
// clients are never counted. It panics on a negative n, which
// RunConfig.ResolvedMaxDCRClients already rejects.
func WithMaxDCRClients(n int) MemoryStorageOption {
        if n < 0 {
                panic(fmt.Sprintf("max DCR clients must not be negative, got %d", n))
        }
        return func(s *MemoryStorage) {
                s.maxDCRClients = n
        }
}

// withMaxCIMDClients overrides defaultMaxMemoryCIMDClients, so tests can
// reach the CIMD cap.
func withMaxCIMDClients(n int) MemoryStorageOption {
        return func(s *MemoryStorage) {
                s.maxCIMDClients = n
        }
}

// withClientClock sets the clock used for client expiry, so tests can move
// past a TTL without sleeping.
func withClientClock(now func() time.Time) MemoryStorageOption {
        return func(s *MemoryStorage) {
                s.now = now
        }
}

// NewMemoryStorage creates a new MemoryStorage instance with initialized maps
// and starts the background cleanup goroutine.
func NewMemoryStorage(opts ...MemoryStorageOption) *MemoryStorage {
        s := &MemoryStorage{
                clients:                    make(map[string]fosite.Client),
                authCodes:                  make(map[string]*timedEntry[fosite.Requester]),
                accessTokens:               make(map[string]*timedEntry[fosite.Requester]),
                refreshTokens:              make(map[string]*timedEntry[fosite.Requester]),
                pkceRequests:               make(map[string]*timedEntry[fosite.Requester]),
                upstreamTokens:             make(map[upstreamKey]*timedEntry[*UpstreamTokens]),
                pendingAuthorizations:      make(map[string]*timedEntry[*PendingAuthorization]),
                consentSessions:            make(map[string]*timedEntry[*ConsentSession]),
                clientApprovals:            make(map[[2]string]*timedEntry[*ClientApproval]),
                deviceRequests:             make(map[string]*timedEntry[*DeviceRequest]),
                deviceRequestsByUserCode:   make(map[string]string),
                pendingDeviceLogins:        make(map[string]*timedEntry[*PendingDeviceLogin]),
                pendingDeviceConfirmations: make(map[string]*timedEntry[*PendingDeviceConfirmation]),
                invalidatedCodes:           make(map[string]*timedEntry[bool]),
                clientAssertionJWTs:        make(map[string]time.Time),
                assertionJWTs:              make(map[assertionJWTKey]time.Time),
                users:                      make(map[string]*User),
                providerIdentities:         make(map[string]*ProviderIdentity),
                dcrCredentials:             make(map[DCRKey]*DCRCredentials),
                cleanupInterval:            DefaultCleanupInterval,
                dcrClientExpiry:            make(map[string]time.Time),
                cimdClientExpiry:           make(map[string]time.Time),
                configured:                 make(map[string]struct{}),
                maxDCRClients:              DefaultMaxDCRClients,
                maxCIMDClients:             defaultMaxMemoryCIMDClients,
                now:                        time.Now,
                stopCleanup:                make(chan struct{}),
                cleanupDone:                make(chan struct{}),
        }

        for _, opt := range opts {
                opt(s)
        }

        // Start background cleanup goroutine
        go s.cleanupLoop()

        return s
}

// Health is a no-op for in-memory storage since it is always available.
func (*MemoryStorage) Health(_ context.Context) error {
        return nil
}

// Close stops the background cleanup goroutine and waits for it to finish.
// This should be called when the storage is no longer needed.
func (s *MemoryStorage) Close() error {
        close(s.stopCleanup)
        <-s.cleanupDone
        return nil
}

// cleanupLoop runs periodic cleanup of expired entries.
func (s *MemoryStorage) cleanupLoop() {
        defer close(s.cleanupDone)

        ticker := time.NewTicker(s.cleanupInterval)
        defer ticker.Stop()

        for {
                select {
                case <-s.stopCleanup:
                        return
                case <-ticker.C:
                        s.cleanupExpired()
                }
        }
}

// cleanupExpired removes all expired entries from storage.
// Uses collect-then-delete pattern: collects expired keys under read lock,
// then deletes under write lock. This minimizes write lock hold time.
//
//nolint:gocyclo // Function is straightforward, just repetitive cleanup loops for each storage type
func (s *MemoryStorage) cleanupExpired() {
        now := time.Now()

        // Phase 1: Collect expired keys under read lock
        s.mu.RLock()

        var expiredAuthCodes []string
        for k, v := range s.authCodes {
                if now.After(v.expiresAt) {
                        expiredAuthCodes = append(expiredAuthCodes, k)
                }
        }

        var expiredInvalidatedCodes []string
        for k, v := range s.invalidatedCodes {
                if now.After(v.expiresAt) {
                        expiredInvalidatedCodes = append(expiredInvalidatedCodes, k)
                }
        }

        var expiredAccessTokens []string
        for k, v := range s.accessTokens {
                if now.After(v.expiresAt) {
                        expiredAccessTokens = append(expiredAccessTokens, k)
                }
        }

        var expiredRefreshTokens []string
        for k, v := range s.refreshTokens {
                if now.After(v.expiresAt) {
                        expiredRefreshTokens = append(expiredRefreshTokens, k)
                }
        }

        var expiredPKCERequests []string
        for k, v := range s.pkceRequests {
                if now.After(v.expiresAt) {
                        expiredPKCERequests = append(expiredPKCERequests, k)
                }
        }

        var expiredUpstreamTokens []upstreamKey
        for k, v := range s.upstreamTokens {
                // Zero expiresAt is the sentinel for "no TTL" (non-expiring token with no session
                // bound). Other entry types never use zero expiresAt, so only upstream tokens need
                // this guard — without it, time.Time{} would compare as before any real time and
                // every non-expiring token would be swept on the next tick.
                if !v.expiresAt.IsZero() && now.After(v.expiresAt) {
                        expiredUpstreamTokens = append(expiredUpstreamTokens, k)
                }
        }

        var expiredPendingAuthorizations []string
        for k, v := range s.pendingAuthorizations {
                if now.After(v.expiresAt) {
                        expiredPendingAuthorizations = append(expiredPendingAuthorizations, k)
                }
        }

        var expiredConsentSessions []string
        var expiredClientApprovals [][2]string
        for k, v := range s.consentSessions {
                if !now.Before(v.expiresAt) {
                        expiredConsentSessions = append(expiredConsentSessions, k)
                }
        }
        for k, v := range s.clientApprovals {
                if !now.Before(v.expiresAt) {
                        expiredClientApprovals = append(expiredClientApprovals, k)
                }
        }

        var expiredDeviceRequests []string
        for k, v := range s.deviceRequests {
                if now.After(v.expiresAt) {
                        expiredDeviceRequests = append(expiredDeviceRequests, k)
                }
        }

        var expiredPendingDeviceLogins []string
        for k, v := range s.pendingDeviceLogins {
                if now.After(v.expiresAt) {
                        expiredPendingDeviceLogins = append(expiredPendingDeviceLogins, k)
                }
        }

        var expiredPendingDeviceConfirmations []string
        for k, v := range s.pendingDeviceConfirmations {
                if now.After(v.expiresAt) {
                        expiredPendingDeviceConfirmations = append(expiredPendingDeviceConfirmations, k)
                }
        }

        var expiredJWTs []string
        for k, v := range s.clientAssertionJWTs {
                if now.After(v) {
                        expiredJWTs = append(expiredJWTs, k)
                }
        }

        var expiredAssertionJWTs []assertionJWTKey
        for k, v := range s.assertionJWTs {
                if now.After(v) {
                        expiredAssertionJWTs = append(expiredAssertionJWTs, k)
                }
        }

        clientNow := s.now()
        var expiredClients []string
        for _, expiry := range []map[string]time.Time{s.dcrClientExpiry, s.cimdClientExpiry} {
                for k, v := range expiry {
                        if !clientNow.Before(v) {
                                expiredClients = append(expiredClients, k)
                        }
                }
        }

        s.mu.RUnlock()

        // Phase 2: Early return if nothing to delete (no write lock needed)
        if len(expiredClients) == 0 &&
                len(expiredAuthCodes) == 0 &&
                len(expiredInvalidatedCodes) == 0 &&
                len(expiredAccessTokens) == 0 &&
                len(expiredRefreshTokens) == 0 &&
                len(expiredPKCERequests) == 0 &&
                len(expiredUpstreamTokens) == 0 &&
                len(expiredPendingAuthorizations) == 0 &&
                len(expiredConsentSessions) == 0 &&
                len(expiredClientApprovals) == 0 &&
                len(expiredDeviceRequests) == 0 &&
                len(expiredPendingDeviceLogins) == 0 &&
                len(expiredPendingDeviceConfirmations) == 0 &&
                len(expiredJWTs) == 0 &&
                len(expiredAssertionJWTs) == 0 {
                return
        }

        // Phase 3: Delete collected keys under write lock
        s.mu.Lock()
        defer s.mu.Unlock()

        for _, k := range expiredClients {
                // RenewClientTTL may have extended the client between the phases.
                if exp, ok := s.clientExpiryLocked(k); ok && !clientNow.Before(exp) {
                        delete(s.clients, k)
                        delete(s.dcrClientExpiry, k)
                        delete(s.cimdClientExpiry, k)
                        delete(s.configured, k)
                }
        }

        for _, k := range expiredAuthCodes {
                delete(s.authCodes, k)
                delete(s.invalidatedCodes, k) // Also clean up associated invalidated code
        }

        for _, k := range expiredInvalidatedCodes {
                delete(s.invalidatedCodes, k)
        }

        for _, k := range expiredAccessTokens {
                delete(s.accessTokens, k)
        }

        for _, k := range expiredRefreshTokens {
                delete(s.refreshTokens, k)
        }

        for _, k := range expiredPKCERequests {
                delete(s.pkceRequests, k)
        }

        for _, k := range expiredUpstreamTokens {
                delete(s.upstreamTokens, k)
        }

        for _, k := range expiredPendingAuthorizations {
                delete(s.pendingAuthorizations, k)
        }

        for _, k := range expiredConsentSessions {
                if v, ok := s.consentSessions[k]; ok && !now.Before(v.expiresAt) {
                        delete(s.consentSessions, k)
                }
        }
        for _, k := range expiredClientApprovals {
                if v, ok := s.clientApprovals[k]; ok && !now.Before(v.expiresAt) {
                        delete(s.clientApprovals, k)
                }
        }

        for _, k := range expiredDeviceRequests {
                if entry, ok := s.deviceRequests[k]; ok {
                        delete(s.deviceRequestsByUserCode, entry.value.UserCode)
                }
                delete(s.deviceRequests, k)
        }

        for _, k := range expiredPendingDeviceLogins {
                delete(s.pendingDeviceLogins, k)
        }

        for _, k := range expiredPendingDeviceConfirmations {
                delete(s.pendingDeviceConfirmations, k)
        }

        for _, k := range expiredJWTs {
                delete(s.clientAssertionJWTs, k)
        }

        for _, k := range expiredAssertionJWTs {
                if exp, ok := s.assertionJWTs[k]; ok && now.After(exp) {
                        delete(s.assertionJWTs, k)
                }
        }
}

// snapshotRequester copies request-owned data while retaining the concrete client reference.
// GetID may initialize an empty ID on the caller's Request during creation.
func snapshotRequester(request fosite.Requester) *fosite.Request {
        form := request.GetRequestForm()
        var formCopy url.Values
        if form != nil {
                formCopy = make(url.Values, len(form))
                for key, values := range form {
                        formCopy[key] = slices.Clone(values)
                }
        }

        snapshot := &fosite.Request{
                ID:                request.GetID(),
                RequestedAt:       request.GetRequestedAt(),
                Client:            request.GetClient(),
                RequestedScope:    slices.Clone(request.GetRequestedScopes()),
                GrantedScope:      slices.Clone(request.GetGrantedScopes()),
                RequestedAudience: slices.Clone(request.GetRequestedAudience()),
                GrantedAudience:   slices.Clone(request.GetGrantedAudience()),
                Form:              formCopy,
        }
        if session := request.GetSession(); session != nil {
                snapshot.Session = session.Clone()
        }
        if langContext, ok := request.(fosite.G11NContext); ok {
                snapshot.Lang = langContext.GetLang()
        }
        // Requester implementations other than fosite.Request may return an empty ID.
        // Initialize the copy before publishing it so Get never writes under RLock.
        snapshot.GetID()
        return snapshot
}

// getExpirationFromRequester extracts expiration time from a fosite.Requester session.
// Returns the provided default if expiration cannot be extracted.
//
// This demonstrates why GetExpiresAt lives on fosite.Session: different token types
// (AccessToken, RefreshToken, AuthorizeCode) have different lifetimes, and Session
// stores per-token-type expiration. The Session is the natural container for token
// metadata, while Requester holds the full authorization context.
func getExpirationFromRequester(request fosite.Requester, tokenType fosite.TokenType, defaultTTL time.Duration) time.Time {
        if request == nil {
                return time.Now().Add(defaultTTL)
        }

        session := request.GetSession()
        if session == nil {
                return time.Now().Add(defaultTTL)
        }

        expTime := session.GetExpiresAt(tokenType)
        if expTime.IsZero() {
                return time.Now().Add(defaultTTL)
        }

        return expTime
}

// RegisterClient creates a client in storage. Always create-only: it returns
// ErrAlreadyExists when a client with the same ID is already registered,
// regardless of the new client's origin. /oauth/register is unauthenticated,
// so this must never clobber an existing client; operator-declared clients
// reconcile through ReconcileConfiguredClient instead.
func (s *MemoryStorage) RegisterClient(_ context.Context, client fosite.Client) error {
        id := client.GetID()
        if registration.IsConfigured(client) {
                return fmt.Errorf("client %q must not carry the configured stamp; use ReconcileConfiguredClient", id)
        }
        if err := ValidateRegisterableClientID(id); err != nil {
                return err
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := s.now()
        if _, exists := s.liveClientLocked(id, now); exists {
                return fmt.Errorf("%w: client %q", ErrAlreadyExists, id)
        }
        if !registration.DCRIssued(client) {
                s.putPermanentClientLocked(id, client)
                return nil
        }
        if s.maxDCRClients > 0 && countsTowardDCRCap(client) {
                if count, oldestExpiry := countLiveLocked(s.dcrClientExpiry, now); count >= s.maxDCRClients {
                        return newClientCapacityError(s.maxDCRClients, oldestExpiry, now)
                }
        }
        s.putExpiringClientLocked(id, client, now.Add(DefaultUnusedDCRClientTTL))
        return nil
}

// UpsertDCRIssuedClient creates or replaces a DCR-issued client at
// client.GetID(). Unlike RegisterClient, an existing row is replaced (and its
// expiry reset to DefaultDCRClientTTL) rather than rejected -- but only when
// the existing row is itself DCR-issued; a configured/SPIFFE-reconciled row at
// the same ID is protected and refuses with ErrAlreadyExists. See the
// ClientRegistry interface doc for the full contract. A client counted
// against the DCR cap is refused, so this path cannot bypass the cap. A new
// CIMD row beyond maxCIMDClients is refused with a *ClientCapacityError;
// replacing an existing one always succeeds.
func (s *MemoryStorage) UpsertDCRIssuedClient(_ context.Context, client fosite.Client) error {
        if !registration.DCRIssued(client) {
                return fmt.Errorf("client %q must carry the DCR-issued marker to use UpsertDCRIssuedClient", client.GetID())
        }
        id := client.GetID()
        if err := ValidateRegisterableClientID(id); err != nil {
                return err
        }
        if countsTowardDCRCap(client) {
                return fmt.Errorf("client %q counts against the DCR client cap and must use RegisterClient", id)
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := s.now()
        existing, exists := s.liveClientLocked(id, now)
        if exists && !registration.DCRIssued(existing) {
                return fmt.Errorf("%w: client %q", ErrAlreadyExists, id)
        }
        if !exists && s.maxCIMDClients > 0 {
                if count, oldestExpiry := countLiveLocked(s.cimdClientExpiry, now); count >= s.maxCIMDClients {
                        return newClientCapacityError(s.maxCIMDClients, oldestExpiry, now)
                }
        }
        s.putExpiringClientLocked(id, client, now.Add(DefaultDCRClientTTL))
        return nil
}

// ReconcileConfiguredClient applies an operator-declared client: creates it
// if absent, otherwise replaces the existing record unless
// configuredClientConflict refuses (or the record is DCR-issued), in which
// case it returns ErrAlreadyExists. See the ClientRegistry interface doc for
// the full contract.
func (s *MemoryStorage) ReconcileConfiguredClient(_ context.Context, client fosite.Client) error {
        if registration.DCRIssued(client) {
                return fmt.Errorf("configured client %q must not carry the DCR-issued marker", client.GetID())
        }
        if err := checkConfiguredClientType(client); err != nil {
                return err
        }
        id := client.GetID()
        if err := ValidateRegisterableClientID(id); err != nil {
                return err
        }

        existingFP, incomingFP, replaced, err := s.reconcileConfiguredClientLocked(id, client)
        if err != nil {
                return err
        }
        if replaced {
                warnConfiguredClientReplaced(id, existingFP, incomingFP)
        }
        return nil
}

// reconcileConfiguredClientLocked stores client under s.mu and reports the
// fingerprints of the row it replaced, or replaced=false when it created a
// new row. The caller logs the replacement after the lock is released.
func (s *MemoryStorage) reconcileConfiguredClientLocked(
        id string, client fosite.Client,
) (existingFP, incomingFP clientFingerprint, replaced bool, err error) {
        s.mu.Lock()
        defer s.mu.Unlock()

        now := s.now()
        existing, exists := s.liveClientLocked(id, now)
        if !exists {
                s.putPermanentClientLocked(id, client)
                s.configured[id] = struct{}{}
                return existingFP, incomingFP, false, nil
        }
        if registration.DCRIssued(existing) {
                return existingFP, incomingFP, false, fmt.Errorf(
                        "%w: client %q is DCR-issued, refusing to overwrite with a configured client",
                        ErrAlreadyExists, id)
        }
        existingFP, incomingFP = fingerprintOfClient(existing), fingerprintOfClient(client)
        if err := configuredClientConflict(id, existingFP, incomingFP); err != nil {
                return existingFP, incomingFP, false, err
        }

        // Replace the stored value so secret rotation and policy changes take
        // effect.
        s.putPermanentClientLocked(id, client)
        s.configured[id] = struct{}{}
        return existingFP, incomingFP, true, nil
}

// liveClientLocked returns the client stored at id, or false when there is
// none or it has expired. Callers must hold s.mu.
func (s *MemoryStorage) liveClientLocked(id string, now time.Time) (fosite.Client, bool) {
        client, ok := s.clients[id]
        if !ok {
                return nil, false
        }
        if exp, ok := s.clientExpiryLocked(id); ok && !now.Before(exp) {
                return nil, false
        }
        return client, true
}

// clientExpiryLocked returns the expiry of the client stored at id, or false
// when it never expires. Callers must hold s.mu.
func (s *MemoryStorage) clientExpiryLocked(id string) (time.Time, bool) {
        if exp, ok := s.dcrClientExpiry[id]; ok {
                return exp, true
        }
        exp, ok := s.cimdClientExpiry[id]
        return exp, ok
}

// expiryMapFor returns the expiry map a DCR-issued client belongs in.
func (s *MemoryStorage) expiryMapFor(client fosite.Client) map[string]time.Time {
        if countsTowardDCRCap(client) {
                return s.dcrClientExpiry
        }
        return s.cimdClientExpiry
}

// putExpiringClientLocked stores a DCR-issued client at id, expiring at
// expiry. Callers must hold s.mu for writing.
func (s *MemoryStorage) putExpiringClientLocked(id string, client fosite.Client, expiry time.Time) {
        s.clients[id] = client
        delete(s.dcrClientExpiry, id)
        delete(s.cimdClientExpiry, id)
        delete(s.configured, id)
        s.expiryMapFor(client)[id] = expiry
}

// putPermanentClientLocked stores a client without the DCR-issued marker at
// id; it never expires. It clears the configured stamp; the caller sets it
// afterwards when the write is a reconcile. Callers must hold s.mu for writing.
func (s *MemoryStorage) putPermanentClientLocked(id string, client fosite.Client) {
        s.clients[id] = client
        delete(s.dcrClientExpiry, id)
        delete(s.cimdClientExpiry, id)
        delete(s.configured, id)
}

// countLiveLocked returns the number of entries in expiry that have not
// expired by now and the earliest expiry among them. Callers must hold s.mu.
// It walks one expiry map, which its cap bounds (plus entries expired since
// the last cleanup), on every write that map's cap checks and on a metric
// collection.
func countLiveLocked(expiry map[string]time.Time, now time.Time) (int, time.Time) {
        var count int
        var oldest time.Time
        for _, exp := range expiry {
                if !now.Before(exp) {
                        continue
                }
                count++
                if oldest.IsZero() || exp.Before(oldest) {
                        oldest = exp
                }
        }
        return count, oldest
}

// CountDCRClients implements DCRClientCounter.
func (s *MemoryStorage) CountDCRClients(_ context.Context) (int64, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()
        count, _ := countLiveLocked(s.dcrClientExpiry, s.now())
        return int64(count), nil
}

// RenewClientTTL sets a live DCR-issued client's expiry to DefaultDCRClientTTL
// from now on proven use (a successful token exchange/refresh), the same
// signal RedisStorage.RenewClientTTL uses. A pre-provisioned client never
// expires, so it needs no renewal. A client that is absent or expired stays
// absent, matching RedisStorage's PEXPIRE no-op on a missing key.
func (s *MemoryStorage) RenewClientTTL(_ context.Context, client fosite.Client) error {
        if client == nil || !registration.DCRIssued(client) {
                return nil
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        id := client.GetID()
        now := s.now()
        stored, exists := s.liveClientLocked(id, now)
        if !exists || !registration.DCRIssued(stored) {
                return nil
        }
        s.expiryMapFor(stored)[id] = now.Add(DefaultDCRClientTTL)
        return nil
}

// -----------------------
// fosite.ClientManager
// -----------------------

// GetClient loads the client by its ID or returns an error if the client does not exist.
func (s *MemoryStorage) GetClient(_ context.Context, id string) (fosite.Client, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        client, ok := s.liveClientLocked(id, s.now())
        if !ok {
                slog.Debug("client not found", "client_id", id)
                return nil, notFoundRFC6749Error("Client not found")
        }
        if dcrErr := registration.ValidatePrivateKeyJWTClientKeys(client); dcrErr != nil {
                slog.Debug("rejecting client whose signing keys no longer meet server policy",
                        "client_id", id, "reason", dcrErr.ErrorDescription)
                return nil, fosite.ErrInvalidClient.WithHint("client signing keys no longer meet server policy")
        }
        if _, stamped := s.configured[id]; stamped && !isReservedPlaceholder(client) {
                marked, ok := registration.MarkConfigured(client)
                if !ok {
                        return nil, notFoundRFC6749Error("Client not found")
                }
                return marked, nil
        }
        return client, nil
}

// assertionReservationTTL bounds how long ClientAssertionJWTValid's atomic
// reservation lives before SetClientAssertionJWT extends it to the
// assertion's real (bounded) expiry. Short and fixed: if a caller wins the
// reservation and then never calls SetClientAssertionJWT (crash, error path),
// the jti becomes reusable again after this window rather than being stuck
// or, worse, permanently reserved.
const assertionReservationTTL = 30 * time.Second

// ClientAssertionJWTValid atomically reserves jti for a short window,
// returning fosite.ErrJTIKnown if it is already reserved or was previously
// consumed and not yet expired. This is the only atomic checkpoint in
// fosite's two-call client-assertion replay protocol (Valid, then
// SetClientAssertionJWT with the real expiry once the rest of the assertion
// has been validated) — without it, two concurrent requests presenting the
// same assertion could both observe jti as unused before either recorded it.
// jti length is bounded here too, before it's ever stored, since fosite
// hands it to us before parsing exp.
func (s *MemoryStorage) ClientAssertionJWTValid(_ context.Context, jti string) error {
        if len(jti) > server.MaxAssertionJTILength {
                return fmt.Errorf("jti exceeds maximum length of %d bytes", server.MaxAssertionJTILength)
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        if exp, ok := s.clientAssertionJWTs[jti]; ok && now.Before(exp) {
                return fosite.ErrJTIKnown
        }

        // Clean up expired JTIs while the lock is already held.
        for k, v := range s.clientAssertionJWTs {
                if now.After(v) {
                        delete(s.clientAssertionJWTs, k)
                }
        }

        s.clientAssertionJWTs[jti] = now.Add(assertionReservationTTL)
        return nil
}

// SetClientAssertionJWT extends jti's replay marker to its real expiry,
// rejecting an exp further than MaxAssertionLifespan in the future. This
// bounds how long a replay-tracking entry survives regardless of what the
// client's own assertion claims — otherwise a registered caller could set
// exp decades out and grow this map indefinitely. Unconditional overwrite,
// not a fresh atomic check: only the caller that already won
// ClientAssertionJWTValid's reservation reaches this call.
func (s *MemoryStorage) SetClientAssertionJWT(_ context.Context, jti string, exp time.Time) error {
        if time.Until(exp) > server.MaxAssertionLifespan {
                return fmt.Errorf("assertion lifetime exceeds maximum of %s", server.MaxAssertionLifespan)
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        s.clientAssertionJWTs[jti] = exp
        return nil
}

// ConsumeAssertionJWT atomically records an assertion JWT as consumed until its
// expiry. Reusing an unexpired (purpose, issuer, jti) tuple returns
// fosite.ErrJTIKnown. Expired assertions are not retained because a caller must
// reject them during JWT validation before attempting replay consumption.
func (s *MemoryStorage) ConsumeAssertionJWT(_ context.Context, purpose, issuer, jti string, exp time.Time) error {
        key := assertionJWTKey{purpose: purpose, issuer: issuer, jti: jti}
        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        if !exp.After(now) {
                return nil
        }

        if storedExp, ok := s.assertionJWTs[key]; ok {
                if storedExp.After(now) {
                        return fosite.ErrJTIKnown
                }
                delete(s.assertionJWTs, key)
        }

        s.assertionJWTs[key] = exp
        return nil
}

// -----------------------
// oauth2.AuthorizeCodeStorage
// -----------------------

// CreateAuthorizeCodeSession stores the authorization request for a given authorization code.
func (s *MemoryStorage) CreateAuthorizeCodeSession(_ context.Context, code string, request fosite.Requester) error {
        if code == "" {
                return fosite.ErrInvalidRequest.WithHint("authorization code cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        stored := snapshotRequester(request)
        expiresAt := getExpirationFromRequester(stored, fosite.AuthorizeCode, DefaultAuthCodeTTL)

        s.authCodes[code] = &timedEntry[fosite.Requester]{
                value:     stored,
                createdAt: now,
                expiresAt: expiresAt,
        }
        return nil
}

// GetAuthorizeCodeSession retrieves the authorization request for a given code.
// If the authorization code has been invalidated, it returns ErrInvalidatedAuthorizeCode
// along with the request (as required by fosite). Expired codes are rejected
// with a not-found error, after the invalidation check so reuse is still detected.
func (s *MemoryStorage) GetAuthorizeCodeSession(_ context.Context, code string, _ fosite.Session) (fosite.Requester, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.authCodes[code]
        if !ok {
                slog.Debug("authorization code not found")
                return nil, notFoundRFC6749Error("Authorization code not found")
        }

        // Check if the code has been invalidated
        if s.invalidatedCodes[code] != nil {
                // Must return the request along with the error as per fosite documentation
                return snapshotRequester(entry.value), fosite.ErrInvalidatedAuthorizeCode
        }

        if time.Now().After(entry.expiresAt) {
                slog.Debug("authorization code expired")
                return nil, notFoundRFC6749Error("Authorization code expired")
        }

        return snapshotRequester(entry.value), nil
}

// InvalidateAuthorizeCodeSession marks an authorization code as used/invalid.
// Subsequent calls to GetAuthorizeCodeSession will return ErrInvalidatedAuthorizeCode.
func (s *MemoryStorage) InvalidateAuthorizeCodeSession(_ context.Context, code string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.authCodes[code]; !ok {
                slog.Debug("authorization code not found for invalidation")
                return notFoundRFC6749Error("Authorization code not found")
        }

        now := time.Now()
        s.invalidatedCodes[code] = &timedEntry[bool]{
                value:     true,
                createdAt: now,
                expiresAt: now.Add(DefaultInvalidatedCodeTTL),
        }
        return nil
}

// -----------------------
// oauth2.AccessTokenStorage
// -----------------------

// CreateAccessTokenSession stores the access token session.
func (s *MemoryStorage) CreateAccessTokenSession(_ context.Context, signature string, request fosite.Requester) error {
        if signature == "" {
                return fosite.ErrInvalidRequest.WithHint("access token signature cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        stored := snapshotRequester(request)
        expiresAt := getExpirationFromRequester(stored, fosite.AccessToken, DefaultAccessTokenTTL)

        s.accessTokens[signature] = &timedEntry[fosite.Requester]{
                value:     stored,
                createdAt: now,
                expiresAt: expiresAt,
        }
        return nil
}

// GetAccessTokenSession retrieves the access token session by its signature.
//
// The session parameter is a prototype for deserialization in persistent backends.
// Our in-memory implementation ignores it because it snapshots Requesters in memory.
// A persistent backend may instead use it to choose a deserialization type.
func (s *MemoryStorage) GetAccessTokenSession(_ context.Context, signature string, _ fosite.Session) (fosite.Requester, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.accessTokens[signature]
        if !ok {
                slog.Debug("access token not found")
                return nil, notFoundRFC6749Error("Access token not found")
        }
        return snapshotRequester(entry.value), nil
}

// DeleteAccessTokenSession removes the access token session.
func (s *MemoryStorage) DeleteAccessTokenSession(_ context.Context, signature string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.accessTokens[signature]; !ok {
                return notFoundRFC6749Error("Access token not found")
        }
        delete(s.accessTokens, signature)
        return nil
}

// -----------------------
// oauth2.RefreshTokenStorage
// -----------------------

// CreateRefreshTokenSession stores the refresh token session.
// The accessSignature parameter is used to link the refresh token to its access token.
// TODO: Store the accessSignature in a refreshToAccess map to enable direct lookup
// during token rotation instead of O(n) scan by request ID in RotateRefreshToken.
func (s *MemoryStorage) CreateRefreshTokenSession(_ context.Context, signature string, _ string, request fosite.Requester) error {
        if signature == "" {
                return fosite.ErrInvalidRequest.WithHint("refresh token signature cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        stored := snapshotRequester(request)
        expiresAt := getExpirationFromRequester(stored, fosite.RefreshToken, DefaultRefreshTokenTTL)

        s.refreshTokens[signature] = &timedEntry[fosite.Requester]{
                value:     stored,
                createdAt: now,
                expiresAt: expiresAt,
        }
        return nil
}

// GetRefreshTokenSession retrieves the refresh token session by its signature.
func (s *MemoryStorage) GetRefreshTokenSession(_ context.Context, signature string, _ fosite.Session) (fosite.Requester, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.refreshTokens[signature]
        if !ok {
                slog.Debug("refresh token not found")
                return nil, notFoundRFC6749Error("Refresh token not found")
        }
        return snapshotRequester(entry.value), nil
}

// DeleteRefreshTokenSession removes the refresh token session.
func (s *MemoryStorage) DeleteRefreshTokenSession(_ context.Context, signature string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.refreshTokens[signature]; !ok {
                return notFoundRFC6749Error("Refresh token not found")
        }
        delete(s.refreshTokens, signature)
        return nil
}

// RotateRefreshToken invalidates a refresh token and all its related token data.
// This is called during token refresh to implement refresh token rotation.
func (s *MemoryStorage) RotateRefreshToken(_ context.Context, requestID string, refreshTokenSignature string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        // Delete the specific refresh token
        delete(s.refreshTokens, refreshTokenSignature)

        // TODO: Use the refreshToAccess map (once implemented) for direct access token lookup
        // instead of O(n) scan by request ID, which may delete unrelated tokens sharing the same ID.
        // Also delete any access tokens associated with this request ID
        for sig, entry := range s.accessTokens {
                if entry.value.GetID() == requestID {
                        delete(s.accessTokens, sig)
                }
        }

        return nil
}

// -----------------------
// oauth2.TokenRevocationStorage
// -----------------------

// RevokeAccessToken marks an access token as revoked.
// This method implements the oauth2.TokenRevocationStorage interface.
//
// Note: This takes requestID, not signature. Per RFC 7009, revoking by request ID
// enables revoking ALL tokens from the same authorization grant. This is why we
// store the full Requester (with its ID) rather than just token values.
//
// The O(n) scan by request ID is acceptable for in-memory storage. Production
// implementations should maintain a reverse index (request_id -> signatures).
func (s *MemoryStorage) RevokeAccessToken(_ context.Context, requestID string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        // Find and remove all access tokens associated with this request ID.
        // Uses Requester.GetID() to match the grant identifier, not the token signature.
        for sig, entry := range s.accessTokens {
                if entry.value.GetID() == requestID {
                        delete(s.accessTokens, sig)
                }
        }

        return nil
}

// RevokeRefreshToken marks a refresh token as revoked.
// This method implements the oauth2.TokenRevocationStorage interface.
//
// Like RevokeAccessToken, this takes requestID to find all refresh tokens from
// the same authorization grant. Per RFC 7009 Section 2.1, implementations SHOULD
// also revoke associated access tokens, which RotateRefreshToken handles.
func (s *MemoryStorage) RevokeRefreshToken(_ context.Context, requestID string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        // Find and remove all refresh tokens associated with this request ID.
        for sig, entry := range s.refreshTokens {
                if entry.value.GetID() == requestID {
                        delete(s.refreshTokens, sig)
                }
        }

        return nil
}

// RevokeRefreshTokenMaybeGracePeriod marks a refresh token as revoked, optionally allowing
// a grace period during which the old token is still valid.
// For this implementation, we don't support grace periods and revoke immediately.
func (s *MemoryStorage) RevokeRefreshTokenMaybeGracePeriod(ctx context.Context, requestID string, _ string) error {
        return s.RevokeRefreshToken(ctx, requestID)
}

// -----------------------
// pkce.PKCERequestStorage
// -----------------------

// CreatePKCERequestSession stores the PKCE request session.
func (s *MemoryStorage) CreatePKCERequestSession(_ context.Context, signature string, request fosite.Requester) error {
        if signature == "" {
                return fosite.ErrInvalidRequest.WithHint("PKCE signature cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        stored := snapshotRequester(request)
        expiresAt := getExpirationFromRequester(stored, fosite.AuthorizeCode, DefaultPKCETTL)

        s.pkceRequests[signature] = &timedEntry[fosite.Requester]{
                value:     stored,
                createdAt: now,
                expiresAt: expiresAt,
        }
        return nil
}

// GetPKCERequestSession retrieves the PKCE request session by its signature.
func (s *MemoryStorage) GetPKCERequestSession(_ context.Context, signature string, _ fosite.Session) (fosite.Requester, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.pkceRequests[signature]
        if !ok {
                slog.Debug("pkce request not found")
                return nil, notFoundRFC6749Error("PKCE request not found")
        }
        return snapshotRequester(entry.value), nil
}

// DeletePKCERequestSession removes the PKCE request session.
func (s *MemoryStorage) DeletePKCERequestSession(_ context.Context, signature string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.pkceRequests[signature]; !ok {
                return notFoundRFC6749Error("PKCE request not found")
        }
        delete(s.pkceRequests, signature)
        return nil
}

// -----------------------
// Upstream Token Storage
// -----------------------

// ResolveUpstreamTokenRowID returns an opaque ID for process-local refresh
// coordination of a logical upstream-token row.
func (*MemoryStorage) ResolveUpstreamTokenRowID(_ context.Context, sessionID, providerName string) (UpstreamTokenRowID, error) {
        if sessionID == "" {
                return "", fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return "", fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }
        return upstreamTokenRowID(sessionID, providerName), nil
}

// StoreUpstreamTokens stores the upstream IDP tokens for a session and provider.
// A defensive copy is made to prevent aliasing issues.
func (s *MemoryStorage) StoreUpstreamTokens(_ context.Context, sessionID, providerName string, tokens *UpstreamTokens) error {
        if sessionID == "" {
                return fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        s.storeUpstreamTokensLocked(upstreamKey{sessionID, providerName}, tokens)
        return nil
}

// CompareAndSwapUpstreamTokens stores tokens for (sessionID, providerName)
// only if the refresh token currently held there equals expectedRefreshToken.
// See the interface doc (UpstreamTokenStorage.CompareAndSwapUpstreamTokens)
// for the coordination contract. The compare-then-write happens under s.mu,
// so it is atomic with respect to every other reader/writer of this backend.
func (s *MemoryStorage) CompareAndSwapUpstreamTokens(
        _ context.Context, sessionID, providerName, expectedRefreshToken string, tokens *UpstreamTokens,
) error {
        if sessionID == "" {
                return fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        key := upstreamKey{sessionID, providerName}
        var currentRefreshToken string
        if entry, ok := s.upstreamTokens[key]; ok && entry.value != nil {
                currentRefreshToken = entry.value.RefreshToken
        }
        if currentRefreshToken != expectedRefreshToken {
                return ErrConcurrentRefresh
        }

        s.storeUpstreamTokensLocked(key, tokens)
        return nil
}

// storeUpstreamTokensLocked writes tokens for key, replacing any existing
// entry. Callers must hold s.mu for writing.
func (s *MemoryStorage) storeUpstreamTokensLocked(key upstreamKey, tokens *UpstreamTokens) {
        now := time.Now()
        // Add DefaultRefreshTokenTTL beyond access token expiry so the refresh token
        // survives in storage for transparent token refresh by the middleware.
        // Zero ExpiresAt means the token never expires; no TTL is applied.
        expiresAt := func() time.Time {
                if tokens == nil {
                        return now.Add(DefaultAccessTokenTTL + DefaultRefreshTokenTTL)
                }
                if !tokens.ExpiresAt.IsZero() {
                        return tokens.ExpiresAt.Add(DefaultRefreshTokenTTL)
                }
                if !tokens.SessionExpiresAt.IsZero() {
                        return tokens.SessionExpiresAt.Add(DefaultRefreshTokenTTL)
                }
                return time.Time{} // non-expiring token with no known session bound
        }()

        s.upstreamTokens[key] = &timedEntry[*UpstreamTokens]{
                value:     cloneUpstreamTokens(tokens),
                createdAt: now,
                expiresAt: expiresAt,
        }
}

// cloneUpstreamTokens returns a field-by-field copy of t, or nil if t is nil.
func cloneUpstreamTokens(t *UpstreamTokens) *UpstreamTokens {
        if t == nil {
                return nil
        }
        return &UpstreamTokens{
                ProviderID:       t.ProviderID,
                AccessToken:      t.AccessToken,
                RefreshToken:     t.RefreshToken,
                IDToken:          t.IDToken,
                ExpiresAt:        t.ExpiresAt,
                SessionExpiresAt: t.SessionExpiresAt,
                UserID:           t.UserID,
                UpstreamSubject:  t.UpstreamSubject,
                ClientID:         t.ClientID,
        }
}

// GetUpstreamTokens retrieves the upstream IDP tokens for a session and provider.
// Returns a defensive copy to prevent aliasing issues.
func (s *MemoryStorage) GetUpstreamTokens(_ context.Context, sessionID, providerName string) (*UpstreamTokens, error) {
        if sessionID == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.upstreamTokens[upstreamKey{sessionID, providerName}]
        if !ok {
                slog.Debug("upstream tokens not found", "session_id", sessionID, "provider_name", providerName)
                return nil, notFoundRFC6749Error("Upstream tokens not found")
        }

        // Return a defensive copy to prevent aliasing issues
        result := cloneUpstreamTokens(entry.value)
        if result == nil {
                return nil, nil
        }

        // Check the token's own ExpiresAt (access token expiry), not the entry's expiresAt
        // (storage TTL which includes DefaultRefreshTokenTTL buffer for refresh token survival).
        // Return tokens along with ErrExpired so callers can use the refresh token.
        if !result.ExpiresAt.IsZero() && time.Now().After(result.ExpiresAt) {
                slog.Debug("upstream tokens expired", "session_id", sessionID, "provider_name", providerName)
                return result, ErrExpired
        }

        return result, nil
}

// GetAllUpstreamTokens retrieves all upstream IDP tokens for a session across all providers.
// Returns a map of providerName -> tokens with defensive copies.
// Returns an empty map (not error) for unknown sessions.
// Includes expired tokens (no expiry filtering at bulk-read level).
func (s *MemoryStorage) GetAllUpstreamTokens(_ context.Context, sessionID string) (map[string]*UpstreamTokens, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        result := make(map[string]*UpstreamTokens)
        for key, entry := range s.upstreamTokens {
                if key.sessionID != sessionID {
                        continue
                }
                // Defensive copy (cloneUpstreamTokens handles nil)
                result[key.providerName] = cloneUpstreamTokens(entry.value)
        }

        return result, nil
}

// DeleteUpstreamTokens removes all upstream IDP tokens for a session (all providers).
func (s *MemoryStorage) DeleteUpstreamTokens(_ context.Context, sessionID string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        found := false
        for key := range s.upstreamTokens {
                if key.sessionID == sessionID {
                        delete(s.upstreamTokens, key)
                        found = true
                }
        }
        if !found {
                return notFoundRFC6749Error("Upstream tokens not found")
        }
        return nil
}

// DeleteUpstreamTokensForProvider removes tokens for a single (sessionID, providerName),
// leaving sibling providers' rows intact. Absent row returns nil (not ErrNotFound).
func (s *MemoryStorage) DeleteUpstreamTokensForProvider(_ context.Context, sessionID, providerName string) error {
        if sessionID == "" {
                return fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        delete(s.upstreamTokens, upstreamKey{sessionID, providerName})
        return nil
}

// compareExpiry orders ExpiresAt values for the GetLatestUpstreamTokensForUser
// tie-breaker. Non-expiring rows (zero ExpiresAt — "alive forever") rank latest;
// among finite expiries, later ranks latest. Mirrors time.Compare but with the
// zero sentinel reinterpreted. Returns -1/0/+1.
func compareExpiry(a, b time.Time) int {
        aZero, bZero := a.IsZero(), b.IsZero()
        switch {
        case aZero && bZero:
                return 0
        case aZero:
                return 1
        case bZero:
                return -1
        }
        return a.Compare(b)
}

// GetLatestUpstreamTokensForUser implements UpstreamTokenStorage.
//
// Expired tokens (past ExpiresAt) are returned so callers can use the refresh
// token; filtering by access-token expiry is the caller's responsibility.
// See the interface declaration in types.go for the full contract.
func (s *MemoryStorage) GetLatestUpstreamTokensForUser(_ context.Context, userID, providerID string) (*UpstreamTokens, error) {
        if userID == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("user ID cannot be empty")
        }
        if providerID == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("provider ID cannot be empty")
        }

        s.mu.RLock()
        defer s.mu.RUnlock()

        var winner *UpstreamTokens
        for _, entry := range s.upstreamTokens {
                if entry.value == nil || entry.value.UserID != userID || entry.value.ProviderID != providerID {
                        continue
                }
                if winner == nil || compareExpiry(entry.value.ExpiresAt, winner.ExpiresAt) > 0 {
                        winner = entry.value
                }
        }

        if winner == nil {
                return nil, notFoundRFC6749Error("Upstream tokens not found")
        }

        return cloneUpstreamTokens(winner), nil
}

// -----------------------
// Pending Authorization Storage
// -----------------------

// StorePendingAuthorization stores a pending authorization request.
// The pending authorization is keyed by the internal state used to correlate
// the upstream IDP callback.
func (s *MemoryStorage) StorePendingAuthorization(_ context.Context, state string, pending *PendingAuthorization) error {
        if state == "" {
                return fosite.ErrInvalidRequest.WithHint("state cannot be empty")
        }
        if pending == nil {
                return fosite.ErrInvalidRequest.WithHint("pending authorization cannot be nil")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        expiresAt := pending.CreatedAt.Add(DefaultPendingAuthorizationTTL)
        if !expiresAt.After(now) {
                return ErrExpired
        }

        // Make a defensive copy to prevent aliasing issues
        pendingCopy := &PendingAuthorization{
                ClientID:                pending.ClientID,
                RedirectURI:             pending.RedirectURI,
                State:                   pending.State,
                PKCEChallenge:           pending.PKCEChallenge,
                PKCEMethod:              pending.PKCEMethod,
                Scopes:                  slices.Clone(pending.Scopes),
                Resource:                pending.Resource,
                ConsentStage:            pending.ConsentStage,
                RememberConsent:         pending.RememberConsent,
                FirstProviderSubject:    pending.FirstProviderSubject,
                ExpectedUserID:          pending.ExpectedUserID,
                ExpectedProviderSubject: pending.ExpectedProviderSubject,
                ConsentSessionDigest:    pending.ConsentSessionDigest,
                InternalState:           pending.InternalState,
                UpstreamPKCEVerifier:    pending.UpstreamPKCEVerifier,
                UpstreamNonce:           pending.UpstreamNonce,
                BrowserBindingHash:      pending.BrowserBindingHash,
                UpstreamProviderName:    pending.UpstreamProviderName,
                SessionID:               pending.SessionID,
                ResolvedUserID:          pending.ResolvedUserID,
                ResolvedUserName:        pending.ResolvedUserName,
                ResolvedUserEmail:       pending.ResolvedUserEmail,
                VerifiedUpstream:        pending.VerifiedUpstream,
                SingleLeg:               pending.SingleLeg,
                ChainUpstreams:          slices.Clone(pending.ChainUpstreams),
                CreatedAt:               pending.CreatedAt,
        }

        s.pendingAuthorizations[state] = &timedEntry[*PendingAuthorization]{
                value:     pendingCopy,
                createdAt: now,
                expiresAt: expiresAt,
        }
        return nil
}

// LoadPendingAuthorization retrieves a pending authorization by internal state.
// Returns a defensive copy to prevent aliasing issues.
func (s *MemoryStorage) LoadPendingAuthorization(_ context.Context, state string) (*PendingAuthorization, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.pendingAuthorizations[state]
        if !ok {
                slog.Debug("pending authorization not found")
                return nil, notFoundRFC6749Error("Pending authorization not found")
        }

        // Check if expired
        if !entry.expiresAt.After(time.Now()) {
                slog.Debug("pending authorization expired")
                return nil, ErrExpired
        }

        // Return a defensive copy to prevent aliasing issues
        pending := entry.value
        if pending == nil {
                return nil, nil
        }
        return &PendingAuthorization{
                ClientID:                pending.ClientID,
                RedirectURI:             pending.RedirectURI,
                State:                   pending.State,
                PKCEChallenge:           pending.PKCEChallenge,
                PKCEMethod:              pending.PKCEMethod,
                Scopes:                  slices.Clone(pending.Scopes),
                Resource:                pending.Resource,
                ConsentStage:            pending.ConsentStage,
                RememberConsent:         pending.RememberConsent,
                FirstProviderSubject:    pending.FirstProviderSubject,
                ExpectedUserID:          pending.ExpectedUserID,
                ExpectedProviderSubject: pending.ExpectedProviderSubject,
                ConsentSessionDigest:    pending.ConsentSessionDigest,
                InternalState:           pending.InternalState,
                UpstreamPKCEVerifier:    pending.UpstreamPKCEVerifier,
                UpstreamNonce:           pending.UpstreamNonce,
                BrowserBindingHash:      pending.BrowserBindingHash,
                UpstreamProviderName:    pending.UpstreamProviderName,
                SessionID:               pending.SessionID,
                ResolvedUserID:          pending.ResolvedUserID,
                ResolvedUserName:        pending.ResolvedUserName,
                ResolvedUserEmail:       pending.ResolvedUserEmail,
                VerifiedUpstream:        pending.VerifiedUpstream,
                SingleLeg:               pending.SingleLeg,
                ChainUpstreams:          slices.Clone(pending.ChainUpstreams),
                CreatedAt:               pending.CreatedAt,
        }, nil
}

// DeletePendingAuthorization removes a pending authorization.
func (s *MemoryStorage) DeletePendingAuthorization(_ context.Context, state string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.pendingAuthorizations[state]; !ok {
                return notFoundRFC6749Error("Pending authorization not found")
        }
        delete(s.pendingAuthorizations, state)
        return nil
}

// -----------------------
// Pending Device Login Storage
// -----------------------

// StorePendingDeviceLogin stores a pending device-flow verification-page
// login, keyed by the internal state used to correlate the upstream IDP
// callback.
func (s *MemoryStorage) StorePendingDeviceLogin(_ context.Context, state string, pending *PendingDeviceLogin) error {
        if state == "" {
                return fosite.ErrInvalidRequest.WithHint("state cannot be empty")
        }
        if pending == nil {
                return fosite.ErrInvalidRequest.WithHint("pending device login cannot be nil")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        pendingCopy := &PendingDeviceLogin{
                DeviceCode:           pending.DeviceCode,
                UserCode:             pending.UserCode,
                UpstreamPKCEVerifier: pending.UpstreamPKCEVerifier,
                UpstreamNonce:        pending.UpstreamNonce,
                BrowserBindingHash:   pending.BrowserBindingHash,
                UpstreamProviderName: pending.UpstreamProviderName,
                CreatedAt:            pending.CreatedAt,
        }

        s.pendingDeviceLogins[state] = &timedEntry[*PendingDeviceLogin]{
                value:     pendingCopy,
                createdAt: now,
                expiresAt: now.Add(DefaultDeviceLoginTTL),
        }
        return nil
}

// LoadPendingDeviceLogin retrieves a pending device login by state. Returns a
// defensive copy to prevent aliasing issues.
func (s *MemoryStorage) LoadPendingDeviceLogin(_ context.Context, state string) (*PendingDeviceLogin, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.pendingDeviceLogins[state]
        if !ok {
                return nil, notFoundRFC6749Error("Pending device login not found")
        }
        if time.Now().After(entry.expiresAt) {
                return nil, ErrExpired
        }

        pending := entry.value
        return &PendingDeviceLogin{
                DeviceCode:           pending.DeviceCode,
                UserCode:             pending.UserCode,
                UpstreamPKCEVerifier: pending.UpstreamPKCEVerifier,
                UpstreamNonce:        pending.UpstreamNonce,
                BrowserBindingHash:   pending.BrowserBindingHash,
                UpstreamProviderName: pending.UpstreamProviderName,
                CreatedAt:            pending.CreatedAt,
        }, nil
}

// DeletePendingDeviceLogin removes a pending device login.
func (s *MemoryStorage) DeletePendingDeviceLogin(_ context.Context, state string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.pendingDeviceLogins[state]; !ok {
                return notFoundRFC6749Error("Pending device login not found")
        }
        delete(s.pendingDeviceLogins, state)
        return nil
}

// -----------------------
// Pending Device Confirmation Storage
// -----------------------

// StorePendingDeviceConfirmation stores a pending device-flow confirmation,
// keyed by an opaque token minted by the caller. See
// PendingDeviceConfirmation's doc comment for why the resolved identity is
// addressed by token rather than round-tripped through the browser.
func (s *MemoryStorage) StorePendingDeviceConfirmation(
        _ context.Context, token string, pending *PendingDeviceConfirmation,
) error {
        if token == "" {
                return fosite.ErrInvalidRequest.WithHint("token cannot be empty")
        }
        if pending == nil {
                return fosite.ErrInvalidRequest.WithHint("pending device confirmation cannot be nil")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        now := time.Now()
        pendingCopy := &PendingDeviceConfirmation{
                DeviceCode:        pending.DeviceCode,
                UserCode:          pending.UserCode,
                ResolvedUserID:    pending.ResolvedUserID,
                ResolvedUserName:  pending.ResolvedUserName,
                ResolvedUserEmail: pending.ResolvedUserEmail,
                VerifiedUpstream:  pending.VerifiedUpstream,
                UpstreamTokens:    cloneUpstreamTokens(pending.UpstreamTokens),
                Synthetic:         pending.Synthetic,
                CreatedAt:         pending.CreatedAt,
        }

        s.pendingDeviceConfirmations[token] = &timedEntry[*PendingDeviceConfirmation]{
                value:     pendingCopy,
                createdAt: now,
                expiresAt: now.Add(DefaultDeviceLoginTTL),
        }
        return nil
}

// LoadPendingDeviceConfirmation retrieves a pending device confirmation by
// token. Returns a defensive copy to prevent aliasing issues.
func (s *MemoryStorage) LoadPendingDeviceConfirmation(
        _ context.Context, token string,
) (*PendingDeviceConfirmation, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.pendingDeviceConfirmations[token]
        if !ok {
                return nil, notFoundRFC6749Error("Pending device confirmation not found")
        }
        if time.Now().After(entry.expiresAt) {
                return nil, ErrExpired
        }

        pending := entry.value
        return &PendingDeviceConfirmation{
                DeviceCode:        pending.DeviceCode,
                UserCode:          pending.UserCode,
                ResolvedUserID:    pending.ResolvedUserID,
                ResolvedUserName:  pending.ResolvedUserName,
                ResolvedUserEmail: pending.ResolvedUserEmail,
                VerifiedUpstream:  pending.VerifiedUpstream,
                UpstreamTokens:    cloneUpstreamTokens(pending.UpstreamTokens),
                Synthetic:         pending.Synthetic,
                CreatedAt:         pending.CreatedAt,
        }, nil
}

// DeletePendingDeviceConfirmation removes a pending device confirmation.
func (s *MemoryStorage) DeletePendingDeviceConfirmation(_ context.Context, token string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.pendingDeviceConfirmations[token]; !ok {
                return notFoundRFC6749Error("Pending device confirmation not found")
        }
        delete(s.pendingDeviceConfirmations, token)
        return nil
}

// -----------------------
// Device Code Storage
// -----------------------

// cloneDeviceRequest returns a defensive copy of device, cloning its slice
// fields so neither the caller nor the store can mutate the other's data
// through a shared backing array.
func cloneDeviceRequest(device *DeviceRequest) *DeviceRequest {
        return &DeviceRequest{
                DeviceCode:        device.DeviceCode,
                UserCode:          device.UserCode,
                ClientID:          device.ClientID,
                Scopes:            slices.Clone(device.Scopes),
                Audience:          slices.Clone(device.Audience),
                Status:            device.Status,
                Interval:          device.Interval,
                LastPolledAt:      device.LastPolledAt,
                ResolvedUserID:    device.ResolvedUserID,
                ResolvedUserName:  device.ResolvedUserName,
                ResolvedUserEmail: device.ResolvedUserEmail,
                VerifiedUpstream:  device.VerifiedUpstream,
                SessionID:         device.SessionID,
                CreatedAt:         device.CreatedAt,
        }
}

// getUnexpiredDeviceRequestEntry looks up the device request entry keyed by
// deviceCode, returning ErrNotFound if absent or ErrExpired if its TTL has
// elapsed. Callers must hold s.mu (read or write lock).
func (s *MemoryStorage) getUnexpiredDeviceRequestEntry(deviceCode string) (*timedEntry[*DeviceRequest], error) {
        entry, ok := s.deviceRequests[deviceCode]
        if !ok {
                return nil, fmt.Errorf("%w: device request not found", ErrNotFound)
        }
        if time.Now().After(entry.expiresAt) {
                return nil, ErrExpired
        }
        return entry, nil
}

// StoreDeviceRequest stores a new pending device request, indexed by both
// DeviceCode and UserCode.
func (s *MemoryStorage) StoreDeviceRequest(_ context.Context, device *DeviceRequest) error {
        if device == nil {
                return fosite.ErrInvalidRequest.WithHint("device request cannot be nil")
        }
        if device.DeviceCode == "" {
                return fosite.ErrInvalidRequest.WithHint("device code cannot be empty")
        }
        if device.UserCode == "" {
                return fosite.ErrInvalidRequest.WithHint("user code cannot be empty")
        }
        if device.Status != DeviceRequestStatusPending {
                return fosite.ErrInvalidRequest.WithHint("device request must be created with pending status")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.deviceRequests[device.DeviceCode]; ok {
                return fmt.Errorf("%w: device code %q", ErrAlreadyExists, device.DeviceCode)
        }
        if _, ok := s.deviceRequestsByUserCode[device.UserCode]; ok {
                return fmt.Errorf("%w: user code %q", ErrAlreadyExists, device.UserCode)
        }

        now := time.Now()
        s.deviceRequests[device.DeviceCode] = &timedEntry[*DeviceRequest]{
                value:     cloneDeviceRequest(device),
                createdAt: now,
                expiresAt: now.Add(DefaultDeviceRequestTTL),
        }
        s.deviceRequestsByUserCode[device.UserCode] = device.DeviceCode
        return nil
}

// LoadDeviceRequestByDeviceCode retrieves a device request by its device_code.
// Returns a defensive copy to prevent aliasing issues.
func (s *MemoryStorage) LoadDeviceRequestByDeviceCode(_ context.Context, deviceCode string) (*DeviceRequest, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, err := s.getUnexpiredDeviceRequestEntry(deviceCode)
        if err != nil {
                return nil, err
        }
        return cloneDeviceRequest(entry.value), nil
}

// LoadDeviceRequestByUserCode retrieves a device request by its user_code,
// for the verification page. Returns a defensive copy to prevent aliasing issues.
func (s *MemoryStorage) LoadDeviceRequestByUserCode(_ context.Context, userCode string) (*DeviceRequest, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        deviceCode, ok := s.deviceRequestsByUserCode[userCode]
        if !ok {
                return nil, fmt.Errorf("%w: device request not found", ErrNotFound)
        }
        entry, err := s.getUnexpiredDeviceRequestEntry(deviceCode)
        if err != nil {
                return nil, err
        }
        return cloneDeviceRequest(entry.value), nil
}

// MarkDeviceRequestAuthorized transitions a pending device request to
// authorized, attaching the resolved identity.
func (s *MemoryStorage) MarkDeviceRequestAuthorized(
        _ context.Context, deviceCode string, auth DeviceAuthorization,
) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        entry, err := s.getUnexpiredDeviceRequestEntry(deviceCode)
        if err != nil {
                return err
        }
        if entry.value.Status != DeviceRequestStatusPending {
                return fmt.Errorf("%w: device request is %q, not pending", ErrInvalidState, entry.value.Status)
        }

        entry.value.Status = DeviceRequestStatusAuthorized
        entry.value.ResolvedUserID = auth.UserID
        entry.value.ResolvedUserName = auth.UserName
        entry.value.ResolvedUserEmail = auth.UserEmail
        entry.value.VerifiedUpstream = auth.VerifiedUpstream
        entry.value.SessionID = auth.SessionID
        return nil
}

// MarkDeviceRequestDenied transitions a pending device request to denied.
func (s *MemoryStorage) MarkDeviceRequestDenied(_ context.Context, deviceCode string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        entry, err := s.getUnexpiredDeviceRequestEntry(deviceCode)
        if err != nil {
                return err
        }
        if entry.value.Status != DeviceRequestStatusPending {
                return fmt.Errorf("%w: device request is %q, not pending", ErrInvalidState, entry.value.Status)
        }

        entry.value.Status = DeviceRequestStatusDenied
        return nil
}

// UpdateDeviceRequestLastPolledAt records the time of the most recent poll.
func (s *MemoryStorage) UpdateDeviceRequestLastPolledAt(_ context.Context, deviceCode string, polledAt time.Time) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        entry, err := s.getUnexpiredDeviceRequestEntry(deviceCode)
        if err != nil {
                return err
        }
        entry.value.LastPolledAt = polledAt
        return nil
}

// DeleteDeviceRequest removes a device request, e.g. once its token has been issued.
func (s *MemoryStorage) DeleteDeviceRequest(_ context.Context, deviceCode string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        entry, ok := s.deviceRequests[deviceCode]
        if !ok {
                return fmt.Errorf("%w: device request not found", ErrNotFound)
        }
        delete(s.deviceRequests, deviceCode)
        delete(s.deviceRequestsByUserCode, entry.value.UserCode)
        return nil
}

// -----------------------
// User Storage
// -----------------------

// providerIdentityKey creates a unique key for a provider identity.
// The key format is "len(providerID):providerID:providerSubject" for O(1) lookup.
// The length prefix ensures collision-free keys even if providerID or providerSubject
// contain colons (which is valid per RFC 7519 StringOrURI semantics for OIDC subjects).
func providerIdentityKey(providerID, providerSubject string) string {
        return fmt.Sprintf("%d:%s:%s", len(providerID), providerID, providerSubject)
}

// CreateUser creates a new user account.
// Returns ErrAlreadyExists if a user with the same ID already exists.
func (s *MemoryStorage) CreateUser(_ context.Context, user *User) error {
        if user == nil {
                return fosite.ErrInvalidRequest.WithHint("user cannot be nil")
        }
        if user.ID == "" {
                return fosite.ErrInvalidRequest.WithHint("user ID cannot be empty")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        if _, exists := s.users[user.ID]; exists {
                return fmt.Errorf("%w: user already exists", ErrAlreadyExists)
        }

        // Make a defensive copy
        s.users[user.ID] = &User{
                ID:        user.ID,
                CreatedAt: user.CreatedAt,
                UpdatedAt: user.UpdatedAt,
        }
        return nil
}

// GetUser retrieves a user by their internal ID.
// Returns ErrNotFound if the user does not exist.
func (s *MemoryStorage) GetUser(_ context.Context, id string) (*User, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        user, ok := s.users[id]
        if !ok {
                return nil, fmt.Errorf("%w: user not found", ErrNotFound)
        }

        // Return a defensive copy
        return &User{
                ID:        user.ID,
                CreatedAt: user.CreatedAt,
                UpdatedAt: user.UpdatedAt,
        }, nil
}

// DeleteUser removes a user account and all associated provider identities.
// Returns ErrNotFound if the user does not exist.
func (s *MemoryStorage) DeleteUser(_ context.Context, id string) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if _, exists := s.users[id]; !exists {
                return fmt.Errorf("%w: user not found", ErrNotFound)
        }

        // Delete all associated provider identities
        for key, identity := range s.providerIdentities {
                if identity.UserID == id {
                        delete(s.providerIdentities, key)
                }
        }

        // Delete all associated upstream tokens
        for key, entry := range s.upstreamTokens {
                if entry.value != nil && entry.value.UserID == id {
                        delete(s.upstreamTokens, key)
                }
        }

        delete(s.users, id)
        return nil
}

// CreateProviderIdentity links a provider identity to a user.
// Returns ErrAlreadyExists if this provider identity is already linked.
func (s *MemoryStorage) CreateProviderIdentity(_ context.Context, identity *ProviderIdentity) error {
        if identity == nil {
                return fosite.ErrInvalidRequest.WithHint("identity cannot be nil")
        }
        if identity.UserID == "" {
                return fosite.ErrInvalidRequest.WithHint("user ID cannot be empty")
        }
        if identity.ProviderID == "" {
                return fosite.ErrInvalidRequest.WithHint("provider ID cannot be empty")
        }
        if identity.ProviderSubject == "" {
                return fosite.ErrInvalidRequest.WithHint("provider subject cannot be empty")
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        // Verify user exists before linking identity
        if _, exists := s.users[identity.UserID]; !exists {
                return fmt.Errorf("%w: user not found", ErrNotFound)
        }

        key := providerIdentityKey(identity.ProviderID, identity.ProviderSubject)
        if _, exists := s.providerIdentities[key]; exists {
                return fmt.Errorf("%w: provider identity already linked", ErrAlreadyExists)
        }

        // Make a defensive copy
        s.providerIdentities[key] = &ProviderIdentity{
                UserID:          identity.UserID,
                ProviderID:      identity.ProviderID,
                ProviderSubject: identity.ProviderSubject,
                LinkedAt:        identity.LinkedAt,
                LastUsedAt:      identity.LastUsedAt,
        }
        return nil
}

// GetProviderIdentity retrieves a provider identity by provider ID and subject.
// This is the primary lookup path during authentication callbacks.
// Returns ErrNotFound if the identity does not exist.
func (s *MemoryStorage) GetProviderIdentity(_ context.Context, providerID, providerSubject string) (*ProviderIdentity, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        key := providerIdentityKey(providerID, providerSubject)
        identity, ok := s.providerIdentities[key]
        if !ok {
                return nil, fmt.Errorf("%w: provider identity not found", ErrNotFound)
        }

        // Return a defensive copy
        return &ProviderIdentity{
                UserID:          identity.UserID,
                ProviderID:      identity.ProviderID,
                ProviderSubject: identity.ProviderSubject,
                LinkedAt:        identity.LinkedAt,
                LastUsedAt:      identity.LastUsedAt,
        }, nil
}

// UpdateProviderIdentityLastUsed updates the LastUsedAt timestamp for a provider identity.
// This should be called after each successful authentication via this identity.
// Returns ErrNotFound if the identity does not exist.
func (s *MemoryStorage) UpdateProviderIdentityLastUsed(
        _ context.Context, providerID, providerSubject string, lastUsedAt time.Time,
) error {
        s.mu.Lock()
        defer s.mu.Unlock()

        key := providerIdentityKey(providerID, providerSubject)
        identity, ok := s.providerIdentities[key]
        if !ok {
                return fmt.Errorf("%w: provider identity not found", ErrNotFound)
        }

        identity.LastUsedAt = lastUsedAt
        return nil
}

// GetUserProviderIdentities returns all provider identities linked to a user.
// Returns an empty slice (not error) if the user exists but has no linked identities.
// Returns ErrNotFound if the user does not exist.
func (s *MemoryStorage) GetUserProviderIdentities(_ context.Context, userID string) ([]*ProviderIdentity, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        // Verify user exists
        if _, exists := s.users[userID]; !exists {
                return nil, fmt.Errorf("%w: user not found", ErrNotFound)
        }

        // Collect all identities for this user
        var identities []*ProviderIdentity
        for _, identity := range s.providerIdentities {
                if identity.UserID == userID {
                        // Return defensive copies
                        identities = append(identities, &ProviderIdentity{
                                UserID:          identity.UserID,
                                ProviderID:      identity.ProviderID,
                                ProviderSubject: identity.ProviderSubject,
                                LinkedAt:        identity.LinkedAt,
                                LastUsedAt:      identity.LastUsedAt,
                        })
                }
        }

        return identities, nil
}

// -----------------------
// DCR Credentials Storage
// -----------------------

// cloneDCRCredentials returns a field-by-field copy of c, or nil if c is nil.
// All fields are values (no slices, maps, or pointers), so a shallow copy is
// sufficient — adding a new reference-typed field requires updating this
// helper to deep-copy that field.
func cloneDCRCredentials(c *DCRCredentials) *DCRCredentials {
        if c == nil {
                return nil
        }
        cp := *c
        return &cp
}

// StoreDCRCredentialsIfAbsent claims creds.Key for creds. The credentials
// are stored under their own Key field; callers must populate it before
// calling. A defensive copy is made so subsequent caller mutations do not
// affect persisted state.
//
// Create-if-absent, not overwrite: a single process's dcrFlight singleflight
// (see pkg/auth/dcr) already prevents concurrent same-key writers within
// this process, so the check below is not closing a live race here — it is
// contract symmetry with RedisStorage.StoreDCRCredentialsIfAbsent, which
// DOES need it to prevent two replicas from independently registering RFC
// 7591 clients for the same Key and racing on which write wins.
//
// The in-memory backend has no native TTL, so "absent" cannot be a plain
// map-presence check: the Redis backend's rows self-evict via TTL derived
// from ClientSecretExpiresAt, so a claim there naturally succeeds again once
// the old row expires. To keep behaviour symmetric, an existing entry whose
// ClientSecretExpiresAt is non-zero and already in the past is treated as
// absent — the claim proceeds and overwrites it — so a secret's expiry does
// not permanently block re-registration on this backend the way a naive
// presence check would.
//
// Validation is delegated to validateDCRCredentialsForStore so the rejection
// set stays in sync with sibling backends.
func (s *MemoryStorage) StoreDCRCredentialsIfAbsent(_ context.Context, creds *DCRCredentials) (*DCRCredentials, error) {
        if err := validateDCRCredentialsForStore(creds); err != nil {
                return nil, err
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        if existing, ok := s.dcrCredentials[creds.Key]; ok {
                expired := !existing.ClientSecretExpiresAt.IsZero() && time.Now().After(existing.ClientSecretExpiresAt)
                if !expired {
                        return cloneDCRCredentials(existing), nil
                }
        }

        s.dcrCredentials[creds.Key] = cloneDCRCredentials(creds)
        return cloneDCRCredentials(creds), nil
}

// UpdateDCRCredentialsIfPresent replaces the entry at creds.Key with creds
// when one physically exists, and returns ErrNotFound (wrapped) otherwise. A
// defensive copy is stored so subsequent caller mutations do not affect
// persisted state, mirroring StoreDCRCredentialsIfAbsent.
//
// Presence is a plain map lookup, deliberately NOT the TTL-aware "absent"
// check StoreDCRCredentialsIfAbsent uses: an expired-but-present entry is
// still updatable here. See the DCRCredentialStore interface docs for why
// Update gates on physical presence rather than liveness. The in-memory
// backend has no native TTL, so ClientSecretExpiresAt is retained verbatim on
// the rewritten entry exactly as StoreDCRCredentialsIfAbsent retains it.
//
// Validation is delegated to validateDCRCredentialsForStore so the rejection
// set stays in sync with sibling backends.
func (s *MemoryStorage) UpdateDCRCredentialsIfPresent(_ context.Context, creds *DCRCredentials) (*DCRCredentials, error) {
        if err := validateDCRCredentialsForStore(creds); err != nil {
                return nil, err
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        if _, ok := s.dcrCredentials[creds.Key]; !ok {
                return nil, notFoundRFC6749Error("DCR credentials not found")
        }

        s.dcrCredentials[creds.Key] = cloneDCRCredentials(creds)
        return cloneDCRCredentials(creds), nil
}

// UpdateDCRCredentialsIfUnchanged replaces the entry at creds.Key with creds
// only when the stored entry still equals expected, returning ErrNotFound
// (wrapped) when no entry exists and ErrDCRCredentialsChanged when it differs.
// The compare and the write happen under s.mu, so no other writer can slip in
// between them. Presence and ClientSecretExpiresAt handling are identical to
// UpdateDCRCredentialsIfPresent.
func (s *MemoryStorage) UpdateDCRCredentialsIfUnchanged(
        _ context.Context, creds, expected *DCRCredentials,
) (*DCRCredentials, error) {
        if err := validateDCRCompareAndSet(creds, expected); err != nil {
                return nil, err
        }

        s.mu.Lock()
        defer s.mu.Unlock()

        existing, ok := s.dcrCredentials[creds.Key]
        if !ok {
                return nil, notFoundRFC6749Error("DCR credentials not found")
        }
        if !dcrCredentialsEqual(existing, expected) {
                return nil, ErrDCRCredentialsChanged
        }

        s.dcrCredentials[creds.Key] = cloneDCRCredentials(creds)
        return cloneDCRCredentials(creds), nil
}

// dcrCredentialsEqual reports whether a and b hold the same persisted values.
// Time fields are compared with time.Time.Equal rather than ==, so two values
// for the same instant compare equal regardless of monotonic-clock reading or
// location.
//
// The final == covers every other field and only compiles while
// DCRCredentials stays comparable. A new time.Time field added to
// DCRCredentials MUST be compared with Equal and zeroed here alongside the
// existing two; == on it would reintroduce the monotonic-clock/location
// mismatch. TestDCRCredentialsEqual_TimeFieldsCompareByInstant fails if one
// is missed.
func dcrCredentialsEqual(a, b *DCRCredentials) bool {
        ac, bc := *a, *b
        if !ac.CreatedAt.Equal(bc.CreatedAt) || !ac.ClientSecretExpiresAt.Equal(bc.ClientSecretExpiresAt) {
                return false
        }
        ac.CreatedAt, bc.CreatedAt = time.Time{}, time.Time{}
        ac.ClientSecretExpiresAt, bc.ClientSecretExpiresAt = time.Time{}, time.Time{}
        return ac == bc
}

// GetDCRCredentials retrieves DCR credentials by key.
// Returns a defensive copy; returns ErrNotFound (wrapped) on miss.
func (s *MemoryStorage) GetDCRCredentials(_ context.Context, key DCRKey) (*DCRCredentials, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()

        entry, ok := s.dcrCredentials[key]
        if !ok {
                slog.Debug("dcr credentials not found",
                        "issuer", key.Issuer,
                        "upstream_id", key.UpstreamID,
                        "redirect_uri", key.RedirectURI,
                )
                return nil, notFoundRFC6749Error("DCR credentials not found")
        }

        return cloneDCRCredentials(entry), nil
}

// -----------------------
// Metrics/Stats (for testing and monitoring)
// -----------------------

// Stats contains statistics about the storage contents.
type Stats struct {
        Clients               int
        AuthCodes             int
        AccessTokens          int
        RefreshTokens         int
        PKCERequests          int
        UpstreamTokens        int
        PendingAuthorizations int
        InvalidatedCodes      int
        ClientAssertionJWTs   int
        Users                 int
        ProviderIdentities    int
        DCRCredentials        int
}

// Stats returns current statistics about storage contents.
// This is useful for testing and monitoring.
func (s *MemoryStorage) Stats() Stats {
        s.mu.RLock()
        defer s.mu.RUnlock()

        return Stats{
                Clients:               len(s.clients),
                AuthCodes:             len(s.authCodes),
                AccessTokens:          len(s.accessTokens),
                RefreshTokens:         len(s.refreshTokens),
                PKCERequests:          len(s.pkceRequests),
                UpstreamTokens:        len(s.upstreamTokens),
                PendingAuthorizations: len(s.pendingAuthorizations),
                InvalidatedCodes:      len(s.invalidatedCodes),
                ClientAssertionJWTs:   len(s.clientAssertionJWTs),
                Users:                 len(s.users),
                ProviderIdentities:    len(s.providerIdentities),
                DCRCredentials:        len(s.dcrCredentials),
        }
}

// Compile-time interface compliance checks
var (
        _ Storage                          = (*MemoryStorage)(nil)
        _ PendingAuthorizationStorage      = (*MemoryStorage)(nil)
        _ DeviceCodeStorage                = (*MemoryStorage)(nil)
        _ PendingDeviceLoginStorage        = (*MemoryStorage)(nil)
        _ PendingDeviceConfirmationStorage = (*MemoryStorage)(nil)
        _ ClientRegistry                   = (*MemoryStorage)(nil)
        _ UpstreamTokenStorage             = (*MemoryStorage)(nil)
        _ UserStorage                      = (*MemoryStorage)(nil)
        _ DCRCredentialStore               = (*MemoryStorage)(nil)
        _ AssertionJWTConsumer             = (*MemoryStorage)(nil)
)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "net/url"
        "slices"
        "strconv"
        "strings"
        "time"

        "github.com/go-jose/go-jose/v3"
        "github.com/ory/fosite"
        "github.com/redis/go-redis/v9"

        "github.com/stacklok/toolhive-core/redisconn"
        "github.com/stacklok/toolhive/pkg/authserver/server"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authserver/server/session"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// nullMarker is used to store nil upstream tokens in Redis.
const nullMarker = "null"

// pastExpiryDCRTTL is the bounded TTL applied when a caller writes DCR
// credentials whose ClientSecretExpiresAt is already in the past. The row is
// still accepted (the resolver decides when to re-register) but it self-evicts
// almost immediately so the store does not hold an already-expired secret
// forever. One second is small enough that no caller can usefully read the
// row, large enough that real Redis applies the TTL reliably (sub-second TTLs
// are technically supported via PEXPIRE but the second-grain TTL command is
// the broadly-tested path), and short enough that operational metrics do not
// confuse this row with a healthy long-lived registration.
const pastExpiryDCRTTL = time.Second

// maxDCRClaimRetries bounds the WATCH/MULTI retry loops of
// StoreDCRCredentialsIfAbsent and UpdateDCRCredentialsIfUnchanged. go-redis
// does not retry Watch internally: a concurrent write to the watched key
// (another replica claiming, refreshing, updating, or evicting the same row)
// aborts the pipelined EXEC with redis.TxFailedErr, which Watch returns to
// the caller unwrapped. Retrying a small, fixed number of times lets a real
// concurrent write during the exact race these methods exist to close
// resolve on its own rather than failing the caller with a spurious error.
// Mirrors maxConfiguredClientReconcileRetries above for the same reason.
const maxDCRClaimRetries = 3

// warnOnCleanupErr logs a warning when a best-effort cleanup operation fails.
//
// Secondary index cleanup in Redis (SRem from reverse-lookup sets, Del of orphaned
// keys) is intentionally best-effort: the primary operation has already succeeded,
// and failing the overall request due to an index cleanup error would be worse than
// leaving a stale entry. Stale index entries are bounded by TTL expiration of the
// underlying keys and are tolerated on read (readers skip missing entries).
//
// TODO: Add a periodic background cleanup job to scan and remove stale entries
// from secondary index sets (KeyTypeReqIDAccess, KeyTypeReqIDRefresh,
// KeyTypeUserUpstream, KeyTypeUserProviders) to prevent unbounded growth.
func warnOnCleanupErr(err error, operation, key string) {
        if err != nil {
                slog.Warn("best-effort index cleanup failed",
                        "operation", operation, "key", key, "error", err)
        }
}

// filterIndexMembersByPrefix splits SMEMBERS results into members that share
// this storage instance's key prefix and those that do not. Multi-key
// operations like MGet and Del fail with CROSSSLOT under Redis Cluster when
// any member hashes to a different slot, and a stray un-prefixed entry
// (legacy data, an external admin op, a test fixture) is the most likely
// source of such drift. Filtering at read time turns a hard cluster failure
// into a logged warning, while remaining a no-op on standalone Redis.
func filterIndexMembersByPrefix(prefix string, members []string) (kept, dropped []string) {
        for _, m := range members {
                if strings.HasPrefix(m, prefix) {
                        kept = append(kept, m)
                } else {
                        dropped = append(dropped, m)
                }
        }
        return kept, dropped
}

// warnDroppedIndexMembers emits a warning for each foreign member returned by
// filterIndexMembersByPrefix so operators can identify the source of cluster-
// incompatible entries.
func warnDroppedIndexMembers(operation, indexKey, expectedPrefix string, dropped []string) {
        for _, m := range dropped {
                slog.Warn("dropping foreign index member to prevent CROSSSLOT",
                        "operation", operation, "indexKey", indexKey,
                        "expectedPrefix", expectedPrefix, "member", m)
        }
}

// RedisStorage implements the Storage interface backed by Redis.
// Supports standalone mode (single endpoint), Sentinel failover mode, and
// Cluster mode. It provides distributed storage for OAuth2 tokens, authorization
// codes, user data, and pending authorizations, enabling horizontal scaling.
type RedisStorage struct {
        client    redis.UniversalClient
        keyPrefix string

        // maxDCRClients caps the live clients counted by countsTowardDCRCap; zero
        // means unlimited. The count lives in the KeyTypeDCRClientIdx sorted set,
        // so every replica sharing this Redis enforces one cap.
        maxDCRClients int

        // now is the clock DCR cap index scores are computed from.
        now func() time.Time
}

// RedisStorageOption configures a RedisStorage.
type RedisStorageOption func(*RedisStorage)

// WithRedisMaxDCRClients sets the cap on live clients counted by
// countsTowardDCRCap; zero means unlimited. At the cap RegisterClient returns
// a *ClientCapacityError and writes nothing. It panics on a negative n, which
// RunConfig.ResolvedMaxDCRClients already rejects.
func WithRedisMaxDCRClients(n int) RedisStorageOption {
        if n < 0 {
                panic(fmt.Sprintf("max DCR clients must not be negative, got %d", n))
        }
        return func(s *RedisStorage) {
                s.maxDCRClients = n
        }
}

// withRedisClock sets the clock DCR cap index scores are computed from, so
// tests can keep it in step with miniredis's key TTLs.
func withRedisClock(now func() time.Time) RedisStorageOption {
        return func(s *RedisStorage) {
                s.now = now
        }
}

func newRedisStorage(client redis.UniversalClient, keyPrefix string, opts []RedisStorageOption) *RedisStorage {
        s := &RedisStorage{
                client:        client,
                keyPrefix:     keyPrefix,
                maxDCRClients: DefaultMaxDCRClients,
                now:           time.Now,
        }
        for _, opt := range opts {
                opt(s)
        }
        return s
}

// storedSession is a serializable wrapper for fosite.Requester.
// This allows storing fosite sessions in Redis as JSON.
type storedSession struct {
        ClientID          string              `json:"client_id"`
        RequestedAt       time.Time           `json:"requested_at"`
        RequestedScopes   []string            `json:"requested_scopes"`
        GrantedScopes     []string            `json:"granted_scopes"`
        RequestedAudience []string            `json:"requested_audience"`
        GrantedAudience   []string            `json:"granted_audience"`
        Form              map[string][]string `json:"form"`
        RequestID         string              `json:"request_id"`
        Session           json.RawMessage     `json:"session"`
}

// NewRedisStorage creates Redis-backed storage. Connection-mode topology,
// timeouts, TLS, and credentials are configured through cfg; keyPrefix is the
// per-tenant key prefix (a single Redis Cluster hash tag combining namespace
// and name — see DeriveKeyPrefix, e.g. "thv:auth:{ns:name}:") and must be
// non-empty.
//
// Connection-mode validation, timeout defaults, client construction (standalone,
// cluster, or sentinel), TLS plumbing, and connectivity verification are
// delegated to the shared toolhive-core redisconn package. cfg.Password may be
// empty when the Redis server does not require authentication (the auth server
// does not mandate ACL auth); the keyPrefix is storage-specific and required.
func NewRedisStorage(
        ctx context.Context, cfg redisconn.Config, keyPrefix string, opts ...RedisStorageOption,
) (*RedisStorage, error) {
        if keyPrefix == "" {
                return nil, errors.New("invalid redis configuration: key prefix is required")
        }

        client, err := redisconn.NewClient(ctx, &cfg)
        if err != nil {
                return nil, err
        }

        return newRedisStorage(client, keyPrefix, opts), nil
}

// NewRedisStorageWithClient creates a RedisStorage with a pre-configured client.
// This is useful for testing with miniredis.
func NewRedisStorageWithClient(
        client redis.UniversalClient, keyPrefix string, opts ...RedisStorageOption,
) *RedisStorage {
        return newRedisStorage(client, keyPrefix, opts)
}

// defaultSessionFactory creates session prototypes for deserialization.
// Name and email are empty because they are preserved in the JWT Extra map
// from the serialized session data during deserialization.
func defaultSessionFactory(subject, idpSessionID, clientID string) fosite.Session {
        return session.New(subject, idpSessionID, clientID, session.UserClaims{})
}

// Health checks Redis connectivity.
func (s *RedisStorage) Health(ctx context.Context) error {
        if err := s.client.Ping(ctx).Err(); err != nil {
                return fmt.Errorf("redis health check failed: %w", err)
        }
        return nil
}

// Close closes the Redis client connection.
func (s *RedisStorage) Close() error {
        return s.client.Close()
}

// -----------------------
// ClientRegistry
// -----------------------

// storedClient is a serializable wrapper for OAuth clients.
type storedClient struct {
        ID            string   `json:"id"`
        ClientName    string   `json:"client_name,omitempty"`
        Secret        []byte   `json:"secret,omitempty"` //nolint:gosec // G117: field legitimately holds sensitive data
        RedirectURIs  []string `json:"redirect_uris"`
        GrantTypes    []string `json:"grant_types"`
        ResponseTypes []string `json:"response_types"`
        Scopes        []string `json:"scopes"`
        Audience      []string `json:"audience"`
        // Resources is the RFC 8707 resource allowlist, populated only for a
        // Reserved client that implements resourceScopedClient (today, always the
        // SPIFFE static-client placeholder). Gated on Reserved rather than the
        // interface check alone so buildStoredClient's write and clientFromStored's
        // read stay symmetric by construction. Nil for every other client shape,
        // and omitted from the JSON encoding in that case.
        Resources           []string `json:"resources,omitempty"`
        IdentityFingerprint string   `json:"identity_fingerprint,omitempty"`
        Public              bool     `json:"public"`
        // TokenEndpointAuthMethod is the auth method registered for the client
        // ("none", "client_secret_basic", "client_secret_post"). Empty means the
        // row predates confidential-client support; see GetClient for how legacy
        // rows are interpreted.
        TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty"`
        // DCRIssued is true when the row was written from a dynamically registered
        // client. It is persisted rather than inferred from the auth method: a
        // pre-provisioned confidential client also carries a method, and inferring
        // from it would hand that permanent client an anti-bloat TTL on its first
        // token request. Rows predating this field unmarshal with DCRIssued=false.
        // DCR mints a new client_id per registration and never rewrites an existing
        // row, so a pre-existing row with no DCRIssued value is itself DCR-issued —
        // it just reads back unmarked. clientFromStored compensates for this on
        // read using the Redis key's TTL as a discriminator: a legacy row with no
        // auth method, Public=true, AND a live TTL on its key is treated as
        // DCR-issued even though the persisted field says otherwise, because
        // RegisterClient only ever wrote that TTL for DCR-issued rows — a
        // pre-provisioned public client (a bare *fosite.DefaultClient) was always
        // written with no TTL. A legacy confidential row (no method, Public=false)
        // is not compensated for — it predates confidential DCR support entirely,
        // so it cannot be DCR-issued.
        DCRIssued bool `json:"dcr_issued,omitempty"`
        // Reserved is true when the row is a SPIFFE static-client durable
        // placeholder (see staticClientPlaceholder in spiffe_decorator.go) —
        // never a real, authenticatable client. It is checked before GrantTypes/
        // ResponseTypes/Secret/Public are trusted on read: clientFromStored
        // re-wraps a Reserved row in the same inert type the placeholder was
        // built with, regardless of what those other fields happen to contain,
        // so the row cannot become issuable via fosite's own zero-value
        // defaulting (see clientFromStored and inertPlaceholderClient).
        Reserved bool `json:"reserved,omitempty"`
        // Configured is true when the row was written by ReconcileConfiguredClient,
        // i.e. it belongs to ToolHive configuration rather than to DCR or an
        // embedder. Only that method sets it; buildStoredClient, RegisterClient and
        // UpsertDCRIssuedClient never do. Precedence on read: reserved > configured
        // > dcr_issued > unmarked. A stamped row is never treated as DCR-issued, so
        // the legacy "public with a TTL" compensation is skipped for it. The field
        // is not part of storedClient.fingerprint, so stamping a row written by an
        // older binary does not read as a policy change. Rows written before this
        // field existed unmarshal with Configured=false.
        Configured bool `json:"configured,omitempty"`
        // JSONWebKeys stores only the inline public keys used by private_key_jwt.
        // JSONWebKeysURI and other assertion material are intentionally not stored.
        JSONWebKeys                       *jose.JSONWebKeySet `json:"jwks,omitempty"`
        TokenEndpointAuthSigningAlgorithm string              `json:"token_endpoint_auth_signing_alg,omitempty"`
}

// publicJSONWebKeySet copies an inline JWKS while removing private and
// symmetric key material before it is persisted.
func publicJSONWebKeySet(jwks *jose.JSONWebKeySet) *jose.JSONWebKeySet {
        if jwks == nil {
                return nil
        }

        publicKeys := make([]jose.JSONWebKey, 0, len(jwks.Keys))
        for _, key := range jwks.Keys {
                publicKey := key
                if !key.IsPublic() {
                        publicKey = key.Public()
                }
                if publicKey.Key != nil && publicKey.IsPublic() {
                        publicKeys = append(publicKeys, publicKey)
                }
        }
        return &jose.JSONWebKeySet{Keys: publicKeys}
}

// clientFromStored rebuilds a fosite.Client from its persisted form. hasTTL
// reports whether the Redis key currently carries an expiry (TTL >= 0); it is
// the discriminator used to compensate for the legacy public shape below, and
// is meaningless for rows that carry an explicit DCRIssued marker.
//
// The read logic is two-way so that drift between the Public flag and the
// auth method can only ever *add* secret verification, never remove it:
//
//   - A row with no method recorded covers two populations, both rebuilt as a
//     bare *fosite.DefaultClient: a pre-provisioned non-OIDC client
//     (permanent — RegisterClient only populates the column for
//     fosite.OpenIDConnectClient implementations, so a non-OIDC client will
//     never carry one, by design, not by omission) and a row written before
//     confidential-client support existed (transitional). Neither shape was
//     ever assigned a pinned method — the public one no more than the
//     confidential one — so neither should acquire method enforcement on
//     read-back: fosite only enforces token_endpoint_auth_method on
//     fosite.OpenIDConnectClient implementations, and a bare
//     *fosite.DefaultClient does not satisfy that interface. IsPublic still
//     reports correctly because it reads the Public field directly
//     (fosite.DefaultClient.IsPublic() returns c.Public, uninfluenced by the
//     auth method).
//   - A row with a method was written by registration.New (the only path
//     that populates the column), so it is rebuilt as a
//     *fosite.DefaultOpenIDConnectClient — fosite enforces the pinned method
//     at the token endpoint. IsPublic is derived as
//     (Public && method == "none"): a row that somehow carries both
//     Public=true and a secret-based method reads back confidential, forcing
//     secret verification rather than dropping it.
//
// DCR-issued marking: a row is treated as DCR-issued when it was persisted as
// such (stored.DCRIssued — the primary signal, always honored regardless of
// TTL), OR when it has the legacy public shape (no method, Public=true) AND
// the key currently carries a TTL. That second clause exists because rows
// written before the DCRIssued field existed have neither field set, and the
// legacy public shape is ambiguous on its own: it matches both a pre-existing
// DCR-issued public client (which must keep renewing its TTL) and a
// pre-provisioned public client registered as a bare *fosite.DefaultClient
// (which must never acquire one — Redis EXPIRE on a key with no TTL creates
// one). The two are distinguishable by how RegisterClient originally wrote
// the key: DCR-issued rows always got DefaultDCRClientTTL, pre-provisioned
// rows always got none. Gating the legacy-shape compensation on "the key
// currently has a TTL" reproduces that distinction without needing to touch
// the stored JSON. A legacy confidential row (no method, Public=false) is
// never marked regardless of TTL: it predates confidential DCR support
// entirely and cannot be DCR-issued.
//
// A row with Configured set is returned through registration.MarkConfigured
// before either DCR rule above (dcr_issued, legacy public+TTL) is applied.
func clientFromStored(stored storedClient, hasTTL bool) fosite.Client {
        // A reserved SPIFFE placeholder is re-wrapped in the exact inert type it
        // was built with, ignoring whatever GrantTypes/ResponseTypes/Secret/
        // Public happen to be on the row: a bare *fosite.DefaultClient would
        // otherwise satisfy fosite's own zero-value defaulting for
        // GetGrantTypes/GetResponseTypes (["authorization_code"] / ["code"]) on
        // this exact reconstruction path, undoing the guarantee the placeholder
        // exists to provide. This check must run before the method-based branch
        // below: a reserved row never carries a TokenEndpointAuthMethod, but even
        // if it somehow did, Reserved must still win.
        if stored.Reserved {
                return inertPlaceholderClient{
                        DefaultClient: &fosite.DefaultClient{
                                ID:       stored.ID,
                                Scopes:   stored.Scopes,
                                Audience: stored.Audience,
                                Public:   false,
                        },
                        resources:           stored.Resources,
                        identityFingerprint: stored.IdentityFingerprint,
                }
        }

        method := stored.TokenEndpointAuthMethod
        if method == "" {
                client := &fosite.DefaultClient{
                        ID:            stored.ID,
                        Secret:        stored.Secret,
                        RedirectURIs:  stored.RedirectURIs,
                        GrantTypes:    stored.GrantTypes,
                        ResponseTypes: stored.ResponseTypes,
                        Scopes:        stored.Scopes,
                        Audience:      stored.Audience,
                        Public:        stored.Public,
                }
                if stored.Configured {
                        return markConfiguredFromStored(client)
                }
                if stored.DCRIssued || (stored.Public && hasTTL) {
                        return registration.MarkDCRIssuedWithName(client, stored.ClientName)
                }
                return client
        }
        oidcClient := &fosite.DefaultOpenIDConnectClient{
                DefaultClient: &fosite.DefaultClient{
                        ID:            stored.ID,
                        Secret:        stored.Secret,
                        RedirectURIs:  stored.RedirectURIs,
                        GrantTypes:    stored.GrantTypes,
                        ResponseTypes: stored.ResponseTypes,
                        Scopes:        stored.Scopes,
                        Audience:      stored.Audience,
                        Public:        stored.Public && method == oauthproto.TokenEndpointAuthMethodNone,
                },
                TokenEndpointAuthMethod:           method,
                JSONWebKeys:                       stored.JSONWebKeys,
                TokenEndpointAuthSigningAlgorithm: stored.TokenEndpointAuthSigningAlgorithm,
        }
        if stored.Configured {
                return markConfiguredFromStored(oidcClient)
        }
        if stored.DCRIssued {
                return registration.MarkDCRIssuedWithName(oidcClient, stored.ClientName)
        }
        return oidcClient
}

// markConfiguredFromStored stamps a client clientFromStored just rebuilt. Both
// shapes it builds are accepted by MarkConfigured; if that ever changes, the
// row reads back as an inert placeholder instead of unstamped.
func markConfiguredFromStored(client fosite.Client) fosite.Client {
        marked, ok := registration.MarkConfigured(client)
        if !ok {
                return inertPlaceholderClient{DefaultClient: &fosite.DefaultClient{ID: client.GetID()}}
        }
        return marked
}

// buildStoredClient converts client to its serializable form. Record the
// registered auth method when the client exposes one. Clients that don't
// implement fosite.OpenIDConnectClient (e.g. pre-provisioned confidential
// clients built as bare *fosite.DefaultClient) leave the field empty on
// purpose: Public alone carries the meaning it already carries, and
// clientFromStored treats the empty method as a legacy row. Do NOT substitute
// a "none" fallback here — that would silently reclassify a confidential row
// as public on read-back.
func buildStoredClient(client fosite.Client) storedClient {
        stored := storedClient{
                ID:            client.GetID(),
                Secret:        client.GetHashedSecret(),
                RedirectURIs:  client.GetRedirectURIs(),
                GrantTypes:    client.GetGrantTypes(),
                ResponseTypes: client.GetResponseTypes(),
                Scopes:        client.GetScopes(),
                Audience:      client.GetAudience(),
                Public:        client.IsPublic(),
        }
        stored.Reserved = isReservedPlaceholder(client)
        // Resources and IdentityFingerprint are only ever restored on read for a
        // Reserved row (see clientFromStored); gating the write the same way
        // keeps the two symmetric by construction instead of relying on every
        // resourceScopedClient/spiffeIdentityClient implementation always being
        // Reserved.
        if stored.Reserved {
                if rc, ok := client.(resourceScopedClient); ok {
                        stored.Resources = rc.Resources()
                }
                if identity, ok := client.(spiffeIdentityClient); ok {
                        stored.IdentityFingerprint = identity.IdentityFingerprint()
                }
        }
        if oidcClient, ok := client.(fosite.OpenIDConnectClient); ok {
                stored.TokenEndpointAuthMethod = oidcClient.GetTokenEndpointAuthMethod()
                stored.JSONWebKeys = publicJSONWebKeySet(oidcClient.GetJSONWebKeys())
                stored.TokenEndpointAuthSigningAlgorithm = oidcClient.GetTokenEndpointAuthSigningAlgorithm()
        }
        stored.DCRIssued = registration.DCRIssued(client)
        if stored.DCRIssued {
                stored.ClientName = registration.ClientName(client)
        }
        return stored
}

// RegisterClient creates a client in storage. Always create-only via SetNX:
// it atomically returns ErrAlreadyExists when a client with the same ID is
// already registered, regardless of the new client's origin. /oauth/register
// is unauthenticated, so this must never clobber an existing client;
// operator-declared clients reconcile through ReconcileConfiguredClient
// instead.
func (s *RedisStorage) RegisterClient(ctx context.Context, client fosite.Client) error {
        if registration.IsConfigured(client) {
                return fmt.Errorf("client %q must not carry the configured stamp; use ReconcileConfiguredClient", client.GetID())
        }
        if err := ValidateRegisterableClientID(client.GetID()); err != nil {
                return err
        }

        key := redisKey(s.keyPrefix, KeyTypeClient, client.GetID())
        stored := buildStoredClient(client)

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal client: %w", err)
        }

        if countsTowardDCRCap(client) {
                return s.registerCountedClient(ctx, key, client.GetID(), data)
        }

        // DCR-issued clients (public and confidential) expire to prevent unbounded
        // growth from unauthenticated dynamic registration; RenewClientTTL raises
        // the TTL on proven use so actively-used clients do not expire. The TTL is
        // keyed on the registration.DCRIssued marker, not IsPublic: a confidential
        // client minted via DCR is exactly the registration an attacker would pile
        // up, while a pre-provisioned confidential client (no marker) must never
        // acquire a TTL — Redis EXPIRE on a key with no TTL creates one, so gating
        // on IsPublic would give an existing permanent client an expiry on its
        // first token request after upgrade.
        ttl := time.Duration(0)
        if registration.DCRIssued(client) {
                ttl = DefaultUnusedDCRClientTTL
        }

        created, err := s.client.SetNX(ctx, key, data, ttl).Result()
        if err != nil {
                return fmt.Errorf("failed to register client: %w", err)
        }
        if !created {
                return fmt.Errorf("%w: client %q", ErrAlreadyExists, client.GetID())
        }
        return nil
}

// DCR cap index scores are expiry times in Unix milliseconds taken from the
// replica's clock (RedisStorage.now) and passed to the scripts as an argument.
// Redis TIME would avoid clock skew between replicas, but it is in the @fast
// ACL category, which the documented auth-server ACL does not grant, so
// using it would break existing ACL-restricted deployments on upgrade. Skew
// under NTP is milliseconds, and its only effect is that an index entry is
// trimmed that much before or after its row expires.

// Status codes returned by registerCountedClientScript.
const (
        registerStatusCreated  int64 = 1
        registerStatusAtCap    int64 = 0
        registerStatusConflict int64 = -1
)

// registerCountedClientScript creates a client counted against the DCR cap.
// KEYS[1] is the client row, KEYS[2] the KeyTypeDCRClientIdx sorted set.
// ARGV: row data, TTL in ms, cap (0 = unlimited), client ID, now in ms.
//
// It reports an existing row first, as the memory store does, then trims
// index entries whose expiry has passed, rejects at the cap, and otherwise
// creates the row and indexes it, all in one atomic step so concurrent
// registrations on any replica cannot overshoot the cap. Returns {1} on
// success, {0, oldest_expiry_ms} at the cap, {-1} when the row already
// exists.
var registerCountedClientScript = redis.NewScript(`
local cap = tonumber(ARGV[3])
local now = tonumber(ARGV[5])
if redis.call('EXISTS', KEYS[1]) == 1 then
        return {-1}
end
redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', now)
if cap > 0 and redis.call('ZCARD', KEYS[2]) >= cap then
        local oldest = redis.call('ZRANGE', KEYS[2], 0, 0, 'WITHSCORES')
        return {0, tonumber(oldest[2])}
end
redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
redis.call('ZADD', KEYS[2], now + tonumber(ARGV[2]), ARGV[4])
return {1}
`)

// registerCountedClient runs registerCountedClientScript for a client counted
// against the DCR cap, giving it DefaultUnusedDCRClientTTL.
func (s *RedisStorage) registerCountedClient(ctx context.Context, key, id string, data []byte) error {
        now := s.now()
        res, err := registerCountedClientScript.Run(ctx, s.client,
                []string{key, s.dcrClientIndexKey()},
                data, DefaultUnusedDCRClientTTL.Milliseconds(), s.maxDCRClients, id, now.UnixMilli(),
        ).Slice()
        if err != nil {
                return fmt.Errorf("failed to register client: %w", err)
        }
        if len(res) == 0 {
                return errors.New("failed to register client: empty script reply")
        }
        status, ok := res[0].(int64)
        if !ok {
                return fmt.Errorf("failed to register client: unexpected script status type %T", res[0])
        }
        switch status {
        case registerStatusCreated:
                return nil
        case registerStatusConflict:
                return fmt.Errorf("%w: client %q", ErrAlreadyExists, id)
        case registerStatusAtCap:
                if len(res) < 2 {
                        return errors.New("failed to register client: at-cap reply has no oldest expiry")
                }
                oldestMs, ok := res[1].(int64)
                if !ok {
                        return fmt.Errorf("failed to register client: unexpected oldest expiry type %T", res[1])
                }
                return newClientCapacityError(s.maxDCRClients, time.UnixMilli(oldestMs), now)
        default:
                return fmt.Errorf("failed to register client: unexpected script status %d", status)
        }
}

// dcrClientIndexKey returns the KeyTypeDCRClientIdx sorted set key.
func (s *RedisStorage) dcrClientIndexKey() string {
        return s.keyPrefix + KeyTypeDCRClientIdx
}

// CountDCRClients implements DCRClientCounter: the number of index entries
// whose expiry has not passed. It only reads, so a metric collection never
// writes.
func (s *RedisStorage) CountDCRClients(ctx context.Context) (int64, error) {
        minScore := "(" + strconv.FormatInt(s.now().UnixMilli(), 10)
        n, err := s.client.ZCount(ctx, s.dcrClientIndexKey(), minScore, "+inf").Result()
        if err != nil {
                return 0, fmt.Errorf("failed to count DCR clients: %w", err)
        }
        return n, nil
}

// UpsertDCRIssuedClient creates or replaces a DCR-issued client at
// client.GetID(). Unlike RegisterClient (create-only via SetNX), an existing
// row is replaced -- and its TTL refreshed to DefaultDCRClientTTL -- when the
// existing row is itself DCR-issued; it refuses with ErrAlreadyExists when the
// existing row is NOT DCR-issued, protecting a configured/SPIFFE-reconciled
// client from being clobbered. See the ClientRegistry interface doc for the
// full contract.
//
// The read-check-write sequence for an existing key runs inside a Redis
// WATCH/MULTI transaction, mirroring ReconcileConfiguredClient, so a
// concurrent writer cannot interleave between the DCR-issued check and the
// write; see maxConfiguredClientReconcileRetries for why this method retries
// redis.TxFailedErr itself, up to the same bounded count. Unlike
// ReconcileConfiguredClient (ttl=0, permanent), the write here always uses
// DefaultDCRClientTTL: this row is TTL-bounded like any other DCR-issued
// client. A client counted against the DCR cap is refused, so this path can
// neither bypass the cap nor leave the index out of step.
func (s *RedisStorage) UpsertDCRIssuedClient(ctx context.Context, client fosite.Client) error {
        if !registration.DCRIssued(client) {
                return fmt.Errorf("client %q must carry the DCR-issued marker to use UpsertDCRIssuedClient", client.GetID())
        }
        if err := ValidateRegisterableClientID(client.GetID()); err != nil {
                return err
        }
        if countsTowardDCRCap(client) {
                return fmt.Errorf("client %q counts against the DCR client cap and must use RegisterClient", client.GetID())
        }

        key := redisKey(s.keyPrefix, KeyTypeClient, client.GetID())
        stored := buildStoredClient(client)
        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal client: %w", err)
        }

        setPipelined := func(tx *redis.Tx) error {
                _, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
                        pipe.Set(ctx, key, data, DefaultDCRClientTTL)
                        return nil
                })
                return err
        }

        txFn := func(tx *redis.Tx) error {
                existingData, getErr := tx.Get(ctx, key).Bytes()
                if errors.Is(getErr, redis.Nil) {
                        return setPipelined(tx)
                }
                if getErr != nil {
                        return fmt.Errorf("failed to get existing client: %w", getErr)
                }

                ttl, ttlErr := tx.TTL(ctx, key).Result()
                if ttlErr != nil {
                        return fmt.Errorf("failed to get existing client TTL: %w", ttlErr)
                }

                var existingStored storedClient
                if unmarshalErr := json.Unmarshal(existingData, &existingStored); unmarshalErr != nil {
                        return fmt.Errorf("failed to unmarshal existing client: %w", unmarshalErr)
                }

                if !registration.DCRIssued(clientFromStored(existingStored, ttl >= 0)) {
                        return fmt.Errorf("%w: client %q", ErrAlreadyExists, client.GetID())
                }

                return setPipelined(tx)
        }

        var watchErr error
        for attempt := 0; attempt < maxConfiguredClientReconcileRetries; attempt++ {
                watchErr = s.client.Watch(ctx, txFn, key)
                if !errors.Is(watchErr, redis.TxFailedErr) {
                        return watchErr
                }
        }
        return watchErr
}

// fingerprint reads the persisted fields directly rather than going through
// clientFromStored. This matters because fosite.DefaultClient.GetGrantTypes
// and GetResponseTypes each silently substitute a single-element default
// (["authorization_code"] / ["code"]) whenever the underlying field is empty
// (see clientFromStored) — comparing through the reconstructed client would
// report a false mismatch between two rows that both deliberately carry no
// grant/response types, such as the SPIFFE static-client placeholder in
// spiffe_decorator.go reconciling against itself. Reading the stored fields
// directly sidesteps that defaulting entirely.
func (s storedClient) fingerprint() clientFingerprint {
        return clientFingerprint{
                scopes:              s.Scopes,
                audience:            s.Audience,
                grantTypes:          s.GrantTypes,
                responseTypes:       s.ResponseTypes,
                resources:           s.Resources,
                identityFingerprint: s.IdentityFingerprint,
                public:              s.Public,
                reserved:            s.Reserved,
        }
}

// maxConfiguredClientReconcileRetries bounds the retry loop
// ReconcileConfiguredClient runs around its WATCH/MULTI transaction. go-redis
// does not retry Watch internally: a concurrent write to the watched key aborts
// the pipelined EXEC with redis.TxFailedErr, which Watch returns to the caller
// unwrapped (see $GOMODCACHE/github.com/redis/go-redis/v9@*/tx.go). Retrying a
// small, fixed number of times lets a real concurrent write during the exact
// race this feature exists to close (see preflightDurableCollisions in
// spiffe_decorator.go) resolve on its own rather than failing the caller's
// startup path with a spurious error.
const maxConfiguredClientReconcileRetries = 5

// ReconcileConfiguredClient applies an operator-declared client: creates it
// if absent, otherwise replaces the existing record unless
// configuredClientConflict refuses (or the record is DCR-issued), in which
// case it returns ErrAlreadyExists. See the ClientRegistry interface doc for
// the full contract.
//
// The read-check-write sequence for an existing key runs inside a Redis
// WATCH/MULTI transaction so a concurrent writer (e.g. a DCR registration on
// another replica racing this reconciliation) cannot interleave between the
// conflict check and the write. Redis aborts the transaction if the
// watched key changed concurrently, and go-redis surfaces that as
// redis.TxFailedErr rather than retrying it — see
// maxConfiguredClientReconcileRetries for why this method retries that error
// itself, up to a bounded count. A configured client is stored with ttl=0
// (operator-declared clients never expire), clearing any TTL the row might
// have inherited.
func (s *RedisStorage) ReconcileConfiguredClient(ctx context.Context, client fosite.Client) error {
        if registration.DCRIssued(client) {
                return fmt.Errorf("configured client %q must not carry the DCR-issued marker", client.GetID())
        }
        if err := checkConfiguredClientType(client); err != nil {
                return err
        }
        if err := ValidateRegisterableClientID(client.GetID()); err != nil {
                return err
        }

        key := redisKey(s.keyPrefix, KeyTypeClient, client.GetID())
        stored := buildStoredClient(client)
        stored.Configured = true
        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal client: %w", err)
        }

        setPipelined := func(tx *redis.Tx) error {
                _, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
                        pipe.Set(ctx, key, data, 0)
                        // A configured client never counts against the DCR cap. The
                        // row it replaces cannot be DCR-issued (refused above), but an
                        // index entry outliving its expired row would otherwise be
                        // counted until its score passes.
                        pipe.ZRem(ctx, s.dcrClientIndexKey(), client.GetID())
                        return nil
                })
                return err
        }

        // replaced holds the fingerprint of the row the successful attempt
        // overwrote, nil when the key was absent.
        var replaced *clientFingerprint
        txFn := func(tx *redis.Tx) error {
                existingData, getErr := tx.Get(ctx, key).Bytes()
                if errors.Is(getErr, redis.Nil) {
                        return setPipelined(tx)
                }
                if getErr != nil {
                        return fmt.Errorf("failed to get existing client: %w", getErr)
                }

                ttl, ttlErr := tx.TTL(ctx, key).Result()
                if ttlErr != nil {
                        return fmt.Errorf("failed to get existing client TTL: %w", ttlErr)
                }

                var existingStored storedClient
                if unmarshalErr := json.Unmarshal(existingData, &existingStored); unmarshalErr != nil {
                        return fmt.Errorf("failed to unmarshal existing client: %w", unmarshalErr)
                }

                // DCR-issued detection goes through the reconstructed client (not the
                // raw existingStored.DCRIssued field) so it inherits clientFromStored's
                // legacy-row compensation for rows written before the DCRIssued field
                // existed. The fingerprint comparison below deliberately does NOT go
                // through the reconstructed client — see storedClient.fingerprint.
                if registration.DCRIssued(clientFromStored(existingStored, ttl >= 0)) {
                        return fmt.Errorf("%w: client %q is DCR-issued, refusing to overwrite with a configured client",
                                ErrAlreadyExists, client.GetID())
                }
                existingFP := existingStored.fingerprint()
                if conflictErr := configuredClientConflict(client.GetID(), existingFP, stored.fingerprint()); conflictErr != nil {
                        return conflictErr
                }

                replaced = &existingFP
                return setPipelined(tx)
        }

        var watchErr error
        for attempt := 0; attempt < maxConfiguredClientReconcileRetries; attempt++ {
                replaced = nil
                watchErr = s.client.Watch(ctx, txFn, key)
                if !errors.Is(watchErr, redis.TxFailedErr) {
                        break
                }
        }
        // replaced is only meaningful when watchErr is nil: the last attempt may
        // have set it before EXEC aborted.
        if watchErr == nil && replaced != nil {
                warnConfiguredClientReplaced(client.GetID(), *replaced, stored.fingerprint())
        }
        return watchErr
}

// GetClient loads the client by its ID.
//
// The value and the key's TTL are fetched together in one pipelined round
// trip: clientFromStored needs the TTL to disambiguate the legacy public
// client shape (see DCRIssued and clientFromStored for why), and issuing GET
// then TTL as two separate round trips would leave a window where a
// concurrent RenewClientTTL or expiry changes the answer between them.
func (s *RedisStorage) GetClient(ctx context.Context, id string) (fosite.Client, error) {
        key := redisKey(s.keyPrefix, KeyTypeClient, id)

        pipe := s.client.Pipeline()
        getCmd := pipe.Get(ctx, key)
        ttlCmd := pipe.TTL(ctx, key)
        if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
                return nil, fmt.Errorf("failed to get client: %w", err)
        }

        data, err := getCmd.Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Client not found")
                }
                return nil, fmt.Errorf("failed to get client: %w", err)
        }

        // TTL returns -2 for a missing key (already excluded above by getCmd
        // succeeding) and -1 for a key with no expiry ("persist"); any value >= 0
        // means the key currently carries an expiry.
        ttl, err := ttlCmd.Result()
        if err != nil {
                return nil, fmt.Errorf("failed to get client TTL: %w", err)
        }
        hasTTL := ttl >= 0

        var stored storedClient
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal client: %w", err)
        }

        client := clientFromStored(stored, hasTTL)
        if dcrErr := registration.ValidatePrivateKeyJWTClientKeys(client); dcrErr != nil {
                slog.DebugContext(ctx, "rejecting client whose signing keys no longer meet server policy",
                        "client_id", id, "reason", dcrErr.ErrorDescription)
                return nil, fosite.ErrInvalidClient.WithHint("client signing keys no longer meet server policy")
        }
        return client, nil
}

// renewClientTTLScript sets a client row's TTL and, when the client counts
// against the DCR cap, its index score, in one atomic step. KEYS[1] is the
// client row, KEYS[2] the KeyTypeDCRClientIdx sorted set. ARGV: TTL in ms,
// client ID, "1" when counted, now in ms. PEXPIRE on a missing row is a no-op, and the
// index is left alone so an expired client is not re-counted.
var renewClientTTLScript = redis.NewScript(`
if redis.call('PEXPIRE', KEYS[1], ARGV[1]) == 1 and ARGV[3] == '1' then
        redis.call('ZADD', KEYS[2], tonumber(ARGV[4]) + tonumber(ARGV[1]), ARGV[2])
end
return 0
`)

// RenewClientTTL extends a DCR-issued client's registration TTL to DefaultDCRClientTTL.
// This also raises a newly registered client off DefaultUnusedDCRClientTTL on
// its first token exchange, and moves its DCR cap index score to match.
//
// Call this on a proven-use signal — a successful token exchange/refresh — not on a
// client read. GetClient is reached from the unauthenticated front-channel
// /oauth/authorize handler before any authentication, and DCR-issued public clients
// have no secret, so renewing there would let any caller who knows a public client_id
// keep its row alive indefinitely and defeat the anti-bloat TTL. Renewing on token
// issuance ties registration survival to actual use.
//
// Only DCR-issued clients carry a TTL; pre-provisioned clients are stored without one
// and are left untouched. The gate is the registration.DCRIssued marker, not IsPublic:
// EXPIRE on a key with no TTL creates one, so keying on IsPublic would hand an
// existing permanent confidential client an expiry on its first token request after
// upgrade. PEXPIRE on a missing key is a no-op, so a client whose row has already
// expired (or a non-persisted CIMD client) is safely ignored.
func (s *RedisStorage) RenewClientTTL(ctx context.Context, client fosite.Client) error {
        if client == nil || !registration.DCRIssued(client) {
                return nil
        }
        key := redisKey(s.keyPrefix, KeyTypeClient, client.GetID())
        counted := "0"
        if countsTowardDCRCap(client) {
                counted = "1"
        }
        return renewClientTTLScript.Run(ctx, s.client,
                []string{key, s.dcrClientIndexKey()},
                DefaultDCRClientTTL.Milliseconds(), client.GetID(), counted, s.now().UnixMilli(),
        ).Err()
}

// ClientAssertionJWTValid atomically reserves jti for a short window via
// Redis SET NX, returning fosite.ErrJTIKnown if it is already reserved or
// was previously consumed and not yet expired. This is the only atomic
// checkpoint in fosite's two-call client-assertion replay protocol (Valid,
// then SetClientAssertionJWT with the real expiry once the rest of the
// assertion has been validated) — an Exists-then-Set pair (the prior
// implementation) lets two concurrent requests presenting the same
// assertion both observe jti as unused before either records it. jti
// length is bounded here too, before it's ever sent to Redis, since fosite
// hands it to us before parsing exp.
func (s *RedisStorage) ClientAssertionJWTValid(ctx context.Context, jti string) error {
        if len(jti) > server.MaxAssertionJTILength {
                return fmt.Errorf("jti exceeds maximum length of %d bytes", server.MaxAssertionJTILength)
        }

        key := redisKey(s.keyPrefix, KeyTypeJWT, jti)
        _, err := s.client.SetArgs(ctx, key, "1", redis.SetArgs{TTL: assertionReservationTTL, Mode: "NX"}).Result()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return fosite.ErrJTIKnown
                }
                return fmt.Errorf("failed to reserve JWT: %w", err)
        }
        return nil
}

// SetClientAssertionJWT extends jti's replay marker to its real expiry,
// rejecting an exp further than MaxAssertionLifespan in the future. This
// bounds how long a replay-tracking entry survives regardless of what the
// client's own assertion claims — otherwise a registered caller could set
// exp decades out and give Redis a key with a correspondingly enormous TTL.
// Unconditional overwrite, not a fresh atomic check: only the caller that
// already won ClientAssertionJWTValid's reservation reaches this call. If
// the JWT has already expired (exp is in the past), this deletes the
// reservation instead: there is no need to keep it for replay detection
// since the assertion will be rejected on expiry checks before reaching the
// JTI lookup.
func (s *RedisStorage) SetClientAssertionJWT(ctx context.Context, jti string, exp time.Time) error {
        if time.Until(exp) > server.MaxAssertionLifespan {
                return fmt.Errorf("assertion lifetime exceeds maximum of %s", server.MaxAssertionLifespan)
        }

        key := redisKey(s.keyPrefix, KeyTypeJWT, jti)

        ttl := time.Until(exp)
        if ttl <= 0 {
                slog.Debug("removing reservation for already-expired client assertion JWT",
                        "jti", jti, "exp", exp)
                return s.client.Del(ctx, key).Err()
        }

        return s.client.Set(ctx, key, "1", ttl).Err()
}

// ConsumeAssertionJWT atomically records an assertion JWT as consumed until its
// expiry. Reusing an unexpired (purpose, issuer, jti) tuple returns
// fosite.ErrJTIKnown. Redis errors are returned so callers fail closed rather
// than accepting an assertion whose replay status is unknown.
func (s *RedisStorage) ConsumeAssertionJWT(ctx context.Context, purpose, issuer, jti string, exp time.Time) error {
        ttl := time.Until(exp)
        if ttl <= 0 {
                return nil
        }

        key := redisAssertionJWTKey(s.keyPrefix, purpose, issuer, jti)
        _, err := s.client.SetArgs(ctx, key, "1", redis.SetArgs{TTL: ttl, Mode: "NX"}).Result()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return fosite.ErrJTIKnown
                }
                return fmt.Errorf("consume assertion JWT replay marker: %w", err)
        }
        return nil
}

// -----------------------
// oauth2.AuthorizeCodeStorage
// -----------------------

// CreateAuthorizeCodeSession stores the authorization request for a given authorization code.
func (s *RedisStorage) CreateAuthorizeCodeSession(ctx context.Context, code string, request fosite.Requester) error {
        if code == "" {
                return fosite.ErrInvalidRequest.WithHint("authorization code cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        key := redisKey(s.keyPrefix, KeyTypeAuthCode, code)
        ttl := getTTLFromRequester(request, fosite.AuthorizeCode, DefaultAuthCodeTTL)

        data, err := marshalRequester(request)
        if err != nil {
                return fmt.Errorf("failed to marshal request: %w", err)
        }

        return s.client.Set(ctx, key, data, ttl).Err()
}

// GetAuthorizeCodeSession retrieves the authorization request for a given code.
// Matches memory.go's pattern: get the auth code data first, then check if
// invalidated. InvalidateAuthorizeCodeSession extends the auth code TTL to match
// the invalidation marker, so the data is always available when the marker exists.
func (s *RedisStorage) GetAuthorizeCodeSession(ctx context.Context, code string, _ fosite.Session) (fosite.Requester, error) {
        key := redisKey(s.keyPrefix, KeyTypeAuthCode, code)
        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Authorization code not found")
                }
                return nil, fmt.Errorf("failed to get authorization code: %w", err)
        }

        request, err := unmarshalRequester(ctx, data, s)
        if err != nil {
                return nil, fmt.Errorf("failed to unmarshal request: %w", err)
        }

        // Check if the code has been invalidated.
        invalidatedKey := redisKey(s.keyPrefix, KeyTypeInvalidated, code)
        invalidated, err := s.client.Exists(ctx, invalidatedKey).Result()
        if err != nil {
                return nil, fmt.Errorf("failed to check invalidation status: %w", err)
        }
        if invalidated > 0 {
                // Must return the request along with the error as per fosite documentation.
                // Fosite needs the Requester to extract the request ID for token revocation
                // during replay attack handling (RFC 6819).
                return request, fosite.ErrInvalidatedAuthorizeCode
        }

        return request, nil
}

// InvalidateAuthorizeCodeSession marks an authorization code as used/invalid.
// It extends the auth code key's TTL to match the invalidation marker, ensuring
// GetAuthorizeCodeSession can always return the Requester alongside
// ErrInvalidatedAuthorizeCode as required by fosite for token revocation.
func (s *RedisStorage) InvalidateAuthorizeCodeSession(ctx context.Context, code string) error {
        key := redisKey(s.keyPrefix, KeyTypeAuthCode, code)

        // Check if the code exists
        exists, err := s.client.Exists(ctx, key).Result()
        if err != nil {
                return fmt.Errorf("failed to check authorization code: %w", err)
        }
        if exists == 0 {
                return notFoundRFC6749Error("Authorization code not found")
        }

        // Atomically: create invalidation marker and extend auth code TTL to match.
        // The auth code data must outlive the invalidation marker so that
        // GetAuthorizeCodeSession can return the Requester for replay detection.
        invalidatedKey := redisKey(s.keyPrefix, KeyTypeInvalidated, code)
        pipe := s.client.TxPipeline()
        pipe.Set(ctx, invalidatedKey, "1", DefaultInvalidatedCodeTTL)
        pipe.Expire(ctx, key, DefaultInvalidatedCodeTTL)
        _, err = pipe.Exec(ctx)
        return err
}

// -----------------------
// oauth2.AccessTokenStorage
// -----------------------

// CreateAccessTokenSession stores the access token session.
func (s *RedisStorage) CreateAccessTokenSession(ctx context.Context, signature string, request fosite.Requester) error {
        if signature == "" {
                return fosite.ErrInvalidRequest.WithHint("access token signature cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        key := redisKey(s.keyPrefix, KeyTypeAccess, signature)
        ttl := getTTLFromRequester(request, fosite.AccessToken, DefaultAccessTokenTTL)

        data, err := marshalRequester(request)
        if err != nil {
                return fmt.Errorf("failed to marshal request: %w", err)
        }

        // Store token, add to request ID index, and set index TTL atomically.
        reqIDKey := redisSetKey(s.keyPrefix, KeyTypeReqIDAccess, request.GetID())
        pipe := s.client.TxPipeline()
        pipe.Set(ctx, key, data, ttl)
        pipe.SAdd(ctx, reqIDKey, signature)
        pipe.Expire(ctx, reqIDKey, ttl)
        _, err = pipe.Exec(ctx)
        return err
}

// GetAccessTokenSession retrieves the access token session by its signature.
func (s *RedisStorage) GetAccessTokenSession(ctx context.Context, signature string, _ fosite.Session) (fosite.Requester, error) {
        key := redisKey(s.keyPrefix, KeyTypeAccess, signature)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Access token not found")
                }
                return nil, fmt.Errorf("failed to get access token: %w", err)
        }

        return unmarshalRequester(ctx, data, s)
}

// DeleteAccessTokenSession removes the access token session.
func (s *RedisStorage) DeleteAccessTokenSession(ctx context.Context, signature string) error {
        key := redisKey(s.keyPrefix, KeyTypeAccess, signature)

        // Get the request first to find the request ID for cleaning up the index
        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return notFoundRFC6749Error("Access token not found")
                }
                return fmt.Errorf("failed to get access token: %w", err)
        }

        // Delete the token
        if err := s.client.Del(ctx, key).Err(); err != nil {
                return fmt.Errorf("failed to delete access token: %w", err)
        }

        // Best-effort secondary index cleanup (see warnOnCleanupErr).
        var stored storedSession
        if err := json.Unmarshal(data, &stored); err == nil && stored.RequestID != "" {
                reqIDKey := redisSetKey(s.keyPrefix, KeyTypeReqIDAccess, stored.RequestID)
                warnOnCleanupErr(s.client.SRem(ctx, reqIDKey, signature).Err(), "SRem", reqIDKey)
        }

        return nil
}

// -----------------------
// oauth2.RefreshTokenStorage
// -----------------------

// CreateRefreshTokenSession stores the refresh token session.
func (s *RedisStorage) CreateRefreshTokenSession(
        ctx context.Context, signature string, _ string, request fosite.Requester,
) error {
        if signature == "" {
                return fosite.ErrInvalidRequest.WithHint("refresh token signature cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        key := redisKey(s.keyPrefix, KeyTypeRefresh, signature)
        ttl := getTTLFromRequester(request, fosite.RefreshToken, DefaultRefreshTokenTTL)

        data, err := marshalRequester(request)
        if err != nil {
                return fmt.Errorf("failed to marshal request: %w", err)
        }

        // Store token, add to request ID index, and set index TTL atomically.
        reqIDKey := redisSetKey(s.keyPrefix, KeyTypeReqIDRefresh, request.GetID())
        pipe := s.client.TxPipeline()
        pipe.Set(ctx, key, data, ttl)
        pipe.SAdd(ctx, reqIDKey, signature)
        pipe.Expire(ctx, reqIDKey, ttl)
        _, err = pipe.Exec(ctx)
        return err
}

// GetRefreshTokenSession retrieves the refresh token session by its signature.
func (s *RedisStorage) GetRefreshTokenSession(ctx context.Context, signature string, _ fosite.Session) (fosite.Requester, error) {
        key := redisKey(s.keyPrefix, KeyTypeRefresh, signature)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Refresh token not found")
                }
                return nil, fmt.Errorf("failed to get refresh token: %w", err)
        }

        return unmarshalRequester(ctx, data, s)
}

// DeleteRefreshTokenSession removes the refresh token session.
func (s *RedisStorage) DeleteRefreshTokenSession(ctx context.Context, signature string) error {
        key := redisKey(s.keyPrefix, KeyTypeRefresh, signature)

        // Get the request first to find the request ID for cleaning up the index
        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return notFoundRFC6749Error("Refresh token not found")
                }
                return fmt.Errorf("failed to get refresh token: %w", err)
        }

        // Delete the token
        if err := s.client.Del(ctx, key).Err(); err != nil {
                return fmt.Errorf("failed to delete refresh token: %w", err)
        }

        // Best-effort secondary index cleanup (see warnOnCleanupErr).
        var stored storedSession
        if err := json.Unmarshal(data, &stored); err == nil && stored.RequestID != "" {
                reqIDKey := redisSetKey(s.keyPrefix, KeyTypeReqIDRefresh, stored.RequestID)
                warnOnCleanupErr(s.client.SRem(ctx, reqIDKey, signature).Err(), "SRem", reqIDKey)
        }

        return nil
}

// RotateRefreshToken invalidates a refresh token and all its related token data.
// This is a no-op if the token does not exist (returns nil), matching the behavior
// of the in-memory implementation. All cleanup operations are best-effort
// (see warnOnCleanupErr); the new refresh token has already been issued by fosite,
// so partial cleanup is acceptable.
func (s *RedisStorage) RotateRefreshToken(ctx context.Context, requestID string, refreshTokenSignature string) error {
        // Delete the specific refresh token. Del returns the number of keys removed;
        // 0 means the token did not exist (already rotated or never created).
        refreshKey := redisKey(s.keyPrefix, KeyTypeRefresh, refreshTokenSignature)
        deleted, err := s.client.Del(ctx, refreshKey).Result()
        if err != nil {
                warnOnCleanupErr(err, "Del", refreshKey)
        }
        if deleted == 0 {
                slog.Debug("refresh token not found during rotation, treating as no-op",
                        "request_id", requestID, "signature", refreshTokenSignature)
        }

        // Remove from the request ID index
        reqIDRefreshKey := redisSetKey(s.keyPrefix, KeyTypeReqIDRefresh, requestID)
        warnOnCleanupErr(s.client.SRem(ctx, reqIDRefreshKey, refreshTokenSignature).Err(), "SRem", reqIDRefreshKey)

        // Delete all access tokens associated with this request ID
        reqIDAccessKey := redisSetKey(s.keyPrefix, KeyTypeReqIDAccess, requestID)
        signatures, err := s.client.SMembers(ctx, reqIDAccessKey).Result()
        if err == nil {
                for _, sig := range signatures {
                        accessKey := redisKey(s.keyPrefix, KeyTypeAccess, sig)
                        warnOnCleanupErr(s.client.Del(ctx, accessKey).Err(), "Del", accessKey)
                }
                warnOnCleanupErr(s.client.Del(ctx, reqIDAccessKey).Err(), "Del", reqIDAccessKey)
        }

        return nil
}

// -----------------------
// oauth2.TokenRevocationStorage
// -----------------------

// RevokeAccessToken marks an access token as revoked by request ID.
func (s *RedisStorage) RevokeAccessToken(ctx context.Context, requestID string) error {
        reqIDKey := redisSetKey(s.keyPrefix, KeyTypeReqIDAccess, requestID)
        signatures, err := s.client.SMembers(ctx, reqIDKey).Result()
        if err != nil && !errors.Is(err, redis.Nil) {
                return fmt.Errorf("failed to get access token signatures: %w", err)
        }

        for _, sig := range signatures {
                accessKey := redisKey(s.keyPrefix, KeyTypeAccess, sig)
                warnOnCleanupErr(s.client.Del(ctx, accessKey).Err(), "Del", accessKey)
        }

        // Clean up the index
        warnOnCleanupErr(s.client.Del(ctx, reqIDKey).Err(), "Del", reqIDKey)

        return nil
}

// RevokeRefreshToken marks a refresh token as revoked by request ID.
func (s *RedisStorage) RevokeRefreshToken(ctx context.Context, requestID string) error {
        reqIDKey := redisSetKey(s.keyPrefix, KeyTypeReqIDRefresh, requestID)
        signatures, err := s.client.SMembers(ctx, reqIDKey).Result()
        if err != nil && !errors.Is(err, redis.Nil) {
                return fmt.Errorf("failed to get refresh token signatures: %w", err)
        }

        for _, sig := range signatures {
                refreshKey := redisKey(s.keyPrefix, KeyTypeRefresh, sig)
                warnOnCleanupErr(s.client.Del(ctx, refreshKey).Err(), "Del", refreshKey)
        }

        // Clean up the index
        warnOnCleanupErr(s.client.Del(ctx, reqIDKey).Err(), "Del", reqIDKey)

        return nil
}

// RevokeRefreshTokenMaybeGracePeriod marks a refresh token as revoked, optionally allowing
// a grace period. For this implementation, we revoke immediately.
func (s *RedisStorage) RevokeRefreshTokenMaybeGracePeriod(ctx context.Context, requestID string, _ string) error {
        return s.RevokeRefreshToken(ctx, requestID)
}

// -----------------------
// pkce.PKCERequestStorage
// -----------------------

// CreatePKCERequestSession stores the PKCE request session.
func (s *RedisStorage) CreatePKCERequestSession(ctx context.Context, signature string, request fosite.Requester) error {
        if signature == "" {
                return fosite.ErrInvalidRequest.WithHint("PKCE signature cannot be empty")
        }
        if request == nil {
                return fosite.ErrInvalidRequest.WithHint("request cannot be nil")
        }

        key := redisKey(s.keyPrefix, KeyTypePKCE, signature)
        ttl := getTTLFromRequester(request, fosite.AuthorizeCode, DefaultPKCETTL)

        data, err := marshalRequester(request)
        if err != nil {
                return fmt.Errorf("failed to marshal request: %w", err)
        }

        return s.client.Set(ctx, key, data, ttl).Err()
}

// GetPKCERequestSession retrieves the PKCE request session by its signature.
func (s *RedisStorage) GetPKCERequestSession(ctx context.Context, signature string, _ fosite.Session) (fosite.Requester, error) {
        key := redisKey(s.keyPrefix, KeyTypePKCE, signature)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("PKCE request not found")
                }
                return nil, fmt.Errorf("failed to get PKCE request: %w", err)
        }

        return unmarshalRequester(ctx, data, s)
}

// DeletePKCERequestSession removes the PKCE request session.
func (s *RedisStorage) DeletePKCERequestSession(ctx context.Context, signature string) error {
        key := redisKey(s.keyPrefix, KeyTypePKCE, signature)

        result, err := s.client.Del(ctx, key).Result()
        if err != nil {
                return fmt.Errorf("failed to delete PKCE request: %w", err)
        }
        if result == 0 {
                return notFoundRFC6749Error("PKCE request not found")
        }

        return nil
}

// -----------------------
// Upstream Token Storage
// -----------------------

// storedUpstreamTokens is a serializable wrapper for UpstreamTokens.
// Time fields use int64 Unix epoch; 0 is the sentinel meaning "not set".
// Neither time field uses omitempty so that 0 is always present and the
// read path can use a consistent != 0 check for both.
type storedUpstreamTokens struct {
        ProviderID       string `json:"provider_id"`
        AccessToken      string `json:"access_token"`  //nolint:gosec // G117: field legitimately holds sensitive data
        RefreshToken     string `json:"refresh_token"` //nolint:gosec // G117: field legitimately holds sensitive data
        IDToken          string `json:"id_token"`
        ExpiresAt        int64  `json:"expires_at"`
        SessionExpiresAt int64  `json:"session_expires_at"`
        UserID           string `json:"user_id"`
        UpstreamSubject  string `json:"upstream_subject"`
        ClientID         string `json:"client_id"`
}

// toUpstreamTokens converts the stored int64 epoch fields to time.Time and returns
// the populated UpstreamTokens. Zero epoch values are preserved as zero time.Time.
func (s *storedUpstreamTokens) toUpstreamTokens() *UpstreamTokens {
        var expiresAt time.Time
        if s.ExpiresAt != 0 {
                expiresAt = time.Unix(s.ExpiresAt, 0)
        }
        var sessionExpiresAt time.Time
        if s.SessionExpiresAt != 0 {
                sessionExpiresAt = time.Unix(s.SessionExpiresAt, 0)
        }
        return &UpstreamTokens{
                ProviderID:       s.ProviderID,
                AccessToken:      s.AccessToken,
                RefreshToken:     s.RefreshToken,
                IDToken:          s.IDToken,
                ExpiresAt:        expiresAt,
                SessionExpiresAt: sessionExpiresAt,
                UserID:           s.UserID,
                UpstreamSubject:  s.UpstreamSubject,
                ClientID:         s.ClientID,
        }
}

// newStoredUpstreamTokens converts tokens to its serializable form, encoding
// zero time.Time fields as epoch 0 -- the "no expiry" sentinel toUpstreamTokens
// decodes back to a zero time.Time (time.Time{}.Unix() itself returns
// -62135596800, not a useful sentinel). Returns nil for nil input.
func newStoredUpstreamTokens(tokens *UpstreamTokens) *storedUpstreamTokens {
        if tokens == nil {
                return nil
        }
        var expiresAtUnix int64
        if !tokens.ExpiresAt.IsZero() {
                expiresAtUnix = tokens.ExpiresAt.Unix()
        }
        var sessionExpiresAtUnix int64
        if !tokens.SessionExpiresAt.IsZero() {
                sessionExpiresAtUnix = tokens.SessionExpiresAt.Unix()
        }
        return &storedUpstreamTokens{
                ProviderID:       tokens.ProviderID,
                AccessToken:      tokens.AccessToken,
                RefreshToken:     tokens.RefreshToken,
                IDToken:          tokens.IDToken,
                ExpiresAt:        expiresAtUnix,
                SessionExpiresAt: sessionExpiresAtUnix,
                UserID:           tokens.UserID,
                UpstreamSubject:  tokens.UpstreamSubject,
                ClientID:         tokens.ClientID,
        }
}

// storeUpstreamTokensScript atomically reads the existing UserID, writes new token
// data, updates the session index set, and updates user reverse-index sets.
// This prevents a race condition where concurrent writes for the same session
// could leave orphaned entries in user sets.
//
// KEYS[1] = per-provider token key (e.g. "thv:auth:{ns:name}:upstream:{sessionID}:{providerName}")
// KEYS[2] = session index set key (e.g. "thv:auth:{ns:name}:upstream:idx:{sessionID}")
// ARGV[1] = new token data (JSON or "null" marker)
// ARGV[2] = TTL in milliseconds
// ARGV[3] = new UserID ("" if no user)
// ARGV[4] = user upstream set key prefix (e.g. "thv:auth:{ns:name}:user:upstream:")
//
// Cluster slot invariant: this script reads oldUserID from KEYS[1] inside the
// script body (atomic with the rest of the work), so the user-set keys are
// constructed dynamically as `ARGV[4] .. userID` rather than being passed as
// declared KEYS. Every dynamically-built key MUST therefore inherit the
// `{ns:name}` hash tag that is baked into ARGV[4] (s.keyPrefix). All callers
// derive ARGV[4] from s.keyPrefix, which DeriveKeyPrefix builds with the
// `{ns:name}` hash tag, so user-set keys land on the same Redis Cluster slot
// as KEYS[1] and KEYS[2]. A future refactor that strips the hash tag — or
// rebuilds setPrefix from raw inputs without the tag — will silently pass on
// standalone Redis and fail with CROSSSLOT under Cluster.
var storeUpstreamTokensScript = redis.NewScript(`
local oldUserID = ""
local existing = redis.call('GET', KEYS[1])
if existing and existing ~= "null" then
    local ok, decoded = pcall(cjson.decode, existing)
    if ok and type(decoded) == "table" and decoded.user_id and decoded.user_id ~= "" then
        oldUserID = decoded.user_id
    end
end
` + upstreamRowWriteAndIndexScriptBody)

// upstreamRowWriteAndIndexScriptBody is the write-and-index Lua fragment
// shared verbatim by storeUpstreamTokensScript (unconditional overwrite) and
// casUpstreamTokensScript (compare-and-swap): it SETs the new value, then
// maintains the session index set's TTL/PERSIST invariants (see the comment
// above), then updates the user reverse-index sets. Both scripts read the
// existing row into oldUserID (and, for the CAS script, existingRefreshToken)
// before this fragment runs, so it is textually identical between the two —
// concatenated in Go rather than duplicated, so a future change to one script's
// write/index behavior cannot silently drift from the other's.
const upstreamRowWriteAndIndexScriptBody = `
local ttlMs = tonumber(ARGV[2])
if ttlMs > 0 then
    redis.call('SET', KEYS[1], ARGV[1], 'PX', ttlMs)
else
    redis.call('SET', KEYS[1], ARGV[1])
end

-- Maintain the session index set's TTL.
--
-- Invariant: the index set must outlive every per-provider key it points to.
--   * If ANY member is non-expiring (ttlMs == 0), the index set must also be
--     persistent. Otherwise the index evicts and we lose the ability to find
--     (and clean up) the per-provider key, leaking it forever.
--   * If ALL members are expiring, the index TTL must be at least the longest
--     member TTL.
--
-- Known trade-off: an in-place rewrite of the same (sessionID, providerName)
-- slot from non-expiring to expiring leaves the index PERSIST'd, even though
-- the rewritten member was the sole persistence anchor. Detecting this would
-- require tracking per-member TTL state in the index, which adds complexity.
-- We accept the trade-off because DeleteUpstreamTokens and session GC clean
-- the index up anyway. This behaviour is pinned by the test
-- "same provider rewrite from non-expiring to expiring keeps PERSIST'd until
-- rewrite" in redis_test.go — a future maintainer who tightens the rule will
-- see that test fail and can find the rationale here.
--
-- We discriminate two cases that look identical AFTER SADD (both have PTTL == -1):
--   1. A fresh set our SADD just created. We own the TTL decision unconditionally.
--   2. An existing set previously PERSIST'd by a non-expiring member. Must stay
--      persistent — applying a TTL here is the bug this script must avoid.
--
-- The trick: read EXISTS BEFORE SADD. SADD creates the set as a side-effect,
-- so post-SADD any "no TTL" state is ambiguous; pre-SADD it is not.
local idxExisted = redis.call('EXISTS', KEYS[2])
redis.call('SADD', KEYS[2], KEYS[1])

if ttlMs == 0 then
    -- This member never expires. Make the index persistent (no-op if already so).
    redis.call('PERSIST', KEYS[2])
elseif idxExisted == 0 then
    -- Fresh set; our SADD created it. Apply our TTL.
    redis.call('PEXPIRE', KEYS[2], ttlMs)
else
    -- Existing set; its TTL summarises prior members.
    local idxTTL = redis.call('PTTL', KEYS[2])
    if idxTTL == -1 then
        -- A previous non-expiring write PERSIST'd it. Leave it alone —
        -- applying a TTL here is the bug we are fixing.
    elseif idxTTL < ttlMs then
        -- Existing TTL is shorter than this member's. Extend.
        redis.call('PEXPIRE', KEYS[2], ttlMs)
    end
    -- else: idxTTL >= ttlMs, index already outlives this member.
end

local newUserID = ARGV[3]
local setPrefix = ARGV[4]

if oldUserID ~= "" and oldUserID ~= newUserID then
    redis.call('SREM', setPrefix .. oldUserID, KEYS[1])
end

if newUserID ~= "" then
    redis.call('SADD', setPrefix .. newUserID, KEYS[1])
end

return 1
`

// casUpstreamTokensScript is the compare-and-swap sibling of
// storeUpstreamTokensScript: it additionally gates the write on the existing
// row's refresh_token matching ARGV[5] before doing anything else, returning
// 0 (no write performed) on a mismatch instead of 1. This is what makes
// CompareAndSwapUpstreamTokens's stored row deterministic when redeeming a
// single-use, rotating upstream refresh token across multiple replicas of an
// application sharing this Redis: a replica whose read is stale by the time
// it tries to write loses the CAS instead of clobbering a winning replica's
// rotated token. (This orders writes to storage; it is not by itself a
// guarantee that the redemption is safe at the upstream provider — see the
// CompareAndSwapUpstreamTokens interface doc.)
//
// KEYS[1] = per-provider token key
// KEYS[2] = session index set key
// ARGV[1] = new token data (JSON or "null" marker)
// ARGV[2] = TTL in milliseconds
// ARGV[3] = new UserID ("" if no user)
// ARGV[4] = user upstream set key prefix
// ARGV[5] = expected refresh token ("" means "no row exists yet, or the
//
//        existing row carries no refresh token")
//
// The write-and-index body is shared verbatim with storeUpstreamTokensScript
// via upstreamRowWriteAndIndexScriptBody — see that constant's doc comment.
var casUpstreamTokensScript = redis.NewScript(`
local oldUserID = ""
local existingRefreshToken = ""
local existing = redis.call('GET', KEYS[1])
if existing and existing ~= "null" then
    local ok, decoded = pcall(cjson.decode, existing)
    if ok and type(decoded) == "table" then
        if decoded.user_id and decoded.user_id ~= "" then
            oldUserID = decoded.user_id
        end
        if decoded.refresh_token then
            existingRefreshToken = decoded.refresh_token
        end
    end
end

if existingRefreshToken ~= ARGV[5] then
    return 0
end
` + upstreamRowWriteAndIndexScriptBody)

// marshalUpstreamTokensWithTTL marshals tokens and calculates TTL.
func marshalUpstreamTokensWithTTL(tokens *UpstreamTokens) ([]byte, time.Duration, error) {
        if tokens == nil {
                return []byte(nullMarker), DefaultAccessTokenTTL, nil
        }

        stored := newStoredUpstreamTokens(tokens)

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return nil, 0, fmt.Errorf("failed to marshal upstream tokens: %w", err)
        }

        // Add DefaultRefreshTokenTTL beyond access token expiry so the refresh token
        // survives in storage for transparent token refresh by the middleware.
        // Zero ExpiresAt means the token never expires; return ttl=0 so the Lua script
        // stores the key without a Redis TTL.
        var ttl time.Duration // zero means no Redis TTL (non-expiring token with no known session bound)
        if !tokens.ExpiresAt.IsZero() {
                ttl = time.Until(tokens.ExpiresAt) + DefaultRefreshTokenTTL
                if ttl <= 0 {
                        // Access token and its refresh grace period have both passed — evict promptly.
                        ttl = time.Second
                }
        } else if !tokens.SessionExpiresAt.IsZero() {
                ttl = time.Until(tokens.SessionExpiresAt) + DefaultRefreshTokenTTL
                if ttl <= 0 {
                        // Session bound and its refresh grace period have both passed — evict promptly.
                        ttl = time.Second
                }
        }

        return data, ttl, nil
}

// StoreUpstreamTokens stores the upstream IDP tokens for a session and provider.
// Uses a Lua script to atomically write token data, update the session index set,
// and update user reverse-index sets, preventing race conditions on concurrent writes.
func (s *RedisStorage) StoreUpstreamTokens(ctx context.Context, sessionID, providerName string, tokens *UpstreamTokens) error {
        if sessionID == "" {
                return fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        key := redisUpstreamKey(s.keyPrefix, sessionID, providerName)
        idxKey := redisSetKey(s.keyPrefix, KeyTypeUpstreamIdx, sessionID)

        data, ttl, err := marshalUpstreamTokensWithTTL(tokens)
        if err != nil {
                return err
        }

        newUserID := ""
        if tokens != nil {
                newUserID = tokens.UserID
        }

        userSetKeyPrefix := s.keyPrefix + KeyTypeUserUpstream + ":"

        _, err = storeUpstreamTokensScript.Run(ctx, s.client,
                []string{key, idxKey},
                string(data),
                ttl.Milliseconds(),
                newUserID,
                userSetKeyPrefix,
        ).Result()
        if err != nil {
                return fmt.Errorf("failed to store upstream tokens: %w", err)
        }

        return nil
}

// CompareAndSwapUpstreamTokens stores tokens for (sessionID, providerName)
// only if the refresh token currently stored there equals
// expectedRefreshToken; see the interface doc
// (UpstreamTokenStorage.CompareAndSwapUpstreamTokens) for the coordination
// contract. Uses casUpstreamTokensScript so the comparison and the write (and
// its index maintenance) happen as one atomic Redis operation.
func (s *RedisStorage) CompareAndSwapUpstreamTokens(
        ctx context.Context, sessionID, providerName, expectedRefreshToken string, tokens *UpstreamTokens,
) error {
        if sessionID == "" {
                return fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        key := redisUpstreamKey(s.keyPrefix, sessionID, providerName)
        idxKey := redisSetKey(s.keyPrefix, KeyTypeUpstreamIdx, sessionID)

        data, ttl, err := marshalUpstreamTokensWithTTL(tokens)
        if err != nil {
                return err
        }

        newUserID := ""
        if tokens != nil {
                newUserID = tokens.UserID
        }

        userSetKeyPrefix := s.keyPrefix + KeyTypeUserUpstream + ":"

        wrote, err := casUpstreamTokensScript.Run(ctx, s.client,
                []string{key, idxKey},
                string(data),
                ttl.Milliseconds(),
                newUserID,
                userSetKeyPrefix,
                expectedRefreshToken,
        ).Int64()
        if err != nil {
                return fmt.Errorf("failed to compare-and-swap upstream tokens: %w", err)
        }
        if wrote == 0 {
                return ErrConcurrentRefresh
        }

        return nil
}

// ResolveUpstreamTokenRowID returns an opaque ID for process-local refresh
// coordination of a logical upstream-token row.
func (s *RedisStorage) ResolveUpstreamTokenRowID(_ context.Context, sessionID, providerName string) (UpstreamTokenRowID, error) {
        if sessionID == "" {
                return "", fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return "", fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }
        return upstreamTokenRowID(redisUpstreamKey(s.keyPrefix, sessionID, providerName)), nil
}

// GetUpstreamTokens retrieves the upstream IDP tokens for a session and provider.
// Returns a new UpstreamTokens struct deserialized from Redis, which acts as
// a defensive copy - callers cannot modify the stored data by mutating the return value.
func (s *RedisStorage) GetUpstreamTokens(ctx context.Context, sessionID, providerName string) (*UpstreamTokens, error) {
        if sessionID == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        key := redisUpstreamKey(s.keyPrefix, sessionID, providerName)
        return s.getUpstreamTokensFromKey(ctx, key)
}

// GetAllUpstreamTokens retrieves all upstream IDP tokens for a session across all providers.
// Uses SMEMBERS on the session index set to find all provider keys, then MGET to fetch them.
// Returns a map of providerName -> tokens. Returns an empty map for unknown sessions.
func (s *RedisStorage) GetAllUpstreamTokens(ctx context.Context, sessionID string) (map[string]*UpstreamTokens, error) {
        idxKey := redisSetKey(s.keyPrefix, KeyTypeUpstreamIdx, sessionID)
        result := make(map[string]*UpstreamTokens)

        // Get all provider keys from the session index set
        providerKeys, err := s.client.SMembers(ctx, idxKey).Result()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return result, nil
                }
                return nil, fmt.Errorf("failed to get upstream token index: %w", err)
        }

        providerKeys, dropped := filterIndexMembersByPrefix(s.keyPrefix, providerKeys)
        warnDroppedIndexMembers("GetAllUpstreamTokens", idxKey, s.keyPrefix, dropped)

        if len(providerKeys) == 0 {
                return result, nil
        }

        // Fetch all provider tokens in a single MGET
        values, err := s.client.MGet(ctx, providerKeys...).Result()
        if err != nil {
                return nil, fmt.Errorf("failed to get upstream tokens: %w", err)
        }

        // The per-provider key format is "{prefix}upstream:{sessionID}:{providerName}"
        // Extract providerName by stripping the prefix + "upstream:{sessionID}:"
        keyPrefix := fmt.Sprintf("%s%s:%s:", s.keyPrefix, KeyTypeUpstream, sessionID)

        for i, val := range values {
                if val == nil {
                        continue
                }

                data, ok := val.(string)
                if !ok {
                        slog.Warn("skipping upstream token entry: unexpected type", "key", providerKeys[i])
                        continue
                }

                tokens, parseErr := unmarshalUpstreamTokens([]byte(data))
                if parseErr != nil && !errors.Is(parseErr, ErrExpired) {
                        slog.Warn("skipping corrupt upstream token entry", "key", providerKeys[i], "error", parseErr)
                        continue
                }

                // Extract provider name from the key
                providerName := ""
                if len(providerKeys[i]) > len(keyPrefix) {
                        providerName = providerKeys[i][len(keyPrefix):]
                }
                if providerName == "" && tokens != nil {
                        providerName = tokens.ProviderID
                }

                if providerName != "" {
                        result[providerName] = tokens
                }
        }

        return result, nil
}

// DeleteUpstreamTokens removes all upstream IDP tokens for a session (all providers).
// Uses the session index set to find all provider keys and deletes them atomically.
func (s *RedisStorage) DeleteUpstreamTokens(ctx context.Context, sessionID string) error {
        idxKey := redisSetKey(s.keyPrefix, KeyTypeUpstreamIdx, sessionID)

        // Get all provider keys from the session index set
        providerKeys, err := s.client.SMembers(ctx, idxKey).Result()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return notFoundRFC6749Error("Upstream tokens not found")
                }
                return fmt.Errorf("failed to get upstream token index: %w", err)
        }

        providerKeys, dropped := filterIndexMembersByPrefix(s.keyPrefix, providerKeys)
        warnDroppedIndexMembers("DeleteUpstreamTokens", idxKey, s.keyPrefix, dropped)

        if len(providerKeys) == 0 {
                return notFoundRFC6749Error("Upstream tokens not found")
        }

        // Collect UserIDs for reverse-index cleanup before deleting
        var userIDs []string
        for _, providerKey := range providerKeys {
                data, getErr := s.client.Get(ctx, providerKey).Bytes()
                if getErr != nil {
                        continue
                }
                if string(data) != nullMarker {
                        var stored storedUpstreamTokens
                        if unmarshalErr := json.Unmarshal(data, &stored); unmarshalErr == nil && stored.UserID != "" {
                                userIDs = append(userIDs, stored.UserID)
                        }
                }
        }

        // Delete all provider keys and the index set
        keysToDelete := append(slices.Clone(providerKeys), idxKey)
        if err := s.client.Del(ctx, keysToDelete...).Err(); err != nil {
                return fmt.Errorf("failed to delete upstream tokens: %w", err)
        }

        // Best-effort secondary index cleanup for user:upstream sets
        for _, userID := range userIDs {
                userUpstreamSetKey := redisSetKey(s.keyPrefix, KeyTypeUserUpstream, userID)
                for _, providerKey := range providerKeys {
                        warnOnCleanupErr(s.client.SRem(ctx, userUpstreamSetKey, providerKey).Err(), "SRem", userUpstreamSetKey)
                }
        }

        return nil
}

// DeleteUpstreamTokensForProvider removes tokens for a single (sessionID, providerName),
// leaving sibling providers' rows intact. Absent row returns nil.
func (s *RedisStorage) DeleteUpstreamTokensForProvider(ctx context.Context, sessionID, providerName string) error {
        if sessionID == "" {
                return fosite.ErrInvalidRequest.WithHint("session ID cannot be empty")
        }
        if providerName == "" {
                return fosite.ErrInvalidRequest.WithHint("provider name cannot be empty")
        }

        key := redisUpstreamKey(s.keyPrefix, sessionID, providerName)
        idxKey := redisSetKey(s.keyPrefix, KeyTypeUpstreamIdx, sessionID)

        // Best-effort: read UserID for reverse-index cleanup before deleting.
        var userID string
        data, err := s.client.Get(ctx, key).Bytes()
        if err == nil && string(data) != nullMarker {
                var stored storedUpstreamTokens
                if unmarshalErr := json.Unmarshal(data, &stored); unmarshalErr == nil {
                        userID = stored.UserID
                }
        }
        // redis.Nil (key absent) or any error on GET: skip user-index cleanup gracefully.

        // Del returns count 0 on absent key — treat as nil (not ErrNotFound).
        if delErr := s.client.Del(ctx, key).Err(); delErr != nil {
                return fmt.Errorf("failed to delete upstream tokens for provider: %w", delErr)
        }

        // Best-effort secondary index cleanup — leave sibling members in idxKey intact.
        // If this was the last provider for the session, the set becomes empty and
        // expires by its own TTL; we do not delete the set key itself (intentional,
        // matching the best-effort cleanup posture of DeleteUpstreamTokens).
        warnOnCleanupErr(s.client.SRem(ctx, idxKey, key).Err(), "SRem", idxKey)
        if userID != "" {
                userUpstreamSetKey := redisSetKey(s.keyPrefix, KeyTypeUserUpstream, userID)
                warnOnCleanupErr(s.client.SRem(ctx, userUpstreamSetKey, key).Err(), "SRem", userUpstreamSetKey)
        }

        return nil
}

// GetLatestUpstreamTokensForUser returns the upstream token row for (userID, providerID)
// with the highest ExpiresAt across all sessions.
//
// Expired tokens (past ExpiresAt) are returned so callers can use the refresh
// token; filtering by access-token expiry is the caller's responsibility.
// See the interface declaration in types.go for the full contract.
func (s *RedisStorage) GetLatestUpstreamTokensForUser(ctx context.Context, userID, providerID string) (*UpstreamTokens, error) {
        if userID == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("user ID cannot be empty")
        }
        if providerID == "" {
                return nil, fosite.ErrInvalidRequest.WithHint("provider ID cannot be empty")
        }

        setKey := redisSetKey(s.keyPrefix, KeyTypeUserUpstream, userID)

        members, err := s.client.SMembers(ctx, setKey).Result()
        if err != nil && !errors.Is(err, redis.Nil) {
                return nil, fmt.Errorf("failed to get user upstream index: %w", err)
        }

        members, dropped := filterIndexMembersByPrefix(s.keyPrefix, members)
        warnDroppedIndexMembers("GetLatestUpstreamTokensForUser", setKey, s.keyPrefix, dropped)

        if len(members) == 0 {
                return nil, notFoundRFC6749Error("Upstream tokens not found")
        }

        values, err := s.client.MGet(ctx, members...).Result()
        if err != nil {
                return nil, fmt.Errorf("failed to get upstream tokens: %w", err)
        }

        var winner *storedUpstreamTokens
        for i, val := range values {
                stored, ok := parseUserUpstreamEntry(val, userID, providerID, members[i])
                if !ok {
                        continue
                }

                // Tie-breaker: prefer non-expiring rows (ExpiresAt == 0, the no-expiry
                // sentinel for providers like Slack and GitHub OAuth Apps). Among finite
                // expiries, the latest wins. This aligns with the rest of the package
                // treating zero ExpiresAt as "alive forever".
                if winner == nil || compareExpiryInt64(stored.ExpiresAt, winner.ExpiresAt) > 0 {
                        winner = stored
                }
        }

        if winner == nil {
                return nil, notFoundRFC6749Error("Upstream tokens not found")
        }

        return winner.toUpstreamTokens(), nil
}

// compareExpiryInt64 is the int64 (Unix-epoch) variant of compareExpiry. Zero
// (no-expiry sentinel) ranks latest; among finite epochs, the larger ranks
// latest. Returns -1/0/+1.
func compareExpiryInt64(a, b int64) int {
        switch {
        case a == 0 && b == 0:
                return 0
        case a == 0:
                return 1
        case b == 0:
                return -1
        case a > b:
                return 1
        case a < b:
                return -1
        }
        return 0
}

// parseUserUpstreamEntry parses one raw Redis value from the user-upstream index
// and returns the decoded storedUpstreamTokens together with a match flag.
// It returns (nil, false) for nil values, type mismatches, deletion tombstones,
// JSON decode errors, and rows whose UserID or ProviderID does not match the
// requested userID / providerID. keyName is used only for warning log messages.
//
// The UserID comparison makes the ownership postcondition hold on the row
// itself rather than on the index that named it. The index is keyed by user and
// every writer files a row under its own UserID, so a foreign row is not
// reachable today; the set has no TTL and is never garbage collected, which
// makes it the one index where dead members accumulate, and the memory backend
// already recomputes the same check while scanning.
func parseUserUpstreamEntry(val any, userID, providerID, keyName string) (*storedUpstreamTokens, bool) {
        if val == nil {
                // Dangling set member: the per-provider key has been TTL-evicted.
                // Skip it; the next write will clean up the index entry (best-effort).
                return nil, false
        }

        data, ok := val.(string)
        if !ok {
                slog.Warn("skipping upstream token entry: unexpected type", "key", keyName)
                return nil, false
        }

        if data == nullMarker {
                // Deletion tombstone — treat as absent.
                return nil, false
        }

        var stored storedUpstreamTokens
        if unmarshalErr := json.Unmarshal([]byte(data), &stored); unmarshalErr != nil {
                slog.Warn("skipping corrupt upstream token entry", "key", keyName, "error", unmarshalErr)
                return nil, false
        }

        if stored.UserID != userID || stored.ProviderID != providerID {
                return nil, false
        }

        return &stored, true
}

// getUpstreamTokensFromKey retrieves and deserializes upstream tokens from a specific Redis key.
func (s *RedisStorage) getUpstreamTokensFromKey(ctx context.Context, key string) (*UpstreamTokens, error) {
        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Upstream tokens not found")
                }
                return nil, fmt.Errorf("failed to get upstream tokens: %w", err)
        }

        return unmarshalUpstreamTokens(data)
}

// unmarshalUpstreamTokens deserializes upstream tokens from JSON bytes.
func unmarshalUpstreamTokens(data []byte) (*UpstreamTokens, error) {
        // Handle null marker
        if string(data) == nullMarker {
                return nil, nil
        }

        var stored storedUpstreamTokens
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal upstream tokens: %w", err)
        }

        tokens := stored.toUpstreamTokens()

        // tokens.ExpiresAt is zero when the stored epoch was 0 (the write path stores
        // 0 for zero time.Time, which toUpstreamTokens leaves as the zero time). Skip
        // the expiry check in that case since Redis TTL handles the actual expiration.
        // Return tokens along with ErrExpired so callers can use the refresh token.
        if !tokens.ExpiresAt.IsZero() && time.Now().After(tokens.ExpiresAt) {
                return tokens, ErrExpired
        }

        return tokens, nil
}

// -----------------------
// DCR Credentials Storage
// -----------------------

// storedDCRCredentials is the on-the-wire JSON representation of DCRCredentials.
// Time fields use int64 Unix epoch; 0 is the sentinel meaning "not set", matching
// the storedUpstreamTokens convention. ClientSecretExpiresAt == 0 specifically
// encodes the RFC 7591 §3.2.1 "client_secret does not expire" semantics, in
// which case StoreDCRCredentialsIfAbsent persists the entry without a Redis TTL.
type storedDCRCredentials struct {
        // Embed the canonical key so a row recovered without its lookup key
        // (e.g. via SCAN during diagnostics) still self-identifies.
        KeyIssuer      string `json:"key_issuer"`
        KeyUpstreamID  string `json:"key_upstream_id"`
        KeyRedirectURI string `json:"key_redirect_uri"`
        KeyScopesHash  string `json:"key_scopes_hash"`
        // Omitted when empty so rows written before the field existed and rows
        // without a rotation key share one wire shape.
        KeyRotationKey string `json:"key_rotation_key,omitempty"`

        ProviderName string `json:"provider_name,omitempty"`

        ClientID                string `json:"client_id"`
        ClientSecret            string `json:"client_secret,omitempty"` //nolint:gosec // G117: field legitimately holds sensitive data
        TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty"`

        // Bearer token for the RFC 7592 management endpoint.
        //nolint:gosec // G117: field legitimately holds sensitive data
        RegistrationAccessToken string `json:"registration_access_token,omitempty"`
        RegistrationClientURI   string `json:"registration_client_uri,omitempty"`

        AuthorizationEndpoint string `json:"authorization_endpoint,omitempty"`
        TokenEndpoint         string `json:"token_endpoint,omitempty"`

        CreatedAt             int64 `json:"created_at"`
        ClientSecretExpiresAt int64 `json:"client_secret_expires_at"`
}

// newStoredDCRCredentials converts creds to its stored form, mirroring
// newStoredUpstreamTokens. Time fields are truncated to Unix seconds; a zero
// time.Time becomes the 0 "not set" sentinel.
func newStoredDCRCredentials(creds *DCRCredentials) storedDCRCredentials {
        stored := storedDCRCredentials{
                KeyIssuer:               creds.Key.Issuer,
                KeyUpstreamID:           creds.Key.UpstreamID,
                KeyRedirectURI:          creds.Key.RedirectURI,
                KeyScopesHash:           creds.Key.ScopesHash,
                KeyRotationKey:          creds.Key.RotationKey,
                ProviderName:            creds.ProviderName,
                ClientID:                creds.ClientID,
                ClientSecret:            creds.ClientSecret,
                TokenEndpointAuthMethod: creds.TokenEndpointAuthMethod,
                RegistrationAccessToken: creds.RegistrationAccessToken,
                RegistrationClientURI:   creds.RegistrationClientURI,
                AuthorizationEndpoint:   creds.AuthorizationEndpoint,
                TokenEndpoint:           creds.TokenEndpoint,
        }
        if !creds.CreatedAt.IsZero() {
                stored.CreatedAt = creds.CreatedAt.Unix()
        }
        if !creds.ClientSecretExpiresAt.IsZero() {
                stored.ClientSecretExpiresAt = creds.ClientSecretExpiresAt.Unix()
        }
        return stored
}

// toDCRCredentials decodes the stored form back into the public type, mirroring
// storedUpstreamTokens.toUpstreamTokens. Zero epoch values become the zero
// time.Time, preserving the "not set" sentinel.
func (s *storedDCRCredentials) toDCRCredentials() *DCRCredentials {
        var createdAt time.Time
        if s.CreatedAt != 0 {
                createdAt = time.Unix(s.CreatedAt, 0)
        }
        var clientSecretExpiresAt time.Time
        if s.ClientSecretExpiresAt != 0 {
                clientSecretExpiresAt = time.Unix(s.ClientSecretExpiresAt, 0)
        }
        return &DCRCredentials{
                Key: DCRKey{
                        Issuer:      s.KeyIssuer,
                        UpstreamID:  s.KeyUpstreamID,
                        RedirectURI: s.KeyRedirectURI,
                        ScopesHash:  s.KeyScopesHash,
                        RotationKey: s.KeyRotationKey,
                },
                ProviderName:            s.ProviderName,
                ClientID:                s.ClientID,
                ClientSecret:            s.ClientSecret,
                TokenEndpointAuthMethod: s.TokenEndpointAuthMethod,
                RegistrationAccessToken: s.RegistrationAccessToken,
                RegistrationClientURI:   s.RegistrationClientURI,
                AuthorizationEndpoint:   s.AuthorizationEndpoint,
                TokenEndpoint:           s.TokenEndpoint,
                CreatedAt:               createdAt,
                ClientSecretExpiresAt:   clientSecretExpiresAt,
        }
}

// StoreDCRCredentialsIfAbsent claims creds.Key for creds via a Redis
// WATCH/MULTI compare-and-set (the same pattern ReconcileConfiguredClient
// uses above), returning the authoritative durable value: creds itself on a
// successful claim, or the existing winner's credentials when the key
// already holds a live (non-expired) entry. Callers MUST use the returned
// value — RFC 7591 dynamic registration mints a unique client_id/client_secret
// on every call, so two replicas racing a cache miss for the same Key
// register genuinely different OAuth clients with the upstream, and only the
// value the durable store agrees to is a credential either replica actually
// holds. Defensive copy is provided implicitly by JSON serialisation —
// caller mutations after the call cannot reach the persisted bytes.
//
// # Expired existing rows are claimable
//
// An existing row whose ClientSecretExpiresAt is non-zero and already in the
// past is treated as absent, mirroring MemoryStorage.StoreDCRCredentialsIfAbsent:
// it is overwritten with creds rather than handed back as if it were still
// the authoritative winner — but only when creds itself is fresher (creds'
// own ClientSecretExpiresAt is zero or still in the future). When both the
// existing row and the incoming creds are already expired — the startup
// case where many replicas independently re-resolve a dead upstream
// registration and each mints equally-stale credentials — the existing row
// is returned as-is instead. Overwriting would gain nothing there, and
// without this guard every concurrent claimant would keep re-observing an
// expired row and re-entering the write path, risking exhausting
// maxDCRClaimRetries under contention instead of converging the way the
// genuinely-absent case does. WATCH still guards the overwrite path: if
// another replica claims or refreshes the same key between our GET and our
// SET, EXEC aborts with redis.TxFailedErr and the loop below retries, so
// this does not reopen the double-registration race the NX-based
// predecessor of this method was built to close — it only changes what
// happens when the existing row is provably expired.
//
// This means a caller CAN receive an already-expired credential back from
// this call — deliberately, per the above, rather than as an oversight. The
// only thing that later re-checks an expiry against ResolveCredentials'
// cache is lookupCachedResolution, which only runs on a *subsequent*
// ResolveCredentials call for the same key; nothing currently re-resolves
// after the one-time boot-time call in embeddedauthserver.go. So a replica
// that receives an expired credential here keeps it for the rest of its
// process lifetime, failing cleanly with an upstream invalid_client/
// invalid_grant on every use until restart. Closing that gap needs
// revisiting credentials after boot, not another guard in this function —
// tracked in #6496.
//
// # TTL
//
// When creds.ClientSecretExpiresAt is non-zero (the upstream advertised an
// RFC 7591 §3.2.1 client_secret_expires_at), the entry is stored with a Redis
// TTL derived from time.Until(ClientSecretExpiresAt) so the row evicts before
// the upstream rejects the secret at the token endpoint. When zero (RFC 7591
// "never"), TTL=0 is used and the entry is long-lived.
//
// If ClientSecretExpiresAt is already in the past at call time, the entry is
// written with the bounded TTL pastExpiryDCRTTL (1 second) rather than rejected
// or stored long-lived. This keeps the store's "already-expired secret" window
// narrow even if the resolver never re-reads the row, and matches the
// fail-loud-but-tolerant posture: the caller's expiry timestamp round-trips so
// a downstream reader can still observe it and trigger re-registration.
//
// Validation is delegated to validateDCRCredentialsForStore so the rejection
// set stays in sync with MemoryStorage and any future backend.
func (s *RedisStorage) StoreDCRCredentialsIfAbsent(ctx context.Context, creds *DCRCredentials) (*DCRCredentials, error) {
        if err := validateDCRCredentialsForStore(creds); err != nil {
                return nil, err
        }

        key := redisDCRKey(s.keyPrefix, creds.Key)

        data, ttl, err := marshalDCRCredentialsForStore(creds)
        if err != nil {
                return nil, err
        }

        // result is set by txFn on every non-error path: either to creds (a
        // successful claim) or to the decoded existing row (a live winner
        // already present). It must not be read until s.client.Watch returns nil.
        var result *DCRCredentials
        txFn := func(tx *redis.Tx) error {
                var txErr error
                result, txErr = dcrClaimOrReturnWinner(ctx, tx, key, data, ttl, creds)
                return txErr
        }

        // Retry bound mirrors ReconcileConfiguredClient's maxConfiguredClientReconcileRetries:
        // go-redis does not retry Watch internally, so a concurrent writer to the
        // watched key aborts EXEC with redis.TxFailedErr, which must be retried
        // here for the race this method exists to close to resolve on its own.
        var watchErr error
        for attempt := 0; attempt < maxDCRClaimRetries; attempt++ {
                watchErr = s.client.Watch(ctx, txFn, key)
                if watchErr == nil {
                        return result, nil
                }
                if !errors.Is(watchErr, redis.TxFailedErr) {
                        return nil, fmt.Errorf("failed to store dcr credentials: %w", watchErr)
                }
        }
        return nil, fmt.Errorf("failed to claim dcr credentials after %d attempts: %w", maxDCRClaimRetries, watchErr)
}

// marshalDCRCredentialsForStore serializes creds to its Redis wire form and
// derives the TTL to store it with. Split out of StoreDCRCredentialsIfAbsent
// purely to keep that method's cyclomatic complexity down; it has no
// transaction-retry concerns of its own.
//
// Derive Redis TTL from ClientSecretExpiresAt:
//   - Zero (unset)  -> TTL=0 (no expiration) per RFC 7591 §3.2.1 "never".
//   - Future expiry -> TTL = time.Until(expiry).
//   - Past expiry   -> TTL = pastExpiryDCRTTL (bounded eviction window).
//
// See StoreDCRCredentialsIfAbsent's docstring for the past-expiry rationale.
func marshalDCRCredentialsForStore(creds *DCRCredentials) ([]byte, time.Duration, error) {
        stored := newStoredDCRCredentials(creds)

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return nil, 0, fmt.Errorf("failed to marshal dcr credentials: %w", err)
        }

        ttl := time.Duration(0)
        if !creds.ClientSecretExpiresAt.IsZero() {
                if until := time.Until(creds.ClientSecretExpiresAt); until > 0 {
                        ttl = until
                } else {
                        ttl = pastExpiryDCRTTL
                }
        }
        return data, ttl, nil
}

// dcrClaimOrReturnWinner runs the read-check-write body of
// StoreDCRCredentialsIfAbsent's WATCH transaction: an absent row, or an
// existing row that is expired while creds is fresher, is claimed with
// (data, ttl) inside MULTI; a live existing row — or one that is expired but
// no fresher than creds — is decoded and returned as-is, with no write. See
// StoreDCRCredentialsIfAbsent's "Expired existing rows" doc section for why.
// Split out purely to keep StoreDCRCredentialsIfAbsent's cyclomatic
// complexity down.
func dcrClaimOrReturnWinner(
        ctx context.Context, tx *redis.Tx, key string, data []byte, ttl time.Duration, creds *DCRCredentials,
) (*DCRCredentials, error) {
        claim := func() error {
                _, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
                        pipe.Set(ctx, key, data, ttl)
                        return nil
                })
                return err
        }

        existingData, getErr := tx.Get(ctx, key).Bytes()
        if errors.Is(getErr, redis.Nil) {
                if err := claim(); err != nil {
                        return nil, err
                }
                return creds, nil
        }
        if getErr != nil {
                return nil, fmt.Errorf("failed to get existing dcr credentials: %w", getErr)
        }

        var existingStored storedDCRCredentials
        if unmarshalErr := json.Unmarshal(existingData, &existingStored); unmarshalErr != nil {
                return nil, fmt.Errorf("failed to unmarshal existing dcr credentials: %w", unmarshalErr)
        }
        existing := existingStored.toDCRCredentials()

        expired := !existing.ClientSecretExpiresAt.IsZero() && time.Now().After(existing.ClientSecretExpiresAt)
        if !expired {
                return existing, nil
        }

        // The existing row is expired. Only overwrite it when the incoming
        // credentials are actually fresher — a genuine new registration
        // reclaiming a dead slot. If the incoming credentials are ALSO already
        // expired (the startup-reconciliation case: many replicas independently
        // re-resolving a dead upstream registration all mint equally-stale
        // credentials), overwriting gains nothing and, under contention, every
        // concurrent claimant would keep re-observing an expired row and
        // re-entering this write path, burning through maxDCRClaimRetries
        // instead of converging. Returning the stable existing row lets every
        // claimant converge without a write, matching the "genuinely absent"
        // case's self-limiting behaviour.
        incomingExpired := !creds.ClientSecretExpiresAt.IsZero() && time.Now().After(creds.ClientSecretExpiresAt)
        if incomingExpired {
                return existing, nil
        }

        // The existing row is provably expired: treat it as absent and overwrite
        // it, matching MemoryStorage's semantics. WATCH still protects this — see
        // StoreDCRCredentialsIfAbsent's "Expired existing rows" doc section.
        if err := claim(); err != nil {
                return nil, err
        }
        return creds, nil
}

// UpdateDCRCredentialsIfPresent replaces the row at creds.Key with creds using
// a single Redis SET with the XX flag, returning ErrNotFound (wrapped) when no
// row physically exists at the key. It never creates: XX makes Redis itself
// refuse the write when the key is absent, so an update racing a concurrent
// delete or TTL eviction fails with ErrNotFound rather than silently
// re-creating the row.
//
// The write needs no WATCH/MULTI (unlike StoreDCRCredentialsIfAbsent, which
// must read the existing row to decide whether it may claim the slot). Nothing
// here depends on the old value: only the key's bare existence gates the
// write, and SET XX evaluates that existence and performs the write in one
// atomic server-side step. There is therefore no read-check-write window to
// guard, no lost-update race, and no possibility of redis.TxFailedErr.
//
// # Presence is physical, not liveness
//
// Unlike StoreDCRCredentialsIfAbsent, this does not treat an expired existing
// row as absent — any row whose key still exists is updatable, including one
// whose ClientSecretExpiresAt has passed but whose Redis key has not yet
// self-evicted. See the DCRCredentialStore interface docs for why Update gates
// on physical presence: it exists so a decorator can rewrite a row it just read
// via GetDCRCredentials, which itself does not filter on expiry.
//
// # TTL
//
// The rewritten row's TTL is derived from the incoming creds by
// marshalDCRCredentialsForStore, identically to StoreDCRCredentialsIfAbsent: a
// future ClientSecretExpiresAt sets that TTL, a zero value clears it (the row
// becomes long-lived, since a SET without KEEPTTL discards any existing TTL),
// and a past value uses the bounded pastExpiryDCRTTL. An update can therefore
// extend, shorten, or clear the row's TTL exactly as an initial store would.
//
// Validation is delegated to validateDCRCredentialsForStore so the rejection
// set stays in sync with MemoryStorage and any future backend.
func (s *RedisStorage) UpdateDCRCredentialsIfPresent(ctx context.Context, creds *DCRCredentials) (*DCRCredentials, error) {
        if err := validateDCRCredentialsForStore(creds); err != nil {
                return nil, err
        }

        key := redisDCRKey(s.keyPrefix, creds.Key)

        data, ttl, err := marshalDCRCredentialsForStore(creds)
        if err != nil {
                return nil, err
        }

        // Mode XX: write only if the key already exists. Redis replies nil when it
        // does not, which go-redis surfaces as redis.Nil.
        err = s.client.SetArgs(ctx, key, data, redis.SetArgs{Mode: "XX", TTL: ttl}).Err()
        if errors.Is(err, redis.Nil) {
                return nil, notFoundRFC6749Error("DCR credentials not found")
        }
        if err != nil {
                return nil, fmt.Errorf("failed to update dcr credentials: %w", err)
        }

        return cloneDCRCredentials(creds), nil
}

// UpdateDCRCredentialsIfUnchanged replaces the row at creds.Key with creds
// only when the stored row still equals expected, returning ErrNotFound
// (wrapped) when no row exists and ErrDCRCredentialsChanged when it differs.
//
// The compare and the write run in a WATCH/MULTI transaction on the single
// row key — the same pattern StoreDCRCredentialsIfAbsent uses — so the
// operation is atomic against concurrent writers and valid in standalone,
// Sentinel and Cluster modes alike (one key means one hash slot). If another
// writer touches the key between the GET and EXEC, EXEC aborts with
// redis.TxFailedErr and the read-compare-write is retried, up to
// maxDCRClaimRetries; the retry re-reads the row and so normally resolves to
// ErrDCRCredentialsChanged. If every attempt aborts, a generic error is
// returned and nothing was written by this call.
//
// The comparison is between decoded stored forms (expected is normalised via
// newStoredDCRCredentials, and storedDCRCredentials holds only strings and
// int64s, so == compares every field by value), not raw JSON bytes, so it is unaffected by field
// ordering or by a row written by an older encoder, and a value obtained from
// GetDCRCredentials always matches the row it was read from despite the
// one-second time precision of the stored form.
//
// Presence and TTL handling are identical to UpdateDCRCredentialsIfPresent:
// an expired-but-present row is updatable, and the rewritten row's TTL is
// derived from creds by marshalDCRCredentialsForStore. The write uses SET XX
// for parity with that method, although WATCH already guarantees the key
// still exists at EXEC.
func (s *RedisStorage) UpdateDCRCredentialsIfUnchanged(
        ctx context.Context, creds, expected *DCRCredentials,
) (*DCRCredentials, error) {
        if err := validateDCRCompareAndSet(creds, expected); err != nil {
                return nil, err
        }

        key := redisDCRKey(s.keyPrefix, creds.Key)

        data, ttl, err := marshalDCRCredentialsForStore(creds)
        if err != nil {
                return nil, err
        }
        want := newStoredDCRCredentials(expected)

        txFn := func(tx *redis.Tx) error {
                existingData, getErr := tx.Get(ctx, key).Bytes()
                if errors.Is(getErr, redis.Nil) {
                        return notFoundRFC6749Error("DCR credentials not found")
                }
                if getErr != nil {
                        return fmt.Errorf("failed to get existing dcr credentials: %w", getErr)
                }

                var existing storedDCRCredentials
                if unmarshalErr := json.Unmarshal(existingData, &existing); unmarshalErr != nil {
                        return fmt.Errorf("failed to unmarshal existing dcr credentials: %w", unmarshalErr)
                }
                if existing != want {
                        return ErrDCRCredentialsChanged
                }

                _, pipeErr := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
                        pipe.SetArgs(ctx, key, data, redis.SetArgs{Mode: "XX", TTL: ttl})
                        return nil
                })
                return pipeErr
        }

        var watchErr error
        for attempt := 0; attempt < maxDCRClaimRetries; attempt++ {
                watchErr = s.client.Watch(ctx, txFn, key)
                if watchErr == nil {
                        return cloneDCRCredentials(creds), nil
                }
                if errors.Is(watchErr, ErrNotFound) || errors.Is(watchErr, ErrDCRCredentialsChanged) {
                        return nil, watchErr
                }
                if !errors.Is(watchErr, redis.TxFailedErr) {
                        return nil, fmt.Errorf("failed to compare-and-set dcr credentials: %w", watchErr)
                }
        }
        return nil, fmt.Errorf("failed to compare-and-set dcr credentials after %d attempts: %w", maxDCRClaimRetries, watchErr)
}

// GetDCRCredentials retrieves the credentials previously persisted under key.
// Returns ErrNotFound (wrapped) when no entry exists. The returned value is a
// fresh struct decoded from JSON, which acts as a defensive copy.
//
// An unpopulated key (empty Issuer, UpstreamID, RedirectURI, or ScopesHash) cannot match
// any stored row because StoreDCRCredentialsIfAbsent rejects such keys, so a Get
// against one is a normal miss — ErrNotFound — matching
// MemoryStorage.GetDCRCredentials and the DCRCredentialStore interface contract.
func (s *RedisStorage) GetDCRCredentials(ctx context.Context, key DCRKey) (*DCRCredentials, error) {
        redisKey := redisDCRKey(s.keyPrefix, key)
        data, err := s.client.Get(ctx, redisKey).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("DCR credentials not found")
                }
                return nil, fmt.Errorf("failed to get dcr credentials: %w", err)
        }

        var stored storedDCRCredentials
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal dcr credentials: %w", err)
        }

        return stored.toDCRCredentials(), nil
}

// -----------------------
// Pending Authorization Storage
// -----------------------

// storedPendingAuthorization is a serializable wrapper for PendingAuthorization.
type storedPendingAuthorization struct {
        ClientID                string                  `json:"client_id"`
        RedirectURI             string                  `json:"redirect_uri"`
        State                   string                  `json:"state"`
        PKCEChallenge           string                  `json:"pkce_challenge"`
        PKCEMethod              string                  `json:"pkce_method"`
        Scopes                  []string                `json:"scopes"`
        Resource                string                  `json:"resource,omitempty"`
        ConsentStage            string                  `json:"consent_stage,omitempty"`
        RememberConsent         bool                    `json:"remember_consent,omitempty"`
        FirstProviderSubject    string                  `json:"first_provider_subject,omitempty"`
        ExpectedUserID          string                  `json:"expected_user_id,omitempty"`
        ExpectedProviderSubject string                  `json:"expected_provider_subject,omitempty"`
        ConsentSessionDigest    string                  `json:"consent_session_digest,omitempty"`
        InternalState           string                  `json:"internal_state"`
        UpstreamPKCEVerifier    string                  `json:"upstream_pkce_verifier"`
        UpstreamNonce           string                  `json:"upstream_nonce"`
        BrowserBindingHash      string                  `json:"browser_binding_hash,omitempty"`
        UpstreamProviderName    string                  `json:"upstream_provider_name,omitempty"`
        SessionID               string                  `json:"session_id,omitempty"`
        ResolvedUserID          string                  `json:"resolved_user_id,omitempty"`
        ResolvedUserName        string                  `json:"resolved_user_name,omitempty"`
        ResolvedUserEmail       string                  `json:"resolved_user_email,omitempty"`
        VerifiedUpstream        session.UpstreamSubject `json:"verified_upstream,omitzero"`
        SingleLeg               bool                    `json:"single_leg,omitempty"`
        ChainUpstreams          []string                `json:"chain_upstreams,omitempty"`
        CreatedAt               int64                   `json:"created_at"`
        CreatedAtNano           int64                   `json:"created_at_nano,omitempty"`
}

// StorePendingAuthorization stores a pending authorization request.
func (s *RedisStorage) StorePendingAuthorization(ctx context.Context, state string, pending *PendingAuthorization) error {
        if state == "" {
                return fosite.ErrInvalidRequest.WithHint("state cannot be empty")
        }
        if pending == nil {
                return fosite.ErrInvalidRequest.WithHint("pending authorization cannot be nil")
        }

        key := redisKey(s.keyPrefix, KeyTypePending, state)

        stored := storedPendingAuthorization{
                ClientID:                pending.ClientID,
                RedirectURI:             pending.RedirectURI,
                State:                   pending.State,
                PKCEChallenge:           pending.PKCEChallenge,
                PKCEMethod:              pending.PKCEMethod,
                Scopes:                  slices.Clone(pending.Scopes),
                Resource:                pending.Resource,
                ConsentStage:            pending.ConsentStage,
                RememberConsent:         pending.RememberConsent,
                FirstProviderSubject:    pending.FirstProviderSubject,
                ExpectedUserID:          pending.ExpectedUserID,
                ExpectedProviderSubject: pending.ExpectedProviderSubject,
                ConsentSessionDigest:    pending.ConsentSessionDigest,
                InternalState:           pending.InternalState,
                UpstreamPKCEVerifier:    pending.UpstreamPKCEVerifier,
                UpstreamNonce:           pending.UpstreamNonce,
                BrowserBindingHash:      pending.BrowserBindingHash,
                UpstreamProviderName:    pending.UpstreamProviderName,
                SessionID:               pending.SessionID,
                ResolvedUserID:          pending.ResolvedUserID,
                ResolvedUserName:        pending.ResolvedUserName,
                ResolvedUserEmail:       pending.ResolvedUserEmail,
                VerifiedUpstream:        pending.VerifiedUpstream,
                SingleLeg:               pending.SingleLeg,
                ChainUpstreams:          slices.Clone(pending.ChainUpstreams),
                CreatedAt:               pending.CreatedAt.Unix(),
                CreatedAtNano:           pending.CreatedAt.UnixNano(),
        }

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal pending authorization: %w", err)
        }

        ttl := time.Until(pending.CreatedAt.Add(DefaultPendingAuthorizationTTL))
        if ttl <= 0 {
                return ErrExpired
        }
        return s.client.Set(ctx, key, data, ttl).Err()
}

// LoadPendingAuthorization retrieves a pending authorization by internal state.
func (s *RedisStorage) LoadPendingAuthorization(ctx context.Context, state string) (*PendingAuthorization, error) {
        key := redisKey(s.keyPrefix, KeyTypePending, state)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Pending authorization not found")
                }
                return nil, fmt.Errorf("failed to get pending authorization: %w", err)
        }

        var stored storedPendingAuthorization
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal pending authorization: %w", err)
        }

        createdAt := time.Unix(stored.CreatedAt, 0)
        if stored.CreatedAtNano != 0 {
                createdAt = time.Unix(0, stored.CreatedAtNano)
        }

        // Check if expired (TTL should handle this, but double-check)
        if !createdAt.Add(DefaultPendingAuthorizationTTL).After(time.Now()) {
                return nil, ErrExpired
        }

        return &PendingAuthorization{
                ClientID:                stored.ClientID,
                RedirectURI:             stored.RedirectURI,
                State:                   stored.State,
                PKCEChallenge:           stored.PKCEChallenge,
                PKCEMethod:              stored.PKCEMethod,
                Scopes:                  slices.Clone(stored.Scopes),
                Resource:                stored.Resource,
                ConsentStage:            stored.ConsentStage,
                RememberConsent:         stored.RememberConsent,
                FirstProviderSubject:    stored.FirstProviderSubject,
                ExpectedUserID:          stored.ExpectedUserID,
                ExpectedProviderSubject: stored.ExpectedProviderSubject,
                ConsentSessionDigest:    stored.ConsentSessionDigest,
                InternalState:           stored.InternalState,
                UpstreamPKCEVerifier:    stored.UpstreamPKCEVerifier,
                UpstreamNonce:           stored.UpstreamNonce,
                BrowserBindingHash:      stored.BrowserBindingHash,
                UpstreamProviderName:    stored.UpstreamProviderName,
                SessionID:               stored.SessionID,
                ResolvedUserID:          stored.ResolvedUserID,
                ResolvedUserName:        stored.ResolvedUserName,
                ResolvedUserEmail:       stored.ResolvedUserEmail,
                VerifiedUpstream:        stored.VerifiedUpstream,
                SingleLeg:               stored.SingleLeg,
                ChainUpstreams:          slices.Clone(stored.ChainUpstreams),
                CreatedAt:               createdAt,
        }, nil
}

// DeletePendingAuthorization removes a pending authorization.
func (s *RedisStorage) DeletePendingAuthorization(ctx context.Context, state string) error {
        key := redisKey(s.keyPrefix, KeyTypePending, state)

        result, err := s.client.Del(ctx, key).Result()
        if err != nil {
                return fmt.Errorf("failed to delete pending authorization: %w", err)
        }
        if result == 0 {
                return notFoundRFC6749Error("Pending authorization not found")
        }

        return nil
}

// -----------------------
// Pending Device Login Storage
// -----------------------

// storedPendingDeviceLogin is a serializable wrapper for PendingDeviceLogin.
type storedPendingDeviceLogin struct {
        DeviceCode           string `json:"device_code"`
        UserCode             string `json:"user_code"`
        UpstreamPKCEVerifier string `json:"upstream_pkce_verifier"`
        UpstreamNonce        string `json:"upstream_nonce"`
        BrowserBindingHash   string `json:"browser_binding_hash,omitempty"`
        UpstreamProviderName string `json:"upstream_provider_name,omitempty"`
        CreatedAt            int64  `json:"created_at"`
}

// StorePendingDeviceLogin stores a pending device-flow verification-page login.
func (s *RedisStorage) StorePendingDeviceLogin(ctx context.Context, state string, pending *PendingDeviceLogin) error {
        if state == "" {
                return fosite.ErrInvalidRequest.WithHint("state cannot be empty")
        }
        if pending == nil {
                return fosite.ErrInvalidRequest.WithHint("pending device login cannot be nil")
        }

        key := redisKey(s.keyPrefix, KeyTypePendingDeviceLogin, state)

        stored := storedPendingDeviceLogin{
                DeviceCode:           pending.DeviceCode,
                UserCode:             pending.UserCode,
                UpstreamPKCEVerifier: pending.UpstreamPKCEVerifier,
                UpstreamNonce:        pending.UpstreamNonce,
                BrowserBindingHash:   pending.BrowserBindingHash,
                UpstreamProviderName: pending.UpstreamProviderName,
                CreatedAt:            pending.CreatedAt.Unix(),
        }

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal pending device login: %w", err)
        }

        return s.client.Set(ctx, key, data, DefaultDeviceLoginTTL).Err()
}

// LoadPendingDeviceLogin retrieves a pending device login by state.
func (s *RedisStorage) LoadPendingDeviceLogin(ctx context.Context, state string) (*PendingDeviceLogin, error) {
        key := redisKey(s.keyPrefix, KeyTypePendingDeviceLogin, state)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Pending device login not found")
                }
                return nil, fmt.Errorf("failed to get pending device login: %w", err)
        }

        var stored storedPendingDeviceLogin
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal pending device login: %w", err)
        }

        createdAt := time.Unix(stored.CreatedAt, 0)
        if time.Since(createdAt) > DefaultDeviceLoginTTL {
                return nil, ErrExpired
        }

        return &PendingDeviceLogin{
                DeviceCode:           stored.DeviceCode,
                UserCode:             stored.UserCode,
                UpstreamPKCEVerifier: stored.UpstreamPKCEVerifier,
                UpstreamNonce:        stored.UpstreamNonce,
                BrowserBindingHash:   stored.BrowserBindingHash,
                UpstreamProviderName: stored.UpstreamProviderName,
                CreatedAt:            createdAt,
        }, nil
}

// DeletePendingDeviceLogin removes a pending device login.
func (s *RedisStorage) DeletePendingDeviceLogin(ctx context.Context, state string) error {
        key := redisKey(s.keyPrefix, KeyTypePendingDeviceLogin, state)

        result, err := s.client.Del(ctx, key).Result()
        if err != nil {
                return fmt.Errorf("failed to delete pending device login: %w", err)
        }
        if result == 0 {
                return notFoundRFC6749Error("Pending device login not found")
        }
        return nil
}

// -----------------------
// Pending Device Confirmation Storage
// -----------------------

// storedPendingDeviceConfirmation is a serializable wrapper for
// PendingDeviceConfirmation.
type storedPendingDeviceConfirmation struct {
        DeviceCode        string                  `json:"device_code"`
        UserCode          string                  `json:"user_code"`
        ResolvedUserID    string                  `json:"resolved_user_id,omitempty"`
        ResolvedUserName  string                  `json:"resolved_user_name,omitempty"`
        ResolvedUserEmail string                  `json:"resolved_user_email,omitempty"`
        VerifiedUpstream  session.UpstreamSubject `json:"verified_upstream,omitzero"`
        UpstreamTokens    *storedUpstreamTokens   `json:"upstream_tokens,omitempty"`
        Synthetic         bool                    `json:"synthetic,omitempty"`
        CreatedAt         int64                   `json:"created_at"`
}

// StorePendingDeviceConfirmation stores a pending device-flow confirmation,
// keyed by an opaque token. See PendingDeviceConfirmation's doc comment for
// why the resolved identity is addressed by token rather than round-tripped
// through the browser.
func (s *RedisStorage) StorePendingDeviceConfirmation(
        ctx context.Context, token string, pending *PendingDeviceConfirmation,
) error {
        if token == "" {
                return fosite.ErrInvalidRequest.WithHint("token cannot be empty")
        }
        if pending == nil {
                return fosite.ErrInvalidRequest.WithHint("pending device confirmation cannot be nil")
        }

        key := redisKey(s.keyPrefix, KeyTypePendingDeviceConfirmation, token)

        stored := storedPendingDeviceConfirmation{
                DeviceCode:        pending.DeviceCode,
                UserCode:          pending.UserCode,
                ResolvedUserID:    pending.ResolvedUserID,
                ResolvedUserName:  pending.ResolvedUserName,
                ResolvedUserEmail: pending.ResolvedUserEmail,
                VerifiedUpstream:  pending.VerifiedUpstream,
                UpstreamTokens:    newStoredUpstreamTokens(pending.UpstreamTokens),
                Synthetic:         pending.Synthetic,
                CreatedAt:         pending.CreatedAt.Unix(),
        }

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal pending device confirmation: %w", err)
        }

        return s.client.Set(ctx, key, data, DefaultDeviceLoginTTL).Err()
}

// LoadPendingDeviceConfirmation retrieves a pending device confirmation by token.
func (s *RedisStorage) LoadPendingDeviceConfirmation(
        ctx context.Context, token string,
) (*PendingDeviceConfirmation, error) {
        key := redisKey(s.keyPrefix, KeyTypePendingDeviceConfirmation, token)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, notFoundRFC6749Error("Pending device confirmation not found")
                }
                return nil, fmt.Errorf("failed to get pending device confirmation: %w", err)
        }

        var stored storedPendingDeviceConfirmation
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal pending device confirmation: %w", err)
        }

        createdAt := time.Unix(stored.CreatedAt, 0)
        if time.Since(createdAt) > DefaultDeviceLoginTTL {
                return nil, ErrExpired
        }

        var upstreamTokens *UpstreamTokens
        if stored.UpstreamTokens != nil {
                upstreamTokens = stored.UpstreamTokens.toUpstreamTokens()
        }

        return &PendingDeviceConfirmation{
                DeviceCode:        stored.DeviceCode,
                UserCode:          stored.UserCode,
                ResolvedUserID:    stored.ResolvedUserID,
                ResolvedUserName:  stored.ResolvedUserName,
                ResolvedUserEmail: stored.ResolvedUserEmail,
                VerifiedUpstream:  stored.VerifiedUpstream,
                UpstreamTokens:    upstreamTokens,
                Synthetic:         stored.Synthetic,
                CreatedAt:         createdAt,
        }, nil
}

// DeletePendingDeviceConfirmation removes a pending device confirmation.
func (s *RedisStorage) DeletePendingDeviceConfirmation(ctx context.Context, token string) error {
        key := redisKey(s.keyPrefix, KeyTypePendingDeviceConfirmation, token)

        result, err := s.client.Del(ctx, key).Result()
        if err != nil {
                return fmt.Errorf("failed to delete pending device confirmation: %w", err)
        }
        if result == 0 {
                return notFoundRFC6749Error("Pending device confirmation not found")
        }
        return nil
}

// -----------------------
// Device Code Storage
// -----------------------

// storedDeviceRequest is a serializable wrapper for DeviceRequest.
// LastPolledAtNano is stored at nanosecond precision (unlike CreatedAt's
// epoch-seconds), because enforcePollInterval in deviceflow/handler.go
// compares it against MinInterval and a poll landing fractionally under the
// required interval must not appear to have waited long enough merely due to
// storage rounding -- see deviceTimeToUnixNano / deviceUnixNanoToTime.
type storedDeviceRequest struct {
        DeviceCode        string                  `json:"device_code"`
        UserCode          string                  `json:"user_code"`
        ClientID          string                  `json:"client_id"`
        Scopes            []string                `json:"scopes"`
        Audience          []string                `json:"audience,omitempty"`
        Status            DeviceRequestStatus     `json:"status"`
        IntervalSeconds   int64                   `json:"interval_seconds,omitempty"`
        LastPolledAtNano  int64                   `json:"last_polled_at_nano,omitempty"`
        ResolvedUserID    string                  `json:"resolved_user_id,omitempty"`
        ResolvedUserName  string                  `json:"resolved_user_name,omitempty"`
        ResolvedUserEmail string                  `json:"resolved_user_email,omitempty"`
        VerifiedUpstream  session.UpstreamSubject `json:"verified_upstream,omitzero"`
        SessionID         string                  `json:"session_id,omitempty"`
        CreatedAt         int64                   `json:"created_at"`
}

// deviceTimeToUnixNano converts t to epoch nanoseconds for JSON storage, using 0 to
// mean the zero time.Time so a never-polled request round-trips through
// LastPolledAt's IsZero() check rather than colliding with the Unix epoch.
func deviceTimeToUnixNano(t time.Time) int64 {
        if t.IsZero() {
                return 0
        }
        return t.UnixNano()
}

// deviceUnixNanoToTime is the inverse of deviceTimeToUnixNano.
func deviceUnixNanoToTime(unixNano int64) time.Time {
        if unixNano == 0 {
                return time.Time{}
        }
        return time.Unix(0, unixNano)
}

// toDeviceRequest converts the wire representation back to a DeviceRequest,
// cloning slice fields so the caller cannot mutate this storage's JSON-decoded
// backing arrays.
func (stored *storedDeviceRequest) toDeviceRequest() *DeviceRequest {
        return &DeviceRequest{
                DeviceCode:        stored.DeviceCode,
                UserCode:          stored.UserCode,
                ClientID:          stored.ClientID,
                Scopes:            slices.Clone(stored.Scopes),
                Audience:          slices.Clone(stored.Audience),
                Status:            stored.Status,
                Interval:          time.Duration(stored.IntervalSeconds) * time.Second,
                LastPolledAt:      deviceUnixNanoToTime(stored.LastPolledAtNano),
                ResolvedUserID:    stored.ResolvedUserID,
                ResolvedUserName:  stored.ResolvedUserName,
                ResolvedUserEmail: stored.ResolvedUserEmail,
                VerifiedUpstream:  stored.VerifiedUpstream,
                SessionID:         stored.SessionID,
                CreatedAt:         time.Unix(stored.CreatedAt, 0),
        }
}

// StoreDeviceRequest stores a new pending device request, indexed by both
// DeviceCode and UserCode.
//
// Both keys are created with SETNX inside a single Redis transaction, so a
// concurrent double-store cannot silently overwrite either index. Because
// MULTI/EXEC applies each queued command unconditionally (there is no
// short-circuit between them), the two SETNX calls can succeed and fail
// independently; when exactly one did, the successful key is deleted so no
// half-written pair is ever left behind.
func (s *RedisStorage) StoreDeviceRequest(ctx context.Context, device *DeviceRequest) error {
        if device == nil {
                return fosite.ErrInvalidRequest.WithHint("device request cannot be nil")
        }
        if device.DeviceCode == "" {
                return fosite.ErrInvalidRequest.WithHint("device code cannot be empty")
        }
        if device.UserCode == "" {
                return fosite.ErrInvalidRequest.WithHint("user code cannot be empty")
        }
        if device.Status != DeviceRequestStatusPending {
                return fosite.ErrInvalidRequest.WithHint("device request must be created with pending status")
        }

        deviceKey := redisKey(s.keyPrefix, KeyTypeDeviceCode, device.DeviceCode)
        userCodeKey := redisKey(s.keyPrefix, KeyTypeDeviceUserCode, device.UserCode)

        stored := storedDeviceRequest{
                DeviceCode:        device.DeviceCode,
                UserCode:          device.UserCode,
                ClientID:          device.ClientID,
                Scopes:            slices.Clone(device.Scopes),
                Audience:          slices.Clone(device.Audience),
                Status:            device.Status,
                IntervalSeconds:   int64(device.Interval / time.Second),
                LastPolledAtNano:  deviceTimeToUnixNano(device.LastPolledAt),
                ResolvedUserID:    device.ResolvedUserID,
                ResolvedUserName:  device.ResolvedUserName,
                ResolvedUserEmail: device.ResolvedUserEmail,
                VerifiedUpstream:  device.VerifiedUpstream,
                SessionID:         device.SessionID,
                CreatedAt:         device.CreatedAt.Unix(),
        }

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal device request: %w", err)
        }

        pipe := s.client.TxPipeline()
        deviceCmd := pipe.SetNX(ctx, deviceKey, data, DefaultDeviceRequestTTL)
        userCodeCmd := pipe.SetNX(ctx, userCodeKey, device.DeviceCode, DefaultDeviceRequestTTL)
        if _, err := pipe.Exec(ctx); err != nil {
                return fmt.Errorf("failed to store device request: %w", err)
        }

        deviceCreated, userCodeCreated := deviceCmd.Val(), userCodeCmd.Val()
        if deviceCreated && userCodeCreated {
                return nil
        }
        if deviceCreated {
                warnOnCleanupErr(s.client.Del(ctx, deviceKey).Err(), "StoreDeviceRequest cleanup", deviceKey)
                return fmt.Errorf("%w: user code %q", ErrAlreadyExists, device.UserCode)
        }
        if userCodeCreated {
                warnOnCleanupErr(s.client.Del(ctx, userCodeKey).Err(), "StoreDeviceRequest cleanup", userCodeKey)
        }
        return fmt.Errorf("%w: device code %q", ErrAlreadyExists, device.DeviceCode)
}

// getDeviceRequestByDeviceCode is the shared get+unmarshal+expiry-check logic
// used by both LoadDeviceRequestByDeviceCode and LoadDeviceRequestByUserCode
// (once the latter has resolved its device_code via the secondary index).
func (s *RedisStorage) getDeviceRequestByDeviceCode(ctx context.Context, deviceCode string) (*DeviceRequest, error) {
        key := redisKey(s.keyPrefix, KeyTypeDeviceCode, deviceCode)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, fmt.Errorf("%w: device request not found", ErrNotFound)
                }
                return nil, fmt.Errorf("failed to get device request: %w", err)
        }

        var stored storedDeviceRequest
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal device request: %w", err)
        }

        // Check if expired (TTL should handle this, but double-check).
        if time.Since(time.Unix(stored.CreatedAt, 0)) > DefaultDeviceRequestTTL {
                return nil, ErrExpired
        }

        return stored.toDeviceRequest(), nil
}

// LoadDeviceRequestByDeviceCode retrieves a device request by its device_code.
func (s *RedisStorage) LoadDeviceRequestByDeviceCode(ctx context.Context, deviceCode string) (*DeviceRequest, error) {
        return s.getDeviceRequestByDeviceCode(ctx, deviceCode)
}

// LoadDeviceRequestByUserCode retrieves a device request by its user_code,
// for the verification page.
func (s *RedisStorage) LoadDeviceRequestByUserCode(ctx context.Context, userCode string) (*DeviceRequest, error) {
        userCodeKey := redisKey(s.keyPrefix, KeyTypeDeviceUserCode, userCode)

        deviceCode, err := s.client.Get(ctx, userCodeKey).Result()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, fmt.Errorf("%w: device request not found", ErrNotFound)
                }
                return nil, fmt.Errorf("failed to get device request user_code index: %w", err)
        }

        return s.getDeviceRequestByDeviceCode(ctx, deviceCode)
}

// updateDeviceRequest performs a read-modify-write on the device_code record
// identified by deviceCode: it loads the current record, applies mutate, and
// writes it back with redis.KeepTTL so the record's remaining TTL is
// preserved. mutate returns an error (e.g. ErrInvalidState) to abort the
// write without touching the stored record.
//
// The read-modify-write runs inside a Redis WATCH/MULTI transaction,
// mirroring DeleteDeviceRequest: two concurrent callers racing the same
// device_code (e.g. a double-click authorize+deny, or two verification-page
// tabs) cannot both observe the pre-mutation status and both commit --
// whichever EXECs second aborts with redis.TxFailedErr and retries against
// the now-updated record, so mutate's pending-status check (or any other
// invariant it enforces) actually holds across the transition rather than
// being a TOCTOU check on a plain read-then-write. See
// maxConfiguredClientReconcileRetries for why the retry count is bounded.
func (s *RedisStorage) updateDeviceRequest(
        ctx context.Context, deviceCode string, mutate func(*storedDeviceRequest) error,
) error {
        key := redisKey(s.keyPrefix, KeyTypeDeviceCode, deviceCode)

        txFn := func(tx *redis.Tx) error {
                data, err := tx.Get(ctx, key).Bytes()
                if err != nil {
                        if errors.Is(err, redis.Nil) {
                                return fmt.Errorf("%w: device request not found", ErrNotFound)
                        }
                        return fmt.Errorf("failed to get device request: %w", err)
                }

                var stored storedDeviceRequest
                if err := json.Unmarshal(data, &stored); err != nil {
                        return fmt.Errorf("failed to unmarshal device request: %w", err)
                }

                if time.Since(time.Unix(stored.CreatedAt, 0)) > DefaultDeviceRequestTTL {
                        return ErrExpired
                }

                if err := mutate(&stored); err != nil {
                        return err
                }

                updated, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
                if err != nil {
                        return fmt.Errorf("failed to marshal device request: %w", err)
                }

                _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
                        pipe.Set(ctx, key, updated, redis.KeepTTL)
                        return nil
                })
                return err
        }

        // Watch returns any non-TxFailedErr error from txFn (not-found, expired,
        // mutate's own validation error, marshal/get failures) immediately without
        // retrying, since a fresh GET would only reproduce it. Only a genuine
        // transaction conflict (another writer changed the watched key between our
        // GET and EXEC) is retried, up to maxConfiguredClientReconcileRetries.
        var watchErr error
        for attempt := 0; attempt < maxConfiguredClientReconcileRetries; attempt++ {
                watchErr = s.client.Watch(ctx, txFn, key)
                if !errors.Is(watchErr, redis.TxFailedErr) {
                        return watchErr
                }
        }
        return fmt.Errorf("failed to update device request: %w", watchErr)
}

// MarkDeviceRequestAuthorized transitions a pending device request to
// authorized, attaching the resolved identity.
func (s *RedisStorage) MarkDeviceRequestAuthorized(
        ctx context.Context, deviceCode string, auth DeviceAuthorization,
) error {
        return s.updateDeviceRequest(ctx, deviceCode, func(stored *storedDeviceRequest) error {
                if stored.Status != DeviceRequestStatusPending {
                        return fmt.Errorf("%w: device request is %q, not pending", ErrInvalidState, stored.Status)
                }
                stored.Status = DeviceRequestStatusAuthorized
                stored.ResolvedUserID = auth.UserID
                stored.ResolvedUserName = auth.UserName
                stored.ResolvedUserEmail = auth.UserEmail
                stored.VerifiedUpstream = auth.VerifiedUpstream
                stored.SessionID = auth.SessionID
                return nil
        })
}

// MarkDeviceRequestDenied transitions a pending device request to denied.
func (s *RedisStorage) MarkDeviceRequestDenied(ctx context.Context, deviceCode string) error {
        return s.updateDeviceRequest(ctx, deviceCode, func(stored *storedDeviceRequest) error {
                if stored.Status != DeviceRequestStatusPending {
                        return fmt.Errorf("%w: device request is %q, not pending", ErrInvalidState, stored.Status)
                }
                stored.Status = DeviceRequestStatusDenied
                return nil
        })
}

// UpdateDeviceRequestLastPolledAt records the time of the most recent poll.
func (s *RedisStorage) UpdateDeviceRequestLastPolledAt(ctx context.Context, deviceCode string, polledAt time.Time) error {
        return s.updateDeviceRequest(ctx, deviceCode, func(stored *storedDeviceRequest) error {
                stored.LastPolledAtNano = deviceTimeToUnixNano(polledAt)
                return nil
        })
}

// DeleteDeviceRequest removes a device request, e.g. once its token has been
// issued so the device_code cannot be redeemed twice. Both the canonical
// record and the user_code secondary index are deleted in one pipeline.
//
// The get-then-delete sequence runs inside a Redis WATCH/MULTI transaction,
// mirroring UpsertDCRIssuedClient: if another caller deletes (or otherwise
// modifies) the key between our GET and EXEC, EXEC aborts with
// redis.TxFailedErr. Retrying then re-runs the GET, which now observes
// redis.Nil and reports ErrNotFound to the loser -- so a second concurrent or
// replayed DeleteDeviceRequest for the same device_code can never also
// observe success, which is what made this method single-use. See
// maxConfiguredClientReconcileRetries for why the retry count is bounded.
func (s *RedisStorage) DeleteDeviceRequest(ctx context.Context, deviceCode string) error {
        key := redisKey(s.keyPrefix, KeyTypeDeviceCode, deviceCode)

        txFn := func(tx *redis.Tx) error {
                data, err := tx.Get(ctx, key).Bytes()
                if err != nil {
                        if errors.Is(err, redis.Nil) {
                                return fmt.Errorf("%w: device request not found", ErrNotFound)
                        }
                        return fmt.Errorf("failed to get device request: %w", err)
                }

                var stored storedDeviceRequest
                if err := json.Unmarshal(data, &stored); err != nil {
                        return fmt.Errorf("failed to unmarshal device request: %w", err)
                }

                userCodeKey := redisKey(s.keyPrefix, KeyTypeDeviceUserCode, stored.UserCode)

                _, err = tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
                        pipe.Del(ctx, key)
                        pipe.Del(ctx, userCodeKey)
                        return nil
                })
                return err
        }

        // Watch returns any non-TxFailedErr error from txFn (not-found, get/unmarshal
        // failures) immediately without retrying, since a fresh GET would only
        // reproduce it. Only a genuine transaction conflict (another caller deleted
        // or modified the watched key between our GET and EXEC) is retried, up to
        // maxConfiguredClientReconcileRetries -- at which point the retry re-runs
        // the GET, which now observes redis.Nil and reports ErrNotFound to the
        // loser, so a second concurrent or replayed delete can never also observe
        // success.
        var watchErr error
        for attempt := 0; attempt < maxConfiguredClientReconcileRetries; attempt++ {
                watchErr = s.client.Watch(ctx, txFn, key)
                if !errors.Is(watchErr, redis.TxFailedErr) {
                        return watchErr
                }
        }
        return fmt.Errorf("failed to delete device request: %w", watchErr)
}

// -----------------------
// User Storage
// -----------------------

// storedUser is a serializable wrapper for User.
type storedUser struct {
        ID        string `json:"id"`
        CreatedAt int64  `json:"created_at"`
        UpdatedAt int64  `json:"updated_at"`
}

// CreateUser creates a new user account.
func (s *RedisStorage) CreateUser(ctx context.Context, user *User) error {
        if user == nil {
                return fosite.ErrInvalidRequest.WithHint("user cannot be nil")
        }
        if user.ID == "" {
                return fosite.ErrInvalidRequest.WithHint("user ID cannot be empty")
        }

        key := redisKey(s.keyPrefix, KeyTypeUser, user.ID)

        stored := storedUser{
                ID:        user.ID,
                CreatedAt: user.CreatedAt.Unix(),
                UpdatedAt: user.UpdatedAt.Unix(),
        }

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal user: %w", err)
        }

        // Use SetNX for atomic check-and-set to prevent race conditions.
        // Users don't expire (TTL=0).
        result, err := s.client.SetNX(ctx, key, data, 0).Result()
        if err != nil {
                return fmt.Errorf("failed to create user: %w", err)
        }
        if !result {
                return fmt.Errorf("%w: user already exists", ErrAlreadyExists)
        }

        return nil
}

// GetUser retrieves a user by their internal ID.
func (s *RedisStorage) GetUser(ctx context.Context, id string) (*User, error) {
        key := redisKey(s.keyPrefix, KeyTypeUser, id)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, fmt.Errorf("%w: user not found", ErrNotFound)
                }
                return nil, fmt.Errorf("failed to get user: %w", err)
        }

        var stored storedUser
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal user: %w", err)
        }

        return &User{
                ID:        stored.ID,
                CreatedAt: time.Unix(stored.CreatedAt, 0),
                UpdatedAt: time.Unix(stored.UpdatedAt, 0),
        }, nil
}

// DeleteUser removes a user account and all associated data.
func (s *RedisStorage) DeleteUser(ctx context.Context, id string) error {
        key := redisKey(s.keyPrefix, KeyTypeUser, id)

        // Check if user exists
        exists, err := s.client.Exists(ctx, key).Result()
        if err != nil {
                return fmt.Errorf("failed to check user existence: %w", err)
        }
        if exists == 0 {
                return fmt.Errorf("%w: user not found", ErrNotFound)
        }

        // Delete the user
        if err := s.client.Del(ctx, key).Err(); err != nil {
                return fmt.Errorf("failed to delete user: %w", err)
        }

        // Best-effort cascade deletion of associated data (see warnOnCleanupErr).
        // The user record is already deleted; orphaned keys will expire via TTL.
        userProviderSetKey := redisSetKey(s.keyPrefix, KeyTypeUserProviders, id)
        providerKeys, err := s.client.SMembers(ctx, userProviderSetKey).Result()
        if err == nil {
                for _, providerKey := range providerKeys {
                        warnOnCleanupErr(s.client.Del(ctx, providerKey).Err(), "Del", providerKey)
                }
                warnOnCleanupErr(s.client.Del(ctx, userProviderSetKey).Err(), "Del", userProviderSetKey)
        }

        // Delete associated upstream tokens
        userUpstreamSetKey := redisSetKey(s.keyPrefix, KeyTypeUserUpstream, id)
        upstreamKeys, err := s.client.SMembers(ctx, userUpstreamSetKey).Result()
        if err == nil {
                for _, upstreamKey := range upstreamKeys {
                        warnOnCleanupErr(s.client.Del(ctx, upstreamKey).Err(), "Del", upstreamKey)
                }
                warnOnCleanupErr(s.client.Del(ctx, userUpstreamSetKey).Err(), "Del", userUpstreamSetKey)
        }

        return nil
}

// -----------------------
// Provider Identity Storage
// -----------------------

// storedProviderIdentity is a serializable wrapper for ProviderIdentity.
type storedProviderIdentity struct {
        UserID          string `json:"user_id"`
        ProviderID      string `json:"provider_id"`
        ProviderSubject string `json:"provider_subject"`
        LinkedAt        int64  `json:"linked_at"`
        LastUsedAt      int64  `json:"last_used_at"`
}

// CreateProviderIdentity links a provider identity to a user.
func (s *RedisStorage) CreateProviderIdentity(ctx context.Context, identity *ProviderIdentity) error {
        if identity == nil {
                return fosite.ErrInvalidRequest.WithHint("identity cannot be nil")
        }
        if identity.UserID == "" {
                return fosite.ErrInvalidRequest.WithHint("user ID cannot be empty")
        }
        if identity.ProviderID == "" {
                return fosite.ErrInvalidRequest.WithHint("provider ID cannot be empty")
        }
        if identity.ProviderSubject == "" {
                return fosite.ErrInvalidRequest.WithHint("provider subject cannot be empty")
        }

        // Verify user exists
        userKey := redisKey(s.keyPrefix, KeyTypeUser, identity.UserID)
        exists, err := s.client.Exists(ctx, userKey).Result()
        if err != nil {
                return fmt.Errorf("failed to check user existence: %w", err)
        }
        if exists == 0 {
                return fmt.Errorf("%w: user not found", ErrNotFound)
        }

        key := redisProviderKey(s.keyPrefix, identity.ProviderID, identity.ProviderSubject)

        stored := storedProviderIdentity{
                UserID:          identity.UserID,
                ProviderID:      identity.ProviderID,
                ProviderSubject: identity.ProviderSubject,
                LinkedAt:        identity.LinkedAt.Unix(),
                LastUsedAt:      identity.LastUsedAt.Unix(),
        }

        data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
        if err != nil {
                return fmt.Errorf("failed to marshal identity: %w", err)
        }

        // Use SetNX for atomic check-and-set to prevent race conditions.
        // Provider identities don't expire (TTL=0).
        result, err := s.client.SetNX(ctx, key, data, 0).Result()
        if err != nil {
                return fmt.Errorf("failed to store identity: %w", err)
        }
        if !result {
                return fmt.Errorf("%w: provider identity already linked", ErrAlreadyExists)
        }

        // Add to user's provider identity set for reverse lookup.
        // Note: This set may contain stale references if identities are deleted independently.
        // Stale entries are cleaned up lazily during GetUserProviderIdentities reads.
        // A future DeleteProviderIdentity method should also clean up this set.
        userProviderSetKey := redisSetKey(s.keyPrefix, KeyTypeUserProviders, identity.UserID)
        return s.client.SAdd(ctx, userProviderSetKey, key).Err()
}

// GetProviderIdentity retrieves a provider identity by provider ID and subject.
func (s *RedisStorage) GetProviderIdentity(ctx context.Context, providerID, providerSubject string) (*ProviderIdentity, error) {
        key := redisProviderKey(s.keyPrefix, providerID, providerSubject)

        data, err := s.client.Get(ctx, key).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, fmt.Errorf("%w: provider identity not found", ErrNotFound)
                }
                return nil, fmt.Errorf("failed to get identity: %w", err)
        }

        var stored storedProviderIdentity
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal identity: %w", err)
        }

        return &ProviderIdentity{
                UserID:          stored.UserID,
                ProviderID:      stored.ProviderID,
                ProviderSubject: stored.ProviderSubject,
                LinkedAt:        time.Unix(stored.LinkedAt, 0),
                LastUsedAt:      time.Unix(stored.LastUsedAt, 0),
        }, nil
}

// updateLastUsedScript is a Lua script that atomically updates the LastUsedAt field
// of a provider identity. This prevents race conditions when multiple requests
// try to update the same identity concurrently.
// Returns 1 on success, 0 if the key doesn't exist.
var updateLastUsedScript = redis.NewScript(`
local data = redis.call('GET', KEYS[1])
if not data then
        return 0
end
local identity = cjson.decode(data)
identity.last_used_at = tonumber(ARGV[1])
redis.call('SET', KEYS[1], cjson.encode(identity))
return 1
`)

// UpdateProviderIdentityLastUsed updates the LastUsedAt timestamp for a provider identity.
// Uses a Lua script to ensure atomicity and prevent race conditions.
func (s *RedisStorage) UpdateProviderIdentityLastUsed(
        ctx context.Context, providerID, providerSubject string, lastUsedAt time.Time,
) error {
        key := redisProviderKey(s.keyPrefix, providerID, providerSubject)

        result, err := updateLastUsedScript.Run(ctx, s.client, []string{key}, lastUsedAt.Unix()).Int()
        if err != nil {
                return fmt.Errorf("failed to update identity: %w", err)
        }

        // Script returns 0 if the key doesn't exist
        if result == 0 {
                return fmt.Errorf("%w: provider identity not found", ErrNotFound)
        }

        return nil
}

// GetUserProviderIdentities returns all provider identities linked to a user.
func (s *RedisStorage) GetUserProviderIdentities(ctx context.Context, userID string) ([]*ProviderIdentity, error) {
        // Verify user exists
        userKey := redisKey(s.keyPrefix, KeyTypeUser, userID)
        exists, err := s.client.Exists(ctx, userKey).Result()
        if err != nil {
                return nil, fmt.Errorf("failed to check user existence: %w", err)
        }
        if exists == 0 {
                return nil, fmt.Errorf("%w: user not found", ErrNotFound)
        }

        // Get user's provider identity keys
        userProviderSetKey := redisSetKey(s.keyPrefix, KeyTypeUserProviders, userID)
        keys, err := s.client.SMembers(ctx, userProviderSetKey).Result()
        if err != nil && !errors.Is(err, redis.Nil) {
                return nil, fmt.Errorf("failed to get provider identity keys: %w", err)
        }

        var identities []*ProviderIdentity
        for _, key := range keys {
                data, err := s.client.Get(ctx, key).Bytes()
                if err != nil {
                        if errors.Is(err, redis.Nil) {
                                // Stale set entry — the identity key has expired or was deleted.
                                // Skip silently; stale entries are cleaned up during write
                                // operations (DeleteUser, CreateProviderIdentity) to avoid
                                // mutation side effects in this read-only method.
                                continue
                        }
                        return nil, fmt.Errorf("failed to get identity: %w", err)
                }

                var stored storedProviderIdentity
                if err := json.Unmarshal(data, &stored); err != nil {
                        return nil, fmt.Errorf("failed to unmarshal identity: %w", err)
                }

                identities = append(identities, &ProviderIdentity{
                        UserID:          stored.UserID,
                        ProviderID:      stored.ProviderID,
                        ProviderSubject: stored.ProviderSubject,
                        LinkedAt:        time.Unix(stored.LinkedAt, 0),
                        LastUsedAt:      time.Unix(stored.LastUsedAt, 0),
                })
        }

        return identities, nil
}

// -----------------------
// Serialization Helpers
// -----------------------

// marshalRequester serializes a fosite.Requester to JSON.
// The full session is serialized as a JSON blob to preserve all session data
// including JWT claims, headers, and upstream session IDs. This is critical
// for token refresh — fosite's JWT strategy requires the session to implement
// oauth2.JWTSessionContainer, which redisSession did not.
func marshalRequester(request fosite.Requester) ([]byte, error) {
        sessionData, err := json.Marshal(request.GetSession())
        if err != nil {
                return nil, fmt.Errorf("failed to marshal session: %w", err)
        }

        // A clientless grant (fosite's CanSkipClientAuth) is expected to attach a
        // synthetic client (see NewSyntheticClient) rather than leave this nil;
        // this fallback only guards against a future clientless grant that
        // forgets to.
        var clientID string
        if client := request.GetClient(); client != nil {
                clientID = client.GetID()
        }

        stored := storedSession{
                ClientID:          clientID,
                RequestedAt:       request.GetRequestedAt(),
                RequestedScopes:   request.GetRequestedScopes(),
                GrantedScopes:     request.GetGrantedScopes(),
                RequestedAudience: request.GetRequestedAudience(),
                GrantedAudience:   request.GetGrantedAudience(),
                Form:              request.GetRequestForm(),
                RequestID:         request.GetID(),
                Session:           sessionData,
        }

        return json.Marshal(stored)
}

// unmarshalRequester deserializes a fosite.Requester from JSON.
// It requires storage access to look up the client and a session factory
// to create the correct session type for deserialization.
//
// The client lookup goes to this backend directly, below ClientAdmissionStorage,
// so it is identity-only: it names the client a stored session belongs to and
// does not prove that client is still admitted. That is safe because fosite
// always compares the session's client ID with the client the request
// authenticated through the admitted store, and a removed client cannot
// authenticate. Do not use a rehydrated client as an authentication result.
func unmarshalRequester(ctx context.Context, data []byte, s *RedisStorage) (fosite.Requester, error) {
        var stored storedSession
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to unmarshal session: %w", err)
        }

        // A synthetic client (see NewSyntheticClient) was never registered via
        // RegisterClient, so its ID alone reconstructs it rather than looking it
        // up through the client registry.
        var client fosite.Client
        if IsSyntheticClientID(stored.ClientID) {
                client = NewSyntheticClient(stored.ClientID)
        } else {
                // This lookup consults only this storage's own client rows: a client
                // that was resolved dynamically at authorize time but never persisted
                // would fail exactly here, surfacing to the client as a bare
                // invalid_grant at the token endpoint — so log the client id and the
                // underlying error once, which otherwise leave no server-side trace
                // (see issue #6187). The returned error deliberately does not repeat
                // the client id: all four call sites propagate it up to fosite, which
                // would otherwise report the same failure to the operator twice.
                fetchedClient, err := s.GetClient(ctx, stored.ClientID)
                if err != nil {
                        slog.WarnContext(ctx, "failed to resolve client for stored session",
                                "client_id", stored.ClientID, "error", err)
                        return nil, fmt.Errorf("failed to get client for session: %w", err)
                }
                client = fetchedClient
        }

        // Create a session prototype via factory, then deserialize the full session
        // blob into it. This preserves JWT claims, headers, and upstream session IDs.
        sess := defaultSessionFactory("", "", "")
        if len(stored.Session) > 0 {
                if err := json.Unmarshal(stored.Session, sess); err != nil {
                        return nil, fmt.Errorf("failed to unmarshal session data: %w", err)
                }
        }

        return &fosite.Request{
                ID:                stored.RequestID,
                RequestedAt:       stored.RequestedAt,
                Client:            client,
                RequestedScope:    stored.RequestedScopes,
                GrantedScope:      stored.GrantedScopes,
                RequestedAudience: stored.RequestedAudience,
                GrantedAudience:   stored.GrantedAudience,
                Form:              url.Values(stored.Form),
                Session:           sess,
        }, nil
}

// getTTLFromRequester extracts the TTL from a fosite.Requester session.
func getTTLFromRequester(request fosite.Requester, tokenType fosite.TokenType, defaultTTL time.Duration) time.Duration {
        if request == nil {
                return defaultTTL
        }

        sess := request.GetSession()
        if sess == nil {
                return defaultTTL
        }

        expTime := sess.GetExpiresAt(tokenType)
        if expTime.IsZero() {
                return defaultTTL
        }

        ttl := time.Until(expTime)
        if ttl <= 0 {
                return defaultTTL
        }

        return ttl
}

// Compile-time interface compliance checks
var (
        _ Storage                          = (*RedisStorage)(nil)
        _ PendingAuthorizationStorage      = (*RedisStorage)(nil)
        _ DeviceCodeStorage                = (*RedisStorage)(nil)
        _ PendingDeviceLoginStorage        = (*RedisStorage)(nil)
        _ PendingDeviceConfirmationStorage = (*RedisStorage)(nil)
        _ ClientRegistry                   = (*RedisStorage)(nil)
        _ UpstreamTokenStorage             = (*RedisStorage)(nil)
        _ UserStorage                      = (*RedisStorage)(nil)
        _ DCRCredentialStore               = (*RedisStorage)(nil)
        _ AssertionJWTConsumer             = (*RedisStorage)(nil)
)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "crypto/sha256"
        "fmt"
)

// Key type constants for Redis storage.
// These define the different types of data stored in Redis.
const (
        // KeyTypeAccess is the key type for access tokens.
        KeyTypeAccess = "access"

        // KeyTypeRefresh is the key type for refresh tokens.
        KeyTypeRefresh = "refresh"

        // KeyTypeAuthCode is the key type for authorization codes.
        KeyTypeAuthCode = "authcode"

        // KeyTypePKCE is the key type for PKCE requests.
        KeyTypePKCE = "pkce"

        // KeyTypeClient is the key type for OAuth clients.
        KeyTypeClient = "client"

        // KeyTypeUser is the key type for users.
        KeyTypeUser = "user"

        // KeyTypeProvider is the key type for provider identities.
        KeyTypeProvider = "provider"

        // KeyTypeUpstream is the key type for upstream tokens.
        KeyTypeUpstream = "upstream"

        // KeyTypePending is the key type for pending authorizations.
        KeyTypePending        = "pending"
        KeyTypeConsentSession = "consent:session"
        KeyTypeClientApproval = "consent:approval"

        // KeyTypeDeviceCode is the key type for RFC 8628 device authorization
        // requests, keyed by device_code.
        KeyTypeDeviceCode = "device"

        // KeyTypeDeviceUserCode is the key type for the user_code -> device_code
        // secondary index used by the verification page.
        KeyTypeDeviceUserCode = "device:usercode"

        // KeyTypePendingDeviceLogin is the key type for pending device-flow
        // verification-page logins awaiting the upstream IDP callback.
        KeyTypePendingDeviceLogin = "pending:device"

        // KeyTypePendingDeviceConfirmation is the key type for resolved
        // device-flow logins awaiting an explicit Approve/Deny decision at the
        // verification page.
        KeyTypePendingDeviceConfirmation = "pending:device:confirm"

        // KeyTypeInvalidated is the key type for invalidated authorization codes.
        KeyTypeInvalidated = "invalidated"

        // KeyTypeJWT is the key type for client assertion JWTs.
        KeyTypeJWT = "jwt"

        // KeyTypeAssertionJWT is the key type for consumed assertion JWTs.
        KeyTypeAssertionJWT = "assertion-jwt"

        // KeyTypeReqIDAccess is the key type for request ID to access token mappings.
        KeyTypeReqIDAccess = "reqid:access"

        // KeyTypeReqIDRefresh is the key type for request ID to refresh token mappings.
        KeyTypeReqIDRefresh = "reqid:refresh"

        // KeyTypeUpstreamIdx is the key type for the session index set — a Redis SET that
        // tracks all per-provider token keys (upstream:{sid}:{provider}) belonging to a session.
        // This enables O(1) enumeration via SMEMBERS without scanning the keyspace.
        // Used by GetAllUpstreamTokens (bulk read) and DeleteUpstreamTokens (bulk delete).
        KeyTypeUpstreamIdx = "upstream:idx"

        // KeyTypeUserUpstream is the key type for user to upstream token reverse lookups.
        KeyTypeUserUpstream = "user:upstream"

        // KeyTypeUserProviders is the key type for user to provider identity reverse lookups.
        KeyTypeUserProviders = "user:providers"

        // KeyTypeDCRClientIdx is the key type for the sorted set of client IDs
        // counted against the DCR client cap, each scored by its expiry in Unix
        // milliseconds. It is a single key per prefix, so it shares the client
        // rows' hash slot and one Lua script can update both on Redis Cluster.
        // The name deliberately does not match the "client:*" scan pattern.
        KeyTypeDCRClientIdx = "client-idx"

        // KeyTypeDCR is the key type for RFC 7591 Dynamic Client Registration credentials
        // persisted by an authserver upstream-DCR resolver. Distinct from KeyTypeClient,
        // which holds the authserver's *own* OAuth clients — DCR entries are credentials
        // that *this* authserver registered against an *upstream* authorization server.
        KeyTypeDCR = "dcr"
)

// DeriveKeyPrefix creates the key prefix from the Kubernetes namespace and MCP server name.
// The format is "thv:auth:{ns:name}:" where {ns:name} is a Redis hash tag.
//
// Note: The hash tag format {ns:name} intentionally combines namespace and name
// into a single tag. In Redis Cluster, only the first hash tag determines slot
// assignment. Using {ns}:{name} would only hash on namespace, potentially
// spreading a single server's keys across multiple slots. The combined format
// ensures all keys for a specific server (namespace+name pair) are placed in
// the same slot, enabling atomic multi-key operations like token revocation.
func DeriveKeyPrefix(namespace, name string) string {
        return fmt.Sprintf("thv:auth:{%s:%s}:", namespace, name)
}

// redisKey generates a Redis key with the given prefix, type, and ID.
// The resulting format is "{prefix}{keyType}:{id}". This assumes the id does not
// contain colons; callers that need colon-safe keys should use redisProviderKey
// which uses a length-prefixed format. In practice, IDs passed here are UUIDs,
// opaque token signatures, or system-generated identifiers that do not contain colons.
func redisKey(prefix, keyType, id string) string {
        return fmt.Sprintf("%s%s:%s", prefix, keyType, id)
}

// redisAssertionJWTKey returns the namespaced replay key for an assertion JWT.
// The tuple is length-prefixed before hashing, so arbitrary purpose, issuer, and
// JTI values cannot create ambiguous encodings or inject Redis key segments. The
// digest also avoids exposing assertion identifiers in Redis key names.
func redisAssertionJWTKey(prefix, purpose, issuer, jti string) string {
        tuple := fmt.Sprintf("%d:%s%d:%s%d:%s", len(purpose), purpose, len(issuer), issuer, len(jti), jti)
        digest := sha256.Sum256([]byte(tuple))
        return fmt.Sprintf("%s%s:%x", prefix, KeyTypeAssertionJWT, digest)
}

// redisProviderKey generates a Redis key for provider identities.
// Uses length-prefixed format to handle colons in provider IDs/subjects.
func redisProviderKey(prefix, providerID, providerSubject string) string {
        return fmt.Sprintf("%s%s:%d:%s:%s", prefix, KeyTypeProvider, len(providerID), providerID, providerSubject)
}

// redisDCRKey generates a Redis key for a DCR credential entry, identifying the
// (Issuer, UpstreamID, RedirectURI, ScopesHash) tuple that DCRKey canonicalises.
//
// Format: "{prefix}dcr:<len(issuer)>:<issuer>:<len(upstream_id)>:<upstream_id>:<len(redirect_uri)>:<redirect_uri>:<scopes_hash>"
//
// When DCRKey.RotationKey is non-empty a length-prefixed
// ":<len(rotation_key)>:<rotation_key>" segment follows the scopes hash. An
// empty RotationKey yields exactly the format above, so rows written before
// the field existed keep resolving; the hex scopes hash contains no colon, so
// the two shapes cannot be confused.
//
// The first three segments are length-prefixed to handle colons in UpstreamID
// and RedirectURI (and, for symmetry, Issuer) without ambiguity, mirroring
// redisProviderKey. ScopesHash is expected to be a SHA-256 hex digest produced
// by storage.ScopesHash — only [0-9a-f] and never colon-bearing — so it is
// appended without a length prefix. The format is robust for that domain;
// validateDCRCredentialsForStore (called by every Store path) already
// rejects an empty ScopesHash, and callers are required to compute the hash
// via storage.ScopesHash. Length-prefix collision-safety is preserved on
// the leading segments either way.
//
// The UpstreamID segment (issue #5823) means two upstreams that share the
// consumer's Issuer and RedirectURI but differ only by their authorization
// server no longer map to the same Redis key. Adding it changes the key
// format, so entries written by an older binary (without the segment) will
// miss on lookup and be harmlessly re-registered under the new key. Orphaned
// old rows self-evict only when the upstream asserted an RFC 7591 §3.2.1
// client_secret_expires_at (StoreDCRCredentialsIfAbsent sets a matching Redis TTL for
// those); entries for non-expiring secrets — the common §3.2.1 "never" case —
// carry no TTL and persist indefinitely, so they require manual cleanup.
// There is no automated one-shot migration for this key-format change today.
//
// The public-vs-confidential client distinction is intentionally NOT
// encoded here — see DCRKey's doc for the rationale. Today's two consumers
// register on disjoint RedirectURI address spaces (AS-origin vs RFC 8252
// loopback) so the persisted key cannot collide across profiles. A future
// consumer that defaults its RedirectURI into either space would need to
// add the distinguishing component back to the key format alongside an
// explicit migration story for existing Redis-cached entries.
func redisDCRKey(prefix string, key DCRKey) string {
        base := fmt.Sprintf("%s%s:%d:%s:%d:%s:%d:%s:%s",
                prefix, KeyTypeDCR,
                len(key.Issuer), key.Issuer,
                len(key.UpstreamID), key.UpstreamID,
                len(key.RedirectURI), key.RedirectURI,
                key.ScopesHash)
        if key.RotationKey == "" {
                return base
        }
        return fmt.Sprintf("%s:%d:%s", base, len(key.RotationKey), key.RotationKey)
}

// redisUpstreamKey generates a Redis key for a per-provider upstream token entry.
// Format: "{prefix}upstream:{sessionID}:{providerName}"
// This enables storing tokens from multiple upstream providers per session.
func redisUpstreamKey(prefix, sessionID, providerName string) string {
        return fmt.Sprintf("%s%s:%s:%s", prefix, KeyTypeUpstream, sessionID, providerName)
}

// redisSetKey generates a Redis key for a set that tracks multiple items.
// Used for secondary indexes like request ID -> token signature mappings.
// Same colon assumption as redisKey: the id must not contain colons.
func redisSetKey(prefix, keyType, id string) string {
        return fmt.Sprintf("%s%s:%s", prefix, keyType, id)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "strings"
        "time"

        "github.com/redis/go-redis/v9"
)

// MigrationResult holds counts of items migrated during bulk migration.
type MigrationResult struct {
        TokensMigrated     int
        TokensSkipped      int
        TokensFailed       int
        IdentitiesMigrated int
        IdentitiesSkipped  int
        IdentitiesFailed   int
}

// isLegacyUpstreamProviderID reports whether id is a legacy protocol-type provider
// ID that may be migrated to a logical provider name. Only tokens stored under a
// legacy protocol-type ID should be claimed by the migration path; tokens that
// already carry a logical name must not be re-labelled.
func isLegacyUpstreamProviderID(id string) bool {
        return id == "" || id == "oidc" || id == "oauth2"
}

// MigrateLegacyUpstreamData performs a one-shot bulk migration of legacy upstream
// token keys and provider identity keys to the new multi-upstream format.
//
// Token key migration: renames "upstream:{sessionID}" keys (legacy, no provider
// suffix) to "upstream:{sessionID}:{providerName}" and patches ProviderID.
//
// Provider identity migration: duplicates identities stored under legacy
// protocol-type IDs (e.g. "oidc", "oauth2") to the new logical provider name.
// Legacy identity keys are NOT deleted to allow safe rollback.
//
// The migration is idempotent and crash-safe: each key is migrated independently,
// and existing new-format keys are never overwritten.
//
// Returns an error if any keys fail to migrate, since the request path no longer
// has inline fallbacks for legacy data.
//
// TODO(migration): Remove once all deployments have upgraded past this version.
func (s *RedisStorage) MigrateLegacyUpstreamData(ctx context.Context, providerName, legacyProviderID string) error {
        result := &MigrationResult{}

        if err := s.migrateUpstreamTokenKeys(ctx, providerName, result); err != nil {
                return fmt.Errorf("token key migration: %w", err)
        }

        if err := s.migrateProviderIdentityKeys(ctx, providerName, legacyProviderID, result); err != nil {
                return fmt.Errorf("provider identity migration: %w", err)
        }

        if result.TokensFailed > 0 || result.IdentitiesFailed > 0 {
                return fmt.Errorf("migration incomplete: %d token(s) and %d identity(ies) failed — "+
                        "the request path has no inline fallback for unmigrated legacy data",
                        result.TokensFailed, result.IdentitiesFailed)
        }

        if result.TokensMigrated > 0 || result.IdentitiesMigrated > 0 {
                slog.Info("legacy data migration complete",
                        "tokens_migrated", result.TokensMigrated,
                        "tokens_skipped", result.TokensSkipped,
                        "identities_migrated", result.IdentitiesMigrated,
                        "identities_skipped", result.IdentitiesSkipped,
                )
        }

        return nil
}

// migrateUpstreamTokenKeys scans for legacy upstream token keys and migrates them
// to the new per-provider key format.
func (s *RedisStorage) migrateUpstreamTokenKeys(ctx context.Context, providerName string, result *MigrationResult) error {
        // Scan for all upstream:* keys under this prefix
        pattern := s.keyPrefix + KeyTypeUpstream + ":*"
        // The upstream key prefix length helps distinguish legacy from new-format keys.
        // Legacy: "{prefix}upstream:{sessionID}" — remainder after prefix+upstream: has NO colon
        // New:    "{prefix}upstream:{sessionID}:{providerName}" — remainder has a colon
        // Index:  "{prefix}upstream:idx:{sessionID}" — starts with "idx:"
        upstreamPrefixLen := len(s.keyPrefix) + len(KeyTypeUpstream) + 1 // +1 for the colon after "upstream"

        var cursor uint64
        for {
                keys, nextCursor, err := s.client.Scan(ctx, cursor, pattern, 100).Result()
                if err != nil {
                        return fmt.Errorf("SCAN failed: %w", err)
                }

                for _, key := range keys {
                        if len(key) <= upstreamPrefixLen {
                                continue
                        }
                        remainder := key[upstreamPrefixLen:]

                        // Skip index keys (upstream:idx:...)
                        if strings.HasPrefix(remainder, "idx:") {
                                continue
                        }

                        // Distinguish legacy keys (no colon in remainder) from new-format keys (have colon)
                        if strings.Contains(remainder, ":") {
                                // Already new-format key, skip
                                continue
                        }

                        // This is a legacy key: upstream:{sessionID}
                        sessionID := remainder
                        if err := s.migrateSingleUpstreamToken(ctx, key, sessionID, providerName, result); err != nil {
                                slog.Warn("failed to migrate legacy upstream token",
                                        "key", key, "session_id", sessionID, "error", err)
                                result.TokensFailed++
                        }
                }

                cursor = nextCursor
                if cursor == 0 {
                        break
                }
        }

        return nil
}

// migrateSingleUpstreamToken migrates one legacy upstream token key to the new format.
func (s *RedisStorage) migrateSingleUpstreamToken(
        ctx context.Context,
        legacyKey, sessionID, providerName string,
        result *MigrationResult,
) error {
        // Check if new-format key already exists (idempotent)
        newKey := redisUpstreamKey(s.keyPrefix, sessionID, providerName)
        exists, err := s.client.Exists(ctx, newKey).Result()
        if err != nil {
                return fmt.Errorf("EXISTS check failed for %s: %w", newKey, err)
        }
        if exists > 0 {
                // New key exists — clean up the legacy key so it doesn't re-appear on next startup.
                warnOnCleanupErr(s.client.Del(ctx, legacyKey).Err(), "Del", legacyKey)
                result.TokensSkipped++
                return nil
        }

        // Read the legacy data
        data, err := s.client.Get(ctx, legacyKey).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        result.TokensSkipped++
                        return nil
                }
                return fmt.Errorf("GET failed for %s: %w", legacyKey, err)
        }

        // Handle null marker
        if string(data) == nullMarker {
                result.TokensSkipped++
                return nil
        }

        // Deserialize to check ProviderID
        var stored storedUpstreamTokens
        if err := json.Unmarshal(data, &stored); err != nil {
                return fmt.Errorf("unmarshal failed for %s: %w", legacyKey, err)
        }

        // Only migrate tokens with legacy protocol-type IDs
        if !isLegacyUpstreamProviderID(stored.ProviderID) {
                slog.Debug("skipping legacy upstream token: has logical provider name",
                        "session_id", sessionID, "provider_id", stored.ProviderID)
                result.TokensSkipped++
                return nil
        }

        // Patch ProviderID to the logical provider name
        stored.ProviderID = providerName
        newData, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization
        if err != nil {
                return fmt.Errorf("marshal failed: %w", err)
        }

        // Preserve TTL from the legacy key
        ttl, err := s.client.PTTL(ctx, legacyKey).Result()
        if err != nil {
                return fmt.Errorf("PTTL failed for %s: %w", legacyKey, err)
        }
        // PTTL returns -1 for no expiry, -2 for key not found
        if ttl == -2*time.Millisecond {
                result.TokensSkipped++
                return nil
        }

        // Build the session index key
        idxKey := redisSetKey(s.keyPrefix, KeyTypeUpstreamIdx, sessionID)

        // Pipeline: SET new key + SADD to session index + SADD to user reverse index + DEL legacy key.
        // The user:upstream set must be updated so DeleteUser cascade deletion includes migrated tokens.
        pipe := s.client.TxPipeline()
        if ttl > 0 {
                pipe.Set(ctx, newKey, newData, ttl)
        } else {
                pipe.Set(ctx, newKey, newData, 0)
        }
        pipe.SAdd(ctx, idxKey, newKey)
        if ttl > 0 {
                pipe.PExpire(ctx, idxKey, ttl)
        }
        if stored.UserID != "" {
                userUpstreamKey := redisSetKey(s.keyPrefix, KeyTypeUserUpstream, stored.UserID)
                pipe.SAdd(ctx, userUpstreamKey, newKey)
        }
        pipe.Del(ctx, legacyKey)
        if _, err := pipe.Exec(ctx); err != nil {
                return fmt.Errorf("pipeline exec failed: %w", err)
        }

        slog.Debug("migrated legacy upstream token",
                "session_id", sessionID, "provider_name", providerName)
        result.TokensMigrated++
        return nil
}

// migrateProviderIdentityKeys scans for provider identity keys stored under the
// legacy protocol-type ID and duplicates them under the new logical provider name.
func (s *RedisStorage) migrateProviderIdentityKeys(
        ctx context.Context,
        providerName, legacyProviderID string,
        result *MigrationResult,
) error {
        // Skip if legacyProviderID is the same as providerName (nothing to migrate)
        if legacyProviderID == providerName {
                return nil
        }

        // Skip if legacyProviderID is empty (no legacy identity to look up)
        if legacyProviderID == "" {
                return nil
        }

        // Scan for provider identity keys under the legacy ID.
        // Provider key format: "{prefix}provider:{len(providerID)}:{providerID}:{providerSubject}"
        pattern := fmt.Sprintf("%s%s:%d:%s:*", s.keyPrefix, KeyTypeProvider, len(legacyProviderID), legacyProviderID)

        var cursor uint64
        for {
                keys, nextCursor, err := s.client.Scan(ctx, cursor, pattern, 100).Result()
                if err != nil {
                        return fmt.Errorf("SCAN failed: %w", err)
                }

                for _, key := range keys {
                        if err := s.migrateSingleProviderIdentity(ctx, key, providerName, legacyProviderID, result); err != nil {
                                slog.Warn("failed to migrate legacy provider identity",
                                        "key", key, "error", err)
                                result.IdentitiesFailed++
                        }
                }

                cursor = nextCursor
                if cursor == 0 {
                        break
                }
        }

        return nil
}

// migrateSingleProviderIdentity duplicates a legacy provider identity under the
// new logical provider name. The legacy key is NOT deleted for safe rollback.
func (s *RedisStorage) migrateSingleProviderIdentity(
        ctx context.Context,
        legacyKey, providerName, legacyProviderID string,
        result *MigrationResult,
) error {
        // Read the legacy identity
        data, err := s.client.Get(ctx, legacyKey).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        result.IdentitiesSkipped++
                        return nil
                }
                return fmt.Errorf("GET failed for %s: %w", legacyKey, err)
        }

        var stored storedProviderIdentity
        if err := json.Unmarshal(data, &stored); err != nil {
                return fmt.Errorf("unmarshal failed for %s: %w", legacyKey, err)
        }

        // Build the new key under the logical provider name
        newKey := redisProviderKey(s.keyPrefix, providerName, stored.ProviderSubject)

        // Duplicate the identity with the new provider ID, using SetNX for idempotency
        newStored := storedProviderIdentity{
                UserID:          stored.UserID,
                ProviderID:      providerName,
                ProviderSubject: stored.ProviderSubject,
                LinkedAt:        stored.LinkedAt,
                LastUsedAt:      stored.LastUsedAt,
        }

        newData, err := json.Marshal(newStored) //nolint:gosec // G117 - internal Redis storage serialization
        if err != nil {
                return fmt.Errorf("marshal failed: %w", err)
        }

        // SetNX: only write if new key does not exist (idempotent)
        created, err := s.client.SetNX(ctx, newKey, newData, 0).Result()
        if err != nil {
                return fmt.Errorf("SetNX failed for %s: %w", newKey, err)
        }

        if !created {
                slog.Debug("skipping legacy provider identity: new key already exists",
                        "legacy_provider_id", legacyProviderID,
                        "provider_name", providerName,
                        "provider_subject", stored.ProviderSubject)
                result.IdentitiesSkipped++
                return nil
        }

        // Update the user's provider set to include the new key
        userProviderSetKey := redisSetKey(s.keyPrefix, KeyTypeUserProviders, stored.UserID)
        warnOnCleanupErr(s.client.SAdd(ctx, userProviderSetKey, newKey).Err(), "SAdd", userProviderSetKey)

        slog.Debug("migrated legacy provider identity",
                "legacy_provider_id", legacyProviderID,
                "provider_name", providerName,
                "provider_subject", stored.ProviderSubject,
                "user_id", stored.UserID)
        result.IdentitiesMigrated++
        return nil
}

// BackfillDCRClientIndex adds every live client row counted against the DCR
// cap to the KeyTypeDCRClientIdx sorted set, scored by the row's current
// expiry. Rows written before the index existed have no entry, so without
// this a deployment already holding more than the cap would keep admitting
// registrations past it.
//
// It is idempotent and runs on every startup, in the style of
// MigrateLegacyUpstreamData: ZADD NX never moves an existing entry, so a row
// indexed by RegisterClient or RenewClientTTL keeps its exact score. Running
// it again also picks up rows an older replica wrote during a rolling
// upgrade. A row without a TTL is not DCR-issued and is skipped.
//
// Reads are batched per SCAN page. A row that cannot be parsed is skipped and
// counted in one WARN line rather than failing startup: it would equally fail
// GetClient, and one bad row must not keep every replica down. Only Redis
// errors (SCAN, reads, writes) are returned.
//
// On success it logs one INFO line with the rows scanned, the entries added
// and the duration, and one WARN line when the live count is already at or
// above the cap, so an operator sees why /oauth/register answers 503 from the
// first request.
//
// TODO(migration): Remove once all deployments have upgraded past this version.
func (s *RedisStorage) BackfillDCRClientIndex(ctx context.Context) error {
        node, err := s.nodeForPrefix(ctx)
        if err != nil {
                return err
        }
        pattern := s.keyPrefix + KeyTypeClient + ":*"
        start := time.Now()

        var scanned, added, skipped int64
        var cursor uint64
        for {
                keys, nextCursor, err := node.Scan(ctx, cursor, pattern, 100).Result()
                if err != nil {
                        return fmt.Errorf("SCAN failed: %w", err)
                }
                pageAdded, pageSkipped, err := s.backfillDCRClientPage(ctx, node, keys)
                if err != nil {
                        return err
                }
                scanned += int64(len(keys))
                added += pageAdded
                skipped += pageSkipped

                cursor = nextCursor
                if cursor == 0 {
                        break
                }
        }

        if skipped > 0 {
                slog.Warn("DCR client index backfill skipped client rows it could not parse",
                        "rows_skipped", skipped)
        }
        slog.Info("DCR client index backfill complete",
                "rows_scanned", scanned, "clients_added", added, "duration", time.Since(start))
        s.warnIfDCRClientsAtCap(ctx)
        return nil
}

// nodeForPrefix returns the client to run keyless commands (SCAN) on for
// keys under s.keyPrefix. On Redis Cluster a keyless command goes to an
// arbitrary node, and a SCAN cursor is only valid on the node that issued it,
// so this resolves the master owning the prefix's hash slot. Every key under
// the prefix shares that slot through its {ns:name} hash tag.
func (s *RedisStorage) nodeForPrefix(ctx context.Context) (redis.Cmdable, error) {
        cc, ok := s.client.(*redis.ClusterClient)
        if !ok {
                return s.client, nil
        }
        node, err := cc.MasterForKey(ctx, s.dcrClientIndexKey())
        if err != nil {
                return nil, fmt.Errorf("failed to resolve Redis Cluster node for key prefix: %w", err)
        }
        return node, nil
}

// warnIfDCRClientsAtCap logs one WARN when the live DCR client count is
// already at or above the cap, which makes every registration fail until
// enough clients expire. It is silent when the cap is unlimited. A failed
// count is logged rather than returned: the backfill itself succeeded.
func (s *RedisStorage) warnIfDCRClientsAtCap(ctx context.Context) {
        if s.maxDCRClients == 0 {
                return
        }
        n, err := s.CountDCRClients(ctx)
        if err != nil {
                slog.Warn("could not check the DCR client count against the cap", "error", err)
                return
        }
        if n >= int64(s.maxDCRClients) {
                slog.Warn("DCR client cap already reached; /oauth/register rejects new clients until enough expire",
                        "dcr_clients", n, "max_dcr_clients", s.maxDCRClients,
                        "setting", "storage.max_dcr_clients")
        }
}

// backfillDCRClientPage indexes the counted client rows among keys, one SCAN
// page, with one pipelined read and one ZADD NX. It returns how many entries
// it added and how many rows it skipped as unparsable.
func (s *RedisStorage) backfillDCRClientPage(ctx context.Context, node redis.Cmdable, keys []string) (int64, int64, error) {
        if len(keys) == 0 {
                return 0, 0, nil
        }

        pipe := node.Pipeline()
        getCmds := make([]*redis.StringCmd, len(keys))
        pttlCmds := make([]*redis.DurationCmd, len(keys))
        for i, key := range keys {
                getCmds[i] = pipe.Get(ctx, key)
                pttlCmds[i] = pipe.PTTL(ctx, key)
        }
        // redis.Nil only means a row expired since SCAN; per-row results below
        // tell the cases apart.
        if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
                return 0, 0, fmt.Errorf("failed to read client rows: %w", err)
        }
        now := s.now()

        var members []redis.Z
        var skipped int64
        for i, key := range keys {
                member, ok, err := s.backfillDCRClientEntry(key, getCmds[i], pttlCmds[i], now)
                if err != nil {
                        return 0, 0, err
                }
                if !ok {
                        skipped++
                } else if member != nil {
                        members = append(members, *member)
                }
        }
        if len(members) == 0 {
                return 0, skipped, nil
        }

        added, err := node.ZAddNX(ctx, s.dcrClientIndexKey(), members...).Result()
        if err != nil {
                return 0, 0, fmt.Errorf("failed to index client rows: %w", err)
        }
        return added, skipped, nil
}

// backfillDCRClientEntry returns the index entry for one client row read by
// backfillDCRClientPage, or nil when the row is gone, is not DCR-issued, or
// does not count against the cap. ok is false when the row cannot be parsed
// and is skipped. now is the time of the read.
func (s *RedisStorage) backfillDCRClientEntry(
        key string, getCmd *redis.StringCmd, pttlCmd *redis.DurationCmd, now time.Time,
) (member *redis.Z, ok bool, err error) {
        data, err := getCmd.Bytes()
        if errors.Is(err, redis.Nil) {
                return nil, true, nil // expired since SCAN
        }
        if err != nil {
                return nil, false, fmt.Errorf("failed to read client row: %w", err)
        }
        // PTTL reports a missing key or a key without an expiry as a negative
        // value. Neither is a DCR-issued row.
        ttl, err := pttlCmd.Result()
        if err != nil {
                return nil, false, fmt.Errorf("failed to read client row TTL: %w", err)
        }
        if ttl <= 0 {
                return nil, true, nil
        }

        id := key[len(s.keyPrefix)+len(KeyTypeClient)+1:]
        var stored storedClient
        if err := json.Unmarshal(data, &stored); err != nil || (stored.ID != "" && stored.ID != id) {
                slog.Debug("DCR client index backfill skipping unparsable client row", "key", key)
                return nil, false, nil
        }
        if !countsTowardDCRCap(clientFromStored(stored, true)) {
                return nil, true, nil
        }
        return &redis.Z{Score: float64(now.Add(ttl).UnixMilli()), Member: id}, true, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "crypto/sha256"
        "encoding/json"
        "errors"
        "fmt"
        "slices"
        "time"

        "github.com/ory/fosite"
        "github.com/redis/go-redis/v9"
)

var _ RememberedConsentStorage = (*RedisStorage)(nil)

func approvalKey(prefix, userID, clientID string) (string, error) {
        pair, err := json.Marshal([2]string{userID, clientID})
        if err != nil {
                return "", err
        }
        return redisKey(prefix, KeyTypeClientApproval, fmt.Sprintf("%x", sha256.Sum256(pair))), nil
}

// StoreConsentSession stores a digest with its remaining fixed lifetime.
func (s *RedisStorage) StoreConsentSession(ctx context.Context, digest string, session *ConsentSession) error {
        if digest == "" {
                return fosite.ErrInvalidRequest.WithHint("consent session digest cannot be empty")
        }
        if session == nil {
                return fosite.ErrInvalidRequest.WithHint("consent session cannot be nil")
        }
        ttl := time.Until(session.ExpiresAt)
        if ttl <= 0 {
                return ErrExpired
        }
        data, err := json.Marshal(session)
        if err != nil {
                return fmt.Errorf("marshal consent session: %w", err)
        }
        if err := s.client.Set(ctx, redisKey(s.keyPrefix, KeyTypeConsentSession, digest), data, ttl).Err(); err != nil {
                return fmt.Errorf("store consent session: %w", err)
        }
        return nil
}

// GetConsentSession returns a session only until its stored absolute expiry.
func (s *RedisStorage) GetConsentSession(ctx context.Context, digest string) (*ConsentSession, error) {
        data, err := s.client.Get(ctx, redisKey(s.keyPrefix, KeyTypeConsentSession, digest)).Bytes()
        if errors.Is(err, redis.Nil) {
                return nil, ErrNotFound
        }
        if err != nil {
                return nil, fmt.Errorf("load consent session: %w", err)
        }
        var session ConsentSession
        if err := json.Unmarshal(data, &session); err != nil {
                return nil, fmt.Errorf("decode consent session: %w", err)
        }
        if !time.Now().Before(session.ExpiresAt) {
                return nil, ErrExpired
        }
        return &session, nil
}

// DeleteConsentSession revokes the originating browser session.
func (s *RedisStorage) DeleteConsentSession(ctx context.Context, digest string) error {
        if err := s.client.Del(ctx, redisKey(s.keyPrefix, KeyTypeConsentSession, digest)).Err(); err != nil {
                return fmt.Errorf("delete consent session: %w", err)
        }
        return nil
}

// StoreClientApproval overwrites one user's approval for one client.
func (s *RedisStorage) StoreClientApproval(ctx context.Context, userID, clientID string, approval *ClientApproval) error {
        if userID == "" {
                return fosite.ErrInvalidRequest.WithHint("client approval user ID cannot be empty")
        }
        if clientID == "" {
                return fosite.ErrInvalidRequest.WithHint("client approval client ID cannot be empty")
        }
        if approval == nil {
                return fosite.ErrInvalidRequest.WithHint("client approval cannot be nil")
        }
        ttl := time.Until(approval.ExpiresAt)
        if ttl <= 0 {
                return ErrExpired
        }
        key, err := approvalKey(s.keyPrefix, userID, clientID)
        if err != nil {
                return err
        }
        clone := *approval
        clone.Scopes = slices.Clone(approval.Scopes)
        data, err := json.Marshal(&clone)
        if err != nil {
                return fmt.Errorf("marshal client approval: %w", err)
        }
        if err := s.client.Set(ctx, key, data, ttl).Err(); err != nil {
                return fmt.Errorf("store client approval: %w", err)
        }
        return nil
}

// GetClientApproval returns a copy only until its stored absolute expiry.
func (s *RedisStorage) GetClientApproval(ctx context.Context, userID, clientID string) (*ClientApproval, error) {
        key, err := approvalKey(s.keyPrefix, userID, clientID)
        if err != nil {
                return nil, err
        }
        data, err := s.client.Get(ctx, key).Bytes()
        if errors.Is(err, redis.Nil) {
                return nil, ErrNotFound
        }
        if err != nil {
                return nil, fmt.Errorf("load client approval: %w", err)
        }
        var approval ClientApproval
        if err := json.Unmarshal(data, &approval); err != nil {
                return nil, fmt.Errorf("decode client approval: %w", err)
        }
        if !time.Now().Before(approval.ExpiresAt) {
                return nil, ErrExpired
        }
        return &approval, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "slices"
        "time"

        "github.com/ory/fosite"
)

var _ RememberedConsentStorage = (*MemoryStorage)(nil)

// StoreConsentSession stores only the digest of a browser's session secret.
func (s *MemoryStorage) StoreConsentSession(_ context.Context, digest string, session *ConsentSession) error {
        if digest == "" {
                return fosite.ErrInvalidRequest.WithHint("consent session digest cannot be empty")
        }
        if session == nil {
                return fosite.ErrInvalidRequest.WithHint("consent session cannot be nil")
        }
        if !time.Now().Before(session.ExpiresAt) {
                return ErrExpired
        }
        s.mu.Lock()
        defer s.mu.Unlock()
        clone := *session
        s.consentSessions[digest] = &timedEntry[*ConsentSession]{value: &clone, expiresAt: clone.ExpiresAt}
        return nil
}

// GetConsentSession loads a session without extending its fixed expiry.
func (s *MemoryStorage) GetConsentSession(_ context.Context, digest string) (*ConsentSession, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()
        entry, ok := s.consentSessions[digest]
        if !ok {
                return nil, ErrNotFound
        }
        if !time.Now().Before(entry.expiresAt) {
                return nil, ErrExpired
        }
        clone := *entry.value
        return &clone, nil
}

// DeleteConsentSession revokes a browser session; absence is harmless.
func (s *MemoryStorage) DeleteConsentSession(_ context.Context, digest string) error {
        s.mu.Lock()
        defer s.mu.Unlock()
        delete(s.consentSessions, digest)
        return nil
}

// StoreClientApproval replaces the approval for this user and client.
func (s *MemoryStorage) StoreClientApproval(_ context.Context, userID, clientID string, approval *ClientApproval) error {
        if userID == "" {
                return fosite.ErrInvalidRequest.WithHint("client approval user ID cannot be empty")
        }
        if clientID == "" {
                return fosite.ErrInvalidRequest.WithHint("client approval client ID cannot be empty")
        }
        if approval == nil {
                return fosite.ErrInvalidRequest.WithHint("client approval cannot be nil")
        }
        if !time.Now().Before(approval.ExpiresAt) {
                return ErrExpired
        }
        s.mu.Lock()
        defer s.mu.Unlock()
        clone := *approval
        clone.Scopes = slices.Clone(approval.Scopes)
        s.clientApprovals[[2]string{userID, clientID}] = &timedEntry[*ClientApproval]{value: &clone, expiresAt: clone.ExpiresAt}
        return nil
}

// GetClientApproval loads a copy without extending the approval's fixed expiry.
func (s *MemoryStorage) GetClientApproval(_ context.Context, userID, clientID string) (*ClientApproval, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()
        entry, ok := s.clientApprovals[[2]string{userID, clientID}]
        if !ok {
                return nil, ErrNotFound
        }
        if !time.Now().Before(entry.expiresAt) {
                return nil, ErrExpired
        }
        clone := *entry.value
        clone.Scopes = slices.Clone(clone.Scopes)
        return &clone, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

import (
        "context"
        "fmt"
        "maps"
        "slices"
        "time"

        "github.com/ory/fosite"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
)

// SPIFFEStorageDecorator resolves immutable configured SPIFFE clients before
// the dynamic DCR/CIMD storage backend. Static clients exist only in this
// overlay; they are never persisted or eligible for registration.
type SPIFFEStorageDecorator struct {
        Storage
        clients map[string]fosite.Client
}

// NewSPIFFEStorageDecorator overlays static SPIFFE clients over base. It
// durably claims each configured client ID in base's durable backend (see
// preflightDurableCollisions) before creating the overlay.
func NewSPIFFEStorageDecorator(ctx context.Context, base Storage, clients map[string]fosite.Client) (Storage, error) {
        if base == nil {
                return nil, fmt.Errorf("storage is required")
        }
        if len(clients) == 0 {
                return base, nil
        }
        privateClients := maps.Clone(clients)
        for clientID, client := range privateClients {
                if client == nil {
                        return nil, fmt.Errorf("static SPIFFE client %q is required", clientID)
                }
                if client.GetID() != clientID {
                        return nil, fmt.Errorf("static SPIFFE client map key %q does not match client ID %q", clientID, client.GetID())
                }
        }
        if err := preflightDurableCollisions(ctx, base, privateClients); err != nil {
                return nil, err
        }
        return &SPIFFEStorageDecorator{Storage: base, clients: privateClients}, nil
}

// PreflightSPIFFEStaticClientCollisions durably claims each configured static
// client ID in base's durable backend before the caller starts serving static
// SPIFFE clients from an in-process overlay. See preflightDurableCollisions.
func PreflightSPIFFEStaticClientCollisions(ctx context.Context, base Storage, clients map[string]fosite.Client) error {
        if base == nil {
                return fmt.Errorf("storage is required")
        }
        return preflightDurableCollisions(ctx, base, clients)
}

// preflightDurableCollisions durably claims each configured static SPIFFE
// client ID in base's durable backend using an inert placeholder (see
// staticClientPlaceholder) — never the real SPIFFE client. Reconciliation is
// idempotent: a restart with unchanged association config reconstructs the
// identical placeholder and succeeds silently; a collision with a DCR-issued
// client, a non-placeholder configured client, or a durably-claimed placeholder
// for a *different* association at the same ID, fails loudly. Unlike delegate
// clients, which are replaced on restart, placeholders must match exactly.
//
// This closes a cross-replica race that a read-only GetClient check cannot:
// with Redis and multiple replicas, an older/still-rolling replica without
// this SPIFFE config could otherwise DCR-register the same client ID after a
// newer replica's read-only check passed, leaving different replicas
// resolving different clients for the same ID. Durably claiming the ID here
// makes the reservation visible to every replica's ReconcileConfiguredClient
// and RegisterClient calls immediately.
func preflightDurableCollisions(ctx context.Context, base Storage, clients map[string]fosite.Client) error {
        underlying := Unwrap(base)
        for clientID, client := range clients {
                placeholder := staticClientPlaceholder(client)
                if err := underlying.ReconcileConfiguredClient(ctx, placeholder); err != nil {
                        return fmt.Errorf("claim durable placeholder for static SPIFFE client ID %q: %w", clientID, err)
                }
        }
        return nil
}

// resourceScopedClient is implemented by fosite.Client types (e.g.
// registration.SPIFFEClient) that maintain an RFC 8707 resource allowlist
// independent of GetAudience's audience allowlist. This mirrors the
// identically-shaped interface in pkg/authserver/server/tokenexchange -- both
// packages match it structurally against the same concrete client types
// without either importing the other's unexported interface.
type resourceScopedClient interface {
        Resources() []string
}

// spiffeIdentityClient is implemented by fosite.Client types (e.g.
// spiffeStaticClient) that carry a durable SPIFFE association identity
// fingerprint independent of OAuth policy. Mirrors resourceScopedClient's
// structural-matching pattern above.
type spiffeIdentityClient interface {
        IdentityFingerprint() string
}

// inertPlaceholderClient wraps a *fosite.DefaultClient to suppress its
// implicit defaulting: fosite.DefaultClient.GetGrantTypes and
// GetResponseTypes each return a single-element default
// (["authorization_code"] / ["code"]) when the underlying field is nil,
// treating "omitted" as "the interactive default" per the OIDC registration
// metadata convention they follow. staticClientPlaceholder needs the
// opposite: a client that can never be issued a token via any grant.
// Overriding both getters at the type — rather than storing an empty
// (non-nil) slice on the embedded struct — is required, not just clearer:
// fosite defaults on len() == 0 regardless of nil vs. non-nil-empty, so an
// empty slice alone would not suppress the default.
//
// This override is only reliable for the live Go value. A round trip through
// Redis serializes the *fields* (buildStoredClient reads GetGrantTypes() /
// GetResponseTypes() into JSON, which correctly capture the override's empty
// result), but clientFromStored ordinarily reconstructs a bare
// *fosite.DefaultClient on read, which reintroduces the very defaulting this
// type exists to suppress. storedClient.Reserved closes that gap: it marks a
// row as this placeholder so clientFromStored re-wraps the reconstructed
// client in inertPlaceholderClient on every read, independent of backend.
// See buildStoredClient and clientFromStored in redis.go.
type inertPlaceholderClient struct {
        registration.BackChannelOnlyMarker
        *fosite.DefaultClient
        resources           []string
        identityFingerprint string
}

func (inertPlaceholderClient) GetGrantTypes() fosite.Arguments    { return nil }
func (inertPlaceholderClient) GetResponseTypes() fosite.Arguments { return nil }

// Resources returns the placeholder's RFC 8707 resource allowlist, satisfying
// resourceScopedClient so the fingerprint comparison (see clientFingerprint)
// distinguishes associations that differ only in their resource allowlist.
// Returns a defensive copy, matching SPIFFEClient.Resources and
// SPIFFEAuthorizationPolicy.Resources.
func (c inertPlaceholderClient) Resources() []string { return slices.Clone(c.resources) }

func (c inertPlaceholderClient) IdentityFingerprint() string { return c.identityFingerprint }

// isReservedPlaceholder reports whether client is a staticClientPlaceholder
// value, letting buildStoredClient (redis.go) mark the persisted row so
// clientFromStored can restore the same inert wrapper on read-back. Mirrors
// registration.DCRIssued's marker-detection pattern.
func isReservedPlaceholder(client fosite.Client) bool {
        _, ok := client.(inertPlaceholderClient)
        return ok
}

// staticClientPlaceholder returns an inert stand-in for a real static SPIFFE
// client, safe to durably persist. It is never returned to a caller of
// GetClient — GetClient always resolves a reserved ID from the in-process
// overlay first (see SPIFFEStorageDecorator.GetClient) — but must still be
// structurally impossible to authenticate as or exchange a token for, in
// case that invariant is ever violated: no grant types and no response types
// mean fosite can never issue it a token via any grant handler this auth
// server registers, and it carries no secret. This guarantee holds on every
// backend, including a bare read-back from Redis by a replica with no SPIFFE
// overlay — see inertPlaceholderClient and storedClient.Reserved. It keeps
// the real client's Scopes, Audience, (RFC 8707) Resources, and SPIFFE
// association identity fingerprint so the fingerprint comparison in
// ReconcileConfiguredClient can distinguish "same association, restarted"
// (idempotent) from "different, colliding association at this ID" (a loud
// failure).
// Scopes and Audience are taken directly from actual.GetScopes() /
// GetAudience() without an extra defensive clone here: the only production
// caller (registration.SPIFFEClient) already returns a fresh slice.Clone
// from both getters (see its doc comments), so there is no live backing
// array to alias. The same applies to Resources() when actual implements
// resourceScopedClient; a client that doesn't (falls back to nil).
func staticClientPlaceholder(actual fosite.Client) fosite.Client {
        var resources []string
        var identityFingerprint string
        if rc, ok := actual.(resourceScopedClient); ok {
                resources = rc.Resources()
        }
        if identity, ok := actual.(spiffeIdentityClient); ok {
                identityFingerprint = identity.IdentityFingerprint()
        }
        return inertPlaceholderClient{
                DefaultClient: &fosite.DefaultClient{
                        ID:       actual.GetID(),
                        Scopes:   actual.GetScopes(),
                        Audience: actual.GetAudience(),
                        Public:   false,
                },
                resources:           resources,
                identityFingerprint: identityFingerprint,
        }
}

// ConsumeAssertionJWT delegates assertion replay consumption to the wrapped
// storage, one level down. Storage intentionally does not include this narrow
// capability, so this decorator fails closed when its wrapped storage does
// not provide it, rather than silently skipping past this layer via Unwrap.
func (d *SPIFFEStorageDecorator) ConsumeAssertionJWT(
        ctx context.Context, purpose, issuer, jti string, exp time.Time,
) error {
        consumer, ok := d.Storage.(AssertionJWTConsumer)
        if !ok {
                return fmt.Errorf("wrapped storage %T does not support assertion JWT replay consumption", d.Storage)
        }
        return consumer.ConsumeAssertionJWT(ctx, purpose, issuer, jti, exp)
}

// GetClient returns a configured SPIFFE client before consulting the dynamic backend.
func (d *SPIFFEStorageDecorator) GetClient(ctx context.Context, id string) (fosite.Client, error) {
        if client, ok := d.clients[id]; ok {
                return client, nil
        }
        return d.Storage.GetClient(ctx, id)
}

// RegisterClient rejects reserved static SPIFFE IDs and delegates other registrations.
func (d *SPIFFEStorageDecorator) RegisterClient(ctx context.Context, client fosite.Client) error {
        if client == nil {
                return fmt.Errorf("register client: client is required")
        }
        if _, reserved := d.clients[client.GetID()]; reserved {
                return fmt.Errorf("%w: client ID %q is reserved for a static SPIFFE client", ErrAlreadyExists, client.GetID())
        }
        return d.Storage.RegisterClient(ctx, client)
}

// ReconcileConfiguredClient rejects reserved static SPIFFE IDs and delegates
// other configured-client reconciliation to the base backend.
func (d *SPIFFEStorageDecorator) ReconcileConfiguredClient(ctx context.Context, client fosite.Client) error {
        if client == nil {
                return fmt.Errorf("reconcile configured client: client is required")
        }
        if _, reserved := d.clients[client.GetID()]; reserved {
                return fmt.Errorf("%w: client ID %q is reserved for a static SPIFFE client", ErrAlreadyExists, client.GetID())
        }
        return d.Storage.ReconcileConfiguredClient(ctx, client)
}

// UpsertDCRIssuedClient rejects reserved static SPIFFE IDs and delegates
// other DCR-issued upsert-on-fetch calls (e.g. CIMDStorageDecorator) to the
// base backend. Without this override the guard would be skipped whenever
// the durable SPIFFE placeholder row is absent (reconciliation hasn't run
// yet, or persistence failed), letting a DCR row get created at a reserved
// SPIFFE ID.
func (d *SPIFFEStorageDecorator) UpsertDCRIssuedClient(ctx context.Context, client fosite.Client) error {
        if client == nil {
                return fmt.Errorf("upsert DCR-issued client: client is required")
        }
        if _, reserved := d.clients[client.GetID()]; reserved {
                return fmt.Errorf("%w: client ID %q is reserved for a static SPIFFE client", ErrAlreadyExists, client.GetID())
        }
        return d.Storage.UpsertDCRIssuedClient(ctx, client)
}

// Unwrap returns the dynamic DCR/CIMD storage backend.
func (d *SPIFFEStorageDecorator) Unwrap() Storage { return d.Storage }

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package storage provides storage interfaces and implementations for the
// OAuth authorization server.
package storage

//go:generate mockgen -destination=mocks/mock_storage.go -package=mocks -source=types.go Storage,PendingAuthorizationStorage,DeviceCodeStorage,PendingDeviceLoginStorage,PendingDeviceConfirmationStorage,RememberedConsentStorage,AssertionJWTConsumer,ClientRegistry,UpstreamTokenStorage,UpstreamTokenRefresher,UserStorage,DCRCredentialStore

import (
        "context"
        "crypto/sha256"
        "encoding/hex"
        "errors"
        "fmt"
        "log/slog"
        "slices"
        "sort"
        "strings"
        "time"

        "github.com/ory/fosite"
        "github.com/ory/fosite/handler/oauth2"
        "github.com/ory/fosite/handler/pkce"

        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        authsession "github.com/stacklok/toolhive/pkg/authserver/server/session"
)

// Sentinel errors for storage operations.
// Use errors.Is() to check for these error types.
var (
        // ErrNotFound is returned when an item is not found in storage.
        ErrNotFound = errors.New("storage: item not found")

        // ErrExpired is returned when an item exists but has expired.
        ErrExpired = errors.New("storage: item expired")

        // ErrAlreadyExists is returned when attempting to create an item that already exists.
        ErrAlreadyExists = errors.New("storage: item already exists")

        // ErrInvalidBinding is returned when token binding validation fails
        // (e.g., subject or client ID mismatch).
        ErrInvalidBinding = errors.New("storage: token binding validation failed")

        // ErrClientCapacity is returned by RegisterClient when the number of live
        // clients counted against the DCR cap has reached it. Nothing is written.
        // The returned error is a *ClientCapacityError, which matches this
        // sentinel under errors.Is and carries the Retry-After hint.
        ErrClientCapacity = errors.New("storage: client capacity reached")

        // ErrReservedClientID is returned when a caller attempts to register a
        // real client whose ID collides with SyntheticClientIDPrefix.
        ErrReservedClientID = errors.New("storage: client id uses reserved synthetic prefix")

        // ErrUserNotProvisioned is returned by a UserStorage.CreateUser implementation
        // that deliberately refuses to auto-provision an upstream identity (e.g. a
        // deployment where user accounts are only created out-of-band, such as via
        // SCIM). It signals the authorization callback to deny the login with
        // fosite.ErrAccessDenied instead of treating the refusal as an internal
        // server error.
        ErrUserNotProvisioned = errors.New("storage: user not provisioned")

        // ErrConcurrentRefresh is returned by
        // UpstreamTokenStorage.CompareAndSwapUpstreamTokens when the currently
        // stored refresh token no longer equals the caller's expected value.
        // This covers two distinct situations: another writer (a concurrent
        // refresh in this process or another replica sharing the same storage)
        // already moved the row past it, or the row no longer exists at all
        // (deleted by logout, evicted by TTL, or never created). The caller's own
        // redemption is stale either way and must not be written; callers that
        // need to tell the two situations apart must re-read the row afterward
        // (ErrNotFound means there was no race to lose — see
        // CompareAndSwapUpstreamTokens for the full coordination contract).
        ErrConcurrentRefresh = errors.New("storage: upstream token row changed concurrently")

        // ErrDCRCredentialsChanged is returned by
        // DCRCredentialStore.UpdateDCRCredentialsIfUnchanged when the stored row
        // no longer equals the caller's expected value: another writer replaced it
        // after the caller read it. Nothing is written. A row that is absent
        // altogether is reported as ErrNotFound instead.
        ErrDCRCredentialsChanged = errors.New("storage: dcr credentials row changed concurrently")

        // ErrInvalidState is returned when an operation requires an item to be in a
        // particular lifecycle state (e.g. a pending device request) but it is not.
        ErrInvalidState = errors.New("storage: item is not in the required state")
)

// notFoundRFC6749Error preserves the storage and Fosite not-found identities.
func notFoundRFC6749Error(hint string) *fosite.RFC6749Error {
        return fosite.ErrNotFound.WithHint(hint).WithWrap(ErrNotFound)
}

// DefaultPendingAuthorizationTTL is the default TTL for pending authorization requests.
const DefaultPendingAuthorizationTTL = 10 * time.Minute

// DefaultDeviceRequestTTL bounds how long a device authorization request
// stays valid before the client must restart the flow, matching RFC 8628's
// recommended default expiry.
const DefaultDeviceRequestTTL = 10 * time.Minute

// UpstreamTokens represents tokens obtained from an upstream Identity Provider.
// These tokens are stored with binding fields for security validation and
// ProviderID for multi-IDP support.
type UpstreamTokens struct {
        // ProviderID identifies which upstream provider issued these tokens.
        // This is the logical provider name matching UpstreamConfig.Name,
        // used as the storage key for multi-upstream token lookups.
        ProviderID string

        // Token values from the upstream IDP
        AccessToken  string //nolint:gosec // G117: field legitimately holds sensitive data
        RefreshToken string //nolint:gosec // G117: field legitimately holds sensitive data
        IDToken      string
        // ExpiresAt is when the access token expires. Zero value means the provider
        // did not assert an expiry; callers must treat it as non-expiring.
        ExpiresAt time.Time
        // SessionExpiresAt is the Fosite session expiry time. The callback handler
        // sets it from RefreshTokenLifespan on every initial write; the refresher
        // carries it forward unchanged on each subsequent refresh, and defensively
        // re-anchors it on legacy rows where both ExpiresAt and SessionExpiresAt
        // are zero. Storage backends use it as a fallback storage lifetime when
        // ExpiresAt is zero (non-expiring upstream token), bounding the row to the
        // Fosite session. Set unconditionally — including for tokens with their own
        // ExpiresAt — so the refresh path is safe even if the upstream provider stops
        // asserting expires_in on a later rotation.
        SessionExpiresAt time.Time

        // Security binding fields - validated on lookup to prevent cross-session attacks

        // UserID is the internal ToolHive user ID (references User.ID).
        // This is NOT the upstream provider's subject - it's our stable internal identifier.
        // In multi-IDP scenarios, the same UserID may have tokens from multiple providers.
        UserID string

        // UpstreamSubject is the resolved subject of the upstream IDP's ID token:
        // the "sub" claim by default, or the provider's configured SubjectClaim
        // (e.g. Entra's "oid"). This enables validation that tokens match the
        // expected upstream identity and supports audit logging of which upstream
        // identity was used.
        UpstreamSubject string

        // ClientID is the OAuth client that initiated the authorization.
        // Tokens should only be accessible to the same client that obtained them.
        ClientID string
}

// IsExpired returns true if the access token has expired.
// Returns true for nil receivers (treating nil tokens as expired).
// Zero ExpiresAt means the provider did not assert an expiry; returns false.
// The now parameter allows for deterministic testing.
//
// This method is intended for storage eviction decisions only — it checks exact
// expiry with no preemptive buffer. For "should I refresh before using?" decisions,
// use upstream.Tokens.IsExpiredAt which includes a preemptive buffer window.
func (t *UpstreamTokens) IsExpired(now time.Time) bool {
        if t == nil {
                return true
        }
        return !t.ExpiresAt.IsZero() && now.After(t.ExpiresAt)
}

// DCRKey is the canonical lookup key for a DCR registration. The tuple is
// designed so that any backend (in-memory or Redis) serialises it identically
// without redefining the canonical form. ScopesHash is used rather than a raw
// scope slice so the key is comparable, fixed-size, and order-insensitive.
//
// The key lives in the storage package because both MemoryStorage and the
// future Redis backend must hash keys identically; keeping the canonical form
// next to the persistence implementations prevents drift.
type DCRKey struct {
        // Issuer is the registration consumer's issuer identifier. The dual
        // semantic depends on the consumer profile:
        //   - Embedded authorization server: its OWN local issuer (the embedded
        //     authserver that performed the registration).
        //   - CLI / direct OAuth flow: the UPSTREAM authorization server's
        //     issuer, because the CLI has no separate local issuer of its own.
        // The cache is keyed by this value because two different consumers
        // registering against the same upstream are distinct OAuth clients and
        // must not share credentials. The (Issuer, UpstreamID, RedirectURI,
        // ScopesHash) tuple keeps the two consumer profiles' entries apart via
        // the RedirectURI component (the embedded authserver registers an
        // AS-origin callback while the CLI registers a loopback callback per
        // RFC 8252 §7.3 — the two address spaces are disjoint), so a collision
        // between profiles is impossible by construction even when the
        // upstream is the same. Public-client vs confidential-client
        // separation rides on that same disjoint-RedirectURI property at both
        // the persistent-cache and in-process singleflight layers; encoding it on
        // the key would invalidate every existing Redis-cached entry across a
        // deployment without buying additional protection. If a future
        // consumer brings the two address spaces into collision the key
        // format must gain a consumer-identifier component alongside an
        // explicit migration story.
        Issuer string

        // UpstreamID identifies the upstream authorization server this
        // registration is bound to, disambiguating upstreams that share the
        // consumer's Issuer, RedirectURI, and scope set. Within a single
        // embedded authserver every OAuth2 upstream shares the authserver's own
        // Issuer and the one defaulted {issuer}/oauth/callback RedirectURI, so
        // before this component existed two upstreams configured with equal
        // scopes collided on one cache entry — the second read back the first's
        // dynamically-registered client_id / client_secret and never registered
        // with its own authorization server (issue #5823).
        //
        // The value is the upstream's registration identity: the upstream issuer
        // recovered from the consumer's discovery URL, or the registration
        // endpoint URL when the consumer configured one directly. It is derived
        // by the dcr resolver from the Request; do not hand-build it at call
        // sites.
        //
        // On the discovery-URL path this is the derived issuer, so it identifies
        // the authorization server rather than the exact discovery URL: two
        // configs that resolve to the same issuer intentionally share one
        // registration. It does not fully disambiguate the nonstandard case of a
        // single issuer exposing distinct registration endpoints under custom
        // (non-well-known) discovery paths — see resolveUpstreamKeyIdentity in
        // pkg/auth/dcr for the full rationale.
        UpstreamID string

        // RedirectURI is the redirect URI registered with the upstream
        // authorization server. Embedded-authserver callers register an
        // AS-origin callback; CLI callers register an RFC 8252 loopback
        // callback. The two address spaces are disjoint, which is what makes
        // the per-consumer cache namespace structurally safe today.
        RedirectURI string

        // ScopesHash is the SHA-256 hex digest of the sorted, deduplicated scope
        // list. Use ScopesHash() to compute this value — do NOT hash scopes by
        // hand at call sites; the canonical form must be a single source of truth
        // so the key matches across processes and backends.
        ScopesHash string

        // RotationKey is an optional operator-chosen label that is part of the
        // cache identity but carries no OAuth meaning and is never sent to the
        // upstream. Changing it makes the resolver miss the cache and register a
        // fresh client. It is how an operator retires a registration the upstream
        // has dropped when the resolver's own registration probe cannot tell
        // (see authserver.DCRUpstreamConfig.RotationKey).
        // Empty for every row written before the field existed; backends MUST
        // encode an empty value exactly as they encoded a key without the field
        // so those rows stay reachable. Rows written under a previous value are
        // not deleted and follow the same lifetime rules as any superseded row.
        RotationKey string
}

// ScopesHash returns the SHA-256 hex digest of the canonical OAuth scope set,
// suitable for use as DCRKey.ScopesHash.
//
// Canonicalisation:
//  1. Sort ascending so the digest is order-insensitive — e.g.
//     []string{"openid", "profile"} and []string{"profile", "openid"} hash to
//     the same value.
//  2. Deduplicate so that []string{"openid"} and []string{"openid", "openid"}
//     hash to the same value. An OAuth scope set is a set, not a multiset
//     (RFC 6749 §3.3), and without deduplication a caller that accidentally
//     duplicated a scope would miss cache entries and trigger redundant
//     RFC 7591 registrations.
//  3. Join with newlines (a character not valid in OAuth scope tokens per
//     RFC 6749 §3.3) to avoid collision between e.g. ["ab", "c"] and
//     ["a", "bc"].
//
// nil and empty slice both canonicalise to the same hash.
func ScopesHash(scopes []string) string {
        sorted := slices.Clone(scopes)
        sort.Strings(sorted)
        sorted = slices.Compact(sorted)

        h := sha256.New()
        for i, s := range sorted {
                if i > 0 {
                        _, _ = h.Write([]byte("\n"))
                }
                _, _ = h.Write([]byte(s))
        }
        return hex.EncodeToString(h.Sum(nil))
}

// validateDCRCredentialsForStore enforces the rejection contract that every
// DCRCredentialStore implementation must apply before persisting. Extracting
// it as a free function called by every backend prevents the validation set
// from drifting across implementations: a record that fails loud against one
// backend cannot silently persist against another.
//
// Rejected inputs: nil creds, an unpopulated Key (empty Issuer, UpstreamID,
// RedirectURI, or ScopesHash), and missing RFC 7591 mandatory response fields
// (ClientID, AuthorizationEndpoint, TokenEndpoint). An empty ScopesHash is
// rejected because the canonical digest of any scope set — including the
// empty-scope set via ScopesHash(nil) — is non-empty, so an empty string can
// only be a caller bug. An empty UpstreamID is rejected because the resolver
// always derives it from the Request (issue #5823); an empty value here would
// re-open the cross-upstream collision it was added to close. ClientSecret is
// left permissive because RFC 7591 §2 public clients (auth method "none")
// legitimately register without a secret.
func validateDCRCredentialsForStore(creds *DCRCredentials) error {
        if creds == nil {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials cannot be nil")
        }
        if creds.Key.Issuer == "" {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials key issuer cannot be empty")
        }
        if creds.Key.UpstreamID == "" {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials key upstream_id cannot be empty")
        }
        if creds.Key.RedirectURI == "" {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials key redirect_uri cannot be empty")
        }
        if creds.Key.ScopesHash == "" {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials key scopes_hash cannot be empty")
        }
        if creds.ClientID == "" {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials client_id cannot be empty")
        }
        if creds.AuthorizationEndpoint == "" {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials authorization_endpoint cannot be empty")
        }
        if creds.TokenEndpoint == "" {
                return fosite.ErrInvalidRequest.WithHint("dcr credentials token_endpoint cannot be empty")
        }
        return nil
}

// validateDCRCompareAndSet enforces the input contract of
// DCRCredentialStore.UpdateDCRCredentialsIfUnchanged on top of
// validateDCRCredentialsForStore: expected must be present and must address
// the same row as creds, so a caller cannot gate a write to one key on the
// contents of another.
func validateDCRCompareAndSet(creds, expected *DCRCredentials) error {
        if err := validateDCRCredentialsForStore(creds); err != nil {
                return err
        }
        if expected == nil {
                return fosite.ErrInvalidRequest.WithHint("expected dcr credentials cannot be nil")
        }
        if expected.Key != creds.Key {
                return fosite.ErrInvalidRequest.WithHint("expected dcr credentials key must match creds key")
        }
        return nil
}

// DCRCredentials is the persisted form of an RFC 7591 Dynamic Client
// Registration result. All fields are populated from the upstream's DCR
// response. The RFC 7592 management fields (RegistrationAccessToken,
// RegistrationClientURI) are preserved verbatim so future rotation /
// management flows can use them.
//
// # Defensive copy
//
// Callers receive a defensive copy from the store. Mutations on the returned
// value do not affect persisted state, and mutations on a value passed to
// StoreDCRCredentialsIfAbsent are not observed by subsequent reads. This matches the
// UpstreamTokens contract.
//
// # Lifetime
//
// Entries are long-lived — RFC 7591 client registrations do not expire unless
// the upstream asserts client_secret_expires_at. The in-memory backend
// retains entries for the process lifetime and is intentionally excluded from
// the periodic cleanup loop. The Redis backend applies TTL via SET with a
// duration when ClientSecretExpiresAt is non-zero.
//
// # Converter contract
//
// MUST update both converters in
// pkg/auth/dcr/store.go (resolutionToCredentials and
// credentialsToResolution) when adding, renaming, or removing a field
// here. The two converters are the only translation seam between this
// persisted type and the dcr-package *dcr.Resolution; a field added here
// without a paired converter update will silently fail to round-trip
// across an authserver restart. The round-trip behaviour is pinned by
// TestResolutionCredentialsRoundTrip in
// pkg/auth/dcr/store_test.go.
type DCRCredentials struct {
        // Key is the canonical cache key: (Issuer, UpstreamID, RedirectURI,
        // ScopesHash) plus the optional RotationKey.
        Key DCRKey

        // ProviderName is the upstream's UpstreamRunConfig.Name. Debug / audit
        // only — never used as a primary key. Two upstreams with different
        // ProviderName but identical Key share one credential record.
        ProviderName string

        ClientID                string
        ClientSecret            string //nolint:gosec // G117: field legitimately holds sensitive data
        TokenEndpointAuthMethod string

        // RegistrationAccessToken and RegistrationClientURI are RFC 7592 fields
        // captured for future management operations (rotation, deletion).
        RegistrationAccessToken string //nolint:gosec // G117: field legitimately holds sensitive data
        RegistrationClientURI   string

        AuthorizationEndpoint string
        TokenEndpoint         string

        // CreatedAt is the wall-clock time at which the registration completed.
        // Used to compute staleness for observability — the cache itself does
        // not expire entries based on age (see ClientSecretExpiresAt for the
        // authoritative expiry signal).
        CreatedAt time.Time

        // ClientSecretExpiresAt is the RFC 7591 §3.2.1 client_secret_expires_at
        // value converted from int64 epoch seconds to time.Time. The wire
        // convention is that 0 means "the secret does not expire"; this struct
        // represents that as the zero time.Time so callers can use IsZero()
        // rather than special-casing 0.
        //
        // When non-zero, this is the authoritative signal a backend uses to TTL
        // the persisted entry: the Redis backend plumbs it through SET with a
        // duration so the row evicts before the upstream rejects the secret at the
        // token endpoint. The in-memory backend ignores this field — entries
        // persist for the process lifetime and the resolver re-checks the
        // expiry on read.
        ClientSecretExpiresAt time.Time
}

// DCRCredentialStore is a narrow, segregated interface for persisting
// dynamic-client-registration credentials. Both MemoryStorage and
// RedisStorage implement it; an authserver backed by Redis shares DCR
// credentials across replicas and restarts.
//
// # Cross-replica limitation
//
// Sharing DCR credentials does NOT imply cross-replica session / token
// delivery. Callers that need that must still route through the proxy runner
// and (if applicable) pin sessions to a replica.
//
// # Defensive copy
//
// Implementations MUST defensively copy on Store, Update, and Get so caller
// mutations cannot reach persisted state and vice versa, mirroring the
// UpstreamTokens contract.
//
// # TTL handling
//
// Implementations SHOULD honor a non-zero DCRCredentials.ClientSecretExpiresAt
// as a backend-level TTL when the underlying store supports one (e.g. Redis
// SET with a duration) so an entry evicts before the upstream rejects the
// secret at the token endpoint. Backends without a native TTL (e.g. the
// in-memory backend) retain the field verbatim and rely on the caller —
// typically the runner's resolver — to re-check expiry on read; see
// MemoryStorage.GetDCRCredentials.
// A zero ClientSecretExpiresAt means the upstream did not assert an expiry
// and no TTL is applied.
//
// # Why the key is embedded in DCRCredentials
//
// Both StoreDCRCredentialsIfAbsent and UpdateDCRCredentialsIfPresent take a
// single (ctx, creds) argument rather than the (ctx, key, value) shape used by
// sibling Store* methods on Storage. The DCRKey is embedded as
// DCRCredentials.Key so the persisted blob is
// self-describing: a Redis SCAN, an admin-tool dump, or a cross-replica
// reconciliation path can identify a record's logical cache slot
// (Issuer, UpstreamID, RedirectURI, ScopesHash) from the value alone, without
// reconstructing it from a separately-passed key. This is a deliberate
// asymmetry with the rest of the package — callers must populate creds.Key
// before Store or Update, and implementations validate it (see MemoryStorage
// docs for the rejected-input list).
type DCRCredentialStore interface {
        // GetDCRCredentials returns the credentials for the given key.
        // Returns ErrNotFound (wrapped) if no entry exists for the key.
        // The returned value is a defensive copy.
        GetDCRCredentials(ctx context.Context, key DCRKey) (*DCRCredentials, error)

        // StoreDCRCredentialsIfAbsent claims creds.Key for creds. Returns the
        // authoritative durable value: the caller's own creds on a successful
        // claim, the concurrent winner's value otherwise. Callers MUST use the
        // returned value, not their input creds — RFC 7591 dynamic registration
        // mints a unique client_id/client_secret on every call, so a caller that
        // lost the race and kept using its own creds would hold credentials the
        // durable store does not agree it owns. See the interface-level "TTL
        // handling" section for the contract on ClientSecretExpiresAt. The
        // returned *DCRCredentials is always non-nil when err is nil.
        StoreDCRCredentialsIfAbsent(ctx context.Context, creds *DCRCredentials) (*DCRCredentials, error)

        // UpdateDCRCredentialsIfPresent replaces the record at creds.Key with
        // creds, iff a record currently exists at that key. Returns ErrNotFound
        // (wrapped) if no entry exists — it never creates, so an update racing a
        // delete or a not-yet-created record fails loudly rather than silently
        // creating, keeping the create/update split explicit (the mirror of why
        // StoreDCRCredentialsIfAbsent never silently updates). On success it
        // returns the stored value (a defensive copy) and the returned
        // *DCRCredentials is non-nil.
        //
        // This is the write path a storage decorator needs to rewrite a record's
        // persisted representation in place — re-encoding, compression, a
        // checksum, or similar — without changing its RFC 7591 identity or values:
        // StoreDCRCredentialsIfAbsent is create-only and would silently discard
        // such a rewrite on an existing key.
        //
        // # Presence is physical, not liveness
        //
        // "Present" means a row physically exists at creds.Key, regardless of
        // whether its ClientSecretExpiresAt has passed. This deliberately does NOT
        // mirror StoreDCRCredentialsIfAbsent's "an expired row counts as absent"
        // treatment: that check exists to let a fresh registration reclaim a dead
        // slot in the concurrent-registration race, whereas Update exists to let a
        // decorator rewrite a row it just read. GetDCRCredentials returns
        // physically-present rows without filtering on expiry, so gating Update on
        // liveness would break the Get→transform→Update round-trip the method is
        // for. An expired-but-present row is therefore updatable.
        //
        // Implementations MUST defensively copy on input (mirroring the Store/Get
        // contract) and MUST apply the same ClientSecretExpiresAt TTL handling
        // documented in the interface-level "TTL handling" section and on
        // StoreDCRCredentialsIfAbsent — the rewritten row's backend TTL is derived
        // from the incoming creds, so an update can extend, shorten, or clear the
        // row's TTL exactly as an initial store would.
        UpdateDCRCredentialsIfPresent(ctx context.Context, creds *DCRCredentials) (*DCRCredentials, error)

        // UpdateDCRCredentialsIfUnchanged replaces the record at creds.Key with
        // creds, iff the record currently stored there still equals expected —
        // the value the caller previously read via GetDCRCredentials. It is the
        // compare-and-set sibling of UpdateDCRCredentialsIfPresent, for callers
        // that coordinate re-registration across replicas and must not clobber a
        // newer row written by a concurrent writer after their read (e.g. a lock
        // holder that stalled past its lease before writing).
        //
        // Returns ErrNotFound (wrapped) if no record exists at the key, and
        // ErrDCRCredentialsChanged if the stored record no longer equals
        // expected. In both cases nothing is written. On success it returns the
        // stored value (a defensive copy) and the returned *DCRCredentials is
        // non-nil.
        //
        // Any other error leaves the row's state to be determined by re-reading
        // it. This includes transport failures, where the write may or may not
        // have been applied. It also includes a backend that bounds its internal
        // retries (the Redis backend) giving up because other writers kept
        // touching the key between its compare and its write; in that case this
        // call wrote nothing, and the caller may re-read and try again.
        //
        // expected must be non-nil and its Key must equal creds.Key. Equality is
        // over every persisted field, compared at the backend's storage precision
        // (the Redis backend persists times at one-second resolution), so a value
        // returned by GetDCRCredentials always compares equal to the row it was
        // read from.
        //
        // A decorator that transforms fields on the way in and out (encryption,
        // compression, …) must pass the inner store the stored form of expected —
        // what the inner GetDCRCredentials returned — not the decoded form it
        // handed to its own caller. When the transform is not reproducible (e.g.
        // encryption with a fresh nonce per seal), the decorator cannot derive
        // that stored form from the caller's decoded expected. Instead it should
        // re-read the inner row, decode it, return ErrDCRCredentialsChanged if
        // the decoded value differs from the caller's expected, and otherwise
        // call the inner UpdateDCRCredentialsIfUnchanged with the raw inner row
        // it just read as expected. That stays atomic: if the row changes between
        // the decorator's re-read and the inner write, the inner compare refuses
        // it. A decorator that recovers a row it cannot decode at all already
        // holds the raw inner row and passes it directly. A consumer outside the
        // decorator that needs to compare-and-set a row it cannot decode would
        // need an opaque version or etag instead; that is not provided today.
        //
        // Presence is physical, exactly as for UpdateDCRCredentialsIfPresent, and
        // the rewritten row's backend TTL follows the same ClientSecretExpiresAt
        // rules. Implementations MUST perform the compare and the write
        // atomically with respect to other writers of the same key, and MUST
        // touch only that one key so the operation stays valid under Redis
        // Cluster.
        UpdateDCRCredentialsIfUnchanged(ctx context.Context, creds, expected *DCRCredentials) (*DCRCredentials, error)
}

// User represents a user account in the authorization server.
// A user can have multiple linked provider identities.
// The User.ID is used as the "sub" claim in JWTs issued by ToolHive,
// providing a stable identity across multiple upstream identity providers.
type User struct {
        // ID is the internal, stable user identifier (UUID format).
        // This value is used as the "sub" claim in ToolHive-issued JWTs.
        ID string

        // CreatedAt is when the user account was created.
        CreatedAt time.Time

        // UpdatedAt is when the user account was last modified.
        UpdatedAt time.Time
}

// ProviderIdentity links a user to an upstream identity provider.
// Multiple identities can be linked to a single user for account linking,
// enabling users to authenticate via different providers (e.g., Google and GitHub)
// while maintaining a single ToolHive identity.
type ProviderIdentity struct {
        // UserID is the internal user ID this identity belongs to.
        UserID string

        // ProviderID identifies the upstream provider (e.g., "google", "github").
        ProviderID string

        // ProviderSubject is the subject identifier from the upstream provider.
        ProviderSubject string

        // LinkedAt is when this identity was linked to the user.
        LinkedAt time.Time

        // LastUsedAt is when this identity was last used to authenticate.
        LastUsedAt time.Time
}

// Consent stages for a pending client authorization. Empty is a legacy record
// and must never grant consent.
const (
        ConsentStageAwaiting = "awaiting"
        ConsentStageApproved = "approved"
)

// ConsentSession binds a browser to a previously verified first-upstream identity.
// JSON names match existing Redis rows written without tags.
type ConsentSession struct {
        UserID          string    `json:"UserID"`
        ProviderID      string    `json:"ProviderID"`
        ProviderSubject string    `json:"ProviderSubject"`
        ExpiresAt       time.Time `json:"ExpiresAt"`
}

// ClientApproval is the exact request envelope approved by a user for a client.
// JSON names match existing Redis rows written without tags.
type ClientApproval struct {
        RedirectURI string    `json:"RedirectURI"`
        Scopes      []string  `json:"Scopes"`
        Resource    string    `json:"Resource"`
        ExpiresAt   time.Time `json:"ExpiresAt"`
}

const (
        // ConsentSessionTTL is the fixed lifetime of a remembered browser session.
        ConsentSessionTTL = 7 * 24 * time.Hour
        // ClientApprovalTTL is the fixed lifetime of a remembered client approval.
        ClientApprovalTTL = 30 * 24 * time.Hour
)

// RememberedConsentStorage is an optional browser-session and client-approval capability.
// Callers supply a digest of the session secret; raw cookie values must never be stored.
// Reads return ErrNotFound for absent rows and ErrExpired for expired rows.
type RememberedConsentStorage interface {
        StoreConsentSession(ctx context.Context, digest string, session *ConsentSession) error
        GetConsentSession(ctx context.Context, digest string) (*ConsentSession, error)
        DeleteConsentSession(ctx context.Context, digest string) error
        StoreClientApproval(ctx context.Context, userID, clientID string, approval *ClientApproval) error
        GetClientApproval(ctx context.Context, userID, clientID string) (*ClientApproval, error)
}

// PendingAuthorization tracks a client's authorization request while they
// authenticate with the upstream IDP.
type PendingAuthorization struct {
        // ClientID is the ID of the OAuth client making the authorization request.
        ClientID string

        // RedirectURI is the client's callback URL where we'll redirect after authentication.
        RedirectURI string

        // State is the client's original state parameter for CSRF protection.
        State string

        // PKCEChallenge is the client's PKCE code challenge.
        PKCEChallenge string

        // PKCEMethod is the PKCE challenge method (must be "S256").
        PKCEMethod string

        // Scopes are the OAuth scopes requested by the client.
        Scopes []string

        // Resource is the single audience selected at authorization, either explicitly
        // by the client or from the sole configured allowed audience.
        Resource string

        // ConsentStage is empty on legacy records, awaiting before approval, and
        // approved only after a browser-bound consent POST.
        ConsentStage string

        // RememberConsent is set only by a browser-bound approval on an HTTPS endpoint.
        RememberConsent bool
        // FirstProviderSubject is the non-synthetic first upstream's verified subject.
        FirstProviderSubject string

        // Expected identity and originating session for a remembered, page-skipping flow.
        ExpectedUserID          string
        ExpectedProviderSubject string
        ConsentSessionDigest    string

        // InternalState is our randomly generated state for correlating upstream callback.
        InternalState string

        // UpstreamPKCEVerifier is the PKCE code_verifier for the upstream IDP authorization.
        // This is generated when redirecting to the upstream IDP and used when exchanging
        // the authorization code. See RFC 7636.
        UpstreamPKCEVerifier string

        // UpstreamNonce is the OIDC nonce parameter sent to the upstream IDP for ID Token
        // replay protection. This is validated against the nonce claim in the returned
        // ID Token. See OIDC Core Section 3.1.2.1.
        UpstreamNonce string

        // BrowserBindingHash is the base64url SHA-256 of the browser-binding cookie
        // value that /oauth/authorize set on the user agent when it created this
        // record (or that the previous chain leg set when it created the next one).
        // /oauth/callback only completes the leg when the calling browser presents a
        // cookie whose hash matches, which binds the upstream callback to the browser
        // that started the flow (RFC 6749 Section 10.12). Only the hash is stored so
        // that storage never holds the raw secret. The callback rejects a loaded
        // record with an empty hash (a row written before this field existed).
        BrowserBindingHash string

        // UpstreamProviderName identifies the upstream provider being authenticated in this
        // authorization chain leg. Used in multi-upstream scenarios to route the callback
        // to the correct provider.
        UpstreamProviderName string

        // SessionID is the TSID for session accumulation across authorization chain legs.
        // In multi-upstream scenarios, the same session accumulates tokens from multiple
        // providers across successive authorization legs.
        SessionID string

        // ResolvedUserID is the internal user ID resolved from the FIRST upstream in
        // the chain (upstreams[0]), which establishes the session identity.
        // Empty on the first leg; populated after the first callback for subsequent legs.
        //
        // "First" here is deliberately not called "primary": Cedar's
        // ConfigOptions.PrimaryUpstreamProvider names the upstream whose claims drive
        // authorization, which defaults to this one but can be pinned to a later leg.
        // The two only coincide by default. ResolvedUserName and ResolvedUserEmail
        // below are mirrored into the issued access token, so a pinned deployment that
        // falls back to request-token claims evaluates THIS upstream's profile claims,
        // not the pinned provider's. See docs/authz.md, "What the fallback claims
        // actually are".
        ResolvedUserID string

        // ResolvedUserName is the user display name from the first upstream in the chain.
        // Empty on the first leg; populated after the first callback for subsequent legs.
        ResolvedUserName string

        // ResolvedUserEmail is the user email from the first upstream in the chain.
        // Empty on the first leg; populated after the first callback for subsequent legs.
        ResolvedUserEmail string

        // VerifiedUpstream is the verified identity of the first upstream OIDC login. It
        // is set only from that login, never replaced by later legs, and zero for
        // OAuth2-only or synthetic identities, and not recorded at new logins when SubjectIdentifierAudiences is empty.
        VerifiedUpstream authsession.UpstreamSubject

        // SingleLeg scopes this authorization to exactly one upstream. When true, the
        // callback issues the authorization code as soon as this leg completes instead
        // of consulting nextMissingUpstream to continue the chain. This lets a caller
        // connect a single specific provider without other configured-but-tokenless
        // upstreams hijacking the flow into a full chain walk. Defaults to false, which
        // preserves the multi-upstream chaining behavior.
        SingleLeg bool

        // ChainUpstreams is the ordered, effective set of upstream provider names this
        // authorization must walk, always led by the required first upstream. It is
        // computed once when the first leg resolves — narrowed by the optional upstream
        // filter when one is configured — and carried forward across subsequent legs so
        // the filter is not re-run per leg. Empty on the first leg's inbound pending;
        // populated on every subsequent leg. A subsequent leg that arrives with this
        // unset (e.g. a pending written before this field existed) is rejected rather
        // than recomputed, so the filter is never re-run against a later leg's context.
        ChainUpstreams []string

        // CreatedAt is when the pending authorization was created. It must be set on
        // the first leg and preserved across subsequent legs; storage expires the
        // request at CreatedAt + DefaultPendingAuthorizationTTL.
        CreatedAt time.Time
}

// PendingAuthorizationStorage provides storage operations for pending authorization requests.
// These track the state of in-flight authorization requests while users authenticate
// with the upstream IDP, correlating the upstream callback with the original client request.
type PendingAuthorizationStorage interface {
        // StorePendingAuthorization stores a pending authorization request.
        // The state is used to correlate the upstream IDP callback.
        StorePendingAuthorization(ctx context.Context, state string, pending *PendingAuthorization) error

        // LoadPendingAuthorization retrieves a pending authorization by internal state.
        // Returns ErrNotFound if the state does not exist.
        // Returns ErrExpired if the pending authorization has expired.
        LoadPendingAuthorization(ctx context.Context, state string) (*PendingAuthorization, error)

        // DeletePendingAuthorization removes a pending authorization.
        // Returns ErrNotFound if the state does not exist.
        DeletePendingAuthorization(ctx context.Context, state string) error
}

// DeviceRequestStatus is the lifecycle state of an RFC 8628 device authorization request.
type DeviceRequestStatus string

const (
        // DeviceRequestStatusPending is the initial state: the user has not yet
        // completed (or denied) verification at the verification URI.
        DeviceRequestStatusPending DeviceRequestStatus = "pending"

        // DeviceRequestStatusAuthorized means the user approved the request at
        // the verification page; the resolved identity fields are populated.
        DeviceRequestStatusAuthorized DeviceRequestStatus = "authorized"

        // DeviceRequestStatusDenied means the user explicitly denied the request
        // at the verification page.
        DeviceRequestStatusDenied DeviceRequestStatus = "denied"
)

// DeviceRequest represents one in-flight RFC 8628 device authorization grant.
type DeviceRequest struct {
        // DeviceCode is the opaque, high-entropy value the polling client holds.
        // Never logged (secret-shaped, like an authorization code).
        DeviceCode string

        // UserCode is the short, human-typeable code the user enters at the
        // verification URI. Also secret-shaped: it is the only thing binding a
        // human's browser session to this device_code, so treat it like a code.
        UserCode string

        ClientID string
        Scopes   []string
        Audience []string
        Status   DeviceRequestStatus

        // Interval is the minimum seconds between polls the client must honor
        // (RFC 8628 §3.2/§3.5). Storage does not enforce it directly; callers
        // use LastPolledAt + Interval to decide slow_down.
        Interval time.Duration

        // LastPolledAt is zero until the first poll.
        LastPolledAt time.Time

        // Populated only once Status == DeviceRequestStatusAuthorized, mirroring
        // the Resolved* fields on PendingAuthorization:
        ResolvedUserID    string
        ResolvedUserName  string
        ResolvedUserEmail string
        VerifiedUpstream  authsession.UpstreamSubject // zero when SubjectIdentifierAudiences is empty
        SessionID         string

        CreatedAt time.Time
}

// DeviceAuthorization is the identity MarkDeviceRequestAuthorized attaches to a
// pending device request. Field-for-field it mirrors the Resolved*, SessionID
// and VerifiedUpstream fields it populates on DeviceRequest.
type DeviceAuthorization struct {
        UserID           string
        UserName         string
        UserEmail        string
        SessionID        string
        VerifiedUpstream authsession.UpstreamSubject
}

// DeviceCodeStorage provides storage operations for RFC 8628 device
// authorization requests. A request is created pending at the device
// authorization endpoint, looked up by user_code at the verification page
// and transitioned to authorized/denied there, and polled/consumed by
// device_code at the token endpoint.
type DeviceCodeStorage interface {
        // StoreDeviceRequest stores a new pending device request, indexed by both
        // DeviceCode and UserCode. Returns fosite.ErrInvalidRequest if DeviceCode
        // or UserCode is empty, or Status is not DeviceRequestStatusPending.
        // Returns ErrAlreadyExists if a request already exists under the same
        // UserCode (the caller must regenerate the user_code and retry) or the
        // same DeviceCode.
        StoreDeviceRequest(ctx context.Context, device *DeviceRequest) error

        // LoadDeviceRequestByDeviceCode retrieves a device request by its
        // device_code. Returns ErrNotFound if it does not exist, ErrExpired if
        // its TTL has elapsed.
        LoadDeviceRequestByDeviceCode(ctx context.Context, deviceCode string) (*DeviceRequest, error)

        // LoadDeviceRequestByUserCode retrieves a device request by its
        // user_code, for the verification page. Same not-found/expired semantics.
        LoadDeviceRequestByUserCode(ctx context.Context, userCode string) (*DeviceRequest, error)

        // MarkDeviceRequestAuthorized transitions a pending device request to
        // authorized, attaching the resolved identity. Returns ErrNotFound if
        // deviceCode does not exist, ErrExpired if its TTL has elapsed, and
        // ErrInvalidState if the request's Status is not currently
        // DeviceRequestStatusPending (already authorized or denied) — this call
        // is not idempotent, so a stale verification-page resubmission can never
        // clobber a request the token endpoint already consumed.
        MarkDeviceRequestAuthorized(ctx context.Context, deviceCode string, auth DeviceAuthorization) error

        // MarkDeviceRequestDenied transitions a pending device request to
        // denied. Same ErrNotFound/ErrExpired/ErrInvalidState semantics as
        // MarkDeviceRequestAuthorized.
        MarkDeviceRequestDenied(ctx context.Context, deviceCode string) error

        // UpdateDeviceRequestLastPolledAt records the time of the most recent
        // poll, so the (future) token-endpoint grant handler can enforce the
        // minimum polling Interval (RFC 8628 §3.5 slow_down). Same
        // ErrNotFound/ErrExpired semantics; does not require Status ==
        // pending (a client may poll after authorization races the response).
        UpdateDeviceRequestLastPolledAt(ctx context.Context, deviceCode string, polledAt time.Time) error

        // DeleteDeviceRequest removes a device request, e.g. once its token has
        // been issued so the device_code cannot be redeemed twice. Returns
        // ErrNotFound if it does not already exist.
        DeleteDeviceRequest(ctx context.Context, deviceCode string) error
}

// DefaultDeviceLoginTTL bounds how long a device-flow verification-page login
// attempt (PendingDeviceLogin) or post-login confirmation (
// PendingDeviceConfirmation) may remain unresolved before it is treated as
// expired. A human completing an upstream login and then confirming/denying
// a device should take at most a few minutes; this mirrors
// DefaultDeviceRequestTTL's window.
const DefaultDeviceLoginTTL = 10 * time.Minute

// PendingDeviceLogin correlates a device-flow verification-page login
// attempt with its upstream IDP callback, the same role PendingAuthorization
// plays for the OAuth client authorization_code flow -- but scoped to the
// verification page, which has no client_id/redirect_uri of its own.
type PendingDeviceLogin struct {
        // DeviceCode identifies the DeviceRequest this login attempt is for.
        DeviceCode string

        // UserCode is carried alongside DeviceCode so the callback handler can
        // re-look-up (and re-validate the pending/unexpired state of) the
        // DeviceRequest without a second index lookup by DeviceCode.
        UserCode string

        // UpstreamPKCEVerifier is the PKCE code_verifier for the upstream IDP
        // authorization. See RFC 7636.
        UpstreamPKCEVerifier string

        // UpstreamNonce is the OIDC nonce for ID token replay protection.
        UpstreamNonce string

        // BrowserBindingHash binds this login to the browser that submitted the
        // user_code, exactly as PendingAuthorization.BrowserBindingHash binds an
        // OAuth-client authorization: the verification page's redirect to the
        // upstream sets a cookie whose hash is stored here, and the callback
        // completes the login only for a browser presenting a matching cookie.
        BrowserBindingHash string

        // UpstreamProviderName is the configured upstream this login was sent
        // to; the callback re-validates against it (IDP mix-up defense), mirroring
        // PendingAuthorization.UpstreamProviderName.
        UpstreamProviderName string

        // CreatedAt is when the pending login was created.
        CreatedAt time.Time
}

// PendingDeviceLoginStorage provides storage operations for in-flight
// device-flow verification-page logins: created when a user submits a
// user_code and is redirected to the upstream IDP, and consumed exactly once
// when that IDP redirects back to the device callback endpoint.
type PendingDeviceLoginStorage interface {
        // StorePendingDeviceLogin stores a pending device login. state is used to
        // correlate the upstream IDP callback, exactly as
        // PendingAuthorizationStorage.StorePendingAuthorization uses it for the
        // client authorization_code flow.
        StorePendingDeviceLogin(ctx context.Context, state string, pending *PendingDeviceLogin) error

        // LoadPendingDeviceLogin retrieves a pending device login by state.
        // Returns ErrNotFound if the state does not exist, ErrExpired if
        // DefaultDeviceLoginTTL has elapsed.
        LoadPendingDeviceLogin(ctx context.Context, state string) (*PendingDeviceLogin, error)

        // DeletePendingDeviceLogin removes a pending device login. Returns
        // ErrNotFound if the state does not exist.
        DeletePendingDeviceLogin(ctx context.Context, state string) error
}

// PendingDeviceConfirmation holds a resolved upstream identity awaiting the
// user's explicit Approve/Deny decision at the device verification page.
//
// The resolved identity is deliberately kept server-side, addressed only by
// an opaque, unguessable token handed to the browser as a hidden form field --
// never as editable request data. A confirm-page design that instead echoed
// ResolvedUserID/Name/Email back as editable hidden fields would let a user
// tamper with them client-side and have a device authorized under a spoofed
// identity; keying the server-side record by an opaque token closes that.
type PendingDeviceConfirmation struct {
        // DeviceCode and UserCode identify the DeviceRequest awaiting decision.
        DeviceCode string
        UserCode   string

        // Resolved* mirror the fields MarkDeviceRequestAuthorized ultimately
        // persists onto the DeviceRequest -- see DeviceCodeStorage's doc comment.
        ResolvedUserID    string
        ResolvedUserName  string
        ResolvedUserEmail string
        VerifiedUpstream  authsession.UpstreamSubject

        // UpstreamTokens carries the tokens already exchanged with the upstream
        // IDP, ready to persist under the final session id once one exists.
        // Device-flow sessions have no session id until
        // MarkDeviceRequestAuthorized mints one at confirm time -- unlike the
        // OAuth-client authorization_code flow, where the session id is minted
        // up front and CallbackHandler can call StoreUpstreamTokens directly --
        // so the tokens must be carried here and written by
        // DeviceVerificationConfirmHandler once that id exists. ProviderID,
        // AccessToken, RefreshToken, IDToken, ExpiresAt, UpstreamSubject, and
        // ClientID are populated at exchange time; UserID and SessionExpiresAt
        // are left zero and filled in at confirm time.
        UpstreamTokens *UpstreamTokens

        // Synthetic mirrors the upstream identity's Synthetic flag: true when
        // the upstream has no userinfo/identity config and its subject was
        // synthesized rather than resolved from a real claim. Confirm-time
        // refresh-token carry-forward uses this to skip its account-linking
        // guard, which can never match a rotating synthetic subject.
        Synthetic bool

        // CreatedAt is when the pending confirmation was created.
        CreatedAt time.Time
}

// PendingDeviceConfirmationStorage provides storage operations for resolved
// device-flow logins awaiting an explicit Approve/Deny decision at the
// verification page. See PendingDeviceConfirmation's doc comment for why the
// resolved identity is addressed by an opaque token rather than round-tripped
// through the browser as editable form data.
type PendingDeviceConfirmationStorage interface {
        // StorePendingDeviceConfirmation stores a pending confirmation, keyed by
        // a fresh, unguessable token minted by the caller (e.g. rand.Text()).
        StorePendingDeviceConfirmation(ctx context.Context, token string, pending *PendingDeviceConfirmation) error

        // LoadPendingDeviceConfirmation retrieves a pending confirmation by
        // token. Returns ErrNotFound if the token does not exist, ErrExpired if
        // DefaultDeviceLoginTTL has elapsed.
        LoadPendingDeviceConfirmation(ctx context.Context, token string) (*PendingDeviceConfirmation, error)

        // DeletePendingDeviceConfirmation removes a pending confirmation.
        // Returns ErrNotFound if the token does not exist.
        DeletePendingDeviceConfirmation(ctx context.Context, token string) error
}

// AssertionJWTConsumer atomically records a validated assertion JWT as consumed.
//
// Implementations must treat (purpose, issuer, jti) as the replay key, retain it
// until exp, and return fosite.ErrJTIKnown when that key is already unexpired.
// Purpose keeps distinct assertion profiles separate, while issuer binds a JTI to
// its JWT issuer. This intentionally stays separate from Storage so only
// assertion-grant composition needs replay-consumption access.
type AssertionJWTConsumer interface {
        ConsumeAssertionJWT(ctx context.Context, purpose, issuer, jti string, exp time.Time) error
}

// SyntheticClientIDPrefix marks a client ID minted by NewSyntheticClient. A
// grant that skips client authentication (fosite's CanSkipClientAuth) has no
// real registered fosite.Client to attach to its request, but every storage
// backend's marshal/unmarshal path calls request.GetClient() unconditionally
// — so the request still needs a non-nil one. A synthetic client ID encodes
// everything needed to reconstruct that client, so unmarshaling never has to
// look it up in the client registry (it was never registered there).
const SyntheticClientIDPrefix = "synthetic:"

// NewSyntheticClient returns a public fosite.Client identified only by id,
// for a clientless grant to attach to its fosite.AccessRequester so no
// storage backend ever has to marshal a nil client. id should carry the
// SyntheticClientIDPrefix so IsSyntheticClientID recognizes it again on
// unmarshal.
func NewSyntheticClient(id string) fosite.Client {
        return &fosite.DefaultClient{ID: id, Public: true}
}

// IsSyntheticClientID reports whether id was minted by NewSyntheticClient,
// letting a storage backend reconstruct the client locally on unmarshal
// instead of looking it up in the client registry.
func IsSyntheticClientID(id string) bool {
        return strings.HasPrefix(id, SyntheticClientIDPrefix)
}

// ValidateRegisterableClientID rejects a real client registration whose ID
// collides with the SyntheticClientIDPrefix namespace. Synthetic clients are
// minted directly by NewSyntheticClient and never go through RegisterClient,
// so any caller reaching this check — DCR, static delegate-client config, or
// a future registration path — is registering a real, externally-visible
// client. Every ClientRegistry.RegisterClient implementation must call this
// before persisting: without it, a colliding ID would round-trip through
// Redis as a bare unregistered synthetic client (see IsSyntheticClientID),
// silently discarding its real secret, scopes, and grants.
func ValidateRegisterableClientID(id string) error {
        if IsSyntheticClientID(id) {
                return fmt.Errorf("%w: %q", ErrReservedClientID, id)
        }
        return nil
}

// checkConfiguredClientType refuses a non-reserved client that
// registration.MarkConfigured cannot wrap, so a stamped row is always readable
// as stamped. The reserved SPIFFE placeholder is exempt: it is never wrapped.
func checkConfiguredClientType(client fosite.Client) error {
        if isReservedPlaceholder(client) {
                return nil
        }
        if registration.IsConfigured(client) {
                return fmt.Errorf("configured client %q is already stamped; pass the unwrapped client", client.GetID())
        }
        if _, ok := registration.MarkConfigured(client); !ok {
                return fmt.Errorf("configured client %q has unsupported type %T", client.GetID(), client)
        }
        return nil
}

// sameStringSet reports whether a and b contain the same elements as sets:
// order and duplicate count don't matter, only membership. Canonicalisation
// (sort, then dedup) mirrors ScopesHash's approach so the two stay consistent.
func sameStringSet(a, b []string) bool {
        as := slices.Clone(a)
        bs := slices.Clone(b)
        sort.Strings(as)
        sort.Strings(bs)
        as = slices.Compact(as)
        bs = slices.Compact(bs)
        return slices.Equal(as, bs)
}

// clientFingerprint holds the fields ReconcileConfiguredClient compares when a
// configured client is declared over an existing row at the same ID. How much
// of it matters depends on the row's kind (see configuredClientConflict):
// reserved SPIFFE placeholders compare in full, other configured clients only
// on public. Client secrets are deliberately excluded -- an operator rotating
// a secret must still be able to reconcile. TokenEndpointAuthMethod is
// excluded for the same reason: both backends fully overwrite the stored
// record on a successful reconcile rather than merging.
//
// Deliberately NOT fosite.Client: fosite.DefaultClient's GetGrantTypes/
// GetResponseTypes substitute an interactive default when the underlying list
// is empty, so a deliberately-empty SPIFFE placeholder compared through that
// interface can mismatch itself. Each backend converts its own representation
// into this type; the comparison exists once, so memory- and Redis-backed
// deployments cannot disagree.
//
// Grant and response types, like resources and identity, only matter for
// reserved rows; for non-reserved rows only public is compared.
type clientFingerprint struct {
        scopes        []string
        audience      []string
        grantTypes    []string
        responseTypes []string
        // resources is the RFC 8707 resource allowlist, distinct from audience
        // (see resourceScopedClient in spiffe_decorator.go). Only reserved rows
        // carry it: buildStoredClient persists it for them alone.
        resources []string
        // identityFingerprint is the SPIFFE association hash. Like resources, it
        // is only populated for reserved rows.
        identityFingerprint string
        public              bool
        // reserved marks a SPIFFE static-client placeholder, as opposed to an
        // ordinary configured client such as a delegate.
        reserved bool
}

// equal reports whether f and o represent the same logical client
// configuration: equal scope, audience, grant-type, response-type and
// resource sets, identity, public/confidential class, and reserved kind.
func (f clientFingerprint) equal(o clientFingerprint) bool {
        return sameStringSet(f.scopes, o.scopes) &&
                sameStringSet(f.audience, o.audience) &&
                sameStringSet(f.grantTypes, o.grantTypes) &&
                sameStringSet(f.responseTypes, o.responseTypes) &&
                sameStringSet(f.resources, o.resources) &&
                f.identityFingerprint == o.identityFingerprint &&
                f.public == o.public &&
                f.reserved == o.reserved
}

// policyChanged reports whether the scopes or audience differ between f and
// o, for the operator-visible warning when a configured client is replaced
// with a different policy. Grant and response types are deliberately not
// compared, so a grant-type change is replaced without the WARN. Delegate
// grant types are fixed internally today, so in practice only scopes and
// audience can change.
func (f clientFingerprint) policyChanged(o clientFingerprint) bool {
        return !sameStringSet(f.scopes, o.scopes) || !sameStringSet(f.audience, o.audience)
}

// configuredClientConflict decides whether incoming may replace existing at
// the same client ID in ReconcileConfiguredClient, after the backend has
// already refused DCR-issued rows. It returns nil to replace, or an
// ErrAlreadyExists-wrapped error to refuse:
//   - a reserved placeholder and a non-reserved client never replace each
//     other, in either direction;
//   - two reserved placeholders must match in full, because replicas rolling
//     a placeholder's policy would otherwise authorize different resource
//     sets (#6473);
//   - two non-reserved clients replace each other freely, so a changed
//     scope or audience rolls out on restart, except that a
//     public/confidential flip is refused in either direction. Confidential
//     to public would let refresh tokens issued to a confidential client be
//     used without a secret; the reverse is refused to stay symmetric, since
//     a class change is a different client.
func configuredClientConflict(id string, existing, incoming clientFingerprint) error {
        switch {
        case existing.reserved != incoming.reserved:
                if existing.reserved {
                        return fmt.Errorf("%w: client %q is already registered as a different configured client "+
                                "(the stored row is a SPIFFE-reserved placeholder)", ErrAlreadyExists, id)
                }
                return fmt.Errorf("%w: client %q is already registered as a different configured client "+
                        "(the incoming client is SPIFFE-reserved and the stored one is not)", ErrAlreadyExists, id)
        case existing.reserved:
                if !existing.equal(incoming) {
                        return fmt.Errorf("%w: client %q is already registered as a different configured client "+
                                "(SPIFFE-reserved clients must match exactly)", ErrAlreadyExists, id)
                }
        case existing.public != incoming.public:
                return fmt.Errorf("%w: client %q is already registered as a different configured client "+
                        "(public and confidential clients cannot replace each other)", ErrAlreadyExists, id)
        }
        return nil
}

// warnConfiguredClientReplaced logs, after a successful replace, that a
// configured client's scopes or audience changed. During a rolling update with
// differing configs the replicas overwrite each other, so this is the
// operator's signal for that flapping. Secrets are never logged.
func warnConfiguredClientReplaced(id string, existing, incoming clientFingerprint) {
        if !existing.policyChanged(incoming) {
                return
        }
        slog.Warn("configured client replaced with a different policy",
                "client_id", id,
                "old_scopes", existing.scopes, "new_scopes", incoming.scopes,
                "old_audiences", existing.audience, "new_audiences", incoming.audience)
}

// fingerprintOfClient reads a live fosite.Client. Safe here because a
// deliberately-empty client is the live inertPlaceholderClient, whose
// overridden getters correctly return nil on both sides of the comparison.
// resources and identity are read only for a reserved placeholder, mirroring
// buildStoredClient, so both backends compare the same inputs.
func fingerprintOfClient(c fosite.Client) clientFingerprint {
        reserved := isReservedPlaceholder(c)
        var resources []string
        var identityFingerprint string
        if reserved {
                if rc, ok := c.(resourceScopedClient); ok {
                        resources = rc.Resources()
                }
                if identity, ok := c.(spiffeIdentityClient); ok {
                        identityFingerprint = identity.IdentityFingerprint()
                }
        }
        return clientFingerprint{
                scopes:              c.GetScopes(),
                audience:            c.GetAudience(),
                grantTypes:          c.GetGrantTypes(),
                responseTypes:       c.GetResponseTypes(),
                resources:           resources,
                identityFingerprint: identityFingerprint,
                public:              c.IsPublic(),
                reserved:            reserved,
        }
}

// ClientRegistry provides client registration and lookup operations.
// It embeds fosite.ClientManager for client lookup (GetClient) and adds
// RegisterClient for dynamic client registration (RFC 7591) and
// ReconcileConfiguredClient for operator-declared clients.
type ClientRegistry interface {
        // ClientManager provides client lookup (GetClient)
        fosite.ClientManager

        // RegisterClient registers a new OAuth client. Always create-only: it
        // returns ErrAlreadyExists if a client with the same ID already exists,
        // regardless of the new client's origin. Used by unauthenticated DCR
        // (RFC 7591) and any other caller that must never silently overwrite an
        // existing registration. A client counted against the DCR cap (see
        // ErrClientCapacity) is refused with a *ClientCapacityError when the cap
        // is reached; decorators must propagate it. An existing ID is reported as
        // ErrAlreadyExists before the cap is checked, on both backends.
        RegisterClient(ctx context.Context, client fosite.Client) error

        // UpsertDCRIssuedClient creates or replaces a DCR-issued client at
        // client.GetID(). Unlike RegisterClient (create-only, for the unauthenticated
        // /oauth/register endpoint), this is for callers that independently
        // re-validate the client's authoritative source on every call -- today only
        // CIMDStorageDecorator, which re-fetches and re-validates the document at
        // client.GetID() before calling this. Creates the row if absent. If a row
        // exists, replaces its data and refreshes its TTL only when the existing row
        // is itself DCR-issued; refuses with ErrAlreadyExists if the existing row is
        // NOT DCR-issued (protects a configured/SPIFFE-reconciled client from being
        // clobbered). client MUST carry registration.DCRIssued -- refuses otherwise
        // (mirrors ReconcileConfiguredClient's inverse check). A client that would
        // count against the DCR cap (DCR-issued with an ID that is not a CIMD URL)
        // is refused too: it must go through RegisterClient, which enforces the cap.
        UpsertDCRIssuedClient(ctx context.Context, client fosite.Client) error

        // ReconcileConfiguredClient applies an operator-declared (configured)
        // client: creates it if no client with that ID exists, or replaces the
        // existing record. It refuses with ErrAlreadyExists when:
        //   - the existing record is DCR-issued (registration.DCRIssued);
        //   - one of the two is a reserved SPIFFE placeholder and the other is
        //     not (they never replace each other);
        //   - both are reserved placeholders and their full fingerprints differ
        //     (scopes, audience, grant and response types, resources, identity).
        //     SPIFFE clients are served from an in-process overlay, so the strict
        //     check keeps replicas from authorizing different policies (#6473);
        //   - both are non-reserved and one is public while the other is
        //     confidential.
        // Otherwise two non-reserved clients replace each other, so a changed
        // scope or audience takes effect on restart. A replace that
        // changes scopes or audience logs a WARN. The passed client
        // must not itself carry the registration.DCRIssued marker;
        // ReconcileConfiguredClient returns an error if it does. It also returns
        // an error, before writing anything, for a non-reserved client whose
        // concrete type registration.MarkConfigured cannot wrap
        // (*fosite.DefaultClient and *fosite.DefaultOpenIDConnectClient are
        // supported), because the stamp could not be read back.
        //
        // Rows written here belong to ToolHive configuration. GetClient returns
        // them stamped (registration.IsConfigured), so a caller can hide a row
        // the running server's configuration no longer declares. Embedders that
        // register their own clients use RegisterClient, which never stamps. The
        // stamp is not part of the fingerprint comparison above, and a stamped
        // row is never treated as DCR-issued, whatever its TTL.
        ReconcileConfiguredClient(ctx context.Context, client fosite.Client) error

        // RenewClientTTL extends the registration TTL of a DCR-issued client (public or
        // confidential, gated on the registration.DCRIssued marker) so an actively-used
        // client does not expire mid-lifecycle and is not forced to re-register. Call it on a
        // proven-use signal (e.g. a successful token exchange), NOT on an unauthenticated
        // client read such as the /oauth/authorize lookup. Pre-provisioned/static clients
        // carry no marker and this is a no-op for them on both backends. Both backends
        // set the client's expiry to DefaultDCRClientTTL from now, which also raises a
        // newly registered client off its shorter DefaultUnusedDCRClientTTL. A client
        // that no longer exists is left absent. A renewal failure is non-fatal to the
        // caller's primary operation.
        // CIMDStorageDecorator skips URL-shaped (CIMD) client ids: each successful
        // fetch already resets the row's TTL, and not renewing keeps that TTL as the
        // age bound for its token-endpoint persisted-row fallback.
        RenewClientTTL(ctx context.Context, client fosite.Client) error
}

// UpstreamTokenRowID is an opaque, process-local coordination key for one
// logical upstream token row. Callers must not persist, log, or expose it.
type UpstreamTokenRowID string

// upstreamTokenRowID returns a SHA-256 digest of the supplied row-key components.
// Each component is length-prefixed so the encoding is self-delimiting: distinct
// inputs cannot share a canonical representation when a component contains the
// delimiter. Hashing keeps the session ID out of a value that is easy to log.
func upstreamTokenRowID(components ...string) UpstreamTokenRowID {
        h := sha256.New()
        for _, component := range components {
                // hash.Hash.Write never returns an error.
                _, _ = fmt.Fprintf(h, "%d:%s", len(component), component)
        }
        return UpstreamTokenRowID(hex.EncodeToString(h.Sum(nil)))
}

// UpstreamTokenStorage provides storage for tokens obtained from upstream identity providers.
// The auth server exposes this interface via Server.UpstreamTokenStorage() for use by
// middleware that needs to retrieve upstream tokens (e.g., token swap middleware that
// replaces JWT auth with upstream IDP tokens for backend requests).
//
// Tokens are keyed primarily by (sessionID, providerName) to support multiple upstream
// providers per session. Each provider's tokens are stored and retrieved independently.
// A secondary lookup by (userID, providerID) is exposed via GetLatestUpstreamTokensForUser;
// see that method for usage and security contract.
type UpstreamTokenStorage interface {
        // ResolveUpstreamTokenRowID returns an opaque non-empty identity for the
        // (sessionID, providerName) row. Equal physical rows must resolve to equal
        // IDs; implementations must NOT alias rows that are not physically the
        // same row (e.g. two distinct sessionIDs), because the singleflight leader
        // reads and writes under its own sessionID, and aliasing distinct rows
        // would cross credentials between sessions. A backend whose storage key is
        // derived from (sessionID, providerName) may legitimately map two
        // different (sessionID, providerName) pairs to the same ID when they are
        // the same physical row under that key scheme.
        //
        // Resolution must be local and side-effect-free (no I/O): it is called
        // before the singleflight joins, so any I/O here gives every concurrent
        // caller its own round-trip and partially defeats the deduplication. This
        // ID is solely for process-local refresh coordination: callers must not
        // persist, log, or expose it.
        ResolveUpstreamTokenRowID(ctx context.Context, sessionID, providerName string) (UpstreamTokenRowID, error)

        // StoreUpstreamTokens stores the upstream IDP tokens for a session and provider.
        // The providerName identifies which upstream provider these tokens belong to.
        //
        // This is an unconditional overwrite: it does not check what is currently
        // stored. A refresher redeeming a single-use, rotating refresh token MUST
        // use CompareAndSwapUpstreamTokens instead, so that a redemption raced by
        // another process cannot silently clobber a winning write with a stale
        // one. StoreUpstreamTokens remains correct for every other writer (initial
        // login, the OAuth callback), which has no prior row to race against.
        StoreUpstreamTokens(ctx context.Context, sessionID, providerName string, tokens *UpstreamTokens) error

        // CompareAndSwapUpstreamTokens stores tokens for (sessionID, providerName)
        // only if the refresh token currently stored there equals
        // expectedRefreshToken. Pass "" to match a row with no refresh token: no
        // row exists yet, an existing row's RefreshToken field is itself empty, or
        // an existing row was stored as an explicit no-token placeholder (a nil
        // UpstreamTokens, persisted as the backend's null marker). Returns
        // ErrConcurrentRefresh, and leaves the stored row untouched, when the
        // comparison fails — whether because another writer moved the row past
        // expectedRefreshToken, or because the row no longer exists at all (e.g.
        // deleted by logout or evicted by TTL between the caller's read and this
        // write). Callers that need to distinguish those two cases must re-read
        // the row afterward: an unexpired row means a genuine race was lost to
        // another writer; ErrNotFound on that re-read means there was no race to
        // lose — the row was simply gone.
        //
        // # Purpose
        //
        // This is the coordination primitive that makes the STORED row
        // deterministic under concurrent writers across MULTIPLE PROCESSES
        // sharing the same storage backend (e.g. several horizontally-scaled
        // replicas of an application embedding this auth server, behind the same
        // Redis): whichever writer's expected value still matches when its write
        // lands wins, and every losing writer fails instead of silently
        // clobbering the winner. ResolveUpstreamTokenRowID's singleflight dedup is
        // process-local only; it prevents redundant redemptions within one
        // process but cannot stop two different processes from redeeming the same
        // refresh token at the same time. Read the row, redeem it with the
        // upstream provider, then write with expectedRefreshToken set to the
        // RefreshToken value that was actually redeemed.
        //
        // This orders writes to storage; it is NOT by itself a guarantee that
        // concurrent redemption is safe at the upstream provider. Both processes
        // still call the provider before either one's write lands here, so for a
        // provider enforcing strict single-use rotation the provider may see two
        // redemptions of the same refresh token regardless of which process wins
        // the write below, and may revoke the grant. This primitive is fully
        // sufficient only where the provider tolerates a grace/leeway window in
        // which more than one redeemed child stays valid; otherwise closing the
        // gap requires serializing the redemption itself (a lock around the whole
        // read-redeem-write sequence), which is a separate mechanism this method
        // does not provide.
        //
        // Implementations must perform the comparison and the write atomically
        // with respect to any other writer of the same row (e.g. a Lua script on
        // Redis, or a mutex-guarded read-modify-write in memory).
        CompareAndSwapUpstreamTokens(
                ctx context.Context, sessionID, providerName, expectedRefreshToken string, tokens *UpstreamTokens,
        ) error

        // GetUpstreamTokens retrieves the upstream IDP tokens for a session and provider.
        // Returns ErrNotFound if the session/provider combination does not exist.
        // Returns ErrExpired if the tokens have expired. When ErrExpired is returned,
        // the token data (including refresh token) is also returned to allow callers
        // to attempt a token refresh.
        // Returns ErrInvalidBinding if binding validation fails.
        GetUpstreamTokens(ctx context.Context, sessionID, providerName string) (*UpstreamTokens, error)

        // GetAllUpstreamTokens retrieves all upstream IDP tokens for a session across all providers.
        // Returns a map of providerName -> tokens. Returns an empty map (not error) for unknown sessions.
        // Includes expired tokens (no expiry filtering at bulk-read level).
        GetAllUpstreamTokens(ctx context.Context, sessionID string) (map[string]*UpstreamTokens, error)

        // DeleteUpstreamTokens removes all upstream IDP tokens for a session (all providers).
        // Returns ErrNotFound if the session does not exist.
        DeleteUpstreamTokens(ctx context.Context, sessionID string) error

        // DeleteUpstreamTokensForProvider removes tokens for a single (sessionID, providerName),
        // leaving sibling providers' rows intact. Deleting an absent row is NOT an error (nil).
        // Used as best-effort cleanup of a refresh token rotated at the IdP but not persisted.
        DeleteUpstreamTokensForProvider(ctx context.Context, sessionID, providerName string) error

        // GetLatestUpstreamTokensForUser returns the most recently stored upstream tokens
        // for (userID, providerID) across any session. The "latest" winner is determined
        // by treating non-expiring rows (zero ExpiresAt — providers like Slack and GitHub
        // OAuth Apps that genuinely never expire) as the strongest candidate, falling
        // back to ExpiresAt descending among finite-expiry rows. This aligns with the
        // rest of the package treating zero ExpiresAt as "alive forever" (see IsExpired,
        // cleanupExpired, marshalUpstreamTokensWithTTL). Both backends use the same
        // rule so the carry-forward decision is consistent regardless of deployment
        // shape.
        //
        // This is a secondary lookup that intentionally bypasses the primary
        // (sessionID, providerName) key. It is used by the OAuth callback to preserve a
        // previously-issued refresh token when the upstream IdP omits refresh_token on
        // re-authorization (e.g. Google without prompt=consent). When a sessionID is
        // available, callers should use GetUpstreamTokens instead. This method mirrors
        // the preservation pattern in upstreamTokenRefresher.RefreshAndStore.
        //
        // # Liveness
        //
        // The returned tokens may be expired. Callers MUST NOT assume liveness; they
        // must handle the expired case (typically by reading only the RefreshToken,
        // which remains usable past the access token's expiry). This method does NOT
        // return ErrExpired and the implementation MUST NOT filter expired rows.
        //
        // # Cross-identity safety
        //
        // The returned record is NOT filtered by upstream subject. The same internal
        // UserID can in principle map to multiple upstream subjects on the same provider
        // (account-linking edge cases or data-integrity issues). Callers MUST verify
        // that prior.UpstreamSubject == currentProviderSubject before reusing any
        // credential from the returned record. The OAuth callback applies this guard;
        // other future callers must do the same.
        //
        // Returns ErrNotFound if no record exists for the (userID, providerID) pair.
        GetLatestUpstreamTokensForUser(ctx context.Context, userID, providerID string) (*UpstreamTokens, error)
}

// UpstreamTokenRefresher can refresh expired upstream tokens using their stored refresh token.
// This is implemented by the auth server and used by the upstreamswap middleware to
// transparently refresh tokens without forcing re-authentication.
type UpstreamTokenRefresher interface {
        // RefreshAndStore refreshes the upstream tokens for the given session using
        // the stored refresh token, stores the new tokens, and returns them.
        // Returns an error if the refresh token is empty, revoked, or the refresh fails.
        RefreshAndStore(ctx context.Context, sessionID string, expired *UpstreamTokens) (*UpstreamTokens, error)
}

// UserStorage provides user and provider identity management operations.
// This interface supports multi-IDP scenarios where a single user can authenticate
// via multiple upstream identity providers (e.g., Google and GitHub).
//
// The User type represents the internal ToolHive identity. Its ID becomes the "sub"
// claim in issued JWTs, providing a stable identity across multiple providers.
//
// ProviderIdentity links a user to a specific upstream provider. The
// (ProviderID, ProviderSubject) pair uniquely identifies an upstream identity.
//
// # Account Linking Security
//
// When implementing account linking (one User with multiple ProviderIdentities),
// callers MUST verify user consent before linking. See OAuth 2.0 Security BCP.
//
// TODO(auth): When implementing double-hop auth (Company IDP -> External IDP),
// add the following to this interface:
//   - DeleteProviderIdentity(providerID, subject) - unlink specific provider
//   - Add Primary field to ProviderIdentity to distinguish Company IDP from linked External IDPs
//   - Add ConsentRecord tracking for external provider linking
type UserStorage interface {
        // CreateUser creates a new user account.
        // Returns ErrAlreadyExists if a user with the same ID already exists.
        // A deployment that provisions users out-of-band (e.g. via SCIM) and never
        // auto-creates them here may return ErrUserNotProvisioned to deny the login
        // instead; the caller (the authorization callback) maps that specifically to
        // an OAuth access_denied response rather than a server error.
        CreateUser(ctx context.Context, user *User) error

        // GetUser retrieves a user by their internal ID.
        // Returns ErrNotFound if the user does not exist.
        GetUser(ctx context.Context, id string) (*User, error)

        // DeleteUser removes a user account.
        // Returns ErrNotFound if the user does not exist.
        DeleteUser(ctx context.Context, id string) error

        // CreateProviderIdentity links a provider identity to a user.
        // For account linking scenarios, caller MUST verify user owns the target User
        // (typically via active authenticated session) before linking a new provider.
        // Returns ErrAlreadyExists if this provider identity is already linked.
        CreateProviderIdentity(ctx context.Context, identity *ProviderIdentity) error

        // GetProviderIdentity retrieves a provider identity by provider ID and subject.
        // This is the primary lookup path during authentication callbacks.
        // Returns ErrNotFound if the identity does not exist.
        GetProviderIdentity(ctx context.Context, providerID, providerSubject string) (*ProviderIdentity, error)

        // UpdateProviderIdentityLastUsed updates the LastUsedAt timestamp for a provider identity.
        // This should be called after each successful authentication via this identity.
        // The timestamp supports OIDC auth_time claim when clients use max_age parameter.
        // Returns ErrNotFound if the identity does not exist.
        UpdateProviderIdentityLastUsed(ctx context.Context, providerID, providerSubject string, lastUsedAt time.Time) error

        // GetUserProviderIdentities returns all provider identities linked to a user.
        // This enables queries like "when did this user last authenticate via any provider"
        // which is needed for OIDC max_age enforcement.
        // Returns an empty slice (not error) if the user exists but has no linked identities.
        // Returns ErrNotFound if the user does not exist.
        GetUserProviderIdentities(ctx context.Context, userID string) ([]*ProviderIdentity, error)
}

// Storage combines fosite storage interfaces with ToolHive-specific storage for
// upstream IDP tokens, pending authorization requests, and client registration.
// The auth server requires a Storage implementation; use NewMemoryStorage() for
// single-instance deployments or NewRedisStorage() for distributed deployments.
//
// # Fosite Interface Segregation
//
// Fosite splits storage into separate interfaces (AuthorizeCodeStorage, AccessTokenStorage,
// RefreshTokenStorage, PKCERequestStorage) following the Interface Segregation Principle.
// This enables:
//   - Feature composition: Only enable OAuth features you need
//   - Testing isolation: Mock specific interfaces for focused tests
//   - Clear contracts: Each interface documents its requirements
//
// # Key Design Patterns
//
// All token storage methods store fosite.Requester (not just token values) because token
// validation requires the full authorization context (client, scopes, session).
//
// Methods use two key types:
//   - Signature: Cryptographic token identifier for token-specific operations
//   - Request ID: Grant identifier for finding all tokens from one authorization
//
// See doc.go for comprehensive documentation of fosite's storage design.
type Storage interface {
        // Embed segregated interfaces for IDP tokens, pending authorizations, client registry,
        // and user management for multi-IDP support.
        //
        // DCRCredentialStore is intentionally NOT embedded here: doing so would
        // promote GetDCRCredentials / StoreDCRCredentialsIfAbsent onto every consumer of
        // storage.Storage (handlers, server, registration, etc.), broadening the
        // surface that can read raw client_secret / registration_access_token even
        // when those consumers have no DCR responsibility. Code that legitimately
        // needs DCR access (the runner and authserver constructors) obtains it
        // via an explicit `stor.(DCRCredentialStore)` type assertion at the
        // boundary; the per-backend `var _ DCRCredentialStore = (*MemoryStorage)(nil)`
        // and `var _ DCRCredentialStore = (*RedisStorage)(nil)` checks in
        // memory.go / redis.go provide the compile-time guarantee that production
        // backends satisfy the interface, so the runtime assertion is provably
        // safe at the boundary while keeping the wider Storage surface narrow.
        //
        // DeviceCodeStorage follows the same segregated pattern for the same
        // reason: it is an optional capability (RFC 8628 device flow is off by
        // default), not every Storage implementation needs to support it, and
        // embedding it here would force every future implementation to add all 6
        // methods even when it never enables device flow. Callers that need it
        // (handlers.NewHandler, the deviceflow factory in server_impl.go) obtain
        // it via `stor.(DeviceCodeStorage)` at construction time; the same
        // per-backend `var _ DeviceCodeStorage = (*MemoryStorage)(nil)` /
        // `(*RedisStorage)(nil)` checks provide the compile-time guarantee.
        UpstreamTokenStorage
        PendingAuthorizationStorage
        PendingDeviceLoginStorage
        PendingDeviceConfirmationStorage
        ClientRegistry
        UserStorage

        // AuthorizeCodeStorage provides authorization code storage for the Authorization Code
        // Grant (RFC 6749 Section 4.1). Authorization codes are one-time-use and short-lived.
        // CreateAuthorizeCodeSession stores by code, GetAuthorizeCodeSession retrieves by code,
        // InvalidateAuthorizeCodeSession marks as used (subsequent Gets return ErrInvalidatedAuthorizeCode).
        oauth2.AuthorizeCodeStorage

        // AccessTokenStorage provides access token session storage. Methods use "signature"
        // (derived from token value) as the key for O(1) lookup when validating tokens.
        // The stored fosite.Requester contains the full authorization context including
        // the Session with expiration times via session.GetExpiresAt(fosite.AccessToken).
        oauth2.AccessTokenStorage

        // RefreshTokenStorage provides refresh token session storage. CreateRefreshTokenSession
        // accepts an accessSignature to link refresh tokens to their access tokens for rotation.
        // RotateRefreshToken uses requestID to invalidate both the refresh token and all
        // related access tokens from the same authorization grant.
        oauth2.RefreshTokenStorage

        // TokenRevocationStorage provides token revocation per RFC 7009. RevokeAccessToken
        // and RevokeRefreshToken take requestID (not signature) because RFC 7009 requires
        // revoking a refresh token SHOULD also invalidate associated access tokens, which
        // requires finding all tokens by their common grant identifier.
        oauth2.TokenRevocationStorage

        // PKCERequestStorage provides PKCE challenge/verifier storage (RFC 7636).
        // Stores the code_challenge during authorization, validates code_verifier during
        // token exchange. Keyed by the same code/signature as the authorization code.
        pkce.PKCERequestStorage

        // Health checks connectivity to the storage backend.
        // Returns nil if the storage is healthy and reachable.
        Health(ctx context.Context) error

        // Close releases any resources held by the storage implementation.
        // This should be called when the storage is no longer needed.
        Close() error
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package storage

// Unwrap returns the innermost storage backend by recursively peeling storage
// decorators that expose Unwrap. It is used at construction boundaries that
// require the durable backend, such as static-client collision checks.
func Unwrap(stor Storage) Storage {
        for {
                unwrapper, ok := stor.(interface{ Unwrap() Storage })
                if !ok {
                        return stor
                }
                stor = unwrapper.Unwrap()
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package upstream

import (
        "encoding/base64"
        "errors"
        "fmt"
        "log/slog"
        "strings"
        "sync"

        "github.com/tidwall/gjson"
)

// partialIdentity holds identity fields extracted from a token response body.
// It is used internally to pass the extracted subject, name, and email
// between extractIdentityFromTokenResponse and the provider layer.
type partialIdentity struct {
        Subject string
        Name    string
        Email   string
}

// IdentityFromTokenConfig is the runtime configuration for extracting user
// identity directly from an OAuth2 token endpoint response body.
//
// Each path is a gjson dot-notation path (e.g. "username" or
// "associated_user.id") into the raw JSON body returned by the token
// endpoint. Path semantics, trust-model warnings, and uniqueness notes are
// documented on the corresponding CRD type
// (cmd/thv-operator/api/v1beta1.IdentityFromTokenConfig).
type IdentityFromTokenConfig struct {
        // SubjectPath is the gjson path to the unique user identifier (required).
        SubjectPath string

        // NamePath is the gjson path to the user's display name (optional).
        // Leave empty to skip name extraction.
        NamePath string

        // EmailPath is the gjson path to the user's email address (optional).
        // Leave empty to skip email extraction.
        EmailPath string
}

// registerOnce ensures gjson modifiers are registered exactly once, even
// when RegisterModifiers is called concurrently from multiple goroutines
// (e.g. parallel auth-server constructors in tests or during thundering-herd
// restarts). gjson.AddModifier writes to a global map without internal
// synchronization, so concurrent calls race without this guard.
var registerOnce sync.Once

// RegisterModifiers registers the gjson custom modifiers used by this
// package's path-based identity extractors. Call once during application
// or test wire-up before invoking any extractor that consumes a
// modifier-bearing path. Repeated calls are safe — the registration is
// protected by a sync.Once and executes at most once per process.
//
// Modifiers registered:
//   - @upstreamjwt: see upstreamJWTModifier.
func RegisterModifiers() {
        registerOnce.Do(func() {
                gjson.AddModifier("upstreamjwt", upstreamJWTModifier)
        })
}

// extractIdentityFromTokenResponse extracts user identity fields from a raw
// OAuth2 token endpoint response body using the paths in cfg.
//
// SubjectPath must resolve to a string or number value; objects, arrays, null,
// and missing paths are rejected with ErrIdentityResolutionFailed. NamePath
// and EmailPath are optional: type mismatches or missing paths produce a
// slog.Warn and leave the respective field empty. Empty NamePath/EmailPath in
// cfg means "do not extract" and are skipped silently.
func extractIdentityFromTokenResponse(body []byte, cfg *IdentityFromTokenConfig) (partialIdentity, error) {
        if cfg == nil {
                return partialIdentity{}, errors.New("identity extraction config is required")
        }

        subjectResult := gjson.GetBytes(body, cfg.SubjectPath)
        if err := validateIdentityField(subjectResult); err != nil {
                return partialIdentity{}, fmt.Errorf("%w: subjectPath %q %s", ErrIdentityResolutionFailed, cfg.SubjectPath, err.Error())
        }

        name := extractOptionalField(body, cfg.NamePath, "namePath")
        email := extractOptionalField(body, cfg.EmailPath, "emailPath")

        return partialIdentity{
                Subject: scalarToString(subjectResult),
                Name:    name,
                Email:   email,
        }, nil
}

// scalarToString returns the string representation of a gjson scalar value.
// For Number, it returns the raw JSON token rather than gjson.Result.String(),
// which formats via float64 and would lose precision for integer IDs larger
// than 2^53 (e.g., some upstream providers return 64-bit numeric subjects).
// For String, gjson.Result.String() correctly strips the surrounding quotes.
// The caller must already have validated the type.
func scalarToString(r gjson.Result) string {
        if r.Type == gjson.Number {
                return r.Raw
        }
        return r.String()
}

// validateIdentityField checks that a gjson result is a non-empty scalar
// (string or number). Returns a descriptive error on failure.
func validateIdentityField(result gjson.Result) error {
        if !result.Exists() {
                return errors.New("path not found in token response")
        }
        switch result.Type {
        case gjson.String:
                if result.String() == "" {
                        return errors.New("resolved to an empty string")
                }
                return nil
        case gjson.Number:
                return nil
        case gjson.JSON:
                return errors.New("resolved to an object or array, expected a scalar")
        case gjson.Null, gjson.False, gjson.True:
                return errors.New("resolved to null or unsupported type")
        }
        // Unreachable: all gjson.Type cases are handled above.
        return errors.New("unrecognised gjson result type")
}

// extractOptionalField extracts an optional identity field from the token body.
// Returns an empty string if the path is not configured, missing, or has an
// unexpected type (with a slog.Warn for unexpected types).
func extractOptionalField(body []byte, path, fieldName string) string {
        if path == "" {
                return ""
        }
        result := gjson.GetBytes(body, path)
        if !result.Exists() {
                slog.Warn("optional identity field not found in token response", "field", fieldName, "path", path)
                return ""
        }
        switch result.Type {
        case gjson.String, gjson.Number:
                return scalarToString(result)
        case gjson.JSON, gjson.Null, gjson.False, gjson.True:
                slog.Warn("optional identity field has unexpected type, skipping", "field", fieldName, "path", path)
                return ""
        }
        // Unreachable: all gjson.Type cases are handled above.
        return ""
}

// upstreamJWTModifier is a gjson modifier that decodes the payload of a
// JWT-shaped string value and returns it as JSON, enabling further gjson
// path drilling (e.g. "access_token|@upstreamjwt|sub").
//
// Trust model: NO signature verification. Use only for JWTs received over
// a TLS-authenticated channel directly from the upstream AS's token
// endpoint, where the channel itself provides integrity. For JWTs that
// have transited an untrusted hop, configure the upstream as OIDC and
// use the existing signed-ID-token path instead.
//
// Failure modes (all return ""):
//   - input is not a JSON string
//   - input does not contain exactly three dot-separated parts
//   - the second part is not valid base64url
//
// Returning "" causes the next pipe stage to resolve to gjson.Null, which
// the caller's validateIdentityField rejects as "path not found".
func upstreamJWTModifier(jsonValue, _ string) string {
        token := gjson.Parse(jsonValue).String()
        if token == "" {
                return ""
        }
        parts := strings.Split(token, ".")
        if len(parts) != 3 {
                return ""
        }
        payload, err := base64.RawURLEncoding.DecodeString(parts[1])
        if err != nil {
                return ""
        }
        return string(payload)
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:generate mockgen -destination=mocks/mock_provider.go -package=mocks -source=oauth2.go OAuth2Provider

package upstream

import (
        "context"
        "crypto/sha256"
        "encoding/hex"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "maps"
        "net/http"
        "net/url"
        "strings"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth/oautherr"
        "github.com/stacklok/toolhive/pkg/authserver/oauthparams"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

const (
        // ProviderTypeOAuth2 is for pure OAuth 2.0 providers with explicit endpoints.
        ProviderTypeOAuth2 ProviderType = "oauth2"

        // maxResponseSize is the maximum allowed UserInfo response size (1MB).
        // This prevents memory exhaustion from malicious or malformed responses.
        maxResponseSize = 1 << 20
)

// AuthorizationOption configures authorization URL generation.
type AuthorizationOption func(*authorizationOptions)

type authorizationOptions struct {
        additionalParams map[string]string
}

// WithAdditionalParams adds custom parameters to the authorization URL.
func WithAdditionalParams(params map[string]string) AuthorizationOption {
        return func(o *authorizationOptions) {
                if o.additionalParams == nil {
                        o.additionalParams = make(map[string]string)
                }
                maps.Copy(o.additionalParams, params)
        }
}

// OAuth2Provider handles communication with an upstream Identity Provider.
// This is the base interface for all provider types.
type OAuth2Provider interface {
        // Type returns the provider type.
        Type() ProviderType

        // AuthorizationURL builds the URL to redirect the user to the upstream IDP.
        // state: our internal state to correlate callback
        // codeChallenge: PKCE challenge to send to upstream (if supported)
        // opts: optional configuration such as nonce or additional parameters
        AuthorizationURL(state, codeChallenge string, opts ...AuthorizationOption) (string, error)

        // ExchangeCodeForIdentity exchanges an authorization code for tokens and resolves
        // the user's identity in a single atomic operation. This ensures that OIDC nonce
        // validation (replay protection) cannot be accidentally skipped.
        // For OIDC providers, the nonce is validated against the ID token.
        // For pure OAuth2 providers, identity is resolved via the UserInfo endpoint
        // and the nonce parameter is ignored.
        ExchangeCodeForIdentity(ctx context.Context, code, codeVerifier, nonce string) (*Identity, error)

        // RefreshTokens refreshes the upstream IDP tokens.
        // expectedSubject is the original sub claim; OIDC providers validate it per
        // Section 12.2 when the response includes a new ID token. Pure OAuth2 providers
        // ignore it.
        RefreshTokens(ctx context.Context, refreshToken, expectedSubject string) (*Tokens, error)
}

// CommonOAuthConfig contains fields shared by all OAuth provider types.
// This provides compile-time type safety by separating OIDC and OAuth2 configuration.
type CommonOAuthConfig struct {
        // ClientID is the OAuth client ID registered with the upstream IDP.
        ClientID string `json:"client_id" yaml:"client_id"`

        // ClientSecret is the OAuth client secret registered with the upstream IDP.
        // Optional for public clients (RFC 6749 Section 2.1) which authenticate using
        // PKCE instead of a client secret. Required for confidential clients.
        //nolint:gosec // G117: field legitimately holds sensitive data
        ClientSecret string `json:"client_secret,omitempty" yaml:"client_secret,omitempty"`

        // Scopes are the OAuth scopes to request from the upstream IDP.
        Scopes []string `json:"scopes,omitempty" yaml:"scopes,omitempty"`

        // RedirectURI is the callback URL where the upstream IDP will redirect
        // after authentication.
        RedirectURI string `json:"redirect_uri" yaml:"redirect_uri"`

        // AdditionalAuthorizationParams are extra query parameters to include in the
        // authorization URL. This is useful for providers that require custom parameters
        // such as Google's access_type=offline for obtaining refresh tokens.
        // Framework-managed parameters (response_type, client_id, redirect_uri, scope,
        // state, code_challenge, code_challenge_method, nonce) are not allowed here
        // and will be rejected during validation.
        //nolint:lll // field tags require full JSON+YAML names
        AdditionalAuthorizationParams map[string]string `json:"additional_authorization_params,omitempty" yaml:"additional_authorization_params,omitempty"`

        // AdditionalTokenParams are extra form-body parameters to include in
        // token requests (authorization code exchange and refresh) sent to the
        // upstream IDP's token endpoint. This is useful for providers that
        // enforce RFC 8707 resource indicators on token requests, where the
        // resource parameter must accompany the code exchange and refresh, not
        // only the authorization request.
        // Framework-managed parameters (grant_type, code, redirect_uri, client_id,
        // client_secret, code_verifier, refresh_token, scope) are not allowed here
        // and will be rejected during validation.
        //nolint:lll // field tags require full JSON+YAML names
        AdditionalTokenParams map[string]string `json:"additional_token_params,omitempty" yaml:"additional_token_params,omitempty"`

        // OmitScopeOnRefresh stops RefreshTokens from sending the `scope` parameter
        // on refresh_token grant requests. Scopes are still sent on the
        // authorization request. Defaults to false, which sends Scopes on refresh.
        // Set it for providers that reject `scope` on refresh (e.g. Salesforce
        // returns "invalid_request - scope parameter not supported"); per
        // RFC 6749 §6 the AS then treats the refresh as covering the originally
        // granted scope.
        OmitScopeOnRefresh bool `json:"omit_scope_on_refresh,omitempty" yaml:"omit_scope_on_refresh,omitempty"`
}

// ValidateWithInsecure validates CommonOAuthConfig, allowing http:// redirect URIs for
// non-loopback hosts when insecureAllowHTTP is true.
func (c *CommonOAuthConfig) ValidateWithInsecure(insecureAllowHTTP bool) error {
        if c.ClientID == "" {
                return errors.New("client_id is required")
        }
        if c.RedirectURI == "" {
                return errors.New("redirect_uri is required")
        }
        if err := oauthparams.Validate(c.AdditionalAuthorizationParams); err != nil {
                return err
        }
        if err := oauthparams.ValidateTokenParams(c.AdditionalTokenParams); err != nil {
                return err
        }
        if insecureAllowHTTP {
                return oauthproto.ValidateRedirectURI(c.RedirectURI, oauthproto.RedirectURIPolicyAllowHTTP)
        }
        return validateRedirectURI(c.RedirectURI)
}

// OAuth2Config contains configuration for pure OAuth 2.0 providers without OIDC discovery.
type OAuth2Config struct {
        CommonOAuthConfig `yaml:",inline"`

        // AuthorizationEndpoint is the URL for the OAuth authorization endpoint.
        AuthorizationEndpoint string `json:"authorization_endpoint" yaml:"authorization_endpoint"`

        // TokenEndpoint is the URL for the OAuth token endpoint.
        TokenEndpoint string `json:"token_endpoint" yaml:"token_endpoint"`

        // TokenEndpointAuthMethod is the RFC 7591 client authentication method used
        // at the token endpoint; see authStyleFromMethod for the mapping to
        // oauth2.AuthStyle and the rationale.
        //
        // When empty, the caller's provider-specific default applies. Static
        // OAuth2UpstreamRunConfig clients default to the historical POST-body method
        // regardless of whether a secret is configured; DCR clients receive the
        // negotiated method via applyResolutionToOAuth2Config.
        //nolint:lll // field tags require full JSON+YAML names
        TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty" yaml:"token_endpoint_auth_method,omitempty"`

        // UserInfo contains configuration for fetching user information (optional).
        // When nil, the provider does not support UserInfo fetching.
        UserInfo *UserInfoConfig `json:"userinfo,omitempty" yaml:"userinfo,omitempty"`

        // TokenResponseMapping configures custom field extraction from non-standard token responses.
        // When set, the provider performs the token exchange HTTP call directly (bypassing
        // golang.org/x/oauth2) and extracts fields using gjson dot-notation paths.
        // When nil, standard OAuth 2.0 token response parsing is used.
        // See also: IdentityFromToken for extracting user identity from the same response.
        TokenResponseMapping *TokenResponseMapping `json:"token_response_mapping,omitempty" yaml:"token_response_mapping,omitempty"`

        // IdentityFromToken extracts user identity from the token-endpoint response
        // body when the upstream provider includes identity claims there (e.g.,
        // Snowflake's `username`, Slack's `authed_user.id`). When set, the embedded
        // auth server skips the userinfo HTTP call entirely. See the CRD type
        // (cmd/thv-operator/api/v1beta1.IdentityFromTokenConfig) for the
        // authoritative trust-model and uniqueness documentation.
        IdentityFromToken *IdentityFromTokenConfig `json:"identity_from_token,omitempty" yaml:"identity_from_token,omitempty"`

        // CAFilePath is the path to a PEM CA bundle added to the system roots.
        CAFilePath string `json:"ca_file_path,omitempty" yaml:"ca_file_path,omitempty"`

        // InsecureAllowHTTP permits http:// authorization_endpoint and token_endpoint
        // URLs for non-localhost hosts. Set only for trusted in-cluster deployments.
        InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"`

        // AllowPrivateIPs permits the upstream provider's HTTP client to connect to
        // private IP ranges (RFC-1918, link-local). Use only when the upstream is
        // hosted inside the same cluster and has no public endpoint. HTTP-scheme
        // restrictions are unchanged — HTTPS is still required for non-localhost hosts.
        // Defaults to false.
        AllowPrivateIPs bool `json:"allow_private_ips,omitempty" yaml:"allow_private_ips,omitempty"`
}

// TokenResponseMapping configures extraction of token fields from non-standard
// OAuth token endpoint responses using gjson dot-notation paths.
type TokenResponseMapping struct {
        // AccessTokenPath is the gjson path to the access token (required).
        AccessTokenPath string

        // ScopePath is the gjson path to the scope. Defaults to "scope".
        ScopePath string

        // RefreshTokenPath is the gjson path to the refresh token. Defaults to "refresh_token".
        RefreshTokenPath string

        // ExpiresInPath is the gjson path to the expires_in value. Defaults to "expires_in".
        ExpiresInPath string
}

// Validate checks that OAuth2Config has all required fields, respecting
// c.InsecureAllowHTTP when set.
func (c *OAuth2Config) Validate() error {
        return c.ValidateWithInsecure(c.InsecureAllowHTTP)
}

// ValidateWithInsecure is like Validate but allows http:// endpoints when
// insecureAllowHTTP is true (for trusted in-cluster deployments).
func (c *OAuth2Config) ValidateWithInsecure(insecureAllowHTTP bool) error {
        if c.AuthorizationEndpoint == "" {
                return errors.New("authorization_endpoint is required for OAuth2 providers")
        }
        if err := networking.ValidateEndpointURLWithInsecure(c.AuthorizationEndpoint, insecureAllowHTTP); err != nil {
                return fmt.Errorf("invalid authorization_endpoint: %w", err)
        }
        if c.TokenEndpoint == "" {
                return errors.New("token_endpoint is required for OAuth2 providers")
        }
        if err := networking.ValidateEndpointURLWithInsecure(c.TokenEndpoint, insecureAllowHTTP); err != nil {
                return fmt.Errorf("invalid token_endpoint: %w", err)
        }
        if c.UserInfo != nil {
                if err := c.UserInfo.Validate(); err != nil {
                        return fmt.Errorf("invalid userinfo config: %w", err)
                }
        }
        if c.TokenResponseMapping != nil {
                if c.TokenResponseMapping.AccessTokenPath == "" {
                        return errors.New("token_response_mapping.access_token_path is required when token_response_mapping is set")
                }
        }
        if c.IdentityFromToken != nil {
                if c.IdentityFromToken.SubjectPath == "" {
                        return errors.New("identity_from_token.subject_path is required when identity_from_token is set")
                }
        }
        return c.CommonOAuthConfig.ValidateWithInsecure(insecureAllowHTTP)
}

// validateRedirectURI validates an OAuth redirect URI per RFC 6749 and RFC 8252.
// This is our own callback URL where upstream IDPs redirect back to us. The upstream
// IDP validates this against their registered redirect URIs, so we only do basic checks.
func validateRedirectURI(uri string) error {
        return oauthproto.ValidateRedirectURI(uri, oauthproto.RedirectURIPolicyStrict)
}

// convertOAuth2Token converts an oauth2.Token to our Tokens type.
// It extracts id_token from token extras and validates the response.
func convertOAuth2Token(token *oauth2.Token) (*Tokens, error) {
        if token.AccessToken == "" {
                return nil, errors.New("token response missing access_token")
        }

        // Validate token_type per RFC 6749 Section 5.1
        // The comparison is case-insensitive per the spec
        if !strings.EqualFold(token.TokenType, "bearer") {
                return nil, fmt.Errorf("unexpected token_type: expected \"Bearer\", got %q", token.TokenType)
        }

        // Extract ID token from extras (OIDC providers include it here)
        var idToken string
        if idTokenVal := token.Extra("id_token"); idTokenVal != nil {
                if s, ok := idTokenVal.(string); ok {
                        idToken = s
                }
        }

        return &Tokens{
                AccessToken:  token.AccessToken,
                RefreshToken: token.RefreshToken,
                IDToken:      idToken,
                ExpiresAt:    token.Expiry,
        }, nil
}

// Compile-time interface compliance check.
var _ OAuth2Provider = (*BaseOAuth2Provider)(nil)

// BaseOAuth2Provider implements OAuth 2.0 flows for pure OAuth 2.0 providers.
// This can be used standalone for OAuth 2.0 providers without OIDC support,
// or embedded by OIDCProvider to share common OAuth 2.0 logic.
type BaseOAuth2Provider struct {
        config       *OAuth2Config
        oauth2Config *oauth2.Config
        httpClient   *http.Client
        // ownsHTTPClient is true only when the provider built httpClient itself;
        // see CloseIdleConnections.
        ownsHTTPClient bool
}

// ErrCABundleWithInjectedClient is returned when a provider config sets
// CAFilePath while the caller also injects an HTTP client via
// WithOAuth2HTTPClient / WithHTTPClient. The injected client carries its own TLS
// trust, so the bundle could never take effect; failing loudly keeps a
// trust-anchor misconfiguration a boot-time error rather than a silent no-op.
var ErrCABundleWithInjectedClient = errors.New("caFilePath cannot be combined with an injected HTTP client")

// ErrGrantRejected is wrapped by a token request error when the upstream
// authorization server rejected the grant with the OAuth 2.0 invalid_grant
// error code (RFC 6749 Section 5.2). The classification is
// oautherr.IsRejectedRefreshGrant, shared with the client-side token source:
// invalid_grant on any response except a 5xx or 429, which are treated as
// transient, and a response with no HTTP status is not counted. It applies to
// both token requests a provider makes: on the authorization code exchange it
// means the code or PKCE verifier was rejected, and on refresh it means the
// presented upstream refresh token is expired, revoked or otherwise no longer
// valid, so retrying it cannot succeed. Network failures, 5xx responses and
// other error codes do not wrap it, so callers can tell a rejection from a
// transient failure with errors.Is. It is distinct from fosite.ErrInvalidGrant,
// which is the error the embedded auth server itself returns to its clients.
//
// A rejected refresh token does not always mean the session's upstream
// credential is dead. With a provider that rotates refresh tokens, the
// replica that loses a cross-process refresh race presents an already-rotated
// token and gets invalid_grant while the winner holds valid tokens. The auth
// server's refresher re-reads storage before surfacing the error, which
// resolves the race once the winner's write has landed, but not before.
var ErrGrantRejected = errors.New("invalid_grant")

// IdleConnectionCloser is an optional capability an OAuth2Provider may implement
// to release the idle keep-alive connections held by its HTTP client. It is the
// provider-level analogue of networking.IdleConnectionCloser, and is optional
// rather than part of OAuth2Provider because doc.go invites external
// implementations of that interface, which widening would break. Consumers
// type-assert against it and treat a provider that does not implement it as
// having nothing to release.
//
// Only idle connections are closed: in-flight requests are unaffected and the
// provider stays usable, dialing again on its next call. An implementation must
// leave a caller-supplied client alone — the caller may be sharing it — which is
// the ownership rule BaseOAuth2Provider.CloseIdleConnections implements.
type IdleConnectionCloser interface {
        CloseIdleConnections()
}

// Compile-time capability check.
var _ IdleConnectionCloser = (*BaseOAuth2Provider)(nil)

// OAuth2ProviderOption configures a BaseOAuth2Provider.
type OAuth2ProviderOption func(*BaseOAuth2Provider)

// WithOAuth2HTTPClient sets a custom HTTP client. See
// BaseOAuth2Provider.CloseIdleConnections for the ownership rule this implies.
//
// The injected client fully supersedes the one the provider would have built, so
// the caller owns its request-time scheme enforcement (the HTTPS check
// networking.ValidatingTransport applies) and its dial-time private-IP blocking.
// Config-level validation is unaffected: InsecureAllowHTTP still governs whether
// a plain-http endpoint is accepted at all, via ValidateWithInsecure. Setting
// CAFilePath alongside an injected client is a contradiction and is rejected —
// see ErrCABundleWithInjectedClient.
func WithOAuth2HTTPClient(client *http.Client) OAuth2ProviderOption {
        return func(p *BaseOAuth2Provider) {
                p.httpClient = client
                p.ownsHTTPClient = false
        }
}

// newBaseOAuth2Provider creates a BaseOAuth2Provider with validated config and HTTP client.
// The hostForClient parameter determines which URL to use for HTTP client configuration
// (e.g., TokenEndpoint for OAuth2, Issuer for OIDC).
//
// Options are applied before the default HTTP client is built, so a caller that
// injects one does not pay for a discarded transport (nor for reading the CA
// bundle) on every construction.
//
// IMPORTANT: Callers must ensure config is non-nil before calling this function.
func newBaseOAuth2Provider(
        config *OAuth2Config,
        hostForClient string,
        opts ...OAuth2ProviderOption,
) (_ *BaseOAuth2Provider, retErr error) {
        if err := config.ValidateWithInsecure(config.InsecureAllowHTTP); err != nil {
                return nil, fmt.Errorf("invalid config: %w", err)
        }

        p := &BaseOAuth2Provider{config: config}
        for _, opt := range opts {
                opt(p)
        }

        if p.httpClient == nil {
                httpClient, err := newHTTPClientForHost(
                        hostForClient, config.AllowPrivateIPs, config.InsecureAllowHTTP, config.CAFilePath)
                if err != nil {
                        return nil, fmt.Errorf("failed to create HTTP client: %w", err)
                }
                p.httpClient = httpClient
                p.ownsHTTPClient = true
                // Symmetric with NewOIDCProvider's drain. Nothing below issues a request
                // today — the only later failure is authStyleFromMethod — so the pool is
                // empty and this is currently a no-op; it is here so a future step that
                // performs I/O before returning an error cannot strand a connection.
                defer func() {
                        if retErr != nil {
                                p.CloseIdleConnections()
                        }
                }()
        } else if err := checkInjectedClientCABundle(config.CAFilePath); err != nil {
                return nil, err
        }

        // AuthStyle is derived from the negotiated token_endpoint_auth_method
        // rather than hardcoded — see authStyleFromMethod.
        authStyle, err := authStyleFromMethod(config.TokenEndpointAuthMethod)
        if err != nil {
                return nil, err
        }

        // Create the oauth2.Config for use with golang.org/x/oauth2 library.
        p.oauth2Config = &oauth2.Config{
                ClientID:     config.ClientID,
                ClientSecret: config.ClientSecret,
                RedirectURL:  config.RedirectURI,
                Scopes:       config.Scopes,
                Endpoint: oauth2.Endpoint{
                        AuthURL:   config.AuthorizationEndpoint,
                        TokenURL:  config.TokenEndpoint,
                        AuthStyle: authStyle,
                },
        }

        return p, nil
}

// CloseIdleConnections releases the idle keep-alive connections pooled by the
// provider's HTTP client. Call it when retiring a provider: otherwise the pool's
// sockets and their servicing goroutines are held until the idle timeout
// elapses, and a process that constructs providers repeatedly accumulates them.
//
// This is the ownership rule the package enforces: a client supplied by the
// caller (WithOAuth2HTTPClient / WithHTTPClient) is left alone, because the
// caller may be sharing it and draining it here would cold-start connections
// another live provider is using. Only a client the provider built itself is
// drained.
func (p *BaseOAuth2Provider) CloseIdleConnections() {
        if p.httpClient == nil || !p.ownsHTTPClient {
                return
        }
        p.httpClient.CloseIdleConnections()
}

// checkInjectedClientCABundle rejects a CA bundle configured alongside a
// caller-injected HTTP client: the injected client carries its own TLS trust, so
// the bundle can never take effect, and a trust-anchor decision must fail at
// boot rather than be dropped. AllowPrivateIPs and InsecureAllowHTTP need no
// equivalent check — the former was already inert with an injected client, the
// latter still gates config-level scheme validation.
func checkInjectedClientCABundle(caFilePath string) error {
        if caFilePath == "" {
                return nil
        }
        return fmt.Errorf("%w: caFilePath %q cannot take effect because the injected client carries "+
                "its own TLS trust; configure the CA bundle on that client and clear caFilePath",
                ErrCABundleWithInjectedClient, caFilePath)
}

// authStyleFromMethod maps an RFC 7591 token_endpoint_auth_method to the
// oauth2.AuthStyle the golang.org/x/oauth2 library uses when presenting client
// credentials at the token endpoint:
//
//   - client_secret_basic → AuthStyleInHeader (HTTP Basic Authorization header)
//   - client_secret_post  → AuthStyleInParams (credentials in the POST body)
//   - none / "" (unset)   → AuthStyleInParams (historical default, no secret)
//
// This is the single home for the auth-method rationale. For DCR-registered
// clients the method is copied from the negotiated
// dcr.Resolution.TokenEndpointAuthMethod so the exchange respects whatever the
// upstream advertised — some authorization servers (e.g. Ory Hydra) default
// confidential clients to client_secret_basic and reject client_secret_post
// outright (issue #5865). The same AuthStyle governs the refresh path, which
// shares this oauth2.Config.
//
// AuthStyleAutoDetect is deliberately not used for the default: its Basic-auth
// probe can consume a single-use authorization code against strict servers that
// reject Basic auth, the same failure pkg/auth/oauth/flow.go sidesteps by
// pinning AuthStyleInParams for public clients.
//
// An unrecognised, non-empty method (e.g. private_key_jwt or tls_client_auth,
// which the DCR resolver can negotiate but this client cannot fulfil) returns
// an error rather than silently degrading to POST-body credentials, per the
// "fail loudly on unrecognised enum values" convention in
// .claude/rules/go-style.md. Silently falling back would send an unusable
// credential and reproduce the opaque invalid_client symptom class issue #5865
// was written to eliminate.
func authStyleFromMethod(method string) (oauth2.AuthStyle, error) {
        switch method {
        case oauthproto.TokenEndpointAuthMethodClientSecretBasic:
                return oauth2.AuthStyleInHeader, nil
        case oauthproto.TokenEndpointAuthMethodClientSecretPost,
                oauthproto.TokenEndpointAuthMethodNone,
                "":
                return oauth2.AuthStyleInParams, nil
        default:
                return oauth2.AuthStyleAutoDetect, fmt.Errorf("unsupported token_endpoint_auth_method %q", method)
        }
}

// NewOAuth2Provider creates a new pure OAuth 2.0 provider.
// Use this for providers that don't support OIDC discovery.
// The config must include AuthorizationEndpoint and TokenEndpoint.
func NewOAuth2Provider(config *OAuth2Config, opts ...OAuth2ProviderOption) (*BaseOAuth2Provider, error) {
        if config == nil {
                return nil, errors.New("config is required")
        }

        slog.Info("creating OAuth2 provider",
                "authorization_endpoint", config.AuthorizationEndpoint,
                "token_endpoint", config.TokenEndpoint,
        )

        tokenURL, err := url.Parse(config.TokenEndpoint)
        if err != nil {
                return nil, fmt.Errorf("invalid token endpoint URL: %w", err)
        }
        p, err := newBaseOAuth2Provider(config, tokenURL.Host, opts...)
        if err != nil {
                return nil, err
        }

        slog.Info("oauth2 provider created successfully",
                "authorization_endpoint", config.AuthorizationEndpoint,
                "token_endpoint", config.TokenEndpoint,
        )

        if config.UserInfo == nil && config.IdentityFromToken == nil {
                // Surface synthesis mode at construction so operators see it once
                // per provider rather than only inferring from missing claims later.
                slog.Warn("oauth2 upstream has no userinfo configured; using synthesis mode "+
                        "(non-PII subject from access token, no Name/Email). Configure "+
                        "userInfo if a real identity endpoint exists.",
                        "authorization_endpoint", config.AuthorizationEndpoint,
                )
        }

        return p, nil
}

// Type returns the provider type.
func (*BaseOAuth2Provider) Type() ProviderType {
        return ProviderTypeOAuth2
}

// authorizationEndpoint returns the authorization endpoint URL.
func (p *BaseOAuth2Provider) authorizationEndpoint() string {
        return p.config.AuthorizationEndpoint
}

// AuthorizationURL builds the URL to redirect the user to the upstream IDP.
func (p *BaseOAuth2Provider) AuthorizationURL(state, codeChallenge string, opts ...AuthorizationOption) (string, error) {
        slog.Debug("building authorization URL",
                "authorization_endpoint", p.authorizationEndpoint(),
                "has_pkce", codeChallenge != "",
        )
        return p.buildAuthorizationURL(
                state,
                codeChallenge,
                opts...,
        )
}

// buildAuthorizationURL builds an authorization URL with the given parameters.
// This is the core implementation used by AuthorizationURL and can be extended by embedding types.
func (p *BaseOAuth2Provider) buildAuthorizationURL(
        state string,
        codeChallenge string,
        opts ...AuthorizationOption,
) (string, error) {
        if p.oauth2Config.Endpoint.AuthURL == "" {
                return "", errors.New("authorization endpoint is required")
        }
        if state == "" {
                return "", errors.New("state parameter is required")
        }

        authOpts := &authorizationOptions{}
        if len(p.config.AdditionalAuthorizationParams) > 0 {
                WithAdditionalParams(p.config.AdditionalAuthorizationParams)(authOpts)
        }
        for _, opt := range opts {
                opt(authOpts)
        }

        // Build oauth2 AuthCodeURL options
        var oauth2Opts []oauth2.AuthCodeOption

        // Add PKCE challenge if provided
        if codeChallenge != "" {
                oauth2Opts = append(oauth2Opts,
                        oauth2.SetAuthURLParam("code_challenge", codeChallenge),
                        oauth2.SetAuthURLParam("code_challenge_method", oauthproto.PKCEMethodS256),
                )
        }

        // Add any additional parameters
        for k, v := range authOpts.additionalParams {
                oauth2Opts = append(oauth2Opts, oauth2.SetAuthURLParam(k, v))
        }

        return p.oauth2Config.AuthCodeURL(state, oauth2Opts...), nil
}

// ExchangeCodeForIdentity exchanges an authorization code for tokens and resolves
// the user's identity in a single atomic operation. For pure OAuth2 providers
// (no ID token) the priority chain is:
//
//  1. IdentityFromToken (operator opt-in): extract identity claims from the
//     token-endpoint response body using gjson paths. The userinfo HTTP call
//     is skipped entirely.
//  2. UserInfo endpoint: fetch identity from the configured userinfo URL.
//  3. Synthesis: when neither is configured, synthesizeIdentity derives a
//     non-PII Subject from the access token (rejects empty tokens to prevent
//     the well-known sha256("") collision). Name and Email are empty;
//     Synthetic=true tells the callback handler to bypass UserResolver
//     because the subject rotates per access token.
//
// The nonce parameter is ignored (no ID token to validate).
func (p *BaseOAuth2Provider) ExchangeCodeForIdentity(ctx context.Context, code, codeVerifier, _ string) (*Identity, error) {
        exchanged, err := p.exchangeCodeForTokens(ctx, code, codeVerifier)
        if err != nil {
                return nil, err
        }

        // Priority 1: identityFromToken (configured by operator).
        if p.config.IdentityFromToken != nil {
                if exchanged.identity == nil {
                        // The rewriter logged the extraction failure at WARN with the
                        // operator-supplied path. Surface the same error here so callers
                        // can report it without requiring WARN-level log access.
                        if exchanged.extractionErr != nil {
                                return nil, exchanged.extractionErr
                        }
                        // Defense-in-depth: the rewriter guarantees one of extractedIdentity or
                        // extractionErr is set when the token endpoint returns 200. This branch
                        // fires only if that invariant is ever relaxed (e.g., a non-200-success
                        // status the oauth2 library accepts, or a future rewriter change).
                        return nil, fmt.Errorf(
                                "%w: identityFromToken configured but extraction failed",
                                ErrIdentityResolutionFailed,
                        )
                }
                return &Identity{
                        Tokens:  exchanged.tokens,
                        Subject: exchanged.identity.Subject,
                        Name:    exchanged.identity.Name,
                        Email:   exchanged.identity.Email,
                }, nil
        }

        // Priority 2: userInfo (existing behavior).
        if p.config.UserInfo != nil {
                userInfo, err := p.fetchUserInfo(ctx, exchanged.tokens.AccessToken)
                if err != nil {
                        return nil, fmt.Errorf("%w: %w", ErrIdentityResolutionFailed, err)
                }
                if userInfo == nil || userInfo.Subject == "" {
                        return nil, ErrIdentityResolutionFailed
                }
                return &Identity{
                        Tokens:  exchanged.tokens,
                        Subject: userInfo.Subject,
                        Name:    userInfo.Name,
                        Email:   userInfo.Email,
                        Claims:  userInfo.Claims,
                }, nil
        }

        // Priority 3: synthesis (PR 5094). Subject derived from access token; rotates
        // per token. The callback handler treats Synthetic=true as opt-out from the
        // user-resolver to avoid creating a fresh users row on every re-auth.
        return synthesizeIdentity(exchanged.tokens)
}

// synthesizedSubjectPrefix tags subjects produced by
// synthesizeSubjectFromAccessToken. The prefix is part of the package's
// externally observable contract; downstream code should recognize
// synthesized subjects via the exported IsSynthesizedSubject predicate
// rather than this constant.
const synthesizedSubjectPrefix = "tk-"

// IsSynthesizedSubject reports whether subject was produced by the
// synthesis-mode fallback (vs. resolved from a userinfo endpoint or ID
// token). Use this for code paths that only see the bare subject string —
// e.g., JWT claim consumers, audit pipelines, status conditions. Callers
// holding an *Identity should prefer Identity.Synthetic, which is set at
// the same source of truth.
//
// Purely structural — checks the prefix only, does not validate the digest.
func IsSynthesizedSubject(subject string) bool {
        return strings.HasPrefix(subject, synthesizedSubjectPrefix)
}

// synthesizeSubjectFromAccessToken returns a stable, opaque identifier
// derived from an access token, for OAuth2 upstreams with no userinfo
// endpoint. Output: synthesizedSubjectPrefix + lowercase hex of the first
// 16 bytes of SHA-256(accessToken) — 35 chars total, e.g.
// "tk-89abcdef0123456789abcdef01234567".
//
// The output is non-PII assuming the upstream issues opaque (non-JWT)
// bearer tokens; the digest reveals nothing about the input beyond what an
// attacker holding a candidate token could confirm by re-hashing. 16 bytes
// is sufficient collision resistance for a session-key role.
//
// Only reached when OAuth2Config.UserInfo is nil. OIDC providers always
// have an ID-token-derived subject. Callers must reject empty access
// tokens — sha256("") is a well-known constant, and synthesizing a
// subject from it would collapse distinct sessions onto a single
// storage bucket. synthesizeIdentity enforces this invariant.
func synthesizeSubjectFromAccessToken(accessToken string) string {
        sum := sha256.Sum256([]byte(accessToken))
        return synthesizedSubjectPrefix + hex.EncodeToString(sum[:16])
}

// synthesizeIdentity builds a synthesized Identity for an OAuth2 upstream
// with no userinfo endpoint. Returns ErrIdentityResolutionFailed when the
// access token is empty: sha256("") is the well-known constant
// e3b0c44298fc1c14…, so synthesizing a subject from an empty token would
// collapse every affected session onto a single (UserID, ProviderID)
// storage bucket — a cross-tenant state-mixing hazard. Defense-in-depth:
// convertOAuth2Token already rejects empty AccessToken at exchange time,
// so this guard is unreachable today through ExchangeCodeForIdentity. It
// exists so a future code path (e.g., a custom token-response mapping
// that drops the field) cannot bypass the invariant.
func synthesizeIdentity(tokens *Tokens) (*Identity, error) {
        if tokens.AccessToken == "" {
                return nil, fmt.Errorf("%w: empty access token, cannot synthesize subject", ErrIdentityResolutionFailed)
        }
        return &Identity{
                Tokens:    tokens,
                Subject:   synthesizeSubjectFromAccessToken(tokens.AccessToken),
                Synthetic: true,
        }, nil
}

// tokenExchangeResult bundles the outputs of a successful token exchange:
// the obtained tokens, any identity extracted from the token response body,
// and any error from the identity extraction step. The exchange-level error
// is returned as the function's own error return value.
//
// extractionErr is populated when IdentityFromToken is configured but the
// extractor could not resolve the subject path. It carries operator-actionable
// diagnostics (path name, type description) and already wraps
// ErrIdentityResolutionFailed. The caller must check extractionErr when
// identity is nil and IdentityFromToken is set.
type tokenExchangeResult struct {
        tokens        *Tokens
        identity      *partialIdentity
        extractionErr error
}

// exchangeCodeForTokens exchanges an authorization code for tokens with the upstream IDP.
// It returns a tokenExchangeResult containing the tokens, any identity extracted from
// the token response body, and any extraction error. The function error is the
// exchange-level error (network, HTTP, token parsing).
func (p *BaseOAuth2Provider) exchangeCodeForTokens(
        ctx context.Context, code, codeVerifier string,
) (*tokenExchangeResult, error) {
        if code == "" {
                return nil, errors.New("authorization code is required")
        }

        slog.Debug("exchanging authorization code for tokens",
                "token_endpoint", p.config.TokenEndpoint,
                "has_pkce_verifier", codeVerifier != "",
        )

        // Wrap HTTP client with token response rewriter if mapping or identity extraction
        // is configured. At auth-code time, both mapping (field normalization) and
        // identityCfg (identity extraction) may be active together. Keep a reference to
        // the rewriter so we can read extractedIdentity and extractionErr after Exchange returns.
        httpClient, rewriter := wrapHTTPClientForTokenExchange(
                p.httpClient, p.config.TokenResponseMapping, p.config.IdentityFromToken, p.config.TokenEndpoint,
        )
        ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient)

        // Build exchange options
        var opts []oauth2.AuthCodeOption
        if codeVerifier != "" {
                opts = append(opts, oauth2.VerifierOption(codeVerifier))
        }
        // AuthCodeOptions passed to Exchange land in the POST form body, so
        // configured additional token params (e.g. an RFC 8707 resource
        // indicator) reach the token endpoint the way such ASes require.
        for k, v := range p.config.AdditionalTokenParams {
                opts = append(opts, oauth2.SetAuthURLParam(k, v))
        }

        token, err := p.oauth2Config.Exchange(ctx, code, opts...)
        if err != nil {
                return nil, formatOAuth2Error(err, "token request failed")
        }

        tokens, err := convertOAuth2Token(token)
        if err != nil {
                return nil, err
        }

        slog.Debug("authorization code exchange successful",
                "has_refresh_token", tokens.RefreshToken != "",
                "expires_at", expiresAtLogValue(tokens.ExpiresAt),
        )

        // Read any identity and extraction error captured by the rewriter during the
        // token round-trip. rewriter is nil when neither mapping nor identityCfg is
        // configured; nil-safe.
        result := &tokenExchangeResult{tokens: tokens}
        if rewriter != nil {
                result.identity = rewriter.extractedIdentity
                result.extractionErr = rewriter.extractionErr
        }

        return result, nil
}

// RefreshTokens refreshes the upstream IDP tokens.
//
// Sends `scope` explicitly by default. RFC 6749 §6 makes the param optional
// and says the AS SHOULD preserve original scopes on omission, but some ASes
// (notably Entra ID v1) silently narrow refreshed tokens to the user's default
// consent set when `scope` is omitted, dropping custom resource scopes. Other
// ASes (e.g. Salesforce) reject `scope` on refresh outright; for those,
// CommonOAuthConfig.OmitScopeOnRefresh leaves it out.
//
// Uses oauth2.Config.Exchange with SetAuthURLParam overrides (mirroring the
// pattern in pkg/auth/oauth/non_caching_refresher.go) because the standard
// library's TokenSource refresh path doesn't expose a way to inject scope.
// The empty code= side-effect is tolerated — ASes dispatch on grant_type first.
func (p *BaseOAuth2Provider) RefreshTokens(ctx context.Context, refreshToken, _ string) (*Tokens, error) {
        if refreshToken == "" {
                return nil, errors.New("refresh token is required")
        }

        slog.Debug("refreshing tokens",
                "token_endpoint", p.config.TokenEndpoint,
                "scope_count", len(p.oauth2Config.Scopes),
        )

        // Wrap HTTP client with token response rewriter if mapping is configured.
        // Identity extraction (identityCfg) is intentionally nil here: per Snowflake's
        // contract and the general design, the username/identity field is only present
        // in the initial auth-code response and is omitted on refresh. Identity is
        // cached at auth-code time and read from session storage on subsequent requests.
        // The rewriter is discarded because refresh does not produce identity.
        httpClient, _ := wrapHTTPClientForTokenExchange(p.httpClient, p.config.TokenResponseMapping, nil, p.config.TokenEndpoint)
        ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient)

        opts := []oauth2.AuthCodeOption{
                oauth2.SetAuthURLParam("grant_type", "refresh_token"),
                oauth2.SetAuthURLParam("refresh_token", refreshToken),
        }
        if len(p.oauth2Config.Scopes) > 0 && !p.config.OmitScopeOnRefresh {
                opts = append(opts, oauth2.SetAuthURLParam("scope", strings.Join(p.oauth2Config.Scopes, " ")))
        }
        // Same rationale as in exchangeCodeForTokens: ASes that enforce
        // RFC 8707 on token requests require these on refresh as well.
        for k, v := range p.config.AdditionalTokenParams {
                opts = append(opts, oauth2.SetAuthURLParam(k, v))
        }

        token, err := p.oauth2Config.Exchange(ctx, "", opts...)
        if err != nil {
                return nil, formatOAuth2Error(err, "token request failed")
        }

        tokens, err := convertOAuth2Token(token)
        if err != nil {
                return nil, err
        }
        // AS may not issue a new refresh token (RFC 6749 §6); preserve the old one.
        if tokens.RefreshToken == "" {
                tokens.RefreshToken = refreshToken
        }

        slog.Debug("token refresh successful",
                // Read from token (pre-§6-fallback) so the log accurately reflects
                // whether the AS rotated the refresh token vs the fallback kicking in.
                "has_new_refresh_token", token.RefreshToken != "",
                "expires_at", expiresAtLogValue(tokens.ExpiresAt),
        )

        return tokens, nil
}

// userInfo contains user information retrieved from the upstream IDP.
type userInfo struct {
        // Subject is the unique identifier for the user (sub claim).
        Subject string `json:"sub"`

        // Email is the user's email address.
        Email string `json:"email,omitempty"`

        // Name is the user's full name.
        Name string `json:"name,omitempty"`

        // Claims contains all claims returned by the userinfo endpoint.
        Claims map[string]any `json:"-"`
}

// fetchUserInfo fetches user information from the configured UserInfo endpoint.
// Returns an error if no UserInfo endpoint is configured.
// The field mapping from UserInfoConfig.FieldMapping is used to extract claims
// from non-standard provider responses.
func (p *BaseOAuth2Provider) fetchUserInfo(ctx context.Context, accessToken string) (*userInfo, error) {
        if p.config.UserInfo == nil {
                return nil, errors.New("userinfo endpoint not configured")
        }

        cfg := p.config.UserInfo

        if accessToken == "" {
                return nil, errors.New("access token is required")
        }

        slog.Debug("fetching user info",
                "userinfo_endpoint", cfg.EndpointURL,
        )

        // Determine HTTP method (default GET per OIDC Core Section 5.3.1)
        method := cfg.HTTPMethod
        if method == "" {
                method = http.MethodGet
        }

        req, err := http.NewRequestWithContext(ctx, method, cfg.EndpointURL, nil)
        if err != nil {
                return nil, fmt.Errorf("failed to create request: %w", err)
        }

        // Set authorization header per RFC 6750 (Bearer Token Usage)
        req.Header.Set("Authorization", "Bearer "+accessToken)
        req.Header.Set("Accept", "application/json")

        // Add any additional headers (useful for non-standard providers like GitHub)
        for k, v := range cfg.AdditionalHeaders {
                req.Header.Set(k, v)
        }

        resp, err := p.httpClient.Do(req) //nolint:gosec // G704: URL is from OIDC discovery, not user input
        if err != nil {
                return nil, fmt.Errorf("userinfo request failed: %w", err)
        }
        defer func() { _ = resp.Body.Close() }()

        if resp.StatusCode != http.StatusOK {
                // Drain response body for connection reuse, but don't log it to avoid
                // potentially exposing sensitive information from the upstream provider.
                _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1024))
                slog.Debug("userinfo request failed", //nolint:gosec // G706: status code is an integer
                        "status", resp.StatusCode)
                return nil, fmt.Errorf("userinfo request failed with status %d", resp.StatusCode)
        }

        body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseSize))
        if err != nil {
                return nil, fmt.Errorf("failed to read userinfo response: %w", err)
        }

        var claims map[string]any
        if err := json.Unmarshal(body, &claims); err != nil {
                return nil, fmt.Errorf("failed to parse userinfo response: %w", err)
        }

        // Use configured field mapping for claim extraction
        mapping := cfg.FieldMapping

        // Extract and validate required subject claim
        sub, err := mapping.ResolveSubject(claims)
        if err != nil {
                return nil, fmt.Errorf("userinfo response missing required subject claim: %w", err)
        }

        userInfo := &userInfo{
                Subject: sub,
                Name:    mapping.ResolveName(claims),
                Email:   mapping.ResolveEmail(claims),
                Claims:  claims,
        }

        slog.Debug("user info retrieved",
                "subject", userInfo.Subject,
                "has_email", userInfo.Email != "",
        )

        return userInfo, nil
}

// formatOAuth2Error extracts error details from oauth2.RetrieveError for better error messages.
func formatOAuth2Error(err error, prefix string) error {
        var retrieveErr *oauth2.RetrieveError
        if errors.As(err, &retrieveErr) {
                // RetrieveError contains the OAuth error response
                if oautherr.IsRejectedRefreshGrant(retrieveErr) {
                        return fmt.Errorf("%s: %w - %s", prefix, ErrGrantRejected, retrieveErr.ErrorDescription)
                }
                if retrieveErr.ErrorCode != "" {
                        return fmt.Errorf("%s: %s - %s", prefix, retrieveErr.ErrorCode, retrieveErr.ErrorDescription)
                }
                // Log full response for debugging, but return sanitized error to prevent information disclosure
                slog.Debug("token request failed",
                        "status", retrieveErr.Response.StatusCode,
                        "body", string(retrieveErr.Body))
                return fmt.Errorf("%s with status %d", prefix, retrieveErr.Response.StatusCode)
        }
        // For other errors (network errors, etc.), wrap with context
        return fmt.Errorf("request failed: %w", err)
}

// newHTTPClientForHost creates an HTTP client configured for the given host.
// It enables HTTP and private IPs only for localhost (development/testing),
// or when INSECURE_DISABLE_URL_VALIDATION is set (e.g. Kubernetes dev environments).
// allowPrivateIPs widens only the private-IP gate: it permits connections to
// RFC-1918/link-local addresses (e.g. in-cluster providers) without enabling
// the HTTP scheme for non-localhost hosts.
//
// The host-scoped guard policy lives in networking.NewHostScopedClientBuilder
// so this provider path and the DCR resolver share one implementation.
//
// Unlike the DCR resolver's guarded client, this one deliberately leaves
// keep-alive enabled: the returned client is stored on BaseOAuth2Provider and
// reused for many token-refresh/userinfo calls over the provider's lifetime
// against a single operator-configured host, so disabling keep-alive here
// would pay a fresh TCP+TLS handshake on every call. This trades a narrower
// window — a DNS change for that fixed host between connection reuses is not
// re-checked mid-lifetime — for avoiding that cost on a hot path; the DCR
// resolver's per-request-host, low-frequency calls don't have the same
// trade-off, hence the difference. If this provider's threat model changes
// (e.g. it starts dialing caller-varying hosts), revisit this decision.
//
// The client also sets no response-header timeout, so the client timeout is the
// budget for each call, as for the shared grant client in pkg/oauthproto (see
// grants.go): corporate IdP chains can take more than 10s to send the first
// byte. That holds for every call made with this client, because discovery,
// JWKS, userinfo and the token endpoint are all served by the same IdP.
func newHTTPClientForHost(host string, allowPrivateIPs, insecureAllowHTTP bool, caFilePath string) (*http.Client, error) {
        builder := networking.NewHostScopedClientBuilder(host, allowPrivateIPs, insecureAllowHTTP).
                WithResponseHeaderTimeout(0)
        if caFilePath != "" {
                builder.WithSystemRootsPlusCABundle(caFilePath)
        }
        return builder.Build()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package upstream

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "net"
        "net/http"
        "net/url"
        "regexp"
        "slices"
        "strings"

        "github.com/coreos/go-oidc/v3/oidc"
        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

const (
        // ProviderTypeOIDC is for OIDC providers that support discovery.
        ProviderTypeOIDC ProviderType = "oidc"

        // idTokenValidationAttempts is how many times RefreshTokens will try to
        // verify a newly issued ID token after a successful token-endpoint
        // exchange. The exchange itself is not retried: rotating refresh tokens
        // are single-use, so a second token request would send a consumed grant.
        idTokenValidationAttempts = 3
)

// OIDCConfig contains configuration for OIDC providers that support discovery.
type OIDCConfig struct {
        CommonOAuthConfig

        // Issuer is the URL of the upstream OIDC provider (e.g., https://accounts.google.com).
        // The provider will fetch endpoints from {Issuer}/.well-known/openid-configuration.
        Issuer string

        // SubjectClaim names the validated ID-token claim to use as the upstream
        // subject. Defaults to "sub" when empty (preserves current behavior). Set for
        // IdPs where "sub" isn't stable per user (e.g. Entra/Azure AD's "oid"). This
        // is the OIDC counterpart to OAuth2's IdentityFromTokenConfig.SubjectPath,
        // but narrower: a verbatim top-level claim name (no nested paths) resolving
        // to a non-empty string (no numeric claims).
        //
        // Effectively immutable once users exist: the claim value becomes the
        // upstream subject that resolves to User.ID, so changing it (or fixing a
        // typo) re-keys every existing user and orphans their stored upstream tokens.
        //
        // The claim must also be present on refreshed ID tokens — RefreshTokens
        // resolves through the same path and fails closed if the IdP drops it.
        SubjectClaim string `json:"subject_claim,omitempty" yaml:"subject_claim,omitempty"`

        // CAFilePath is the path to a PEM CA bundle added to the system roots.
        CAFilePath string `json:"ca_file_path,omitempty" yaml:"ca_file_path,omitempty"`

        // AllowPrivateIPs permits the OIDC discovery and token HTTP clients to
        // connect to private IP ranges (RFC-1918, link-local). Use only when the
        // upstream is hosted inside the same cluster and has no public endpoint.
        // Defaults to false.
        AllowPrivateIPs bool `json:"allow_private_ips,omitempty" yaml:"allow_private_ips,omitempty"`

        // InsecureAllowHTTP permits an http:// issuer URL for non-localhost hosts.
        // Only set this for trusted in-cluster Kubernetes deployments.
        // Production deployments reachable outside the cluster MUST use https://.
        InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"`
}

// subjectClaimPattern is the allowed shape for SubjectClaim: a claim-name token
// (letter/underscore start, then letters/digits/underscores). It must stay in
// sync with the CEL rule on the operator CRD field
// (cmd/thv-operator/api/v1beta1 OIDCUpstreamConfig.SubjectClaim) so the Go and
// CRD layers reject the same values.
var subjectClaimPattern = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)

// Validate checks that OIDCConfig has all required fields and valid values,
// respecting c.InsecureAllowHTTP when set.
func (c *OIDCConfig) Validate() error {
        return c.ValidateWithInsecure(c.InsecureAllowHTTP)
}

// ValidateWithInsecure is like Validate but allows http:// issuer URLs for
// non-localhost hosts when insecureAllowHTTP is true (trusted in-cluster deployments).
func (c *OIDCConfig) ValidateWithInsecure(insecureAllowHTTP bool) error {
        if c.Issuer == "" {
                return errors.New("issuer is required for OIDC providers")
        }
        if err := networking.ValidateEndpointURLWithInsecure(c.Issuer, insecureAllowHTTP); err != nil {
                return fmt.Errorf("invalid issuer URL: %w", err)
        }
        // SubjectClaim is optional (empty defaults to "sub"). When set, require a
        // claim-name shape: the claim is looked up verbatim, so values like "oid.sub"
        // or "oid " can never match a real top-level claim. Rejecting at startup here
        // (mirroring the CRD's CEL rule) avoids a silent miss at first login.
        if c.SubjectClaim != "" && !subjectClaimPattern.MatchString(c.SubjectClaim) {
                return fmt.Errorf(
                        "subject_claim %q must be a claim name: start with a letter or "+
                                "underscore and use only letters, digits, and underscores",
                        c.SubjectClaim)
        }
        return c.CommonOAuthConfig.ValidateWithInsecure(insecureAllowHTTP)
}

// ErrNonceMismatch is returned when the nonce claim in the ID token does not match
// the expected nonce from the authorization request.
var ErrNonceMismatch = errors.New("ID token nonce does not match expected value")

// ErrSubjectMismatch is returned when the resolved subject of a refreshed ID
// token (the configured SubjectClaim, or "sub" by default) does not match the
// expected subject from the original token response.
// Per OIDC Core Section 12.2, the sub claim MUST be identical across refreshes.
var ErrSubjectMismatch = errors.New("ID token subject does not match expected value")

// ErrNonceMissing is returned when the ID token does not contain a nonce claim
// but one was expected (because a nonce was sent in the authorization request).
var ErrNonceMissing = errors.New("ID token missing nonce claim when nonce was expected")

// OIDCProviderImpl implements OAuth2Provider for OIDC-compliant identity providers.
// It embeds BaseOAuth2Provider to share common OAuth 2.0 logic while adding
// OIDC-specific functionality like discovery and ID token validation.
// The ResolveIdentity method is overridden to validate ID tokens per OIDC spec.
type OIDCProviderImpl struct {
        *BaseOAuth2Provider                                   // Embed for shared OAuth 2.0 logic
        oidcConfig          *OIDCConfig                       // Store original OIDC config (Issuer + common OAuth fields)
        endpoints           *oauthproto.OIDCDiscoveryDocument // Discovered endpoints for security validation
        forceConsentScreen  bool                              // Force consent screen on auth requests
        verifier            *oidc.IDTokenVerifier             // ID token verifier from go-oidc
}

// Compile-time check that the embedded *BaseOAuth2Provider keeps promoting
// CloseIdleConnections, which (*server).CloseIdleConnections reaches by type
// assertion. It does not pin the behavior — a method declared here would shadow
// the promoted one and still satisfy this — so
// TestServer_Close_DrainsRealOIDCUpstream is the actual guard.
var _ IdleConnectionCloser = (*OIDCProviderImpl)(nil)

// OIDCProviderOption configures an OIDCProvider.
type OIDCProviderOption func(*OIDCProviderImpl)

// WithHTTPClient sets a custom HTTP client for the provider. See
// WithOAuth2HTTPClient, which carries the same ownership and config-supersession
// rules.
func WithHTTPClient(client *http.Client) OIDCProviderOption {
        return func(p *OIDCProviderImpl) {
                p.httpClient = client
                p.ownsHTTPClient = false
        }
}

// WithNonce adds an OIDC nonce parameter to the authorization request.
// The nonce is used to associate a client session with an ID Token and to
// prevent replay attacks. See OIDC Core Section 3.1.2.1.
func WithNonce(nonce string) AuthorizationOption {
        return WithAdditionalParams(map[string]string{"nonce": nonce})
}

// WithForceConsentScreen configures the provider to always request the consent screen
// from the identity provider. When enabled, the "prompt=consent" parameter is added
// to authorization requests, forcing the user to re-consent even if they have
// previously authorized the application.
//
// This is useful for:
//   - Testing OAuth flows to verify consent screen behavior
//   - Obtaining a new refresh token when the original has been lost or revoked
//   - Ensuring the user explicitly confirms permissions after scope changes
//   - Applications that require explicit user consent on every authentication
func WithForceConsentScreen(force bool) OIDCProviderOption {
        return func(p *OIDCProviderImpl) {
                p.forceConsentScreen = force
        }
}

// NewOIDCProvider creates a new OIDC provider.
// It performs OIDC discovery to fetch endpoints and then constructs the
// underlying OAuth2 configuration from the discovered endpoints.
func NewOIDCProvider(
        ctx context.Context,
        config *OIDCConfig,
        opts ...OIDCProviderOption,
) (_ *OIDCProviderImpl, retErr error) {
        if config == nil {
                return nil, errors.New("config is required")
        }

        // Validate OIDC config
        if err := config.ValidateWithInsecure(config.InsecureAllowHTTP); err != nil {
                return nil, fmt.Errorf("invalid config: %w", err)
        }

        slog.Debug("creating OIDC provider",
                "issuer", config.Issuer,
                "client_id", config.ClientID,
        )

        p := &OIDCProviderImpl{
                oidcConfig: config,
                // config will be set after discovery
                BaseOAuth2Provider: &BaseOAuth2Provider{},
        }

        // Apply OIDC-specific options before building the default HTTP client, so a
        // caller that injects one does not pay for a discarded transport (nor for
        // reading the CA bundle) on every construction.
        for _, opt := range opts {
                opt(p)
        }

        if p.httpClient == nil {
                // Create HTTP client for the issuer host
                issuerURL, _ := url.Parse(config.Issuer) // Error already checked in ValidateWithInsecure()
                httpClient, err := newHTTPClientForHost(
                        issuerURL.Host, config.AllowPrivateIPs, config.InsecureAllowHTTP, config.CAFilePath)
                if err != nil {
                        return nil, fmt.Errorf("failed to create HTTP client: %w", err)
                }
                p.httpClient = httpClient
                p.ownsHTTPClient = true
                // Discovery below runs before any later failure (bad claims, an invalid
                // discovery document, a missing openid scope), so a failed construction
                // otherwise abandons a client with a pooled connection that no caller
                // can reach — nothing is returned to call CloseIdleConnections on.
                defer func() {
                        if retErr != nil {
                                p.CloseIdleConnections()
                        }
                }()
        } else if err := checkInjectedClientCABundle(config.CAFilePath); err != nil {
                return nil, err
        }

        // Use go-oidc for discovery - inject custom HTTP client via context
        ctx = oidc.ClientContext(ctx, p.httpClient)
        oidcProvider, err := oidc.NewProvider(ctx, config.Issuer)
        if err != nil {
                return nil, fmt.Errorf("failed to discover OIDC endpoints: %w", err)
        }

        // Extract endpoints from provider claims for security validation.
        // go-oidc validates issuer but doesn't check endpoint origins.
        endpoints := &oauthproto.OIDCDiscoveryDocument{}
        if err := oidcProvider.Claims(endpoints); err != nil {
                return nil, fmt.Errorf("failed to extract provider claims: %w", err)
        }

        // Security validation - go-oidc doesn't check endpoint origins
        if err := validateDiscoveryDocument(endpoints, config.Issuer, config.InsecureAllowHTTP); err != nil {
                return nil, fmt.Errorf("invalid discovery document: %w", err)
        }

        p.endpoints = endpoints

        // Determine scopes: use configured or OIDC defaults
        scopes := config.Scopes
        if len(scopes) == 0 {
                scopes = []string{"openid", "profile", "email"}
        }

        // Validate that openid scope is present for OIDC provider.
        // Per OIDC Core, openid scope is mandatory for ID tokens. Without it, the IDP
        // won't return an ID token, but OIDCProviderImpl requires one for identity resolution.
        if !slices.Contains(scopes, "openid") {
                return nil, errors.New("openid scope is required for OIDC provider; use BaseOAuth2Provider for pure OAuth 2.0 flows")
        }

        // Now create OAuth2Config from discovered endpoints + OIDC config.
        // This allows the embedded BaseOAuth2Provider to use the discovered endpoints
        // for token requests while preserving the original OIDC config.
        // Note: UserInfoEndpoint is stored in p.endpoints, not in OAuth2Config.
        // Copy the full CommonOAuthConfig so that all fields (including
        // AdditionalAuthorizationParams and any future additions) propagate
        // to the embedded BaseOAuth2Provider. Override Scopes since OIDC
        // applies default scope logic above.
        commonCfg := config.CommonOAuthConfig
        commonCfg.Scopes = scopes
        oauth2Config := &OAuth2Config{
                CommonOAuthConfig:     commonCfg,
                AuthorizationEndpoint: p.endpoints.AuthorizationEndpoint,
                TokenEndpoint:         p.endpoints.TokenEndpoint,
                AllowPrivateIPs:       config.AllowPrivateIPs,
        }
        p.config = oauth2Config

        // Create the oauth2.Config for use with golang.org/x/oauth2 library
        // Use go-oidc's endpoint which handles discovery, but explicitly set AuthStyle
        // to ensure client credentials are sent in the request body (not Basic auth header)
        // for consistent behavior across different IDP implementations.
        providerEndpoint := oidcProvider.Endpoint()
        p.oauth2Config = &oauth2.Config{
                ClientID:     config.ClientID,
                ClientSecret: config.ClientSecret,
                RedirectURL:  config.RedirectURI,
                Scopes:       scopes,
                Endpoint: oauth2.Endpoint{
                        AuthURL:   providerEndpoint.AuthURL,
                        TokenURL:  providerEndpoint.TokenURL,
                        AuthStyle: oauth2.AuthStyleInParams,
                },
        }

        // Use go-oidc's built-in verifier for ID token validation
        p.verifier = oidcProvider.Verifier(&oidc.Config{
                ClientID: config.ClientID,
        })

        slog.Debug("oidc provider created successfully",
                "issuer", p.endpoints.Issuer,
                "pkce_supported", p.supportsPKCE(),
                "id_token_validation_enabled", p.verifier != nil,
        )

        return p, nil
}

// Type returns the provider type.
func (*OIDCProviderImpl) Type() ProviderType {
        return ProviderTypeOIDC
}

// ExchangeCodeForIdentity exchanges an authorization code for tokens and validates
// the ID token (including nonce) in a single atomic operation.
// Per OIDC Core Section 3.1.3.3, the ID token MUST be present. The nonce is validated
// against the ID token to prevent replay attacks (Section 3.1.3.7).
func (p *OIDCProviderImpl) ExchangeCodeForIdentity(
        ctx context.Context, code, codeVerifier, nonce string,
) (*Identity, error) {
        if p.endpoints == nil {
                return nil, errors.New("OIDC endpoints not discovered")
        }

        // OIDC resolves identity from the validated ID token, not the token response
        // body, so the extracted identity return value is intentionally discarded.
        // Defense-in-depth: warn if the OAuth2 base config carries a non-nil
        // IdentityFromToken — the field is structurally absent on the OIDC CRD type
        // today, but a future config-loader bug or hand-built runtime config could
        // silently drop the operator's intent without this signal.
        if p.config.IdentityFromToken != nil {
                slog.Debug("OIDC provider ignoring IdentityFromToken; identity is resolved from the validated ID token")
        }
        exchanged, err := p.exchangeCodeForTokens(ctx, code, codeVerifier)
        if err != nil {
                return nil, err
        }

        // OIDC-specific: ID token MUST be present per Section 3.1.3.3.
        if exchanged.tokens.IDToken == "" {
                return nil, fmt.Errorf("%w: ID token required for OIDC provider", ErrIdentityResolutionFailed)
        }

        // Validate ID token with nonce in a single pass — no double-validation.
        validatedToken, err := p.validateIDToken(ctx, exchanged.tokens.IDToken, nonce)
        if err != nil {
                slog.Debug("id token validation failed", "error", err)
                return nil, fmt.Errorf("%w: %w", ErrIdentityResolutionFailed, err)
        }

        slog.Debug("authorization code exchange successful",
                "has_refresh_token", exchanged.tokens.RefreshToken != "",
                "has_id_token", exchanged.tokens.IDToken != "",
                "expires_at", expiresAtLogValue(exchanged.tokens.ExpiresAt),
        )

        // Extract optional standard claims (name, email) from ID token
        var idClaims struct {
                Name  string `json:"name"`
                Email string `json:"email"`
        }
        // Best-effort: if claims extraction fails, we still have the subject
        if err := validatedToken.Claims(&idClaims); err != nil {
                slog.Warn("failed to extract optional claims from ID token",
                        "error", err,
                )
        }

        subject, err := p.resolveSubject(validatedToken)
        if err != nil {
                return nil, fmt.Errorf("%w: %w", ErrIdentityResolutionFailed, err)
        }

        // Capture all ID-token claims for downstream authorization inputs. Best-effort:
        // the subject/name/email above are already resolved, so a claims-extraction
        // failure only loses the enrichment, it does not fail the login.
        var allClaims map[string]any
        if err := validatedToken.Claims(&allClaims); err != nil {
                slog.Warn("failed to extract full claims from ID token for authorization inputs",
                        "error", err,
                )
        }

        return &Identity{
                Tokens:              exchanged.tokens,
                Subject:             subject,
                Name:                idClaims.Name,
                Email:               idClaims.Email,
                Claims:              allClaims,
                VerifiedIssuer:      validatedToken.Issuer,
                VerifiedSubject:     validatedToken.Subject,
                VerifiedStableClaim: p.subjectClaimName(),
                VerifiedStableValue: subject,
        }, nil
}

// subjectClaimName returns the ID-token claim this upstream keys users on.
func (p *OIDCProviderImpl) subjectClaimName() string {
        if p.oidcConfig.SubjectClaim == "" {
                return "sub"
        }
        return p.oidcConfig.SubjectClaim
}

// resolveSubject returns the upstream subject for the validated ID token.
// When SubjectClaim is unset or "sub", it uses the standard OIDC subject.
// Otherwise it extracts the named claim as a non-empty string. It fails loud
// if the claim is missing, empty, or not a string rather than falling back to
// "sub" — a silent fallback would key the user under the wrong identifier.
func (p *OIDCProviderImpl) resolveSubject(token *oidc.IDToken) (string, error) {
        claim := p.subjectClaimName()
        if claim == "sub" {
                return token.Subject, nil
        }

        var allClaims map[string]any
        if err := token.Claims(&allClaims); err != nil {
                return "", fmt.Errorf("extracting claims for subject claim %q: %w", claim, err)
        }

        raw, ok := allClaims[claim]
        if !ok {
                return "", fmt.Errorf("configured subject claim %q not present in ID token", claim)
        }
        value, ok := raw.(string)
        if !ok {
                return "", fmt.Errorf("configured subject claim %q is not a string", claim)
        }
        if value == "" {
                return "", fmt.Errorf("configured subject claim %q is empty", claim)
        }
        return value, nil
}

// validateRefreshedIDToken verifies the ID token returned by a successful
// refresh exchange. Transient JWKS/network failures are retried in place so
// the already-consumed refresh token is not replayed. If every attempt is
// still a transient fetch failure, the unvalidated ID token is dropped and
// (nil, nil) is returned so the caller can keep the new access/refresh tokens
// (the storage layer already carries forward the previous ID token when the
// new one is empty). Permanent verification failures still fail closed.
func (p *OIDCProviderImpl) validateRefreshedIDToken(ctx context.Context, tokens *Tokens) (*oidc.IDToken, error) {
        var lastErr error
        for attempt := 1; attempt <= idTokenValidationAttempts; attempt++ {
                token, err := p.validateIDToken(ctx, tokens.IDToken, "")
                if err == nil {
                        return token, nil
                }
                lastErr = err
                if !isTransientIDTokenValidationError(err) {
                        return nil, fmt.Errorf("ID token validation failed: %w", err)
                }
                slog.Warn("transient ID token validation failure after token refresh; retrying without re-exchanging",
                        "attempt", attempt,
                        "error", err,
                )
        }

        // The token endpoint already succeeded and, for rotating refresh tokens,
        // consumed the grant. Returning an error here would make the caller retry
        // the exchange with the old token and get invalid_grant (#6194). Drop the
        // unvalidated ID token rather than using it; access and refresh tokens
        // still came from the trusted token endpoint.
        slog.Warn("ID token validation failed after token refresh; dropping unvalidated ID token to preserve rotated refresh token",
                "error", lastErr,
        )
        tokens.IDToken = ""
        return nil, nil
}

// isTransientIDTokenValidationError reports whether ID-token verification
// failed because keys could not be fetched, not because the token itself was
// rejected. Permanent failures (bad signature, missing/mismatched nonce) must
// not be retried or bypassed.
func isTransientIDTokenValidationError(err error) bool {
        if err == nil {
                return false
        }
        if errors.Is(err, ErrNonceMismatch) || errors.Is(err, ErrNonceMissing) || errors.Is(err, ErrSubjectMismatch) {
                return false
        }
        if errors.Is(err, context.DeadlineExceeded) {
                return true
        }
        var netErr net.Error
        if errors.As(err, &netErr) && netErr.Timeout() {
                return true
        }
        // go-oidc wraps a JWKS HTTP failure as "fetching keys oidc: get keys failed".
        msg := err.Error()
        return strings.Contains(msg, "fetching keys") || strings.Contains(msg, "get keys failed")
}

// validateIDToken validates an ID token and returns the parsed token.
func (p *OIDCProviderImpl) validateIDToken(ctx context.Context, idToken, nonce string) (*oidc.IDToken, error) {
        if p.verifier == nil {
                return nil, errors.New("ID token verifier not initialized")
        }

        token, err := p.verifier.Verify(ctx, idToken)
        if err != nil {
                return nil, fmt.Errorf("failed to verify ID token: %w", err)
        }

        // Validate nonce if expected (was sent in authorization request).
        // This ensures that when a nonce is provided, the token MUST contain it
        // and it MUST match, preventing replay attacks.
        if nonce != "" {
                if token.Nonce == "" {
                        return nil, ErrNonceMissing
                }
                if token.Nonce != nonce {
                        return nil, ErrNonceMismatch
                }
        }

        return token, nil
}

// supportsPKCE checks if the provider advertises S256 PKCE support.
func (p *OIDCProviderImpl) supportsPKCE() bool {
        if p.endpoints == nil {
                return false
        }
        return p.endpoints.SupportsPKCE()
}

// AuthorizationURL builds the URL to redirect the user to the upstream IDP.
// This overrides the base implementation to add OIDC-specific parameters (nonce, prompt)
// and use discovered endpoints.
func (p *OIDCProviderImpl) AuthorizationURL(state, codeChallenge string, opts ...AuthorizationOption) (string, error) {
        if p.endpoints == nil {
                return "", errors.New("OIDC endpoints not discovered")
        }

        // Apply authorization options to extract nonce for logging
        authOpts := &authorizationOptions{}
        for _, opt := range opts {
                opt(authOpts)
        }

        // Extract nonce from additionalParams if present
        nonce := ""
        if authOpts.additionalParams != nil {
                nonce = authOpts.additionalParams["nonce"]
        }

        slog.Debug("building authorization URL",
                "authorization_endpoint", p.endpoints.AuthorizationEndpoint,
                "has_pkce", codeChallenge != "",
                "has_nonce", nonce != "",
        )

        // PKCE: Per RFC 7636 Section 5, clients SHOULD send PKCE parameters to all
        // servers regardless of whether they advertise support. Servers that don't
        // support PKCE will simply ignore the parameters.
        if codeChallenge != "" && !p.supportsPKCE() {
                slog.Debug("sending PKCE to provider that does not advertise S256 support (per RFC 7636 Section 5)")
        }

        // Merge caller's opts with OIDC-specific params
        allOpts := append(opts, WithAdditionalParams(p.buildOIDCParams())) //nolint:gocritic // intentionally appending single element

        // Use the base implementation which uses oauth2Config (scopes already configured)
        return p.buildAuthorizationURL(state, codeChallenge, allOpts...)
}

// buildOIDCParams builds the OIDC-specific authorization parameters.
func (p *OIDCProviderImpl) buildOIDCParams() map[string]string {
        params := make(map[string]string)

        // Add prompt=consent if configured to force the consent screen
        if p.forceConsentScreen {
                params["prompt"] = "consent"
        }

        return params
}

// RefreshTokens refreshes the upstream IDP tokens.
// This overrides the base implementation to add OIDC-specific ID token validation.
// After a successful token-endpoint exchange, ID-token verification is retried
// on transient JWKS/network failures so a rotating refresh token is not
// replayed. If verification still cannot fetch keys, the unvalidated ID token
// is dropped and the new access/refresh tokens are returned; the storage layer
// carries forward the previous ID token when the new one is empty.
func (p *OIDCProviderImpl) RefreshTokens(ctx context.Context, refreshToken, expectedSubject string) (*Tokens, error) {
        if p.endpoints == nil {
                return nil, errors.New("OIDC endpoints not discovered")
        }

        slog.Debug("refreshing tokens",
                "token_endpoint", p.endpoints.TokenEndpoint,
        )

        // Use base provider's implementation for token refresh
        tokens, err := p.BaseOAuth2Provider.RefreshTokens(ctx, refreshToken, expectedSubject)
        if err != nil {
                return nil, err
        }

        // OIDC-specific: Validate ID token if present.
        // Per OIDC Core Section 12.2, refresh responses MAY include a new ID token
        // (unlike ExchangeCodeForIdentity where it's required per Section 3.1.3.3).
        // Nonce validation is intentionally omitted: Section 12.2 states that
        // refreshed ID tokens SHOULD NOT contain a nonce claim, and no new
        // authorization request exists to provide an expected nonce value.
        // Full nonce validation occurs in ExchangeCodeForIdentity during the initial auth flow.
        if tokens.IDToken != "" && p.verifier != nil {
                token, err := p.validateRefreshedIDToken(ctx, tokens)
                if err != nil {
                        return nil, err
                }
                if token != nil {
                        // The stored expectedSubject is the resolved subject (SubjectClaim, or
                        // "sub" by default). Resolve the refreshed token through the same path
                        // before comparing — comparing the raw "sub" would wrongly reject a
                        // refresh whenever a non-"sub" SubjectClaim is configured. OIDC Core
                        // Section 12.2 still holds: for the default "sub" this is identical to
                        // the original, and a custom claim must likewise be identical.
                        refreshedSubject, err := p.resolveSubject(token)
                        if err != nil {
                                return nil, fmt.Errorf("%w: %w", ErrIdentityResolutionFailed, err)
                        }
                        if expectedSubject != "" && refreshedSubject != expectedSubject {
                                return nil, ErrSubjectMismatch
                        }
                }
        }

        slog.Debug("token refresh successful",
                "has_new_refresh_token", tokens.RefreshToken != "",
                "expires_at", expiresAtLogValue(tokens.ExpiresAt),
        )

        return tokens, nil
}

// validateDiscoveryDocument validates the OIDC discovery document.
//
// It first delegates to OIDCDiscoveryDocument.Validate() for spec-compliant field
// validation (issuer, endpoints, jwks_uri, response_types_supported), then adds
// security validation for endpoint origins.
//
// Note: Issuer match validation (exact match per OIDC spec) is performed by go-oidc's
// NewProvider() before this function is called.
func validateDiscoveryDocument(doc *oauthproto.OIDCDiscoveryDocument, expectedIssuer string, insecureAllowHTTP bool) error {
        // Validate required OIDC fields per spec
        if err := doc.Validate(true); err != nil {
                return err
        }

        // Security: validate that discovered endpoints use secure schemes.
        // This prevents a malicious discovery document from redirecting requests to attacker-controlled servers.
        if err := validateEndpointOrigin(doc.AuthorizationEndpoint, expectedIssuer, insecureAllowHTTP); err != nil {
                return fmt.Errorf("authorization_endpoint origin mismatch: %w", err)
        }

        if err := validateEndpointOrigin(doc.TokenEndpoint, expectedIssuer, insecureAllowHTTP); err != nil {
                return fmt.Errorf("token_endpoint origin mismatch: %w", err)
        }

        // Optional endpoints - only validate if present
        if doc.UserinfoEndpoint != "" {
                if err := validateEndpointOrigin(doc.UserinfoEndpoint, expectedIssuer, insecureAllowHTTP); err != nil {
                        return fmt.Errorf("userinfo_endpoint origin mismatch: %w", err)
                }
        }

        if doc.JWKSURI != "" {
                if err := validateEndpointOrigin(doc.JWKSURI, expectedIssuer, insecureAllowHTTP); err != nil {
                        return fmt.Errorf("jwks_uri origin mismatch: %w", err)
                }
        }

        return nil
}

// validateEndpointOrigin validates that an endpoint URL uses a secure scheme relative to the issuer.
//
// This function enforces scheme consistency (HTTPS for production, HTTP allowed for localhost testing)
// but does NOT enforce host matching. Major identity providers like Google, Microsoft, and others
// commonly use different hosts/domains for their OAuth endpoints:
//   - Google: issuer=accounts.google.com, token_endpoint=oauth2.googleapis.com
//   - Microsoft: issuer=login.microsoftonline.com, various endpoint hosts
//
// The OIDC Discovery spec (https://openid.net/specs/openid-connect-discovery-1_0.html) and
// RFC 8414 (OAuth Authorization Server Metadata) do not require endpoints to be on the same
// host as the issuer. The security model relies on:
//  1. The discovery document being fetched over HTTPS from the configured issuer
//  2. TLS certificate validation ensuring we're talking to the real issuer
//  3. The issuer being a trusted party that controls its own discovery document
//
// If an attacker could compromise the HTTPS connection to the issuer or the issuer itself,
// host validation would provide no additional protection since the attacker controls the
// discovery document contents.
//
// When insecureAllowHTTP is true, HTTP endpoints are permitted for non-localhost issuers.
// This is intended for in-cluster development environments (e.g. a Dex instance in a
// kind cluster) where both issuer and endpoints are served over plain HTTP. Never set
// this in production.
func validateEndpointOrigin(endpoint, issuer string, insecureAllowHTTP bool) error {
        endpointURL, err := url.Parse(endpoint)
        if err != nil {
                return fmt.Errorf("invalid endpoint URL: %w", err)
        }

        issuerURL, err := url.Parse(issuer)
        if err != nil {
                return fmt.Errorf("invalid issuer URL: %w", err)
        }

        // For localhost issuers (development/testing), allow HTTP schemes and any localhost endpoint
        if networking.IsLocalhost(issuerURL.Host) {
                // Endpoint must also be localhost when issuer is localhost
                if !networking.IsLocalhost(endpointURL.Host) {
                        return fmt.Errorf("host mismatch: issuer is localhost but endpoint host is %q", endpointURL.Host)
                }
                // For localhost, we allow both HTTP and HTTPS, no further validation needed
                return nil
        }

        // insecureAllowHTTP permits HTTP for trusted in-cluster deployments (e.g. Dex in kind).
        if insecureAllowHTTP {
                return nil
        }

        // For production issuers, enforce HTTPS on endpoints
        // This prevents protocol downgrade attacks where a malicious discovery document
        // could redirect token requests to an HTTP endpoint, exposing credentials
        if endpointURL.Scheme != networking.HttpsScheme {
                return fmt.Errorf(
                        "scheme mismatch: issuer uses HTTPS but endpoint uses %q "+
                                "(all endpoints must use HTTPS for non-localhost issuers)",
                        endpointURL.Scheme)
        }

        // No host validation - the discovery document comes from a trusted HTTPS source
        // and major providers legitimately use different hosts for different endpoints
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package upstream

import (
        "bytes"
        "encoding/json"
        "io"
        "log/slog"
        "net/http"

        "github.com/tidwall/gjson"
)

// tokenResponseRewriter is an http.RoundTripper that normalizes non-standard
// OAuth token responses before the golang.org/x/oauth2 library parses them,
// and optionally extracts user identity claims from the raw response body.
//
// Some providers (e.g., GovSlack) nest token fields under non-standard paths
// like "authed_user.access_token" instead of the top-level "access_token".
// This RoundTripper intercepts the response, extracts fields using gjson
// dot-notation paths, and rewrites the response body with standard top-level
// field names so the oauth2 library can parse them normally.
//
// When identityCfg is set, identity extraction runs on the RAW pre-rewrite
// body before any field relocation occurs. The extracted identity is stored
// in extractedIdentity and consumed by the caller after Exchange returns.
// If extraction fails, extractionErr holds the cause so callers can surface
// operator-actionable diagnostics (path names, type descriptions) without
// re-reading the response body.
//
// Single-use contract: each instance handles exactly one token exchange and
// must not be shared across goroutines or reused for multiple exchanges. The
// caller reads extractedIdentity and extractionErr after Exchange returns —
// that return is the synchronization edge, so no mutex is needed for this
// one-shot usage pattern.
type tokenResponseRewriter struct {
        base              http.RoundTripper
        mapping           *TokenResponseMapping
        identityCfg       *IdentityFromTokenConfig
        tokenURL          string
        extractedIdentity *partialIdentity
        extractionErr     error
}

// RoundTrip intercepts HTTP responses from the token endpoint and rewrites
// the JSON body to place mapped fields at the top level. Non-token-endpoint
// requests (e.g., userInfo) pass through unchanged.
//
// When identityCfg is set, identity is extracted from the raw response body
// BEFORE the rewrite step so that identity paths are resolved against the
// original provider response, not the normalized form.
func (t *tokenResponseRewriter) RoundTrip(req *http.Request) (*http.Response, error) {
        resp, err := t.base.RoundTrip(req)
        if err != nil {
                return resp, err
        }

        // Only rewrite responses from the token endpoint
        if req.URL.String() != t.tokenURL {
                return resp, nil
        }

        // Only rewrite successful responses (errors should pass through for proper error handling)
        if resp.StatusCode != http.StatusOK {
                return resp, nil
        }

        body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseSize))
        _ = resp.Body.Close()
        if err != nil {
                return nil, err
        }

        // Extract from the raw body before rewriteTokenResponse runs. The rewrite
        // never touches identity-shaped fields today, but extracting first makes
        // the ordering invariant independent of that assumption.
        if t.identityCfg != nil {
                result, extractErr := extractIdentityFromTokenResponse(body, t.identityCfg)
                if extractErr != nil {
                        // WARN so an operator misconfiguration (e.g., wrong subjectPath) is
                        // visible without enabling DEBUG. The error is safe to log: it
                        // contains operator-supplied paths and type descriptions, never
                        // any portion of the response body.
                        slog.Warn("identity extraction from token response failed", "error", extractErr)
                        t.extractionErr = extractErr
                } else {
                        t.extractedIdentity = &result
                }
        }

        // Only run the field-rewrite step when a mapping is configured.
        // When mapping is nil (e.g. only identityCfg is set), pass the body through unchanged.
        var rewritten []byte
        if t.mapping != nil {
                rewritten = rewriteTokenResponse(body, t.mapping)
        } else {
                rewritten = body
        }

        resp.Body = io.NopCloser(bytes.NewReader(rewritten))
        resp.ContentLength = int64(len(rewritten))
        resp.Header.Del("Content-Length")
        return resp, nil
}

// rewriteTokenResponse extracts fields from the raw JSON using gjson paths
// and produces a new JSON object with standard OAuth 2.0 top-level field names.
// Fields that already exist at the top level and aren't overridden by the
// mapping are preserved.
func rewriteTokenResponse(body []byte, mapping *TokenResponseMapping) []byte {
        // Start with the original response to preserve any extra fields
        var original map[string]any
        if err := json.Unmarshal(body, &original); err != nil {
                // If we can't parse, return as-is and let oauth2 library handle the error
                return body
        }

        // Extract and set mapped fields at the top level
        if v := gjson.GetBytes(body, mapping.AccessTokenPath); v.Exists() {
                original["access_token"] = v.String()
        }

        // Always set token_type to "Bearer" for the oauth2 library.
        // Non-standard providers may use different values (e.g., GovSlack uses "user")
        // that the oauth2 library rejects. Since we're already using a custom mapping,
        // the original token_type value is not meaningful for standard validation.
        original["token_type"] = "Bearer"

        if path := pathOrDefault(mapping.RefreshTokenPath, "refresh_token"); path != "" {
                if v := gjson.GetBytes(body, path); v.Exists() {
                        original["refresh_token"] = v.String()
                }
        }

        if path := pathOrDefault(mapping.ExpiresInPath, "expires_in"); path != "" {
                if v := gjson.GetBytes(body, path); v.Exists() && v.Int() > 0 {
                        original["expires_in"] = v.Int()
                }
        }

        if path := pathOrDefault(mapping.ScopePath, "scope"); path != "" {
                if v := gjson.GetBytes(body, path); v.Exists() {
                        original["scope"] = v.String()
                }
        }

        rewritten, err := json.Marshal(original)
        if err != nil {
                return body
        }
        return rewritten
}

// wrapHTTPClientForTokenExchange wraps an HTTP client's transport with a
// tokenResponseRewriter when either a TokenResponseMapping or an
// IdentityFromTokenConfig is configured (or both). Returns the original client
// and nil rewriter when both are nil, so the standard oauth2 library path is
// used. The returned rewriter (when non-nil) can be read after Exchange returns
// to retrieve any identity extracted during the token round-trip.
func wrapHTTPClientForTokenExchange(
        client *http.Client,
        mapping *TokenResponseMapping,
        identityCfg *IdentityFromTokenConfig,
        tokenURL string,
) (*http.Client, *tokenResponseRewriter) {
        if mapping == nil && identityCfg == nil {
                return client, nil
        }

        base := client.Transport
        if base == nil {
                base = http.DefaultTransport
        }

        rewriter := &tokenResponseRewriter{
                base:        base,
                mapping:     mapping,
                identityCfg: identityCfg,
                tokenURL:    tokenURL,
        }

        // Create a shallow copy to avoid mutating the original client
        wrapped := *client
        wrapped.Transport = rewriter
        return &wrapped, rewriter
}

// pathOrDefault returns the path if non-empty, otherwise returns the default.
func pathOrDefault(path, defaultPath string) string {
        if path != "" {
                return path
        }
        return defaultPath
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package upstream

import (
        "log/slog"
        "time"
)

// tokenExpirationBuffer is the time buffer before actual expiration to consider a token expired.
// This accounts for clock skew and network latency.
const tokenExpirationBuffer = 30 * time.Second

// Tokens represents the tokens obtained from an upstream Identity Provider.
// This type is used for token exchange with the IDP, but stored separately
// (see storage.IDPTokens for the storage representation).
type Tokens struct {
        // AccessToken is the access token from the upstream IDP.
        AccessToken string //nolint:gosec // G117: field legitimately holds sensitive data

        // RefreshToken is the refresh token from the upstream IDP (if provided).
        RefreshToken string //nolint:gosec // G117: field legitimately holds sensitive data

        // IDToken is the ID token from the upstream IDP (for OIDC).
        IDToken string

        // ExpiresAt is when the access token expires. Zero value means the provider
        // did not assert an expiry; callers must treat it as non-expiring.
        ExpiresAt time.Time
}

// IsExpired returns true if the access token has expired or will expire within the buffer period.
// Returns true for nil receivers (treating nil tokens as expired).
func (t *Tokens) IsExpired() bool {
        return t.IsExpiredAt(time.Now())
}

// IsExpiredAt returns true if the access token has expired or will expire within the buffer period
// at the given time. This method is primarily for testing to avoid time-based race conditions.
// Returns true for nil receivers (treating nil tokens as expired).
func (t *Tokens) IsExpiredAt(now time.Time) bool {
        if t == nil {
                return true
        }
        if t.ExpiresAt.IsZero() {
                return false
        }
        // Token is expired if it expires at or before (now + buffer)
        // Using !After to include the equality case (expires exactly at boundary)
        return !t.ExpiresAt.After(now.Add(tokenExpirationBuffer))
}

// expiresAtLogValue is a slog.LogValuer wrapper for an ExpiresAt time that
// renders zero time as "none" rather than the misleading year-0001 timestamp
// slog would otherwise produce. As a LogValuer, formatting is deferred until
// the log record is actually emitted, so DEBUG logs do no work when the
// handler level filters them out.
type expiresAtLogValue time.Time

// LogValue implements slog.LogValuer.
func (e expiresAtLogValue) LogValue() slog.Value {
        t := time.Time(e)
        if t.IsZero() {
                return slog.StringValue("none")
        }
        return slog.StringValue(t.Format(time.RFC3339))
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package upstream

import (
        "errors"
        "fmt"
        "net/http"
        "net/url"
        "os"
        "strings"

        "github.com/stacklok/toolhive/pkg/networking"
)

// UserInfoFieldMapping maps provider-specific field names to standard UserInfo fields.
// This allows adapting non-standard provider responses to the canonical UserInfo structure.
//
// Each field supports an ordered list of claim names to try. The first non-empty value
// found will be used. This enables compatibility with providers that use non-standard
// claim names (e.g., GitHub uses "id" instead of "sub", "login" instead of "preferred_username").
//
// Example for GitHub:
//
//        mapping := &UserInfoFieldMapping{
//            SubjectFields: []string{"id", "login"},
//            NameFields:    []string{"name", "login"},
//            EmailFields:   []string{"email"},
//        }
type UserInfoFieldMapping struct {
        // SubjectFields is an ordered list of field names to try for the user ID.
        // The first non-empty value found will be used.
        // Default: ["sub"]
        SubjectFields []string `json:"subject_fields,omitempty" yaml:"subject_fields,omitempty"`

        // NameFields is an ordered list of field names to try for the display name.
        // The first non-empty value found will be used.
        // Default: ["name"]
        NameFields []string `json:"name_fields,omitempty" yaml:"name_fields,omitempty"`

        // EmailFields is an ordered list of field names to try for the email address.
        // The first non-empty value found will be used.
        // Default: ["email"]
        EmailFields []string `json:"email_fields,omitempty" yaml:"email_fields,omitempty"`
}

// Default field names for standard OIDC claims.
var (
        // DefaultSubjectFields is the default field name for subject resolution.
        DefaultSubjectFields = []string{"sub"}

        // DefaultNameFields is the default field name for name resolution.
        DefaultNameFields = []string{"name"}

        // DefaultEmailFields is the default field name for email resolution.
        DefaultEmailFields = []string{"email"}
)

// GetSubjectFields returns the configured subject fields or the default.
func (m *UserInfoFieldMapping) GetSubjectFields() []string {
        if m != nil && len(m.SubjectFields) > 0 {
                return m.SubjectFields
        }
        return DefaultSubjectFields
}

// GetNameFields returns the configured name fields or the default.
func (m *UserInfoFieldMapping) GetNameFields() []string {
        if m != nil && len(m.NameFields) > 0 {
                return m.NameFields
        }
        return DefaultNameFields
}

// GetEmailFields returns the configured email fields or the default.
func (m *UserInfoFieldMapping) GetEmailFields() []string {
        if m != nil && len(m.EmailFields) > 0 {
                return m.EmailFields
        }
        return DefaultEmailFields
}

// ResolveField attempts to extract a string value from claims using an ordered list of field names.
// Returns the first non-empty string value found, or an empty string if none found.
// This function handles type conversion gracefully - non-string values are skipped.
func ResolveField(claims map[string]any, fields []string) string {
        for _, field := range fields {
                if val, ok := claims[field]; ok {
                        switch v := val.(type) {
                        case string:
                                if v != "" {
                                        return v
                                }
                        case float64:
                                // Handle numeric IDs (e.g., GitHub returns numeric "id")
                                return fmt.Sprintf("%.0f", v)
                        case int:
                                return fmt.Sprintf("%d", v)
                        case int64:
                                return fmt.Sprintf("%d", v)
                        }
                }
        }
        return ""
}

// ResolveSubject extracts the subject (user ID) from claims using the configured mapping.
// Returns an error if no subject can be resolved (subject is required).
func (m *UserInfoFieldMapping) ResolveSubject(claims map[string]any) (string, error) {
        fields := m.GetSubjectFields()
        sub := ResolveField(claims, fields)
        if sub == "" {
                return "", fmt.Errorf("subject claim not found (tried fields: %v)", fields)
        }
        return sub, nil
}

// ResolveName extracts the display name from claims using the configured mapping.
// Returns an empty string if no name is found (name is optional).
func (m *UserInfoFieldMapping) ResolveName(claims map[string]any) string {
        return ResolveField(claims, m.GetNameFields())
}

// ResolveEmail extracts the email from claims using the configured mapping.
// Returns an empty string if no email is found (email is optional).
func (m *UserInfoFieldMapping) ResolveEmail(claims map[string]any) string {
        return ResolveField(claims, m.GetEmailFields())
}

// UserInfoConfig contains configuration for fetching user information from an upstream provider.
// This supports both standard OIDC UserInfo endpoints and custom provider-specific endpoints.
// Authentication is always performed using Bearer token in the Authorization header.
type UserInfoConfig struct {
        // EndpointURL is the URL of the userinfo endpoint (required).
        EndpointURL string `json:"endpoint_url" yaml:"endpoint_url"`

        // HTTPMethod is the HTTP method to use (default: GET).
        HTTPMethod string `json:"http_method,omitempty" yaml:"http_method,omitempty"`

        // AdditionalHeaders contains extra headers to include in the request.
        AdditionalHeaders map[string]string `json:"additional_headers,omitempty" yaml:"additional_headers,omitempty"`

        // FieldMapping contains custom field mapping configuration.
        // If nil, standard OIDC field names are used ("sub", "name", "email").
        FieldMapping *UserInfoFieldMapping `json:"field_mapping,omitempty" yaml:"field_mapping,omitempty"`
}

// Validate checks that UserInfoConfig has all required fields and valid values.
func (c *UserInfoConfig) Validate() error {
        if c.EndpointURL == "" {
                return errors.New("endpoint_url is required")
        }

        parsed, err := url.Parse(c.EndpointURL)
        if err != nil {
                return errors.New("endpoint_url must be a valid URL")
        }

        if parsed.Scheme == "" || parsed.Host == "" {
                return errors.New("endpoint_url must be an absolute URL with scheme and host")
        }

        if parsed.Scheme != networking.HttpScheme && parsed.Scheme != networking.HttpsScheme {
                return errors.New("endpoint_url must use http or https scheme")
        }

        // HTTP scheme is only allowed for loopback addresses unless URL validation is disabled.
        // This is consistent with networking.ValidateEndpointURL which checks the same env var.
        if parsed.Scheme == networking.HttpScheme && !networking.IsLocalhost(parsed.Host) {
                if !strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true") {
                        return errors.New("endpoint_url with http scheme requires loopback address (127.0.0.1, ::1, or localhost)")
                }
        }

        // Validate HTTP method if specified (OIDC Core Section 5.3.1 allows GET and POST)
        if c.HTTPMethod != "" && c.HTTPMethod != http.MethodGet && c.HTTPMethod != http.MethodPost {
                return errors.New("http_method must be GET or POST")
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package upstreamfilter

import (
        "context"
        "errors"
        "fmt"

        "github.com/stacklok/toolhive/pkg/authserver/server/handlers"
        "github.com/stacklok/toolhive/pkg/oauthproto/claims"
)

const defaultClaimName = "groups"

// GroupBasedFilter implements handlers.UpstreamFilter by matching the
// principal's group claim against configured rules and returning the union of
// the matching rules' upstream provider sets. See package documentation and
// GroupBasedFilterConfig for semantics.
type GroupBasedFilter struct {
        claim            string
        rules            []GroupRule
        defaultUpstreams []string
}

// Compile-time assertion that GroupBasedFilter satisfies the UpstreamFilter
// contract.
var _ handlers.UpstreamFilter = (*GroupBasedFilter)(nil)

// NewGroupBasedFilter constructs and validates a GroupBasedFilter.
//
// configuredUpstreams is the full ordered list of upstream provider names
// declared on the auth server, including the mandatory first upstream at
// index 0. It is required at construction time to validate that every
// referenced provider exists and that no rule or default references the
// mandatory first upstream (which is always included and cannot be filtered).
//
// It returns an error for any of: fewer than two configured upstreams; no
// rules; empty or duplicate group values within a rule; empty or duplicate
// upstream references within a rule; unknown upstream references; references
// to the mandatory first upstream; unknown or first-upstream references in
// DefaultUpstreams.
func NewGroupBasedFilter(cfg GroupBasedFilterConfig, configuredUpstreams []string) (*GroupBasedFilter, error) {
        if len(configuredUpstreams) < 2 {
                return nil, fmt.Errorf(
                        "group-based upstream filter requires at least two configured upstreams, got %d",
                        len(configuredUpstreams),
                )
        }
        if len(cfg.Rules) == 0 {
                return nil, errors.New("group-based upstream filter requires at least one rule")
        }

        claim := cfg.Claim
        if claim == "" {
                claim = defaultClaimName
        }

        firstUpstream := configuredUpstreams[0]
        validProviders := make(map[string]struct{}, len(configuredUpstreams)-1)
        for _, name := range configuredUpstreams[1:] {
                validProviders[name] = struct{}{}
        }

        for i, rule := range cfg.Rules {
                if err := validateRule(rule, i, firstUpstream, validProviders); err != nil {
                        return nil, err
                }
        }
        if err := validateProviderList("defaultUpstreams", cfg.DefaultUpstreams, firstUpstream, validProviders); err != nil {
                return nil, err
        }

        return &GroupBasedFilter{
                claim:            claim,
                rules:            cloneRules(cfg.Rules),
                defaultUpstreams: append([]string(nil), cfg.DefaultUpstreams...),
        }, nil
}

// FilterUpstreams implements handlers.UpstreamFilter.
//
// It reads principal.GetClaims()[filter.claim], decodes it as a []string via the
// strict OAuth claims decoder, and returns the union of every matching rule's
// UpstreamProviders. The claim MUST be present and MUST be a JSON string
// array — a missing, nil, non-array, or mixed-type claim value fails the
// authorization per the filter contract's fail-closed guarantee for
// security-relevant filters.
//
// A valid empty array or a valid array with no matching rule falls back to
// the configured DefaultUpstreams. Duplicate names across matching rules and
// defaults are deduplicated in the returned slice; the handler layer preserves
// configured order and drops any name that is not a non-first configured
// upstream.
//
// principal (including the map returned by GetClaims()) is treated as read-only. Claim values
// are not logged.
//
// The configured argument is unused: the constructor validated every rule and
// default reference against the full configured upstream list, so runtime
// reconciliation of the returned set against the non-first upstreams is
// redundant with the handler layer's own drop-unknown behavior.
func (f *GroupBasedFilter) FilterUpstreams(
        _ context.Context,
        principal handlers.FilterPrincipal,
        _ []string,
) ([]string, error) {
        principalClaims := principal.GetClaims()
        if principalClaims == nil {
                return nil, errors.New("upstream filter: principal claims are absent")
        }
        raw, present := principalClaims[f.claim]
        if !present {
                return nil, fmt.Errorf("upstream filter: claim %q not present in principal claims", f.claim)
        }
        groups, err := claims.DecodeStringArray(raw)
        if err != nil {
                return nil, fmt.Errorf("upstream filter: claim %q: %w", f.claim, err)
        }

        if len(groups) == 0 {
                return append([]string(nil), f.defaultUpstreams...), nil
        }

        principalGroupSet := make(map[string]struct{}, len(groups))
        for _, g := range groups {
                principalGroupSet[g] = struct{}{}
        }

        effective := make([]string, 0)
        seen := make(map[string]struct{})
        matched := false
        for _, rule := range f.rules {
                if !ruleMatches(rule, principalGroupSet) {
                        continue
                }
                matched = true
                for _, name := range rule.UpstreamProviders {
                        if _, dup := seen[name]; dup {
                                continue
                        }
                        seen[name] = struct{}{}
                        effective = append(effective, name)
                }
        }
        if !matched {
                return append([]string(nil), f.defaultUpstreams...), nil
        }
        return effective, nil
}

// ruleMatches returns true if any of the rule's Groups values is in the
// principal's group set (any-of / OR semantics).
func ruleMatches(rule GroupRule, principalGroupSet map[string]struct{}) bool {
        for _, g := range rule.Groups {
                if _, ok := principalGroupSet[g]; ok {
                        return true
                }
        }
        return false
}

func validateRule(rule GroupRule, index int, firstUpstream string, validProviders map[string]struct{}) error {
        if len(rule.Groups) == 0 {
                return fmt.Errorf("rules[%d]: groups must not be empty", index)
        }
        if hasEmpty(rule.Groups) {
                return fmt.Errorf("rules[%d]: groups must not contain empty values", index)
        }
        if hasDuplicates(rule.Groups) {
                return fmt.Errorf("rules[%d]: groups must not contain duplicate values", index)
        }
        if len(rule.UpstreamProviders) == 0 {
                return fmt.Errorf("rules[%d]: upstreamProviders must not be empty", index)
        }
        return validateProviderList(
                fmt.Sprintf("rules[%d].upstreamProviders", index),
                rule.UpstreamProviders,
                firstUpstream,
                validProviders,
        )
}

func validateProviderList(field string, names []string, firstUpstream string, validProviders map[string]struct{}) error {
        if hasEmpty(names) {
                return fmt.Errorf("%s: must not contain empty values", field)
        }
        if hasDuplicates(names) {
                return fmt.Errorf("%s: must not contain duplicate values", field)
        }
        for _, name := range names {
                if name == firstUpstream {
                        return fmt.Errorf("%s: must not reference the mandatory first upstream %q", field, firstUpstream)
                }
                if _, ok := validProviders[name]; !ok {
                        return fmt.Errorf("%s: %q is not a configured upstream", field, name)
                }
        }
        return nil
}

func hasEmpty(values []string) bool {
        for _, v := range values {
                if v == "" {
                        return true
                }
        }
        return false
}

func hasDuplicates(values []string) bool {
        seen := make(map[string]struct{}, len(values))
        for _, v := range values {
                if _, dup := seen[v]; dup {
                        return true
                }
                seen[v] = struct{}{}
        }
        return false
}

func cloneRules(in []GroupRule) []GroupRule {
        out := make([]GroupRule, len(in))
        for i, r := range in {
                out[i] = GroupRule{
                        Groups:            append([]string(nil), r.Groups...),
                        UpstreamProviders: append([]string(nil), r.UpstreamProviders...),
                }
        }
        return out
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authz

import (
        "sync"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive/pkg/authz/authorizers"
)

// AnnotationCache stores tool annotations indexed by tool name.
// It is safe for concurrent use and nil-safe: calling methods on a nil
// *AnnotationCache is a no-op that returns zero values.
//
// The cache is populated from tools/list responses (via SetFromToolsList)
// and read during tools/call authorization (via Get). This bridges the gap
// between the two separate HTTP requests: annotations are visible in the
// tools/list response but absent from the tools/call request body.
type AnnotationCache struct {
        mu    sync.RWMutex
        tools map[string]*authorizers.ToolAnnotations
}

// NewAnnotationCache creates a new empty annotation cache.
func NewAnnotationCache() *AnnotationCache {
        return &AnnotationCache{
                tools: make(map[string]*authorizers.ToolAnnotations),
        }
}

// Get retrieves annotations for a tool by name. Returns nil if the tool
// is not cached or if the cache itself is nil.
func (c *AnnotationCache) Get(toolName string) *authorizers.ToolAnnotations {
        if c == nil {
                return nil
        }
        c.mu.RLock()
        defer c.mu.RUnlock()
        return c.tools[toolName]
}

// Set stores annotations for a single tool. It is a no-op if the cache
// is nil.
func (c *AnnotationCache) Set(toolName string, annotations *authorizers.ToolAnnotations) {
        if c == nil {
                return
        }
        c.mu.Lock()
        defer c.mu.Unlock()
        c.tools[toolName] = annotations
}

// SetFromToolsList extracts annotations from a tools/list response and
// replaces the entire cache contents. The full replacement ensures that
// tools whose annotations were removed in a subsequent tools/list response
// do not retain stale cached entries.
//
// Only tools that have at least one non-nil annotation hint are cached;
// tools with all-nil hints (the zero value) are skipped to avoid
// unnecessary memory consumption.
//
// This method is a no-op if the cache is nil.
func (c *AnnotationCache) SetFromToolsList(tools []mcp.Tool) {
        if c == nil {
                return
        }
        newTools := make(map[string]*authorizers.ToolAnnotations, len(tools))
        for i := range tools {
                ann := &tools[i].Annotations
                if !hasAnyHint(ann) {
                        continue
                }
                newTools[tools[i].Name] = convertMCPAnnotation(ann)
        }
        c.mu.Lock()
        defer c.mu.Unlock()
        c.tools = newTools
}

// hasAnyHint reports whether the MCP tool annotation has at least one
// non-nil hint field.
func hasAnyHint(ann *mcp.ToolAnnotation) bool {
        return ann.ReadOnlyHint != nil ||
                ann.DestructiveHint != nil ||
                ann.IdempotentHint != nil ||
                ann.OpenWorldHint != nil
}

// convertMCPAnnotation converts an mcp-go ToolAnnotation to the authz
// ToolAnnotations type used by authorizers. Only hint fields are copied;
// the Title field is intentionally omitted because authorizers only use
// hints for policy decisions.
func convertMCPAnnotation(ann *mcp.ToolAnnotation) *authorizers.ToolAnnotations {
        return &authorizers.ToolAnnotations{
                ReadOnlyHint:    ann.ReadOnlyHint,
                DestructiveHint: ann.DestructiveHint,
                IdempotentHint:  ann.IdempotentHint,
                OpenWorldHint:   ann.OpenWorldHint,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authorizers

import "context"

// ToolAnnotations holds MCP tool annotation hints that inform authorization
// decisions. The fields match the MCP specification's tool annotation schema.
// Pointer types are used so callers can distinguish "not set" (nil) from
// an explicit false value.
//
// # Trust Boundary
//
// Annotations MUST be sourced from the server-side tool registry (the MCP
// tools/list response), NOT from the client's tools/call request body.
// Allowing clients to supply their own annotations would let a malicious
// caller set readOnlyHint=true on a destructive tool and bypass policies
// that rely on annotations.
//
// # Authorizer Exposure Paths
//
// The two authorizer implementations expose annotations at different
// locations so that policy authors can reference them:
//
//   - Cedar authorizer: flat on the resource entity — e.g. resource.readOnlyHint
//   - HTTP PDP authorizer: nested in the PORC context — context.mcp.annotations.readOnlyHint
type ToolAnnotations struct {
        ReadOnlyHint    *bool `json:"readOnlyHint,omitempty"`
        DestructiveHint *bool `json:"destructiveHint,omitempty"`
        IdempotentHint  *bool `json:"idempotentHint,omitempty"`
        OpenWorldHint   *bool `json:"openWorldHint,omitempty"`
}

// toolAnnotationsKey is the unexported context key used by
// WithToolAnnotations / ToolAnnotationsFromContext.
type toolAnnotationsKey struct{}

// WithToolAnnotations stores tool annotations in the given context.
func WithToolAnnotations(ctx context.Context, annotations *ToolAnnotations) context.Context {
        return context.WithValue(ctx, toolAnnotationsKey{}, annotations)
}

// ToolAnnotationsFromContext retrieves tool annotations previously stored with
// WithToolAnnotations or as part of ResourceMetadata. It returns nil when no
// annotations are present.
func ToolAnnotationsFromContext(ctx context.Context) *ToolAnnotations {
        return effectiveToolAnnotationsFromContext(ctx)
}

// effectiveToolAnnotationsFromContext keeps the legacy and unified metadata
// readers consistent when both context channels are present. Unified resource
// metadata is authoritative; the legacy value is a compatibility fallback.
func effectiveToolAnnotationsFromContext(ctx context.Context) *ToolAnnotations {
        if metadata, ok := ctx.Value(resourceMetadataKey{}).(ResourceMetadata); ok && metadata.Annotations != nil {
                return metadata.Annotations
        }
        annotations, _ := ctx.Value(toolAnnotationsKey{}).(*ToolAnnotations)
        return annotations
}

// AnnotationsToMap converts non-nil annotation fields to a flat map suitable
// for merging into Cedar resource attributes or HTTP PDP context. Returns nil
// when annotations is nil. Returns an empty (non-nil) map when annotations is
// non-nil but all fields are nil pointers.
func AnnotationsToMap(annotations *ToolAnnotations) map[string]interface{} {
        if annotations == nil {
                return nil
        }

        m := make(map[string]interface{})
        if annotations.ReadOnlyHint != nil {
                m["readOnlyHint"] = *annotations.ReadOnlyHint
        }
        if annotations.DestructiveHint != nil {
                m["destructiveHint"] = *annotations.DestructiveHint
        }
        if annotations.IdempotentHint != nil {
                m["idempotentHint"] = *annotations.IdempotentHint
        }
        if annotations.OpenWorldHint != nil {
                m["openWorldHint"] = *annotations.OpenWorldHint
        }
        return m
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package cedar provides authorization utilities using Cedar policies.
package cedar

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "slices"
        "strings"
        "sync"
        "time"

        cedar "github.com/cedar-policy/cedar-go"
        "github.com/golang-jwt/jwt/v5"

        "github.com/stacklok/toolhive/pkg/auth"
        authclaims "github.com/stacklok/toolhive/pkg/auth/claims"
        "github.com/stacklok/toolhive/pkg/authz/authorizers"
        "github.com/stacklok/toolhive/pkg/syncutil"
)

// ConfigType is the configuration type identifier for Cedar authorization.
const ConfigType = "cedarv1"

func init() {
        // Register the Cedar authorizer factory with the authorizers registry.
        authorizers.Register(ConfigType, &Factory{})
}

// Config represents the complete authorization configuration file structure
// for Cedar authorization. This includes the common version/type fields plus
// the Cedar-specific "cedar" field. This maintains backwards compatibility
// with the v1.0 configuration schema.
type Config struct {
        Version string         `json:"version"`
        Type    string         `json:"type"`
        Options *ConfigOptions `json:"cedar"`
}

// ExtractConfig extracts the Cedar configuration from an authorizers.Config.
// This is useful for tests and other code that needs to inspect the Cedar configuration
// after it has been loaded into the generic Config structure.
// To access the Cedar-specific options (policies, entities), use the returned Config's Cedar field.
func ExtractConfig(authzConfig *authorizers.Config) (*Config, error) {
        if authzConfig == nil {
                return nil, fmt.Errorf("config is nil")
        }
        rawConfig := authzConfig.RawConfig()
        if len(rawConfig) == 0 {
                return nil, fmt.Errorf("config has no raw data")
        }

        var config Config
        if err := json.Unmarshal(rawConfig, &config); err != nil {
                return nil, fmt.Errorf("failed to unmarshal config: %w", err)
        }
        if config.Options == nil {
                return nil, fmt.Errorf("cedar config is nil")
        }
        return &config, nil
}

// InjectUpstreamProvider returns a new authorizers.Config that is identical to
// src except that the Cedar options' PrimaryUpstreamProvider field is set to
// providerName. Any existing PrimaryUpstreamProvider value is overwritten; if
// the Cedar config file already contains a non-empty PrimaryUpstreamProvider
// that differs from providerName, the file value is silently replaced. This is
// intentional: the embedded auth server config is the authoritative source of
// the upstream provider name at runtime. This is used by the runner middleware
// when the embedded auth server is active to wire the upstream provider into
// Cedar evaluation.
//
// If src is not a Cedar config, providerName is empty, or src is nil, src is
// returned unchanged with a nil error. This makes the function safe to call
// unconditionally whenever the embedded auth server is active.
func InjectUpstreamProvider(src *authorizers.Config, providerName string) (*authorizers.Config, error) {
        if src == nil || providerName == "" {
                return src, nil
        }

        cedarCfg, err := ExtractConfig(src)
        if err != nil {
                // src is not a Cedar config (e.g. a future HTTP authorizer); treat as a
                // no-op so callers can apply this unconditionally without needing to
                // know the authorizer type ahead of time.
                slog.Debug("skipping upstream provider injection for non-Cedar config",
                        "provider", providerName, "type", src.Type)
                return src, nil
        }

        cedarCfg.Options.PrimaryUpstreamProvider = providerName
        return authorizers.NewConfig(cedarCfg)
}

// Factory implements the authorizers.AuthorizerFactory interface for Cedar.
type Factory struct{}

// ConfigKey returns the JSON key for Cedar-specific configuration ("cedar").
func (*Factory) ConfigKey() string { return "cedar" }

// ValidateConfig validates the Cedar-specific configuration.
// It receives the full raw config and extracts the Cedar-specific portion.
func (*Factory) ValidateConfig(rawConfig json.RawMessage) error {
        var config Config
        if err := json.Unmarshal(rawConfig, &config); err != nil {
                return fmt.Errorf("failed to parse configuration: %w", err)
        }

        if config.Options == nil {
                return fmt.Errorf("cedar configuration is required (missing 'cedar' field)")
        }

        if len(config.Options.Policies) == 0 {
                return fmt.Errorf("at least one policy is required for Cedar authorization")
        }

        return nil
}

// CreateAuthorizer creates a Cedar Authorizer from the configuration.
// It receives the full raw config and extracts the Cedar-specific portion.
func (*Factory) CreateAuthorizer(rawConfig json.RawMessage, serverName string) (authorizers.Authorizer, error) {
        var config Config
        if err := json.Unmarshal(rawConfig, &config); err != nil {
                return nil, fmt.Errorf("failed to parse configuration: %w", err)
        }

        if config.Options == nil {
                return nil, fmt.Errorf("cedar configuration is required (missing 'cedar' field)")
        }

        return NewCedarAuthorizer(*config.Options, serverName)
}

// Common errors for Cedar authorization
var (
        ErrNoPolicies           = errors.New("no policies loaded")
        ErrInvalidPolicy        = errors.New("invalid policy")
        ErrUnauthorized         = errors.New("unauthorized")
        ErrMissingPrincipal     = errors.New("missing principal")
        ErrMissingAction        = errors.New("missing action")
        ErrMissingResource      = errors.New("missing resource")
        ErrFailedToLoadEntities = errors.New("failed to load entities")
)

// ClientIDContextKey is the key used to store client ID in the context.
type ClientIDContextKey struct{}

// claimPrefix namespaces every JWT claim in the Cedar context/principal
// attributes. claimSetPrefix namespaces the synthetic multi-valued-claim Sets.
// They are deliberately disjoint: because every JWT claim is emitted under
// claimPrefix, a token can never produce a claimSetPrefix key, so the synthetic
// Sets cannot be shadowed or spoofed by claim content.
const (
        claimPrefix    = "claim_"
        claimSetPrefix = "claimset_"
)

// claimSourceAttr is the principal attribute and context key naming which trust
// root asserted the claims this request was evaluated against. Policies read it
// to require upstream provenance, e.g.
//
//        when { principal.thv_claim_source == "upstream:github" }
//
// The name deliberately does not start with claimPrefix or claimSetPrefix: every
// JWT claim is emitted under claimPrefix, so a token cannot produce this key and
// cannot spoof its own provenance. It is written after preprocessClaims for the
// same reason.
const claimSourceAttr = "thv_claim_source"

// Claim-source values. Each names the trust root behind the claims Cedar saw, so
// a policy can tell an upstream assertion from the request token's own word for
// it. The upstream case carries the provider name, since which upstream asserted
// a claim is exactly what a pinned deployment cares about.
//
// The three request-token values are not collapsed into one: they answer
// different questions. claimSourceRequest means no provenance was ever demanded.
// claimSourceNoUpstreamSession means one was demanded and the token affirmatively
// said it has none. claimSourceUpstreamOpaque means one was demanded, a
// credential exists, and it simply cannot be read — a degradation that until now
// no policy could see.
const (
        claimSourceRequest           = "request"
        claimSourceNoUpstreamSession = "request:no-upstream-session"
        claimSourceUpstreamOpaque    = "request:upstream-opaque"
        claimSourceUpstreamPrefix    = "upstream:"
)

// Authorizer authorizes MCP operations using Cedar policies.
type Authorizer struct {
        // Cedar policy set
        policySet *cedar.PolicySet
        // Cedar entities
        entities cedar.EntityMap
        // Entity factory for creating entities
        entityFactory *EntityFactory
        // Mutex for thread safety
        mu sync.RWMutex
        // primaryUpstreamProvider names the upstream IDP provider whose access token
        // is the source of JWT claims for Cedar evaluation, aside from the narrow
        // user-profile supplement described in resolveClaims.
        // When empty, claims from the token on the original client request are used,
        // which may be a ToolHive-issued token or any other bearer token.
        primaryUpstreamProvider string
        // groupClaimName is the JWT claim key that contains group membership.
        // When empty, the well-known defaults are checked ("groups", "roles", etc.).
        groupClaimName string
        // roleClaimName is the JWT claim key that contains role membership.
        // When empty, no role extraction is performed (backward compatible).
        roleClaimName string
        // serverName is the identity of the MCP server this authorizer is scoped to.
        // Used by downstream enterprise features for server-scoped Cedar policies
        // (e.g. resource in MCP::"<server>"). When empty (standalone Cedar usage
        // with no enterprise controller), the authorizer behaves identically to
        // the unscoped case.
        serverName string
        // claimKeyLog rate-limits the diagnostic log of resolved JWT claim keys
        // so it emits at most once per 30 seconds instead of once per authorization check.
        claimKeyLog *syncutil.AtMost
        // supplementLog rate-limits the warning that the upstream access token lacked
        // profile claims. It is deliberately separate from claimKeyLog: a shared
        // timer would let whichever fires first suppress the other for the rest of
        // the window, hiding the claim-key dump on exactly the deployments that need it.
        supplementLog *syncutil.AtMost
        // missingIDTokenLog rate-limits the warning that the primary provider has no
        // usable stored id_token, so profile claims cannot be supplemented at all.
        // Separate from supplementLog because the two are mutually exclusive per
        // request and describe opposite conditions.
        missingIDTokenLog *syncutil.AtMost
        // noSessionLog rate-limits the warning that a request has no upstream session
        // and Cedar is using request-token claims instead. It is separate from
        // claimKeyLog so the warning does not suppress the claim-key diagnostic.
        noSessionLog *syncutil.AtMost
        // multiValuedClaims lists JWT claim names normalized to a canonical unpadded
        // space-delimited string, plus a companion Cedar Set, before Cedar evaluation.
        // See ConfigOptions.MultiValuedClaims.
        multiValuedClaims []string
}

// ConfigOptions represents the Cedar-specific authorization configuration options.
type ConfigOptions struct {
        // Policies is a list of Cedar policy strings
        Policies []string `json:"policies" yaml:"policies"`

        // EntitiesJSON is the JSON string representing Cedar entities
        EntitiesJSON string `json:"entities_json" yaml:"entities_json"`

        // PrimaryUpstreamProvider names the upstream IDP provider whose access
        // token is the primary source of JWT claims for Cedar evaluation.
        // When empty, claims from the original client-request token are used, whether
        // it is ToolHive-issued or another bearer token.
        // When set, a non-nil identity.UpstreamTokens map must contain a non-empty
        // token for this provider; otherwise authorization fails closed. The sole
        // exception is a nil map with identity.SessionlessGrant set by TokenValidator
        // after validation for the configured embedded authorization server; then
        // request-token claims are used.
        //
        // The profile claims in profileClaimsFromIDToken fall back to the SAME
        // provider's id_token when its access token omits them (many OIDC providers
        // assert `email` only in the id_token), so `principal has claim_email` keeps
        // meaning "this upstream asserted an email". Every other claim, including all
        // group/role/scope claims, comes from the upstream access token or not at all.
        // The sole nil-map exception uses only request-token claims because no upstream
        // provider credential exists, never as a substitute for another provider's
        // credentials.
        //
        // Separately from that exception: when the provider's access token is present
        // but opaque rather than JWT-shaped (Google's ya29.*, GitHub's gho_*), its
        // claims cannot be read at all and evaluation degrades to the request token's
        // claims, labelled claimSourceUpstreamOpaque. That is a fallback, not an
        // exception to the fail-closed rule above — the credential is there, it just
        // cannot be parsed. See resolveClaims for the full contract.
        PrimaryUpstreamProvider string `json:"primary_upstream_provider,omitempty" yaml:"primary_upstream_provider,omitempty"`

        // GroupClaimName is the JWT claim key that contains group membership for the
        // principal. When set, it takes priority over the well-known defaults
        // ("groups", "roles", "cognito:groups"). Use this for IDPs that place groups
        // under a URI-style claim (e.g. "https://example.com/groups" in Auth0/Okta).
        // When empty, only the well-known claim names are checked.
        GroupClaimName string `json:"group_claim_name,omitempty" yaml:"group_claim_name,omitempty"`

        // RoleClaimName is the JWT claim key that contains role membership for the
        // principal. When set, the claim is extracted separately from GroupClaimName
        // and both are mapped to the configured group entity type (default "THVGroup").
        // When empty, no role extraction is performed (backward compatible).
        RoleClaimName string `json:"role_claim_name,omitempty" yaml:"role_claim_name,omitempty"`

        // GroupEntityType is the Cedar entity type name used for Client parent UIDs
        // synthesised from JWT group/role claims. Defaults to "THVGroup" when empty,
        // preserving the original behaviour. Must be a valid Cedar identifier — namespaced
        // names (e.g. "Platform::Group") are not yet supported and are rejected at
        // construction. See issue #5072.
        GroupEntityType string `json:"group_entity_type,omitempty" yaml:"group_entity_type,omitempty"`

        // MultiValuedClaims lists JWT claim NAMES (bare, e.g. "scp" or "scope") whose
        // value is exposed to Cedar in two normalized forms, regardless of whether the
        // IdP emits the claim as a space-delimited string (Entra `scp`, Keycloak `scope`)
        // or a JSON array (Okta `scp`):
        //
        //  1. "claim_<name>" (both principal attribute and context) is a space-delimited
        //     string: an array is joined with single spaces; a string value is passed
        //     through VERBATIM — its spacing and order are preserved, not
        //     re-canonicalized. Use this with `like`/`==` string policies.
        //  2. "claimset_<name>" (bare key, both principal attribute and context) is a Cedar
        //     Set of the claim's elements. Use this with `.contains`/`.containsAll`/
        //     `.containsAny` for exact-element matching that cannot be fooled by a
        //     substring (e.g. "Mail.Read" will not match "Mail.ReadWrite"). The "claimset_"
        //     prefix is reserved and never collides with a JWT claim, since all JWT
        //     claims are surfaced under the "claim_" prefix.
        //
        // Only for claims whose ELEMENTS are space-free (OAuth scopes, per RFC 6749 §3.3).
        // Do NOT list group/role claims or any claim whose elements can contain spaces
        // (e.g. group display names) — element boundaries would become ambiguous.
        //
        // Backward compatibility: opt-in. When empty (default), "claim_<name>" is
        // byte-identical to today and no "claimset_<name>" key is added. Listing a claim
        // leaves OTHER claims untouched.
        //
        // Migration hazard: listing a claim that an IdP emits as an ARRAY changes its
        // "claim_<name>" form from a Cedar Set to a String. A pre-existing policy that
        // tested it with `like`/substring — which errored on the Set and so failed
        // closed (deny) — will then match against the joined string and can
        // substring-match (e.g. `context.claim_scp like "*Mail.Read*"` also matches
        // "Mail.ReadWrite"), an unintended grant. Audit existing `like` policies on a
        // claim before listing it, and prefer "claimset_<name>" with `.contains`/
        // `.containsAll`/`.containsAny` for exact-element membership.
        //
        // Note: this unifies claim SHAPE, not NAME — `scp` and `scope` remain distinct
        // claim names, so a policy still references one specific name.
        MultiValuedClaims []string `json:"multi_valued_claims,omitempty" yaml:"multi_valued_claims,omitempty"`
}

// validateGroupEntityType validates a GroupEntityType value. Empty string is
// valid — it means "use the default" and is resolved by NewEntityFactory.
// Non-empty values must:
//   - not contain Cedar's "::" namespace separator (out of scope per #5072), and
//   - parse cleanly as a Cedar identifier when used as an entity type in a
//     synthetic policy. We delegate the identifier-grammar check to cedar-go's
//     policy parser so that future grammar refinements in upstream cedar-go are
//     picked up automatically — this is the source of truth for Cedar identifier
//     validity. Hand-rolling the grammar (reserved words, ANYIDENT regex,
//     __cedar prefix) duplicates rules cedar-go already enforces.
func validateGroupEntityType(s string) error {
        if s == "" {
                return nil
        }

        // Check for namespace separator first: namespaced types are out of scope.
        // This must run before the cedar-go round-trip because the Cedar parser
        // accepts "Foo::Bar" as a valid namespaced type, but we reject it for
        // project-specific reasons.
        if strings.Contains(s, "::") {
                return fmt.Errorf("group_entity_type %q contains \"::\": namespaced entity types are not yet supported", s)
        }

        // Round-trip through cedar-go's policy parser. If the synthesized policy
        // text fails to parse, the type name violates Cedar's identifier grammar
        // (reserved word, invalid character, leading digit, etc.).
        synth := fmt.Sprintf(`permit(principal in %s::"x", action, resource);`, s)
        var p cedar.Policy
        if err := p.UnmarshalCedar([]byte(synth)); err != nil {
                return fmt.Errorf("group_entity_type %q is not a valid Cedar identifier: %w", s, err)
        }
        return nil
}

// NewCedarAuthorizer creates a new Cedar authorizer.
// serverName is a runtime-injected value (not user-authored config) that
// identifies which MCP server this authorizer is scoped to.
// If a second runtime-injected value is needed, bundle both into a
// RuntimeContext struct to keep the factory interface stable.
func NewCedarAuthorizer(options ConfigOptions, serverName string) (authorizers.Authorizer, error) {
        if err := validateGroupEntityType(options.GroupEntityType); err != nil {
                return nil, err
        }

        authorizer := &Authorizer{
                policySet:               cedar.NewPolicySet(),
                entities:                cedar.EntityMap{},
                entityFactory:           NewEntityFactory(cedar.EntityType(options.GroupEntityType)),
                primaryUpstreamProvider: options.PrimaryUpstreamProvider,
                groupClaimName:          options.GroupClaimName,
                roleClaimName:           options.RoleClaimName,
                serverName:              serverName,
                claimKeyLog:             syncutil.NewAtMost(30 * time.Second),
                supplementLog:           syncutil.NewAtMost(30 * time.Second),
                missingIDTokenLog:       syncutil.NewAtMost(30 * time.Second),
                noSessionLog:            syncutil.NewAtMost(30 * time.Second),
                multiValuedClaims:       options.MultiValuedClaims,
        }

        // Load policies
        if len(options.Policies) == 0 {
                return nil, ErrNoPolicies
        }

        for i, policyStr := range options.Policies {
                var policy cedar.Policy
                if err := policy.UnmarshalCedar([]byte(policyStr)); err != nil {
                        return nil, fmt.Errorf("failed to parse policy %d: %w", i, err)
                }

                policyID := cedar.PolicyID(fmt.Sprintf("policy%d", i))
                authorizer.policySet.Add(policyID, &policy)
        }

        // Load entities if provided
        if options.EntitiesJSON != "" {
                if err := json.Unmarshal([]byte(options.EntitiesJSON), &authorizer.entities); err != nil {
                        return nil, fmt.Errorf("failed to parse entities JSON: %w", err)
                }

                // Warn once if entities_json contains stale THVGroup entities while
                // GroupEntityType is configured to a different type. Cedar's `in` operator
                // compares entity UIDs by type name, so the pre-loaded THVGroup entities will
                // never match the synthesised parents and any policy referencing them will
                // silently deny every request.
                if options.GroupEntityType != "" && options.GroupEntityType != string(EntityTypeTHVGroup) {
                        for uid := range authorizer.entities {
                                if uid.Type == EntityTypeTHVGroup {
                                        slog.Warn("Cedar entities_json contains THVGroup entities but GroupEntityType is set to a different value; "+
                                                "synthesised group parents will not match these pre-loaded entities and policies that reference them will silently deny",
                                                "configured_group_entity_type", options.GroupEntityType,
                                                "stale_entity_uid", uid.String())
                                        break
                                }
                        }
                }
        }

        return authorizer, nil
}

// UpdatePolicies updates the Cedar policies.
func (a *Authorizer) UpdatePolicies(policies []string) error {
        a.mu.Lock()
        defer a.mu.Unlock()

        if len(policies) == 0 {
                return ErrNoPolicies
        }

        newPolicySet := cedar.NewPolicySet()

        for i, policyStr := range policies {
                var policy cedar.Policy
                if err := policy.UnmarshalCedar([]byte(policyStr)); err != nil {
                        return fmt.Errorf("failed to parse policy %d: %w", i, err)
                }

                policyID := cedar.PolicyID(fmt.Sprintf("policy%d", i))
                newPolicySet.Add(policyID, &policy)
        }

        a.policySet = newPolicySet
        return nil
}

// UpdateEntities updates the Cedar entities.
func (a *Authorizer) UpdateEntities(entitiesJSON string) error {
        a.mu.Lock()
        defer a.mu.Unlock()

        var newEntities cedar.EntityMap
        if err := json.Unmarshal([]byte(entitiesJSON), &newEntities); err != nil {
                return fmt.Errorf("failed to parse entities JSON: %w", err)
        }

        a.entities = newEntities
        return nil
}

// AddEntity adds or updates an entity in the authorizer's entity store.
func (a *Authorizer) AddEntity(entity cedar.Entity) {
        a.mu.Lock()
        defer a.mu.Unlock()

        a.entities[entity.UID] = entity
}

// RemoveEntity removes an entity from the authorizer's entity store.
func (a *Authorizer) RemoveEntity(uid cedar.EntityUID) {
        a.mu.Lock()
        defer a.mu.Unlock()

        delete(a.entities, uid)
}

// GetEntity retrieves an entity from the authorizer's entity store.
func (a *Authorizer) GetEntity(uid cedar.EntityUID) (cedar.Entity, bool) {
        a.mu.RLock()
        defer a.mu.RUnlock()

        entity, found := a.entities[uid]
        return entity, found
}

// GetEntityFactory returns the entity factory associated with this authorizer.
func (a *Authorizer) GetEntityFactory() *EntityFactory {
        return a.entityFactory
}

// IsAuthorized checks if a request is authorized.
// This is the core authorization method that all other authorization methods use.
// It takes:
// - principal: The entity making the request (e.g., "Client::vscode_extension_123")
// - action: The operation being performed (e.g., "Action::call_tool")
// - resource: The object being accessed (e.g., "Tool::weather")
// - context: Additional information about the request
//
// Note: group-based Cedar policies (e.g. "principal in THVGroup::\"eng\"" with the
// default group entity type — see ConfigOptions.GroupEntityType) require that
// group parent entities are included in the entity map. See #4768 for the group
// parent wiring that will set these up via CreatePrincipalEntity.
// - entities: Optional Cedar entity map with attributes
func (a *Authorizer) IsAuthorized(
        principal, action, resource string,
        contextMap map[string]interface{},
        entities ...cedar.EntityMap,
) (bool, error) {
        a.mu.RLock()
        defer a.mu.RUnlock()

        if principal == "" {
                return false, ErrMissingPrincipal
        }

        if action == "" {
                return false, ErrMissingAction
        }

        if resource == "" {
                return false, ErrMissingResource
        }

        // Parse principal, action, and resource
        principalType, principalID, err := parseCedarEntityID(principal)
        if err != nil {
                return false, err
        }

        actionType, actionID, err := parseCedarEntityID(action)
        if err != nil {
                return false, err
        }

        resourceType, resourceID, err := parseCedarEntityID(resource)
        if err != nil {
                return false, err
        }

        // Create context record
        contextRecord := convertMapToCedarRecord(contextMap)

        // Create Cedar request
        req := cedar.Request{
                Principal: cedar.NewEntityUID(cedar.EntityType(principalType), cedar.String(principalID)),
                Action:    cedar.NewEntityUID(cedar.EntityType(actionType), cedar.String(actionID)),
                Resource:  cedar.NewEntityUID(cedar.EntityType(resourceType), cedar.String(resourceID)),
                Context:   contextRecord,
        }

        // Use the provided entities if available, otherwise use the default entities
        entityMap := a.entities
        if len(entities) > 0 && entities[0] != nil {
                // Merge the request entities with the default entities
                // This allows policies to reference both the request-specific entities
                // and any global entities defined in the authorizer
                mergedEntities := make(cedar.EntityMap)
                for k, v := range a.entities {
                        mergedEntities[k] = v
                }
                for k, v := range entities[0] {
                        mergedEntities[k] = v
                }

                entityMap = mergedEntities
        }

        // Debug logging for authorization
        slog.Debug("cedar authorization check",
                "principal", req.Principal, "action", req.Action, "resource", req.Resource)
        slog.Debug("cedar context", "context", req.Context)

        // Check authorization
        decision, diagnostic := cedar.Authorize(a.policySet, entityMap, req)

        // Log the decision
        slog.Debug("cedar decision", "decision", decision, "diagnostic", diagnostic)

        // Cedar's Authorize returns a Decision and a Diagnostic
        // Check if the Diagnostic contains any errors
        if len(diagnostic.Errors) > 0 {
                return false, fmt.Errorf("authorization error: %v", diagnostic.Errors)
        }
        return decision == cedar.Allow, nil
}

// resolveClaims determines which JWT claims to use for Cedar policy evaluation,
// and returns alongside them the claim-source label naming the trust root that
// asserted them (see claimSourceAttr and the claimSource* values). Every path
// that returns claims returns a label, so a policy can always tell whether it is
// looking at an upstream assertion or the request token's own word.
//
// When primaryUpstreamProvider is empty (the default), the claims of the token on
// the original client request are used as-is. That may be a ToolHive-issued token
// or any other bearer token.
//
// When primaryUpstreamProvider is set (the embedded auth server path), the claim
// source is that provider's upstream credentials — and only that provider's. Its
// access token is primary; the profile claims listed in profileClaimsFromIDToken
// fall back to the same provider's id_token when the access token does not carry
// them, because many OIDC providers assert `email` only in the id_token and the
// access token alone was therefore a deny-all trap (#5916).
//
// Both tokens belong to the same upstream provider and session: the auth
// middleware splits a single credential bundle into Identity.UpstreamTokens and
// Identity.UpstreamIDTokens (see TokenValidator.Middleware in pkg/auth/token.go),
// so a given key denotes the same provider and session in both maps. That is what
// makes a supplemented claim carry the same provenance as one read from the access
// token: `principal has claim_email` means "this upstream asserted an email". (The
// two maps' key SETS may differ — a provider can have an access token but no
// id_token, which is why the fallback is conditional. That is not a case of one key
// meaning two different identities.)
//
// The sole exception is a session-less grant proven by
// Identity.SessionlessGrant. TokenValidator sets that internal runtime fact only
// after validating a token for the configured embedded authorization-server
// issuer. Any non-nil map without a non-empty token for the configured provider,
// including an empty map or a map for another provider, remains an error so request
// claims never substitute for a missing provider's credentials.
//
// The ToolHive-issued token the client presented is otherwise deliberately NOT a
// claim source here, even though it mirrors a name/email — in a multi-upstream chain
// those belong to the first configured upstream, which need not be the pinned
// provider, so using them could attribute one IdP's email to another.
//
// The supplement is restricted to profileClaimsFromIDToken rather than merging
// the two token bodies. Read this before widening it:
//
//   - An id_token's registered claims (`iss`, `aud`, `exp`, `nonce`, `at_hash`,
//     `azp`) describe that token, not the user. Merging them would overwrite the
//     access token's own `aud`/`exp` with values that mean something different.
//   - Authorization-bearing claims (`groups`, `roles`, `scope`) are left alone for
//     now. Taking them from this provider's id_token would be provenance-correct —
//     unlike taking them from the ToolHive-issued token, which must never happen —
//     but it would newly grant requests that deny today, so it belongs in its own
//     change rather than riding along with a deny-all fix.
//   - `sub` is excluded because it becomes the Cedar principal entity ID rather
//     than an attribute, and Cedar has no `has`-style guard in the principal
//     position. An access token without `sub` violates RFC 9068; failing closed
//     with ErrMissingPrincipal beats silently choosing a principal.
//
// An access-token value always wins, so a claim the access token does assert can
// never be shadowed. A claim absent from both tokens stays absent, so `has`-guarded
// policies still fail closed.
//
// The id_token is used without checking `exp`, deliberately: it is read as a record
// of what the upstream asserted at login, not presented as a credential, and
// rejecting it once expired would flip a policy from permit to deny mid-session.
// See Identity.UpstreamIDTokens for that contract and its two consumers.
func (a *Authorizer) resolveClaims(identity *auth.Identity) (jwt.MapClaims, string, error) {
        requestClaims := jwt.MapClaims(identity.Claims)

        if a.primaryUpstreamProvider == "" {
                // Default path: use claims from the original client request's token.
                a.logClaimKeys(claimSourceRequest, requestClaims)
                return requestClaims, claimSourceRequest, nil
        }

        // Embedded auth server path: the upstream IDP token is the primary claim source.
        if identity.UpstreamTokens == nil && identity.SessionlessGrant {
                // TokenValidator proved this token was minted without an upstream login,
                // so no upstream credential can exist for it. Evaluate its own claims.
                a.noSessionLog.Do(func() {
                        slog.Info("token minted without an upstream session; using request-token claims for Cedar evaluation",
                                "provider", a.primaryUpstreamProvider)
                })
                a.logClaimKeys(claimSourceNoUpstreamSession, requestClaims)
                return requestClaims, claimSourceNoUpstreamSession, nil
        }

        if identity.UpstreamTokens == nil {
                // Nil without the validator-established fact: this identity was never
                // enriched with upstream credentials, and nothing authorizes a fallback.
                // It may be anonymous, local, or validated for another issuer.
                return nil, "", fmt.Errorf("no upstream credentials for provider %q and no session-less grant provenance",
                        a.primaryUpstreamProvider)
        }

        upstreamToken, tokenFound := identity.UpstreamTokens[a.primaryUpstreamProvider]
        if !tokenFound || upstreamToken == "" {
                // A session exists but lacks credentials for the configured provider.
                // Do not substitute request claims for another provider's credentials.
                return nil, "", fmt.Errorf("upstream token for provider %q not found in identity",
                        a.primaryUpstreamProvider)
        }

        upstreamClaims, err := parseUpstreamJWTClaims(upstreamToken)
        if err != nil {
                // Distinguish "not JWT-shaped" (opaque OAuth 2.0 access token —
                // Google's ya29.*, GitHub's gho_*, etc.) from "JWT-shaped but
                // malformed/tampered" (a JWT with three segments that fails to
                // parse). Only fall back for the former; preserve the deny for
                // the latter so a tampered upstream JWT cannot bypass policy.
                //
                // Policies that reference upstream-only claims (groups, hd, custom
                // namespaced claims) see those attributes as absent on this branch and
                // must be authored defensively (`principal has claim_groups && ...`).
                if !looksLikeJWT(upstreamToken) {
                        // The Warn shares claimKeyLog with logClaimKeys below so a busy
                        // Google/GitHub deployment does not emit one line per tool call.
                        a.claimKeyLog.Do(func() {
                                slog.Warn("upstream token is not a JWT; falling back to request-token claims for Cedar evaluation",
                                        "provider", a.primaryUpstreamProvider)
                        })
                        a.logClaimKeys(claimSourceUpstreamOpaque, requestClaims)
                        return requestClaims, claimSourceUpstreamOpaque, nil
                }
                return nil, "", fmt.Errorf("failed to parse upstream token for provider %q: %w",
                        a.primaryUpstreamProvider, err)
        }

        merged := a.supplementFromIDToken(identity, upstreamClaims)
        source := claimSourceUpstreamPrefix + a.primaryUpstreamProvider
        a.logClaimKeys(source, merged)
        return merged, source, nil
}

// supplementFromIDToken returns upstreamClaims with the profileClaimsFromIDToken
// it lacks filled in from the primary provider's stored id_token, logging what it
// did. upstreamClaims is not mutated.
//
// A provider with no stored id_token (an OAuth 2.0 upstream that was never asked
// for `openid`, so nothing was captured at login) yields no supplement: the claims
// stay absent and `has`-guarded policies deny. OIDC upstreams always have one —
// the provider rejects a login without it (upstream/oidc.go) and a refresh that
// omits a rotated id_token carries the original forward (upstreamtoken/service.go),
// so it does not vanish mid-session.
func (a *Authorizer) supplementFromIDToken(identity *auth.Identity, upstreamClaims jwt.MapClaims) jwt.MapClaims {
        idToken := identity.UpstreamIDTokens[a.primaryUpstreamProvider] // nil map safe in Go
        if idToken == "" {
                a.missingIDTokenLog.Do(func() {
                        slog.Warn("no upstream ID token stored for provider; policies referencing profile claims "+
                                "the access token omits will deny",
                                "provider", a.primaryUpstreamProvider,
                                "profile_claims", profileClaimsFromIDToken)
                })
                return upstreamClaims
        }

        idTokenClaims, err := parseUpstreamJWTClaims(idToken)
        if err != nil {
                // Do not fail the request: the access token parsed, so policies that only
                // reference its claims must keep working. An unparsable id_token means no
                // supplement, which denies exactly the policies that needed it.
                a.missingIDTokenLog.Do(func() {
                        slog.Warn("upstream ID token for provider is not a parsable JWT; not supplementing profile claims",
                                "provider", a.primaryUpstreamProvider,
                                "error", err)
                })
                return upstreamClaims
        }

        merged, filled := supplementProfileClaims(upstreamClaims, idTokenClaims)
        if len(filled) > 0 {
                // Rate-limited on its own timer, not claimKeyLog: sharing one would let
                // this Warn consume the window and permanently starve the claim-key dump
                // in resolveClaims, which is what you want when debugging a denying policy.
                //
                // Claim keys only — never claim values, which hold user PII.
                a.supplementLog.Do(func() {
                        slog.Warn("upstream access token lacks profile claims; using the upstream ID token's values "+
                                "for Cedar evaluation",
                                "provider", a.primaryUpstreamProvider,
                                "claims", filled)
                })
        }
        return merged
}

// OIDC Core 1.0 §5.1 standard claim names, declared here rather than borrowed
// from another package. What this code reads is an upstream provider's id_token,
// so these are wire names fixed by the OIDC spec — not names ToolHive chooses.
// Anchoring them to a ToolHive constant would invert the dependency: renaming that
// constant would silently change which claim Cedar reads out of a third party's
// token. Two literals also do not justify pulling an authorization-server
// dependency tree into the authorizer.
const (
        oidcNameClaim  = "name"
        oidcEmailClaim = "email"
)

// profileClaimsFromIDToken are the OIDC Core §5.1 profile claims that may be read
// from the primary provider's id_token when its access token omits them. They
// coincide with the two claims the embedded auth server extracts from that same
// id_token: pkg/authserver/upstream/oidc.go reads them by their OIDC names, and the
// callback handler then passes them to session.New
// (pkg/authserver/server/handlers/callback.go). Both sides key off the wire names
// independently — a shared reliance on the spec, not a coupling. If the auth server
// ever mirrors more, adding it here is a separate, deliberate decision.
//
// Widening this list widens what can satisfy a policy — see resolveClaims first.
var profileClaimsFromIDToken = []string{oidcNameClaim, oidcEmailClaim}

// supplementProfileClaims returns a fresh claim set holding every access-token
// claim plus, for each name in profileClaimsFromIDToken that the access token does
// not carry, the id_token's value for that name. It also returns the sorted list of
// names actually supplied by the id_token, for logging.
//
// Neither input is mutated and the result aliases neither, so the caller may hand
// it to code that adds synthetic keys. Absent claims are never fabricated: a name
// missing from both tokens is missing from the result, which keeps `has`-guarded
// policies failing closed.
func supplementProfileClaims(accessTokenClaims, idTokenClaims jwt.MapClaims) (merged jwt.MapClaims, filled []string) {
        merged = make(jwt.MapClaims, len(accessTokenClaims)+len(profileClaimsFromIDToken))
        for k, v := range accessTokenClaims {
                merged[k] = v
        }

        for _, name := range profileClaimsFromIDToken {
                if _, ok := merged[name]; ok {
                        // The access token asserted this claim; it always wins.
                        continue
                }
                v, ok := idTokenClaims[name]
                if !ok {
                        continue
                }
                merged[name] = v
                filled = append(filled, name)
        }

        // Sorted for a canonical, easily-greppable log line — profileClaimsFromIDToken
        // is ordered for readability (name, email), not alphabetically.
        slices.Sort(filled)
        return merged, filled
}

// looksLikeJWT returns true when the token has the three-segment shape of a
// JOSE-compact-serialized JWT (`header.payload.signature`). It does not
// validate the contents; the parser handles that.
func looksLikeJWT(tokenStr string) bool {
        return strings.Count(tokenStr, ".") == 2
}

// logClaimKeys emits a rate-limited DEBUG log listing the JWT claim keys
// available for Cedar policy evaluation.
func (a *Authorizer) logClaimKeys(source string, claims jwt.MapClaims) {
        a.claimKeyLog.Do(func() {
                keys := make([]string, 0, len(claims))
                for k := range claims {
                        keys = append(keys, k)
                }
                slog.Debug("Resolved JWT claim keys for Cedar evaluation",
                        "source", source,
                        "keys", keys)
        })
}

// parseUpstreamJWTClaims parses JWT claims from an upstream access token without
// verifying the signature. The token was already validated by the upstream IDP
// during the OAuth 2.0 code exchange; we only need its claims for Cedar evaluation.
// Returns an error if the token is not a parseable JWT (e.g. opaque token).
func parseUpstreamJWTClaims(tokenStr string) (jwt.MapClaims, error) {
        parser := jwt.NewParser()
        token, _, err := parser.ParseUnverified(tokenStr, jwt.MapClaims{})
        if err != nil {
                return nil, fmt.Errorf("upstream token is not a parseable JWT: %w", err)
        }
        claims, ok := token.Claims.(jwt.MapClaims)
        if !ok {
                return nil, fmt.Errorf("upstream token has unexpected claims type")
        }
        return claims, nil
}

// extractClientIDFromClaims extracts the client ID from JWT claims.
// By default, it uses the "sub" (subject) claim as the client ID.
// This can be customized based on your JWT token structure.
func extractClientIDFromClaims(claims jwt.MapClaims) (string, bool) {
        // Use the GetSubject method to safely extract the "sub" claim
        sub, err := claims.GetSubject()
        if err != nil || sub == "" {
                return "", false
        }

        return sub, true
}

// preprocessClaims adds a "claim_" prefix to all claim keys.
// This makes it clear which values are from the JWT claims.
func preprocessClaims(claims jwt.MapClaims) map[string]interface{} {
        preprocessed := make(map[string]interface{})
        for k, v := range claims {
                claimKey := claimPrefix + k
                preprocessed[claimKey] = v
        }
        return preprocessed
}

// normalizeMultiValuedClaims returns claims with every claim named in names
// and present in claims rewritten to a canonical unpadded space-delimited
// string. This lets a single Cedar `like`/`==` policy match a claim regardless
// of whether the IdP emitted it as a space-delimited string or a JSON array.
// See ConfigOptions.MultiValuedClaims for the full rationale.
//
// When names is empty the input is returned directly (no copy); otherwise a
// shallow copy is returned and the input is never mutated.
//
//   - array ([]interface{} or []string): elements joined with single spaces.
//     Non-string elements are stringified with fmt.Sprint rather than
//     dropped; nested array/object elements are skipped with a slog.Debug
//     since they have no scalar string form.
//   - string: passthrough, completely unchanged (not trimmed).
//   - empty array, or an array with no usable elements: "" (empty string —
//     present but empty).
//   - claim absent from claims, or names empty: left untouched (no key
//     fabricated).
//
// Only top-level claim keys are matched; there is no dot-notation support.
// claims is not mutated.
func normalizeMultiValuedClaims(claims jwt.MapClaims, names []string) jwt.MapClaims {
        if len(names) == 0 {
                // Nothing to normalize. Return the input directly — the sole caller
                // (preprocessClaims) only reads it, so no defensive copy is needed and
                // the opt-in-empty path pays no allocation.
                return claims
        }

        out := make(jwt.MapClaims, len(claims))
        for k, v := range claims {
                out[k] = v
        }

        for _, name := range names {
                v, ok := claims[name]
                if !ok {
                        continue
                }

                switch val := v.(type) {
                case string:
                        // Verbatim passthrough: out[name] already equals val (out is a copy
                        // of claims), so this is intentionally a no-op. The explicit branch
                        // documents the string contract and, importantly, keeps a string
                        // value out of the default "unrecognized type" log below. This is the
                        // one place the string handling diverges from addMultiValuedClaimSets,
                        // which collapses whitespace via strings.Fields for the Set form.
                        out[name] = val
                case []interface{}:
                        out[name] = strings.Join(multiValuedTokens(name, val), " ")
                case []string:
                        out[name] = strings.Join(val, " ")
                default:
                        slog.Debug("multi-valued claim has unrecognized type, leaving unchanged",
                                "claim", name, "type", fmt.Sprintf("%T", v))
                }
        }

        return out
}

// addMultiValuedClaimSets adds a bare "claimset_<name>" key to processed for every
// claim named in names and present in raw, holding the claim's elements as a
// []string (converted to a Cedar Set by convertToCedarValueAtDepth). This is
// the companion to normalizeMultiValuedClaims's "claim_<name>" string form:
// where "claim_<name>" supports substring-style `like` matching, "claimset_<name>"
// supports exact-element `.contains`/`.containsAll`/`.containsAny` matching.
// The "claimset_" prefix is reserved and cannot collide with a JWT claim, since all
// JWT claims are "claim_"-prefixed by preprocessClaims.
//
// processed must already have the "claim_" prefix applied (i.e. this must run
// after preprocessClaims) so the added key stays bare. raw is not mutated.
func addMultiValuedClaimSets(processed map[string]interface{}, raw jwt.MapClaims, names []string) {
        for _, name := range names {
                v, ok := raw[name]
                if !ok {
                        continue
                }

                // The array path shares multiValuedTokens with normalizeMultiValuedClaims
                // (string-only elements) so the Set and the joined string stay in sync.
                // The string path deliberately differs: normalize passes the string
                // through verbatim, whereas here strings.Fields splits it into exact
                // elements (collapsing irregular whitespace) — the two cannot share one
                // coercion for that reason.
                switch val := v.(type) {
                case string:
                        processed[claimSetPrefix+name] = strings.Fields(val)
                case []interface{}:
                        processed[claimSetPrefix+name] = multiValuedTokens(name, val)
                case []string:
                        processed[claimSetPrefix+name] = slices.Clone(val)
                default:
                        slog.Debug("multi-valued claim has unrecognized type, omitting claimset set",
                                "claim", name, "type", fmt.Sprintf("%T", v))
                }
        }
}

// multiValuedTokens extracts the string elements of a []interface{} claim value
// for normalizeMultiValuedClaims and addMultiValuedClaimSets. Only string
// elements are kept: OAuth scope tokens are strings (RFC 6749 §3.3), so a
// non-string element (number, bool, null, or a nested array/object) is
// malformed and is skipped with a slog.Debug rather than stringified. Skipping
// (a) avoids spurious tokens like "<nil>" or "1e+06" polluting both claim
// forms, and (b) keeps the array shape consistent with the string shape, whose
// strings.Fields path likewise yields only real string tokens.
func multiValuedTokens(claimName string, elems []interface{}) []string {
        tokens := make([]string, 0, len(elems))
        for _, elem := range elems {
                s, ok := elem.(string)
                if !ok {
                        slog.Debug("multi-valued claim element is not a string, skipping",
                                "claim", claimName, "type", fmt.Sprintf("%T", elem))
                        continue
                }
                tokens = append(tokens, s)
        }
        return tokens
}

// preprocessArguments adds an "arg_" prefix to all argument keys.
// For complex types, it just notes their presence with an "_present" suffix.
func preprocessArguments(arguments map[string]interface{}) map[string]interface{} {
        if arguments == nil {
                return nil
        }

        preprocessed := make(map[string]interface{})
        for k, v := range arguments {
                argKey := fmt.Sprintf("arg_%s", k)
                switch val := v.(type) {
                case string, bool, int, int64, float64:
                        preprocessed[argKey] = val
                default:
                        // For complex types, just note their presence
                        preprocessed[argKey+"_present"] = true
                }
        }
        return preprocessed
}

// mergeContexts merges multiple context maps into a single map.
// Later maps override earlier maps if there are key conflicts.
func mergeContexts(contextMaps ...map[string]interface{}) map[string]interface{} {
        merged := make(map[string]interface{})
        for _, ctxMap := range contextMaps {
                if ctxMap == nil {
                        continue
                }
                for k, v := range ctxMap {
                        merged[k] = v
                }
        }
        return merged
}

// authorizeToolCall authorizes a tool call operation.
// This method is used when a client tries to call a specific tool.
// It checks if the client is authorized to call the tool with the given context.
// Tool annotations from the context (if present) are included as resource entity
// attributes so Cedar policies can reference them (e.g. resource.readOnlyHint).
func (a *Authorizer) authorizeToolCall(
        ctx context.Context,
        clientID, toolName string,
        claimsMap map[string]interface{},
        attrsMap map[string]interface{},
        groups []string,
) (bool, error) {
        // Extract principal from client ID
        principal := fmt.Sprintf("Client::%s", clientID)

        // Action is to call a tool
        action := "Action::call_tool"

        // Resource is the tool being called
        resource := fmt.Sprintf("Tool::%s", toolName)

        // Read tool annotations from context and include in resource attributes.
        // Annotations are merged first so that the standard attributes ("name",
        // "operation", "feature") always take precedence and cannot be overwritten
        // by annotation keys — intentionally or accidentally.
        resourceMetadata, _ := authorizers.ResourceMetadataFromContext(ctx)
        annotationAttrs := authorizers.AnnotationsToMap(resourceMetadata.Annotations)

        // Create attributes for the entities
        attributes := mergeContexts(annotationAttrs, attrsMap, map[string]interface{}{
                "name":      toolName,
                "operation": "call",
                "feature":   "tool",
        })
        resourceParents := func() []cedar.EntityUID {
                if resourceMetadata.BackendID == "" {
                        return nil
                }
                return []cedar.EntityUID{
                        cedar.NewEntityUID(EntityTypeBackend, cedar.String(resourceMetadata.BackendID)),
                }
        }()

        // Create Cedar entities
        entities, err := a.entityFactory.CreateEntitiesForRequest(
                principal, action, resource, claimsMap, attributes, groups, a.serverName, resourceParents...,
        )
        if err != nil {
                return false, fmt.Errorf("failed to create Cedar entities: %w", err)
        }

        contextMap := mergeContexts(claimsMap, attrsMap)

        // Check authorization with entities
        return a.IsAuthorized(principal, action, resource, contextMap, entities)
}

// authorizePromptGet authorizes a prompt get operation.
// This method is used when a client tries to get a specific prompt.
// It checks if the client is authorized to access the prompt with the given context.
func (a *Authorizer) authorizePromptGet(
        clientID, promptName string,
        claimsMap map[string]interface{},
        attrsMap map[string]interface{},
        groups []string,
) (bool, error) {
        // Extract principal from client ID
        principal := fmt.Sprintf("Client::%s", clientID)

        // Action is to get a prompt
        action := "Action::get_prompt"

        // Resource is the prompt being accessed
        resource := fmt.Sprintf("Prompt::%s", promptName)

        // Create attributes for the entities
        attributes := mergeContexts(map[string]interface{}{
                "name":      promptName,
                "operation": "get",
                "feature":   "prompt",
        }, attrsMap)

        // Create Cedar entities
        entities, err := a.entityFactory.CreateEntitiesForRequest(
                principal, action, resource, claimsMap, attributes, groups, a.serverName,
        )
        if err != nil {
                return false, fmt.Errorf("failed to create Cedar entities: %w", err)
        }

        contextMap := mergeContexts(claimsMap, attrsMap)

        // Check authorization with entities
        return a.IsAuthorized(principal, action, resource, contextMap, entities)
}

// authorizeResourceRead authorizes a resource read operation.
// This method is used when a client tries to read a specific resource.
// It checks if the client is authorized to read the resource.
func (a *Authorizer) authorizeResourceRead(
        clientID, resourceURI string,
        claimsMap map[string]interface{},
        attrsMap map[string]interface{},
        groups []string,
) (bool, error) {
        // Extract principal from client ID
        principal := fmt.Sprintf("Client::%s", clientID)

        // Action is to read a resource
        action := "Action::read_resource"

        // Resource is the resource being accessed. Use the exact URI as the
        // entity ID: entities are built programmatically via cedar.NewEntityUID,
        // which accepts any string on the request side, so no character rewriting
        // is needed. (Policy authors still write the ID as a Cedar string literal,
        // so " and \ need escaping in policy source; see docs/authz.md.) A lossy
        // mapping would let distinct URIs collide onto one entity ID, and a
        // policy grant on one URI would then silently cover every colliding URI.
        resource := fmt.Sprintf("Resource::%s", resourceURI)

        // Create attributes for the entities
        attributes := mergeContexts(map[string]interface{}{
                "name":      resourceURI,
                "uri":       resourceURI,
                "operation": "read",
                "feature":   "resource",
        }, attrsMap)

        // Create Cedar entities
        entities, err := a.entityFactory.CreateEntitiesForRequest(
                principal, action, resource, claimsMap, attributes, groups, a.serverName,
        )
        if err != nil {
                return false, fmt.Errorf("failed to create Cedar entities: %w", err)
        }

        contextMap := mergeContexts(claimsMap, attrsMap)

        // Check authorization with entities
        return a.IsAuthorized(principal, action, resource, contextMap, entities)
}

// authorizeSkillGet authorizes a skills/get operation using its exact URI.
func (a *Authorizer) authorizeSkillGet(
        clientID, skillURI string,
        claimsMap map[string]interface{},
        attrsMap map[string]interface{},
        groups []string,
) (bool, error) {
        principal := fmt.Sprintf("Client::%s", clientID)
        attributes := mergeContexts(map[string]interface{}{
                "name":      skillURI,
                "uri":       skillURI,
                "operation": "get",
                "feature":   "skill",
        }, attrsMap)
        entities, err := a.entityFactory.CreateEntitiesForRequest(
                principal, "Action::get_skill", fmt.Sprintf("Skill::%s", skillURI),
                claimsMap, attributes, groups, a.serverName,
        )
        if err != nil {
                return false, fmt.Errorf("failed to create Cedar entities: %w", err)
        }
        return a.IsAuthorized(principal, "Action::get_skill", fmt.Sprintf("Skill::%s", skillURI),
                mergeContexts(claimsMap, attrsMap), entities)
}

// authorizeFeatureList authorizes a list operation for a feature.
// This method is used when a client tries to list available tools, prompts, or resources.
// It checks if the client is authorized to list the specified feature type.
func (a *Authorizer) authorizeFeatureList(
        clientID string,
        feature authorizers.MCPFeature,
        claimsMap map[string]interface{},
        attrsMap map[string]interface{},
        groups []string,
) (bool, error) {
        // Extract principal from client ID
        principal := fmt.Sprintf("Client::%s", clientID)

        // Action is to list a feature
        action := fmt.Sprintf("Action::list_%ss", feature)

        // Resource is the feature type
        resource := fmt.Sprintf("FeatureType::%s", feature)

        // Create attributes for the entities
        attributes := mergeContexts(map[string]interface{}{
                "type":      string(feature),
                "operation": "list",
                "feature":   string(feature),
        }, attrsMap)

        // Create Cedar entities
        entities, err := a.entityFactory.CreateEntitiesForRequest(
                principal, action, resource, claimsMap, attributes, groups, a.serverName,
        )
        if err != nil {
                return false, fmt.Errorf("failed to create Cedar entities: %w", err)
        }

        contextMap := mergeContexts(claimsMap, attrsMap)

        // Check authorization with entities
        return a.IsAuthorized(principal, action, resource, contextMap, entities)
}

// parseCedarEntityID parses a Cedar entity ID in the format "Type::ID".
// It returns the type and ID parts, or an error if the format is invalid.
func parseCedarEntityID(entityID string) (string, string, error) {
        parts := strings.SplitN(entityID, "::", 2)
        if len(parts) != 2 {
                return "", "", fmt.Errorf("invalid entity ID format: %s", entityID)
        }
        return parts[0], parts[1], nil
}

// AuthorizeWithJWTClaims demonstrates how to use JWT claims with the Cedar authorization middleware.
// This method:
// 1. Extracts JWT claims from the context
// 2. Extracts the client ID from the claims
// 3. Includes the JWT claims in the Cedar context
// 4. Creates entities with appropriate attributes
// 5. Authorizes the operation using the client ID and claims
func (a *Authorizer) AuthorizeWithJWTClaims(
        ctx context.Context,
        feature authorizers.MCPFeature,
        operation authorizers.MCPOperation,
        resourceID string,
        arguments map[string]interface{},
) (bool, error) {
        // Extract Identity from the context
        identity, ok := auth.IdentityFromContext(ctx)
        if !ok {
                return false, ErrMissingPrincipal
        }

        // Resolve the claims source: upstream IDP token or the original request's token.
        resolvedClaims, claimSource, err := a.resolveClaims(identity)
        if err != nil {
                return false, err
        }

        // Extract client ID from the resolved claims.
        clientID, ok := extractClientIDFromClaims(resolvedClaims)
        if !ok {
                return false, ErrMissingPrincipal
        }

        // Extract groups from the group claim (or well-known defaults) and the
        // role claim, merge, and dedup. Both claim sources map to the configured
        // group entity type (default "THVGroup"). Extraction runs BEFORE
        // preprocessClaims so the raw claim
        // values are available.
        // The identity pointer is not mutated here because Identity MUST NOT be
        // modified after it is placed in the request context (concurrent reads).
        groupClaims := extractGroups(resolvedClaims, a.groupClaimName)
        if groupClaims == nil {
                // Fall back to well-known claim names. This covers two cases:
                // 1. No GroupClaimName configured — backward compatible default.
                // 2. GroupClaimName configured but absent from the token — the
                //    documented contract says the custom name takes *priority*
                //    over defaults, not that it replaces them.
                for _, name := range defaultGroupClaimNames {
                        if groupClaims = extractGroups(resolvedClaims, name); groupClaims != nil {
                                break
                        }
                }
        }
        groups := dedup(append(
                groupClaims,
                extractGroups(resolvedClaims, a.roleClaimName)...,
        ))

        // Preprocess claims and arguments. Multi-valued claim normalization runs
        // after group/role extraction (which needs the raw claim shapes) and
        // before the "claim_" prefix is applied.
        normalizedClaims := normalizeMultiValuedClaims(resolvedClaims, a.multiValuedClaims)
        processedClaims := preprocessClaims(normalizedClaims)
        addMultiValuedClaimSets(processedClaims, resolvedClaims, a.multiValuedClaims)
        // Record which trust root asserted these claims. Written after prefixing so
        // no token claim can land on this key; see claimSourceAttr. processedClaims
        // becomes both the principal's attributes and the Cedar context, so the value
        // is readable as principal.thv_claim_source and context.thv_claim_source.
        processedClaims[claimSourceAttr] = claimSource
        processedArgs := preprocessArguments(arguments)

        // Authorize based on the feature and operation.
        switch operation {
        case authorizers.MCPOperationCall:
                if feature == authorizers.MCPFeatureTool {
                        return a.authorizeToolCall(ctx, clientID, resourceID, processedClaims, processedArgs, groups)
                }
        case authorizers.MCPOperationGet:
                return a.authorizeGet(clientID, feature, resourceID, processedClaims, processedArgs, groups)
        case authorizers.MCPOperationRead:
                if feature == authorizers.MCPFeatureResource {
                        return a.authorizeResourceRead(clientID, resourceID, processedClaims, processedArgs, groups)
                }
        case authorizers.MCPOperationList:
                return a.authorizeFeatureList(clientID, feature, processedClaims, processedArgs, groups)
        }
        return false, fmt.Errorf("unsupported feature/operation combination: %s/%s", feature, operation)
}

// authorizeGet dispatches get operations to their feature-specific Cedar mappings.
func (a *Authorizer) authorizeGet(
        clientID string,
        feature authorizers.MCPFeature,
        resourceID string,
        claimsMap map[string]interface{},
        attrsMap map[string]interface{},
        groups []string,
) (bool, error) {
        switch feature {
        case authorizers.MCPFeaturePrompt:
                return a.authorizePromptGet(clientID, resourceID, claimsMap, attrsMap, groups)
        case authorizers.MCPFeatureSkill:
                return a.authorizeSkillGet(clientID, resourceID, claimsMap, attrsMap, groups)
        case authorizers.MCPFeatureTool, authorizers.MCPFeatureResource:
                fallthrough
        default:
                return false, fmt.Errorf("unsupported get feature: %s", feature)
        }
}

// defaultGroupClaimNames lists common group claim names across popular identity
// providers. They are checked in order; the first non-empty match is returned.
//
// Sources:
//   - "groups"         — Microsoft Entra ID, Okta, Auth0, PingIdentity.
//   - "roles"          — Keycloak (realm_access.roles flattened to top-level).
//   - "cognito:groups" — AWS Cognito user pools.
var defaultGroupClaimNames = []string{"groups", "roles", "cognito:groups"}

// resolveNestedClaim resolves a claim value from JWT claims, supporting both
// top-level keys and dot-separated nested paths.
//
// Resolution order:
//  1. Exact top-level match — handles Auth0 / Okta URL-style claim names
//     (e.g. "https://myapp.example.com/roles") that contain dots but are
//     top-level keys in the JWT.
//  2. Dot-notation traversal — handles Keycloak-style nested claims
//     (e.g. "realm_access.roles" → claims["realm_access"]["roles"]).
//
// Returns nil when the claim is absent or traversal hits a non-map value.
func resolveNestedClaim(claims jwt.MapClaims, path string) interface{} {
        if path == "" {
                return nil
        }

        // 1. Exact top-level match (handles Auth0 URL claims with dots).
        if val, ok := claims[path]; ok {
                return val
        }

        // 2. Dot-notation traversal.
        parts := strings.Split(path, ".")
        if len(parts) < 2 {
                return nil // single segment already tried above
        }

        var current interface{} = map[string]interface{}(claims)
        for _, segment := range parts {
                m, ok := current.(map[string]interface{})
                if !ok {
                        return nil
                }
                current, ok = m[segment]
                if !ok {
                        return nil
                }
        }
        return current
}

// extractGroups extracts group/role names from a specific JWT claim.
// It resolves the claim via resolveNestedClaim (supporting both flat and
// dot-notation paths) and decodes the value with authclaims.DecodeStringArrayLenient,
// Cedar's host-owned lenient decoder. The auth-server GroupBasedFilter uses its
// separate strict OAuth claims decoder.
//
// Return value distinguishes "claim absent" from "claim present but empty"
// so callers can decide whether to fall back to defaults:
//   - nil: claimName is empty, the claim is absent, the value has an
//     unsupported scalar/object type (e.g. string, number), or the recognized
//     value is a nil []string. The latter preserves its historical nil-ness,
//     allowing callers to fall back.
//   - non-nil, possibly empty: the claim has a recognized array shape other than
//     a nil []string. Non-string elements are silently dropped, so an []any with
//     no strings (including []any(nil)) yields a non-nil empty slice. A non-nil
//     empty []string{} does too. These cases mean "the IdP said this claim exists
//     with no usable group names" and suppress fallback.
//
// The returned slice never aliases the claim value, so the caller may append
// to it (as authorizeWithClaims does when merging role claims) without writing
// into the token's own backing array.
func extractGroups(claims jwt.MapClaims, claimName string) []string {
        if claimName == "" {
                return nil
        }

        val := resolveNestedClaim(claims, claimName)
        if val == nil {
                return nil
        }

        groups, recognized := authclaims.DecodeStringArrayLenient(val)
        if !recognized {
                slog.Warn("group/role claim has unrecognized type, ignoring",
                        "claim", claimName, "type", fmt.Sprintf("%T", val))
                return nil
        }
        return groups
}

// dedup removes duplicate strings while preserving first-occurrence order.
// Returns nil when the input is nil (not an empty slice) so callers can
// distinguish "no groups" from "empty groups".
func dedup(groups []string) []string {
        if groups == nil {
                return nil
        }

        seen := make(map[string]struct{}, len(groups))
        result := make([]string, 0, len(groups))
        for _, g := range groups {
                if _, exists := seen[g]; exists {
                        continue
                }
                seen[g] = struct{}{}
                result = append(result, g)
        }
        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package cedar provides authorization utilities using Cedar policies.
package cedar

import (
        "log/slog"

        cedar "github.com/cedar-policy/cedar-go"
)

// maxClaimNestingDepth caps how deeply nested maps/arrays in JWT claims are
// converted into Cedar records/sets. Claims can originate from an unverified
// upstream token (see parseUpstreamJWTClaims in core.go, which uses
// jwt.ParseUnverified), so recursion here has no other bound. Matches
// maxSchemaDepth in pkg/vmcp/composer/elicitation_handler.go for consistency.
const maxClaimNestingDepth = 10

const (
        // EntityTypeTHVGroup is the default Cedar entity type representing group membership.
        // It is used when ConfigOptions.GroupEntityType is empty. Principals are added as
        // children of group entities so that Cedar's `in` operator can evaluate
        // group-based policies (e.g. `principal in THVGroup::"engineering"`).
        EntityTypeTHVGroup cedar.EntityType = "THVGroup"
        // EntityTypeBackend represents a logical vMCP backend. Tools are added as
        // children of their originating backend for backend-scoped policies.
        EntityTypeBackend cedar.EntityType = "Backend"
)

// EntityFactory creates Cedar entities for authorization.
type EntityFactory struct {
        // groupEntityType is the Cedar entity type used for Client parent UIDs
        // synthesised from JWT group/role claims. Resolved at construction —
        // defaults to EntityTypeTHVGroup when the caller passes an empty string.
        groupEntityType cedar.EntityType
}

// NewEntityFactory creates a new entity factory. An empty groupEntityType
// falls back to EntityTypeTHVGroup so existing callers stay backward compatible.
func NewEntityFactory(groupEntityType cedar.EntityType) *EntityFactory {
        if groupEntityType == "" {
                groupEntityType = EntityTypeTHVGroup
        }
        return &EntityFactory{groupEntityType: groupEntityType}
}

// CreatePrincipalEntity creates a principal entity with the given ID, attributes,
// and optional parent entity UIDs.
// When no parents are provided, the entity has an empty parent set (backward compatible).
// NOTE: This replaces the previous groups []string parameter from 5c258a11.
// Callers are now responsible for building parent UIDs (see #4768).
func (*EntityFactory) CreatePrincipalEntity(
        principalType, principalID string,
        attributes map[string]interface{},
        parents ...cedar.EntityUID,
) (cedar.EntityUID, cedar.Entity) {
        uid := cedar.NewEntityUID(cedar.EntityType(principalType), cedar.String(principalID))
        attrs := convertMapToCedarRecord(attributes)

        entity := cedar.Entity{
                UID:        uid,
                Parents:    cedar.NewEntityUIDSet(parents...),
                Attributes: attrs,
                Tags:       cedar.NewRecord(cedar.RecordMap{}),
        }

        return uid, entity
}

// CreateActionEntity creates an action entity with the given ID and attributes.
func (*EntityFactory) CreateActionEntity(
        actionType, actionID string,
        attributes map[string]interface{},
) (cedar.EntityUID, cedar.Entity) {
        uid := cedar.NewEntityUID(cedar.EntityType(actionType), cedar.String(actionID))

        // Ensure operation attribute is set
        if attributes == nil {
                attributes = make(map[string]interface{})
        }
        attributes["operation"] = actionID

        attrs := convertMapToCedarRecord(attributes)

        entity := cedar.Entity{
                UID:        uid,
                Parents:    cedar.NewEntityUIDSet(),
                Attributes: attrs,
                Tags:       cedar.NewRecord(cedar.RecordMap{}),
        }

        return uid, entity
}

// CreateResourceEntity creates a resource entity with the given ID, attributes,
// and optional parent entity UIDs.
// When no parents are provided, the entity has an empty parent set (backward compatible).
func (*EntityFactory) CreateResourceEntity(
        resourceType, resourceID string,
        attributes map[string]interface{},
        parents ...cedar.EntityUID,
) (cedar.EntityUID, cedar.Entity) {
        uid := cedar.NewEntityUID(cedar.EntityType(resourceType), cedar.String(resourceID))

        // Ensure name attribute is set — but don't overwrite if the caller
        // already provided one (e.g. authorizeResourceRead sets name to the
        // original URI before sanitization).
        if attributes == nil {
                attributes = make(map[string]interface{})
        }
        if _, exists := attributes["name"]; !exists {
                attributes["name"] = resourceID
        }

        attrs := convertMapToCedarRecord(attributes)

        entity := cedar.Entity{
                UID:        uid,
                Parents:    cedar.NewEntityUIDSet(parents...),
                Attributes: attrs,
                Tags:       cedar.NewRecord(cedar.RecordMap{}),
        }

        return uid, entity
}

// CreateEntitiesForRequest creates entities for a specific authorization request.
// Groups are converted to parent UIDs (using the configured group entity type,
// default "THVGroup") on the principal entity so that Cedar's `in` operator works
// for group-based policies. Unlike the pre-refactor code, no separate group
// entities are inserted into the entity map — those must come from entities_json
// to preserve the role hierarchy.
//
// When serverName is non-empty, resource entities include an MCP parent UID so
// that server-scoped Cedar policies (e.g. `resource in MCP::"github"`) evaluate
// correctly via Cedar's `in` operator. When serverName is empty, resource
// entities have no MCP parent, preserving backward compatibility. Additional
// resource parent UIDs may be supplied for narrower scopes such as vMCP
// backends; their corresponding entities, if needed for hierarchy traversal,
// must come from entities_json.
func (f *EntityFactory) CreateEntitiesForRequest(
        principal, action, resource string,
        claimsMap map[string]interface{},
        attributes map[string]interface{},
        groups []string,
        serverName string,
        extraResourceParents ...cedar.EntityUID,
) (cedar.EntityMap, error) {
        // Parse principal, action, and resource
        principalType, principalID, err := parseCedarEntityID(principal)
        if err != nil {
                return nil, err
        }

        actionType, actionID, err := parseCedarEntityID(action)
        if err != nil {
                return nil, err
        }

        resourceType, resourceID, err := parseCedarEntityID(resource)
        if err != nil {
                return nil, err
        }

        // Create Cedar entities
        entities := make(cedar.EntityMap)

        // Build parent UIDs from groups so the principal's Parents set contains
        // references using the configured group entity type (default "THVGroup"),
        // needed for Cedar's `in` operator. Unlike the pre-refactor code, we do
        // NOT insert separate group entities into the entity map — those come from
        // entities_json and must not be overwritten (see merge-order hazard
        // described in the RFC). #4768 will restructure this further for full role
        // hierarchy support.
        parentUIDs := make([]cedar.EntityUID, 0, len(groups))
        for _, g := range groups {
                parentUIDs = append(parentUIDs, cedar.NewEntityUID(f.groupEntityType, cedar.String(g)))
        }

        principalUID, principalEntity := f.CreatePrincipalEntity(principalType, principalID, claimsMap, parentUIDs...)
        entities[principalUID] = principalEntity

        // Create action entity
        actionUID, actionEntity := f.CreateActionEntity(actionType, actionID, nil)
        entities[actionUID] = actionEntity

        // Build MCP parent for resource entity when serverName is set so that
        // server-scoped policies (e.g. resource in MCP::"github") can match.
        resourceParents := make([]cedar.EntityUID, 0, 1+len(extraResourceParents))
        if serverName != "" {
                resourceParents = append(resourceParents, cedar.NewEntityUID("MCP", cedar.String(serverName)))
        }
        resourceParents = append(resourceParents, extraResourceParents...)

        // Create resource entity
        resourceUID, resourceEntity := f.CreateResourceEntity(resourceType, resourceID, attributes, resourceParents...)
        entities[resourceUID] = resourceEntity

        return entities, nil
}

// convertMapToCedarRecord converts a Go map to a Cedar record. Nesting beyond
// maxClaimNestingDepth is dropped (see convertToCedarValueAtDepth).
func convertMapToCedarRecord(data map[string]interface{}) cedar.Record {
        return convertMapToCedarRecordAtDepth(data, 0)
}

// convertMapToCedarRecordAtDepth is the depth-tracking implementation behind
// convertMapToCedarRecord.
func convertMapToCedarRecordAtDepth(data map[string]interface{}, depth int) cedar.Record {
        if data == nil {
                return cedar.NewRecord(cedar.RecordMap{})
        }

        recordMap := make(cedar.RecordMap)

        for k, v := range data {
                // Convert Go values to Cedar values
                cedarValue := convertToCedarValueAtDepth(v, depth)
                if cedarValue != nil {
                        recordMap[cedar.String(k)] = cedarValue
                }
        }

        return cedar.NewRecord(recordMap)
}

// convertToCedarValueAtDepth converts a Go value to a Cedar value. Maps and
// arrays nested deeper than maxClaimNestingDepth are dropped rather than
// causing unbounded recursion, since claim values can originate from an
// unverified upstream JWT (see parseUpstreamJWTClaims in core.go). depth
// counts how many maps/arrays have already been unwrapped to reach v; once
// it exceeds maxClaimNestingDepth, v is dropped (returns nil) and a warning
// is logged instead of recursing further.
func convertToCedarValueAtDepth(v interface{}, depth int) cedar.Value {
        if depth > maxClaimNestingDepth {
                slog.Warn("claim attribute nested too deeply, dropping",
                        "max_depth", maxClaimNestingDepth)
                return nil
        }

        switch val := v.(type) {
        case bool:
                return convertBoolToCedar(val)
        case string:
                return cedar.String(val)
        case int:
                return cedar.Long(val)
        case int64:
                return cedar.Long(val)
        case float64:
                return convertFloatToCedar(val)
        case []interface{}:
                return convertInterfaceArrayToCedar(val, depth)
        case []string:
                return convertStringArrayToCedar(val)
        case map[string]interface{}:
                return convertMapToCedarRecordAtDepth(val, depth+1)
        default:
                // Skip unsupported types
                return nil
        }
}

// convertBoolToCedar converts a bool to a Cedar value.
func convertBoolToCedar(val bool) cedar.Value {
        if val {
                return cedar.True
        }
        return cedar.False
}

// convertFloatToCedar converts a float64 to a Cedar decimal value.
func convertFloatToCedar(val float64) cedar.Value {
        decimalVal, err := cedar.NewDecimalFromFloat(val)
        if err != nil {
                return nil
        }
        return decimalVal
}

// convertInterfaceArrayToCedar converts an array of interfaces to a Cedar set.
func convertInterfaceArrayToCedar(val []interface{}, depth int) cedar.Value {
        values := make([]cedar.Value, 0, len(val))
        for _, item := range val {
                cedarItem := convertToCedarValueAtDepth(item, depth+1)
                if cedarItem != nil {
                        values = append(values, cedarItem)
                }
        }
        return cedar.NewSet(values...)
}

// convertStringArrayToCedar converts a string array to a Cedar set.
func convertStringArrayToCedar(val []string) cedar.Value {
        values := make([]cedar.Value, 0, len(val))
        for _, item := range val {
                values = append(values, cedar.String(item))
        }
        return cedar.NewSet(values...)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package authorizers provides the authorization framework and abstractions for ToolHive.
// It defines interfaces for authorization decisions and configuration handling.
package authorizers

import (
        "encoding/json"
        "fmt"
        "os"
        "path/filepath"
        "strings"

        "sigs.k8s.io/yaml"
)

// ConfigType represents the type of authorization configuration.
type ConfigType string

// Config represents the authorization configuration.
// This struct contains the common fields (version/type) needed to identify
// which authorizer factory to use. The full raw configuration is preserved
// so that each authorizer implementation can parse it with domain-specific
// knowledge (e.g., Cedar configs have a "cedar" field at the top level).
type Config struct {
        // Version is the version of the configuration format.
        Version string `json:"version" yaml:"version"`

        // Type is the type of authorization configuration (e.g., "cedarv1").
        Type ConfigType `json:"type" yaml:"type"`

        // rawConfig stores the original raw configuration bytes for re-parsing
        // by the authorizer factory with domain-specific knowledge.
        rawConfig json.RawMessage
}

// UnmarshalJSON implements custom JSON unmarshaling that preserves the raw config
// while extracting the version and type fields.
func (c *Config) UnmarshalJSON(data []byte) error {
        // First, extract just version and type
        var header struct {
                Version string     `json:"version"`
                Type    ConfigType `json:"type"`
        }
        if err := json.Unmarshal(data, &header); err != nil {
                return err
        }

        c.Version = header.Version
        c.Type = header.Type
        c.rawConfig = data

        return nil
}

// MarshalJSON implements custom JSON marshaling.
// If we have the original raw config, use that to preserve all fields.
// Otherwise, just marshal version and type.
func (c *Config) MarshalJSON() ([]byte, error) {
        if len(c.rawConfig) > 0 {
                return c.rawConfig, nil
        }

        // Fallback: just marshal version and type
        type alias struct {
                Version string     `json:"version"`
                Type    ConfigType `json:"type"`
        }
        return json.Marshal(&alias{
                Version: c.Version,
                Type:    c.Type,
        })
}

// RawConfig returns the raw configuration bytes for the authorizer factory
// to parse with domain-specific knowledge.
func (c *Config) RawConfig() json.RawMessage {
        return c.rawConfig
}

// LoadConfig loads the authorization configuration from a file.
// It supports both JSON and YAML formats, detected by file extension.
func LoadConfig(path string) (*Config, error) {
        // Validate and clean the path to prevent directory traversal attacks
        cleanPath := filepath.Clean(path)
        if strings.Contains(cleanPath, "..") {
                return nil, fmt.Errorf("path contains directory traversal elements: %s", path)
        }

        // Read the file
        data, err := os.ReadFile(cleanPath)
        if err != nil {
                return nil, fmt.Errorf("failed to read authorization configuration file: %w", err)
        }

        // Determine the file format based on extension
        var config Config
        ext := strings.ToLower(filepath.Ext(cleanPath))

        // Parse the file based on its format
        switch ext {
        case ".yaml", ".yml":
                // Parse YAML - first convert to JSON for consistent handling
                jsonData, err := yaml.YAMLToJSON(data)
                if err != nil {
                        return nil, fmt.Errorf("failed to parse YAML authorization configuration file: %w", err)
                }
                if err := json.Unmarshal(jsonData, &config); err != nil {
                        return nil, fmt.Errorf("failed to parse authorization configuration: %w", err)
                }
        case ".json", "":
                // Parse JSON (default if no extension)
                if err := json.Unmarshal(data, &config); err != nil {
                        return nil, fmt.Errorf("failed to parse JSON authorization configuration file: %w", err)
                }
        default:
                return nil, fmt.Errorf("unsupported file format: %s (supported formats: .json, .yaml, .yml)", ext)
        }

        // Validate the configuration
        if err := config.Validate(); err != nil {
                return nil, err
        }

        return &config, nil
}

// Validate validates the authorization configuration.
func (c *Config) Validate() error {
        // Check if the version is provided
        if c.Version == "" {
                return fmt.Errorf("version is required")
        }

        // Check if the type is provided
        if c.Type == "" {
                return fmt.Errorf("type is required")
        }

        // Get the factory for this config type
        factory := GetFactory(string(c.Type))
        if factory == nil {
                return fmt.Errorf("unsupported configuration type: %s (registered types: %v)",
                        c.Type, RegisteredTypes())
        }

        // Check if we have raw config to validate
        if len(c.rawConfig) == 0 {
                return fmt.Errorf("configuration data is required for type %s", c.Type)
        }

        // Delegate validation to the authorizer factory, passing the full raw config
        if err := factory.ValidateConfig(c.rawConfig); err != nil {
                return fmt.Errorf("invalid %s configuration: %w", c.Type, err)
        }

        return nil
}

// NewConfig creates a new Config from a full configuration structure.
// The fullConfig parameter should be the complete configuration including
// version, type, and authorizer-specific fields (e.g., "cedar" field for Cedar configs).
// This maintains backwards compatibility with the v1.0 configuration schema.
func NewConfig(fullConfig interface{}) (*Config, error) {
        rawConfig, err := json.Marshal(fullConfig)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal configuration: %w", err)
        }

        // Parse the raw config to extract version and type
        var config Config
        if err := json.Unmarshal(rawConfig, &config); err != nil {
                return nil, fmt.Errorf("failed to parse configuration: %w", err)
        }

        return &config, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package http

// ClaimMapper defines the interface for mapping JWT claims to principal attributes.
// Different PDP implementations may have different conventions for claim names
// (e.g., MPE uses m-prefixed claims like "mroles", while OIDC uses standard claims like "roles").
type ClaimMapper interface {
        // MapClaims maps JWT claims to principal attributes suitable for the target PDP.
        // The input is the raw JWT claims map, and the output is a map with claims
        // transformed according to the mapper's conventions.
        MapClaims(claims map[string]any) map[string]any
}

// MPEClaimMapper implements ClaimMapper for Manetu PolicyEngine (MPE).
// MPE uses m-prefixed claims (mroles, mgroups, mclearance, mannotations)
// and also accepts standard OIDC claims (roles, groups, clearance, annotations)
// which are mapped to their m-prefixed equivalents.
type MPEClaimMapper struct{}

// MapClaims maps JWT claims to MPE-compatible principal attributes.
// It maps standard JWT claims to MPE-specific principal attributes:
//   - sub -> sub (subject identifier)
//   - roles/mroles -> mroles (roles)
//   - groups/mgroups -> mgroups (groups)
//   - scope/scopes -> scopes (access scopes)
//   - clearance/mclearance -> mclearance (clearance level)
//   - annotations/mannotations -> mannotations (additional annotations)
//
// Returns map[string]any to ensure the PDP can properly unmarshal
// the PORC structure for identity phase evaluation.
func (*MPEClaimMapper) MapClaims(claims map[string]any) map[string]any {
        principal := make(map[string]any)

        if claims == nil {
                return principal
        }

        // Map standard JWT claims
        if sub, ok := claims["sub"]; ok {
                principal["sub"] = sub
        }

        // Map roles (check both 'roles' and 'mroles')
        if roles, ok := claims["mroles"]; ok {
                principal["mroles"] = roles
        } else if roles, ok := claims["roles"]; ok {
                principal["mroles"] = roles
        }

        // Map groups (check both 'groups' and 'mgroups')
        if groups, ok := claims["mgroups"]; ok {
                principal["mgroups"] = groups
        } else if groups, ok := claims["groups"]; ok {
                principal["mgroups"] = groups
        }

        // Map scopes (check both 'scope' and 'scopes')
        if scopes, ok := claims["scopes"]; ok {
                principal["scopes"] = scopes
        } else if scope, ok := claims["scope"]; ok {
                principal["scopes"] = scope
        }

        // Map clearance level
        if clearance, ok := claims["mclearance"]; ok {
                principal["mclearance"] = clearance
        } else if clearance, ok := claims["clearance"]; ok {
                principal["mclearance"] = clearance
        }

        // Map annotations (initialize empty if not present for identity phase)
        if annotations, ok := claims["mannotations"]; ok {
                principal["mannotations"] = annotations
        } else if annotations, ok := claims["annotations"]; ok {
                principal["mannotations"] = annotations
        } else {
                // Some PDPs require mannotations to be present for identity phase evaluation
                principal["mannotations"] = make(map[string]any)
        }

        return principal
}

// StandardClaimMapper implements ClaimMapper for standard OIDC claims.
// This mapper passes through standard OIDC claim names without modification
// and can be used with PDPs that expect standard OIDC conventions.
type StandardClaimMapper struct{}

// MapClaims maps JWT claims using standard OIDC conventions.
// It preserves standard claim names and normalizes common variations:
//   - sub -> sub (subject identifier)
//   - roles -> roles (roles, preserving standard name)
//   - groups -> groups (groups, preserving standard name)
//   - scope/scopes -> scopes (access scopes, normalized to plural)
func (*StandardClaimMapper) MapClaims(claims map[string]any) map[string]any {
        principal := make(map[string]any)

        if claims == nil {
                return principal
        }

        // Map standard JWT claims
        if sub, ok := claims["sub"]; ok {
                principal["sub"] = sub
        }

        // Map roles (preserve standard name)
        if roles, ok := claims["roles"]; ok {
                principal["roles"] = roles
        }

        // Map groups (preserve standard name)
        if groups, ok := claims["groups"]; ok {
                principal["groups"] = groups
        }

        // Map scopes (normalize to plural form)
        if scopes, ok := claims["scopes"]; ok {
                principal["scopes"] = scopes
        } else if scope, ok := claims["scope"]; ok {
                principal["scopes"] = scope
        }

        return principal
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package http provides authorization using HTTP-based Policy Decision Points (PDPs).
package http

import (
        "encoding/json"
        "fmt"
)

// ConfigType is the configuration type identifier for HTTP-based PDP authorization.
const ConfigType = "httpv1"

// Config represents the complete authorization configuration file structure
// for HTTP-based PDP authorization. This includes the common version/type fields
// plus the PDP-specific "pdp" field.
type Config struct {
        Version string         `json:"version"`
        Type    string         `json:"type"`
        Options *ConfigOptions `json:"pdp"`
}

// ConfigOptions represents the HTTP PDP authorization configuration options.
type ConfigOptions struct {
        // HTTP contains the HTTP connection configuration.
        HTTP *ConnectionConfig `json:"http,omitempty" yaml:"http,omitempty"`

        // Context configures what context information is included in the PORC.
        // By default, no MCP context is included in the PORC.
        Context *ContextConfig `json:"context,omitempty" yaml:"context,omitempty"`

        // ClaimMapping specifies which claim mapper to use for mapping JWT claims
        // to principal attributes. This field is required. Valid values: "mpe", "standard".
        // - "mpe": Uses MPE-specific m-prefixed claims (mroles, mgroups, mclearance, mannotations)
        // - "standard": Uses standard OIDC claim names (roles, groups)
        ClaimMapping string `json:"claim_mapping,omitempty" yaml:"claim_mapping,omitempty"`
}

// ContextConfig configures what context information is included in the PORC.
// All options default to false, meaning no MCP context is included by default.
type ContextConfig struct {
        // IncludeArgs enables inclusion of tool/prompt arguments in context.mcp.args.
        // Default is false.
        IncludeArgs bool `json:"include_args,omitempty" yaml:"include_args,omitempty"`

        // IncludeOperation enables inclusion of MCP operation metadata in context.mcp:
        // feature, operation, and resource_id fields.
        // Default is false.
        IncludeOperation bool `json:"include_operation,omitempty" yaml:"include_operation,omitempty"`
}

// ConnectionConfig contains configuration for the HTTP connection to the PDP.
type ConnectionConfig struct {
        // URL is the base URL of the PDP server (e.g., "http://localhost:9000").
        URL string `json:"url" yaml:"url"`

        // Timeout is the HTTP request timeout in seconds. Default is 30.
        Timeout int `json:"timeout,omitempty" yaml:"timeout,omitempty"`

        // InsecureSkipVerify skips TLS certificate verification. Use only for testing.
        InsecureSkipVerify bool `json:"insecure_skip_verify,omitempty" yaml:"insecure_skip_verify,omitempty"`
}

// parseConfig parses the raw JSON configuration into a Config struct.
func parseConfig(rawConfig json.RawMessage) (*Config, error) {
        var config Config
        if err := json.Unmarshal(rawConfig, &config); err != nil {
                return nil, fmt.Errorf("failed to parse HTTP PDP configuration: %w", err)
        }
        return &config, nil
}

// Validate validates the HTTP PDP configuration options.
func (c *ConfigOptions) Validate() error {
        if c == nil {
                return fmt.Errorf("pdp configuration is required (missing 'pdp' field)")
        }

        // Validate HTTP configuration
        if c.HTTP == nil {
                return fmt.Errorf("http configuration is required")
        }
        if c.HTTP.URL == "" {
                return fmt.Errorf("http.url is required")
        }

        // Validate claim_mapping is specified
        if c.ClaimMapping == "" {
                return fmt.Errorf("claim_mapping is required (valid values: mpe, standard)")
        }

        // Validate claim_mapping value
        if c.ClaimMapping != "mpe" && c.ClaimMapping != "standard" {
                return fmt.Errorf("invalid claim_mapping %q (valid values: mpe, standard)", c.ClaimMapping)
        }

        return nil
}

// GetContextConfig returns the context configuration, or a default empty config if nil.
func (c *ConfigOptions) GetContextConfig() ContextConfig {
        if c.Context == nil {
                return ContextConfig{}
        }
        return *c.Context
}

// GetClaimMapping returns the configured claim mapping type.
// The claim_mapping field is required and validated, so this will never return an empty string.
func (c *ConfigOptions) GetClaimMapping() string {
        return c.ClaimMapping
}

// CreateClaimMapper creates a ClaimMapper based on the configured claim mapping type.
func (c *ConfigOptions) CreateClaimMapper() (ClaimMapper, error) {
        switch c.GetClaimMapping() {
        case "mpe":
                return &MPEClaimMapper{}, nil
        case "standard":
                return &StandardClaimMapper{}, nil
        default:
                return nil, fmt.Errorf("unknown claim mapping type: %s (valid values: mpe, standard)", c.ClaimMapping)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package http

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/authz/authorizers"
)

func init() {
        // Register the HTTP PDP authorizer factory with the authorizers registry.
        authorizers.Register(ConfigType, &Factory{})
}

// Factory implements the authorizers.AuthorizerFactory interface for HTTP PDPs.
type Factory struct{}

// ConfigKey returns the JSON key for HTTP PDP-specific configuration ("pdp").
func (*Factory) ConfigKey() string { return "pdp" }

// ValidateConfig validates the HTTP PDP configuration.
func (*Factory) ValidateConfig(rawConfig json.RawMessage) error {
        config, err := parseConfig(rawConfig)
        if err != nil {
                return err
        }

        if config.Options == nil {
                return fmt.Errorf("pdp configuration is required (missing 'pdp' field)")
        }

        return config.Options.Validate()
}

// CreateAuthorizer creates an HTTP PDP Authorizer from the configuration.
func (*Factory) CreateAuthorizer(rawConfig json.RawMessage, serverName string) (authorizers.Authorizer, error) {
        config, err := parseConfig(rawConfig)
        if err != nil {
                return nil, err
        }

        if config.Options == nil {
                return nil, fmt.Errorf("pdp configuration is required (missing 'pdp' field)")
        }

        // Validate configuration before creating the authorizer
        if err := config.Options.Validate(); err != nil {
                return nil, err
        }

        return NewAuthorizer(*config.Options, serverName)
}

// pdp defines the interface for Policy Decision Point implementations.
type pdp interface {
        Authorize(ctx context.Context, porc PORC, probe bool) (bool, error)
        Close() error
}

// Authorizer implements the authorizers.Authorizer interface using an HTTP PDP.
type Authorizer struct {
        config      ConfigOptions
        pdp         pdp
        porcBuilder *PORCBuilder
}

// NewAuthorizer creates a new HTTP PDP authorizer from the provided configuration.
// Note: This function validates the config as a defensive measure, even though the
// factory also validates. This protects against direct calls to NewAuthorizer that
// bypass the factory.
func NewAuthorizer(config ConfigOptions, serverName string) (*Authorizer, error) {
        if err := config.Validate(); err != nil {
                return nil, err
        }

        slog.Debug("creating new HTTP PDP authorizer", "config", config)
        p, err := NewClient(config.HTTP)
        if err != nil {
                return nil, fmt.Errorf("failed to create HTTP client: %w", err)
        }

        // Create the claim mapper based on configuration
        claimMapper, err := config.CreateClaimMapper()
        if err != nil {
                return nil, fmt.Errorf("failed to create claim mapper: %w", err)
        }

        return &Authorizer{
                config:      config,
                pdp:         p,
                porcBuilder: NewPORCBuilder(serverName, config.GetContextConfig(), claimMapper),
        }, nil
}

// AuthorizeWithJWTClaims implements the authorizers.Authorizer interface.
// It extracts JWT claims from the context, builds a PORC expression,
// and delegates the authorization decision to the configured PDP.
func (a *Authorizer) AuthorizeWithJWTClaims(
        ctx context.Context,
        feature authorizers.MCPFeature,
        operation authorizers.MCPOperation,
        resourceID string,
        arguments map[string]interface{},
) (bool, error) {
        // Extract Identity from the context
        identity, ok := auth.IdentityFromContext(ctx)
        if !ok {
                return false, fmt.Errorf("missing principal: no identity in context")
        }

        // Build PORC expression using identity claims
        porc := a.porcBuilder.Build(feature, operation, resourceID, identity.Claims, arguments)

        // Enrich PORC context with tool annotations for tool-call operations only.
        // Annotations are MCP tool-specific metadata and should not be injected
        // for prompt or resource operations.
        if feature == authorizers.MCPFeatureTool {
                annotations := authorizers.ToolAnnotationsFromContext(ctx)
                if annotationMap := authorizers.AnnotationsToMap(annotations); annotationMap != nil {
                        enrichPORCWithAnnotations(porc, annotationMap)
                }
        }

        // Log the authorization request
        slog.Debug("HTTP PDP authorization check",
                "operation", porc["operation"], "resource", porc["resource"])

        // Delegate to PDP (not in probe mode for actual authorization)
        allowed, err := a.pdp.Authorize(ctx, porc, false)
        if err != nil {
                slog.Debug("HTTP PDP authorization check failed", "error", err)
                return false, fmt.Errorf("authorization failed: %w", err)
        }

        slog.Debug("HTTP PDP authorization result", "allowed", allowed)

        return allowed, nil
}

// Close releases resources used by the authorizer.
func (a *Authorizer) Close() error {
        if a.pdp != nil {
                return a.pdp.Close()
        }
        return nil
}

// enrichPORCWithAnnotations adds tool annotation data into the PORC context
// under the "context.mcp.annotations" path. It navigates and creates the
// nested map structure as needed.
//
// PORCBuilder.Build always sets porc["context"] to a map[string]interface{}
// and the "mcp" sub-key (when present) to the same type, so the type
// assertions below are expected to succeed. If they don't (e.g. a future
// refactor changes the types), we log a warning and create fresh maps rather
// than silently dropping existing context data.
func enrichPORCWithAnnotations(porc PORC, annotationMap map[string]interface{}) {
        porcCtx, ok := porc["context"].(map[string]interface{})
        if !ok {
                if porc["context"] != nil {
                        slog.Warn("enrichPORCWithAnnotations: porc[\"context\"] has unexpected type, creating fresh context")
                }
                porc["context"] = map[string]interface{}{
                        "mcp": map[string]interface{}{"annotations": annotationMap},
                }
                return
        }

        mcpCtx, ok := porcCtx["mcp"].(map[string]interface{})
        if !ok {
                if porcCtx["mcp"] != nil {
                        slog.Warn("enrichPORCWithAnnotations: porc[\"context\"][\"mcp\"] has unexpected type, creating fresh mcp context")
                }
                porcCtx["mcp"] = map[string]interface{}{"annotations": annotationMap}
                return
        }

        mcpCtx["annotations"] = annotationMap
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package http

import (
        "bytes"
        "context"
        "crypto/tls"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        nethttp "net/http"
        "net/url"
        "time"
)

const (
        // defaultTimeout is the default HTTP request timeout in seconds.
        defaultTimeout = 30

        // decisionPath is the PDP decision endpoint path.
        decisionPath = "/decision"
)

// DecisionResponse represents the response from the PDP decision endpoint.
type DecisionResponse struct {
        Allow bool `json:"allow"`
}

// Client handles HTTP communication with the PDP server.
type Client struct {
        baseURL    string
        httpClient *nethttp.Client
}

// NewClient creates a new HTTP client for PDP communication.
func NewClient(config *ConnectionConfig) (*Client, error) {
        slog.Debug("creating new HTTP client", "config", config)

        if config == nil {
                return nil, fmt.Errorf("HTTP configuration is required")
        }

        if config.URL == "" {
                return nil, fmt.Errorf("HTTP URL is required")
        }

        // Validate URL
        parsedURL, err := url.Parse(config.URL)
        if err != nil {
                return nil, fmt.Errorf("invalid URL: %w", err)
        }

        if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" {
                return nil, fmt.Errorf("URL scheme must be http or https, got: %s", parsedURL.Scheme)
        }

        // Determine timeout
        timeout := config.Timeout
        if timeout <= 0 {
                timeout = defaultTimeout
        }

        // Create HTTP client
        httpClient := &nethttp.Client{
                Timeout: time.Duration(timeout) * time.Second,
        }

        // Only set custom transport when we need to override TLS settings
        // Otherwise use http.DefaultTransport which includes proxy support and proper defaults
        if config.InsecureSkipVerify {
                // Clone default transport and override TLS config
                transport := nethttp.DefaultTransport.(*nethttp.Transport).Clone()
                transport.TLSClientConfig = &tls.Config{
                        InsecureSkipVerify: true, //nolint:gosec // User explicitly requested insecure mode
                }
                httpClient.Transport = transport
        }

        return &Client{
                baseURL:    config.URL,
                httpClient: httpClient,
        }, nil
}

// Authorize sends an authorization request to the PDP server.
// It returns true if the request is authorized, false otherwise.
func (c *Client) Authorize(ctx context.Context, porc PORC, probe bool) (bool, error) {
        // Build the decision URL
        decisionURL, err := url.JoinPath(c.baseURL, decisionPath)
        if err != nil {
                return false, fmt.Errorf("failed to build decision URL: %w", err)
        }

        // Add probe parameter if specified (for PDPs that support debugging mode)
        if probe {
                u, err := url.Parse(decisionURL)
                if err != nil {
                        return false, fmt.Errorf("failed to parse decision URL: %w", err)
                }
                q := u.Query()
                q.Set("probe", "true")
                u.RawQuery = q.Encode()
                decisionURL = u.String()
        }

        // Log authorization request without sensitive data
        // PORC can contain sensitive principal attributes and tool arguments, so we only log
        // high-level fields: operation and resource (subject if available)
        logSubject := "unknown"
        if principal, ok := porc["principal"].(map[string]interface{}); ok {
                if sub, ok := principal["sub"].(string); ok {
                        logSubject = sub
                }
        }
        //nolint:gosec // G706: authorization metadata is safe for debug logging
        slog.Debug("HTTP PDP authorization",
                "url", decisionURL, "subject", logSubject,
                "operation", porc["operation"], "resource", porc["resource"])

        // Marshal PORC to JSON
        body, err := json.Marshal(porc)
        if err != nil {
                return false, fmt.Errorf("failed to marshal PORC: %w", err)
        }

        // Create HTTP request
        req, err := nethttp.NewRequestWithContext(ctx, nethttp.MethodPost, decisionURL, bytes.NewReader(body))
        if err != nil {
                return false, fmt.Errorf("failed to create HTTP request: %w", err)
        }

        req.Header.Set("Content-Type", "application/json")
        req.Header.Set("Accept", "application/json")

        // Send request
        resp, err := c.httpClient.Do(req) //nolint:gosec // G704: URL is from server configuration, not user input
        if err != nil {
                return false, fmt.Errorf("HTTP request failed: %w", err)
        }
        defer func() { _ = resp.Body.Close() }()

        // Read response body
        respBody, err := io.ReadAll(resp.Body)
        if err != nil {
                return false, fmt.Errorf("failed to read response body: %w", err)
        }

        // Check HTTP status
        if resp.StatusCode != nethttp.StatusOK {
                return false, fmt.Errorf("PDP server returned status %d: %s", resp.StatusCode, string(respBody))
        }

        // Parse response
        var decision DecisionResponse
        if err := json.Unmarshal(respBody, &decision); err != nil {
                return false, fmt.Errorf("failed to parse decision response: %w", err)
        }

        return decision.Allow, nil
}

// Close closes the HTTP client and releases resources.
func (c *Client) Close() error {
        c.httpClient.CloseIdleConnections()
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package http

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/authz/authorizers"
)

// PORC represents a Principal-Operation-Resource-Context authorization request.
// This is a common input format for Policy Decision Points (PDPs).
type PORC map[string]interface{}

// Principal represents the principal (subject) making the request.
type Principal map[string]interface{}

// Context represents additional context for the authorization decision.
type Context map[string]interface{}

// PORCBuilder builds PORC (Principal-Operation-Resource-Context) expressions
// from MCP authorization parameters for use with HTTP-based PDPs.
type PORCBuilder struct {
        serverID      string
        contextConfig ContextConfig
        claimMapper   ClaimMapper
}

// NewPORCBuilder creates a new PORC builder with the given server ID, context configuration,
// and claim mapper.
func NewPORCBuilder(serverID string, contextConfig ContextConfig, claimMapper ClaimMapper) *PORCBuilder {
        return &PORCBuilder{
                serverID:      serverID,
                contextConfig: contextConfig,
                claimMapper:   claimMapper,
        }
}

// Build creates a PORC expression from MCP authorization parameters.
// It maps MCP concepts to the PORC model:
//   - principal: Built from JWT claims (sub, roles, groups, scopes, etc.)
//   - operation: Derived from MCP feature and operation (e.g., "mcp:tool:call")
//   - resource: The MCP resource identifier (tool name, prompt name, resource URI)
//   - context: Additional context including tool arguments
func (b *PORCBuilder) Build(
        feature authorizers.MCPFeature,
        operation authorizers.MCPOperation,
        resourceID string,
        claims map[string]interface{},
        arguments map[string]interface{},
) PORC {
        porc := make(PORC)

        // Build principal from JWT claims
        porc["principal"] = b.buildPrincipal(claims)

        // Build operation string from MCP feature and operation
        porc["operation"] = b.buildOperation(feature, operation)

        // Set resource identifier
        porc["resource"] = b.buildResource(feature, resourceID)

        // Build context from arguments and additional metadata
        porc["context"] = b.buildContext(feature, operation, resourceID, arguments)

        return porc
}

// buildPORC is a test helper that creates a PORC with all context options enabled.
// This function is only intended for use in tests within this package.
func buildPORC(
        feature authorizers.MCPFeature,
        operation authorizers.MCPOperation,
        resourceID string,
        claims map[string]interface{},
        arguments map[string]interface{},
) PORC {
        // Use a config with all context enabled and MPE claim mapper for testing
        contextConfig := ContextConfig{
                IncludeArgs:      true,
                IncludeOperation: true,
        }
        claimMapper := &MPEClaimMapper{}
        return NewPORCBuilder("test", contextConfig, claimMapper).Build(feature, operation, resourceID, claims, arguments)
}

// buildPrincipal constructs the principal object from JWT claims using the configured ClaimMapper.
// The claim mapping is delegated to the ClaimMapper, which can be configured to use different
// conventions (e.g., MPE-specific m-prefixed claims or standard OIDC claims).
//
// Note: Returns map[string]interface{} (the concrete type) rather than the Principal type alias
// for clarity and to match the ClaimMapper interface. Both types are equivalent for JSON
// marshaling, but using the concrete type avoids unnecessary type assertions and makes the
// return value explicit.
func (b *PORCBuilder) buildPrincipal(claims map[string]interface{}) map[string]interface{} {
        return b.claimMapper.MapClaims(claims)
}

// buildOperation constructs the operation string from MCP feature and operation.
// Format: "mcp:<feature>:<operation>"
// Examples:
//   - "mcp:tool:call" for calling a tool
//   - "mcp:prompt:get" for getting a prompt
//   - "mcp:resource:read" for reading a resource
//   - "mcp:tool:list" for listing tools
func (*PORCBuilder) buildOperation(feature authorizers.MCPFeature, operation authorizers.MCPOperation) string {
        return fmt.Sprintf("mcp:%s:%s", feature, operation)
}

// buildResource constructs the resource identifier.
// Format: "mrn:mcp:<serverID>:<feature>:<resourceID>"
// Examples:
//   - "mrn:mcp:myserver:tool:weather" for the weather tool
//   - "mrn:mcp:myserver:prompt:greeting" for the greeting prompt
//   - "mrn:mcp:myserver:resource:file://data.json" for a resource
func (b *PORCBuilder) buildResource(feature authorizers.MCPFeature, resourceID string) string {
        return fmt.Sprintf("mrn:mcp:%s:%s:%s", b.serverID, feature, resourceID)
}

// buildContext constructs the context object with additional metadata.
// Note: Returns map[string]interface{} (the concrete type) rather than the Context type alias
// for clarity. Both types are equivalent for JSON marshaling, but using the concrete type
// makes the return value explicit and avoids unnecessary type assertions.
//
// The context may include an "mcp" object with the following fields, depending
// on the ContextConfig settings:
//   - feature: The MCP feature type (tool, prompt, resource) - if IncludeOperation is true
//   - operation: The MCP operation (call, get, read, list) - if IncludeOperation is true
//   - resource_id: The resource identifier - if IncludeOperation is true
//   - args: Tool/prompt arguments (if present) - if IncludeArgs is true
//
// Important: The "mcp" object is only included in the context if it would contain
// at least one field. This means:
//   - If both IncludeOperation and IncludeArgs are false, returns an empty context {}
//   - If only IncludeArgs is true but arguments is nil/empty, returns an empty context {}
//   - If only IncludeOperation is true, returns context with mcp object containing operation fields
//   - If both are true and arguments exist, returns context with all enabled fields
//
// This prevents empty mcp objects from being included in the PORC, keeping it minimal.
func (b *PORCBuilder) buildContext(
        feature authorizers.MCPFeature,
        operation authorizers.MCPOperation,
        resourceID string,
        arguments map[string]interface{},
) map[string]interface{} {
        ctx := make(map[string]interface{})

        // Only build the mcp object if at least one context option is enabled
        if !b.contextConfig.IncludeArgs && !b.contextConfig.IncludeOperation {
                return ctx
        }

        // Build nested MCP object with metadata based on configuration
        mcp := make(map[string]interface{})

        // Include operation metadata if enabled
        if b.contextConfig.IncludeOperation {
                mcp["feature"] = string(feature)
                mcp["operation"] = string(operation)
                mcp["resource_id"] = resourceID
        }

        // Include tool/prompt arguments if enabled and present
        if b.contextConfig.IncludeArgs && len(arguments) > 0 {
                mcp["args"] = arguments
        }

        // Only add mcp to context if it has any fields
        if len(mcp) > 0 {
                ctx["mcp"] = mcp
        }

        return ctx
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authorizers

import (
        "encoding/json"
        "fmt"
        "sync"
)

// AuthorizerFactory is the interface that authorizer implementations must satisfy
// to register themselves with the authorizers registry. Each authorizer type
// (e.g., Cedar, OPA) implements this interface to provide validation and
// instantiation of authorizers from their specific configuration format.
type AuthorizerFactory interface {
        // ConfigKey returns the JSON key under which the backend-specific
        // configuration is nested in the full authorizer config blob.
        // For example, Cedar returns "cedar" and HTTP PDP returns "pdp".
        ConfigKey() string

        // ValidateConfig validates the authorizer-specific configuration.
        // The rawConfig is the JSON-encoded authorizer configuration.
        ValidateConfig(rawConfig json.RawMessage) error

        // CreateAuthorizer creates an Authorizer instance from the configuration.
        // The rawConfig is the JSON-encoded authorizer configuration.
        CreateAuthorizer(rawConfig json.RawMessage, serverName string) (Authorizer, error)
}

// registry holds the registered authorizer factories, keyed by config type.
var (
        registryMu sync.RWMutex
        registry   = make(map[string]AuthorizerFactory)
)

// reservedConfigKeys are envelope keys that a factory's ConfigKey() must not
// shadow. The MCPAuthzConfig controller assembles the runtime config as a
// JSON object with reserved top-level keys ("version", "type") plus the
// factory's ConfigKey(); a factory returning one of these would silently
// overwrite the envelope metadata in the map literal. Empty string is
// reserved too — a factory must declare a non-empty key.
var reservedConfigKeys = map[string]struct{}{
        "":        {},
        "version": {},
        "type":    {},
}

// Register registers an AuthorizerFactory for the given config type.
// This is typically called from an init() function in the authorizer package.
// It panics if a factory is already registered for the given type, or if the
// factory's ConfigKey() returns a reserved value ("", "version", "type").
func Register(configType string, factory AuthorizerFactory) {
        registryMu.Lock()
        defer registryMu.Unlock()

        if _, exists := registry[configType]; exists {
                panic(fmt.Sprintf("authorizer factory already registered for type: %s", configType))
        }
        if _, reserved := reservedConfigKeys[factory.ConfigKey()]; reserved {
                panic(fmt.Sprintf(
                        "authorizer factory for type %q uses reserved ConfigKey() %q; "+
                                "reserved keys are: \"\", \"version\", \"type\"",
                        configType, factory.ConfigKey(),
                ))
        }
        registry[configType] = factory
}

// GetFactory returns the AuthorizerFactory for the given config type.
// Returns nil if no factory is registered for the type.
func GetFactory(configType string) AuthorizerFactory {
        registryMu.RLock()
        defer registryMu.RUnlock()

        return registry[configType]
}

// IsRegistered returns true if a factory is registered for the given config type.
func IsRegistered(configType string) bool {
        registryMu.RLock()
        defer registryMu.RUnlock()

        _, exists := registry[configType]
        return exists
}

// RegisteredTypes returns a list of all registered config types.
func RegisteredTypes() []string {
        registryMu.RLock()
        defer registryMu.RUnlock()

        types := make([]string, 0, len(registry))
        for t := range registry {
                types = append(types, t)
        }
        return types
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authorizers

import "context"

// ResourceMetadata carries trusted server-side facts about the resource being
// authorized that are not part of the public Authorizer method signature.
//
// BackendID is the logical vMCP backend identifier. It MUST be sourced from the
// aggregated capability, never from client-supplied request data such as tool
// arguments or an advertised-name prefix. Annotations carry the tool hints from
// that same trusted capability.
type ResourceMetadata struct {
        BackendID   string
        Annotations *ToolAnnotations
}

// resourceMetadataKey is the unexported context key used by
// WithResourceMetadata and ResourceMetadataFromContext.
type resourceMetadataKey struct{}

// WithResourceMetadata stores trusted resource metadata in ctx.
func WithResourceMetadata(ctx context.Context, metadata ResourceMetadata) context.Context {
        return context.WithValue(ctx, resourceMetadataKey{}, metadata)
}

// ResourceMetadataFromContext retrieves trusted resource metadata previously
// stored with WithResourceMetadata. For compatibility, annotations stored with
// WithToolAnnotations are returned when unified metadata has no annotations.
// The second return value reports whether either metadata channel is present.
func ResourceMetadataFromContext(ctx context.Context) (ResourceMetadata, bool) {
        metadata, ok := ctx.Value(resourceMetadataKey{}).(ResourceMetadata)
        if metadata.Annotations == nil {
                metadata.Annotations = effectiveToolAnnotationsFromContext(ctx)
                ok = ok || metadata.Annotations != nil
        }
        return metadata, ok
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package authz provides authorization utilities for MCP servers.
// It supports a pluggable authorizer architecture where different authorization
// backends (e.g., Cedar, OPA) can be registered and used based on configuration.
package authz

import (
        "fmt"
        "net/http"

        "github.com/stacklok/toolhive/pkg/authz/authorizers"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// ConfigType is an alias for authorizers.ConfigType for backward compatibility.
type ConfigType = authorizers.ConfigType

// Config is an alias for authorizers.Config for backward compatibility.
type Config = authorizers.Config

// LoadConfig is an alias for authorizers.LoadConfig for backward compatibility.
var LoadConfig = authorizers.LoadConfig

// NewConfig is an alias for authorizers.NewConfig for backward compatibility.
var NewConfig = authorizers.NewConfig

// CreateMiddlewareFromConfig creates an HTTP middleware from the configuration.
// The passThroughTools parameter is optional (pass nil for none). Tool names in
// this set bypass the response filter's policy check in tools/list responses.
func CreateMiddlewareFromConfig(
        c *Config, serverName string, passThroughTools map[string]struct{},
) (types.MiddlewareFunction, error) {
        // Get the factory for this config type
        factory := authorizers.GetFactory(string(c.Type))
        if factory == nil {
                return nil, fmt.Errorf("unsupported configuration type: %s", c.Type)
        }

        // Create the authorizer using the factory, passing the full raw config
        authz, err := factory.CreateAuthorizer(c.RawConfig(), serverName)
        if err != nil {
                return nil, fmt.Errorf("failed to create %s authorizer: %w", c.Type, err)
        }

        // Return the middleware
        return func(handler http.Handler) http.Handler { return Middleware(authz, handler, passThroughTools) }, nil
}

// GetMiddlewareFromFile loads the authorization configuration from a file and creates an HTTP middleware.
// The passThroughTools parameter is optional (pass nil for none). Tool names in
// this set bypass the response filter's policy check in tools/list responses.
func GetMiddlewareFromFile(serverName, path string, passThroughTools map[string]struct{}) (types.MiddlewareFunction, error) {
        config, err := LoadConfig(path)
        if err != nil {
                return nil, err
        }

        // Create the middleware
        return CreateMiddlewareFromConfig(config, serverName, passThroughTools)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authz

import (
        "bytes"
        "encoding/json"
        "errors"
        "fmt"

        "github.com/stacklok/toolhive/pkg/authz/authorizers"
)

// Some MCP methods name a capability that policy can already talk about, but not
// in a shape MCPMethodToFeatureOperation can express: the feature depends on the
// request body (completion/complete) or a single request names many resources
// (subscriptions/listen). Both used to sit in that map with an empty
// feature/operation pair, which made them always-allowed and let an identity
// denied a prompt or resource reach it anyway through the side door.
//
// These methods are classified here instead. A derived method must NOT also
// appear in MCPMethodToFeatureOperation: one method, one vocabulary, so the
// always-allowed early return can never shadow the checks below
// (TestDerivedMethodsAreNotStaticallyClassified enforces this).

// MCP completion reference discriminators, per the completion/complete schema.
const (
        completionRefTypePrompt   = "ref/prompt"
        completionRefTypeResource = "ref/resource"
)

// maxResourceSubscriptions bounds how many resource URIs one subscriptions/listen
// request may name. Each distinct URI costs one authorization decision, an
// amplification this check introduces: before it, no decision was made at all.
//
// The bound is sized for the SLOWEST supported authorizer, not the fastest.
// Cedar evaluates in-process in microseconds, but the HTTP PDP (authorizers/http)
// issues one outbound POST per decision, sequentially, with no batching or
// caching and a 30s per-call timeout. A caller covered by a wildcard permit
// therefore converts one inbound request into one outbound PDP request per URI.
// A limit tuned to Cedar would be a denial-of-service amplifier against an
// httpv1 deployment.
//
// The limit counts entries as sent, before de-duplication, so it also bounds the
// work of de-duplicating them. Duplicates within the limit collapse to a single
// decision (see subscriptionsListenChecks).
//
// A request over the cap is denied rather than truncated: silently dropping the
// tail would subscribe a caller to less than it asked for while reporting success,
// and would mean the backend registers a different set than was authorized.
const maxResourceSubscriptions = 50

// resourceCheck is one authorization decision: the same feature/operation pair
// the static map carries, bound to the identifier it applies to. featureOperation
// is embedded rather than restated so the two cannot drift apart.
type resourceCheck struct {
        featureOperation
        resourceID string
}

// promptGet and resourceRead name the two decisions every derived method reduces
// to, so a classifier cannot accidentally pair a prompt with a read or authorize
// a resource as a prompt.
func promptGet(name string) resourceCheck {
        return resourceCheck{
                featureOperation: featureOperation{
                        Feature:   authorizers.MCPFeaturePrompt,
                        Operation: authorizers.MCPOperationGet,
                },
                resourceID: name,
        }
}

func resourceRead(uri string) resourceCheck {
        return resourceCheck{
                featureOperation: featureOperation{
                        Feature:   authorizers.MCPFeatureResource,
                        Operation: authorizers.MCPOperationRead,
                },
                resourceID: uri,
        }
}

// deriveChecks turns a method's raw params into the decisions that must all pass
// before the request may be dispatched.
//
// Returning an error means "this request has no resolvable authorization target"
// and is always a denial. It is deliberately distinct from returning zero checks,
// which means "this request names no protected capability" and is an allow: an
// error must never be mistaken for an empty check list, or a malformed request
// would sail through as if it had asked for nothing.
type deriveChecks func(params json.RawMessage) ([]resourceCheck, error)

// derivedAuthorizers maps each derived method to its classifier.
var derivedAuthorizers = map[string]deriveChecks{
        "completion/complete":  completionChecks,
        "subscriptions/listen": subscriptionsListenChecks,
}

// completionRef is the ref object identifying what a completion targets. Name and
// URI are pointers so an absent field is distinguishable from an empty one: an
// empty identifier must be rejected, not passed to a policy that might match a
// broad rule on it.
type completionRef struct {
        Type string  `json:"type"`
        Name *string `json:"name"`
        URI  *string `json:"uri"`
}

type completionParams struct {
        Ref *completionRef `json:"ref"`
}

// completionChecks authorizes completion/complete as the capability its ref names:
// a prompt ref needs the same get decision prompts/get needs, a resource-template
// ref the same read decision resources/read needs. Completion candidates are
// derived from the referenced object, so reaching them without that decision
// discloses exactly what the decision exists to withhold.
//
// Both the operation and the identifier are read from the same branch below.
// pkg/mcp's parser falls back from name to uri independently of type, so a ref
// declaring one kind while carrying the other's field would otherwise let the
// authorized identifier and the dispatched one disagree.
func completionChecks(params json.RawMessage) ([]resourceCheck, error) {
        var p completionParams
        if err := json.Unmarshal(params, &p); err != nil {
                return nil, fmt.Errorf("decoding completion/complete params: %w", err)
        }
        if p.Ref == nil {
                return nil, errors.New("completion/complete: missing ref")
        }
        // No reference kind in the schema carries both fields. One that does is
        // ambiguous about which identifier the backend will act on, so it is refused
        // rather than authorized on the reporter's choice of the two.
        if p.Ref.Name != nil && p.Ref.URI != nil {
                return nil, errors.New("completion/complete: ambiguous ref carries both name and uri")
        }

        switch p.Ref.Type {
        case completionRefTypePrompt:
                if p.Ref.Name == nil || *p.Ref.Name == "" {
                        return nil, errors.New("completion/complete: prompt ref without a name")
                }
                return []resourceCheck{promptGet(*p.Ref.Name)}, nil

        case completionRefTypeResource:
                if p.Ref.URI == nil || *p.Ref.URI == "" {
                        return nil, errors.New("completion/complete: resource ref without a uri")
                }
                return []resourceCheck{resourceRead(*p.Ref.URI)}, nil

        default:
                // Includes the absent discriminator and pkg/mcp's legacy bare-string ref,
                // neither of which says which capability to authorize.
                // Truncated: ref.type is caller-chosen, and an untruncated %q would let a
                // caller decide how many bytes each denied request writes to the log.
                return nil, fmt.Errorf("completion/complete: unsupported ref type %q", truncateForLog(p.Ref.Type))
        }
}

// notificationSubscriptionFields is the complete subscribable set of the
// 2026-07-28 subscriptions/listen params, mirroring pkg/vmcp/server's
// notificationSubscriptions field-for-field. Only resourceSubscriptions names
// concrete resources; the three list-changed flags are server-wide and carry no
// identifier a policy could match.
//
// The set is enforced, not merely documented: a member outside it is refused
// rather than ignored. Silently ignoring unknown members would make this
// classifier safe only by coincidence of today's schema -- a future revision or
// vendor extension adding another URI-carrying field would reopen the bypass with
// no test failing. Refusing is the same default-deny posture
// MCPMethodToFeatureOperation already takes for unknown methods, and it forces a
// deliberate update here when the schema grows.
const fieldResourceSubscriptions = "resourceSubscriptions"

var notificationSubscriptionFields = map[string]struct{}{
        "toolsListChanged":         {},
        "promptsListChanged":       {},
        "resourcesListChanged":     {},
        fieldResourceSubscriptions: {},
}

// listenParams decodes only far enough to reach the members by name. Both levels
// stay raw on purpose: the member map makes an unrecognised member visible, where
// a typed struct would silently drop it, and the raw notifications value keeps an
// explicit JSON null distinguishable from an absent member.
type listenParams struct {
        Notifications json.RawMessage `json:"notifications"`
}

// errNotAnArray and errNotAnObject describe a member whose JSON type contradicts
// the schema. Such a member is refused rather than coerced: this middleware and
// the backend would each have to guess the same meaning for it, and a guess that
// diverges is how an unauthorized URI reaches a subscription. Absent and empty
// are NOT type violations -- both unambiguously name no resource, so both pass
// with zero checks.
var (
        errNotAnArray  = errors.New("subscriptions/listen: resourceSubscriptions must be an array")
        errNotAnObject = errors.New("subscriptions/listen: notifications must be an object")
)

// isJSONNull reports whether raw is the literal JSON null.
func isJSONNull(raw json.RawMessage) bool {
        return string(bytes.TrimSpace(raw)) == "null"
}

// subscriptionsListenChecks authorizes every URI a subscriptions/listen request
// asks to be notified about, as the same read decision resources/subscribe uses.
// A registered subscription streams notifications/resources/updated for its URI,
// which reveals the resource's existence and change timing even though its
// contents stay behind resources/read.
//
// A denied URI fails the whole request rather than being filtered out of it.
// Filtering would mean rewriting the forwarded body, the acknowledgement, and the
// notification stream, and any of the three drifting reopens the hole; refusing
// the request keeps one decision, made once, in one place.
func subscriptionsListenChecks(params json.RawMessage) ([]resourceCheck, error) {
        var p listenParams
        if err := json.Unmarshal(params, &p); err != nil {
                // A body this middleware cannot decode may still carry a URI the backend
                // decodes and honours, so it cannot be forwarded unchecked.
                return nil, fmt.Errorf("decoding subscriptions/listen params: %w", err)
        }
        if len(p.Notifications) == 0 {
                // The member is absent (e.g. "params":{}). Nothing names a resource, so
                // there is nothing to authorize; the backend rejects the missing required
                // field as invalid params and this middleware does not duplicate its schema
                // validation.
                //
                // A request with NO params member at all does not reach here: Params is nil,
                // the Unmarshal above fails, and the request is denied.
                return nil, nil
        }
        if isJSONNull(p.Notifications) {
                return nil, errNotAnObject
        }

        members := map[string]json.RawMessage{}
        if err := json.Unmarshal(p.Notifications, &members); err != nil {
                return nil, fmt.Errorf("decoding notifications: %w", err)
        }

        for member := range members {
                if _, known := notificationSubscriptionFields[member]; !known {
                        return nil, fmt.Errorf("subscriptions/listen: unrecognised notifications member %q",
                                truncateForLog(member))
                }
        }

        var uris []string
        if raw, ok := members[fieldResourceSubscriptions]; ok {
                // Decoded through a pointer so an explicit null is distinguishable from an
                // empty array: unmarshalling null straight into a []string yields a nil
                // slice with err == nil, which would read as "no URIs" and pass with zero
                // checks. An empty array leaves a non-nil pointer and is allowed.
                var decoded *[]string
                if err := json.Unmarshal(raw, &decoded); err != nil {
                        return nil, fmt.Errorf("decoding %s: %w", fieldResourceSubscriptions, err)
                }
                if decoded == nil {
                        return nil, errNotAnArray
                }
                uris = *decoded
        }

        if len(uris) > maxResourceSubscriptions {
                return nil, fmt.Errorf("subscriptions/listen: %d resource subscriptions exceeds the limit of %d",
                        len(uris), maxResourceSubscriptions)
        }

        // De-duplicated, preserving first-seen order. Repeated URIs yield byte-identical
        // decisions, so asking the authorizer again is pure amplification -- against the
        // HTTP PDP it is one wasted network round trip each. Order is preserved so the
        // decision sequence, and therefore which URI a denial is attributed to, is
        // deterministic.
        //
        // This changes only which decisions are made, never what is forwarded: the
        // request body is proxied verbatim, duplicates included.
        seen := make(map[string]struct{}, len(uris))
        checks := make([]resourceCheck, 0, len(uris))
        for _, uri := range uris {
                if uri == "" {
                        return nil, errors.New("subscriptions/listen: empty resource uri")
                }
                if _, duplicate := seen[uri]; duplicate {
                        continue
                }
                seen[uri] = struct{}{}
                checks = append(checks, resourceRead(uri))
        }
        return checks, nil
}

// truncateForLog bounds a caller-supplied string before it reaches a log line.
func truncateForLog(s string) string {
        const maxLoggedLen = 64
        if len(s) <= maxLoggedLen {
                return s
        }
        return s[:maxLoggedLen] + "..."
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package authz provides authorization utilities for MCP servers.
// It supports a pluggable authorizer architecture where different authorization
// backends (e.g., Cedar, OPA) can be registered and used based on configuration.
package authz

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "net/http"
        "time"

        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/authz/authorizers"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer"
        "github.com/stacklok/toolhive/pkg/vmcp/schema"
        "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec"
)

// featureOperation pairs an MCP feature with an operation for authorization checks.
type featureOperation struct {
        Feature   authorizers.MCPFeature
        Operation authorizers.MCPOperation
}

// MCPMethodToFeatureOperation maps MCP method names to feature and operation pairs.
// Methods with empty Feature and Operation are always allowed (protocol-level).
// Methods not in this map are denied by default for security.
var MCPMethodToFeatureOperation = map[string]featureOperation{
        // Core protocol methods - always allowed
        "initialize": {Feature: "", Operation: ""}, // Protocol initialization
        "ping":       {Feature: "", Operation: ""}, // Health check
        // Tool operations - require authorization
        "tools/call": {Feature: authorizers.MCPFeatureTool, Operation: authorizers.MCPOperationCall},
        "tools/list": {Feature: authorizers.MCPFeatureTool, Operation: authorizers.MCPOperationList},

        // Prompt operations - require authorization
        "prompts/get":  {Feature: authorizers.MCPFeaturePrompt, Operation: authorizers.MCPOperationGet},
        "prompts/list": {Feature: authorizers.MCPFeaturePrompt, Operation: authorizers.MCPOperationList},

        // Resource operations - require authorization
        "resources/read":           {Feature: authorizers.MCPFeatureResource, Operation: authorizers.MCPOperationRead},
        "resources/list":           {Feature: authorizers.MCPFeatureResource, Operation: authorizers.MCPOperationList},
        "resources/templates/list": {Feature: authorizers.MCPFeatureResource, Operation: authorizers.MCPOperationList},
        "resources/subscribe":      {Feature: authorizers.MCPFeatureResource, Operation: authorizers.MCPOperationRead},
        "resources/unsubscribe":    {Feature: authorizers.MCPFeatureResource, Operation: authorizers.MCPOperationRead},

        // Skill operations - list responses are filtered by get authorization.
        "skills/get":  {Feature: authorizers.MCPFeatureSkill, Operation: authorizers.MCPOperationGet},
        "skills/list": {Feature: authorizers.MCPFeatureSkill, Operation: authorizers.MCPOperationList},

        // Discovery and capability methods - always allowed
        "features/list": {Feature: "", Operation: authorizers.MCPOperationList}, // Capability discovery
        "roots/list":    {Feature: "", Operation: ""},                           // Root directory discovery

        // server/discover, Modern's (2026-07-28) replacement for initialize+capability
        // negotiation, is always-allowed on THIS path (the single-server pkg/runner HTTP
        // authz Middleware -- vMCP's Modern dispatcher never consults this map at all, it
        // re-homes admission through core.Check*/core.List* directly). The always-allowed
        // choice rests on initialize parity, not on any per-request filtering this map
        // enforces: DiscoverResult carries the exact same Capabilities *ServerCapabilities
        // (+ Instructions) shape InitializeResult does, and "initialize" above has always
        // been always-allowed in this map. discover therefore adds no new exposure class --
        // note ServerCapabilities.Experimental/.Extensions (arbitrary backend-authored maps)
        // and Instructions (free text) are already freeform fields a backend can populate on
        // the always-allowed initialize response today, so "no descriptors" is a property of
        // how vMCP's dispatcher happens to build the value, not a guarantee this wire shape
        // makes on its own. Classifying it as MCPOperationList with an empty Feature
        // would be safe too: response_filter.go's authoritative classifier does not
        // assign server/discover a filter, so protocol-only list methods pass through
        // unchanged. The always-allowed classification is simpler and equally safe here.
        "server/discover": {Feature: "", Operation: ""},

        // Logging and client preferences - always allowed
        "logging/setLevel": {Feature: "", Operation: ""}, // Client preference for server logging

        // NOTE: completion/complete and subscriptions/listen are deliberately absent.
        // Both name a capability policy can already decide on, but not as one static
        // feature/operation pair, so they are classified by derivedAuthorizers
        // (derived_authz.go) instead. Middleware consults that map first, so a
        // duplicate static entry here would not shadow derived checks. The
        // regression to avoid is removing those classifiers and restoring an empty
        // feature/operation pair in this map — that is the original always-allow
        // bypass.

        // Notifications (server-to-client, informational) - always allowed
        "notifications/message":                {Feature: "", Operation: ""}, // General notifications
        "notifications/initialized":            {Feature: "", Operation: ""}, // Initialization complete
        "notifications/progress":               {Feature: "", Operation: ""}, // Progress updates
        "notifications/cancelled":              {Feature: "", Operation: ""}, // Request cancellation
        "notifications/roots/list_changed":     {Feature: "", Operation: ""}, // Roots changed
        "notifications/tools/list_changed":     {Feature: "", Operation: ""}, // Tools changed
        "notifications/prompts/list_changed":   {Feature: "", Operation: ""}, // Prompts changed
        "notifications/resources/list_changed": {Feature: "", Operation: ""}, // Resources changed
        "notifications/resources/updated":      {Feature: "", Operation: ""}, // Resource updated
        "notifications/tasks/status":           {Feature: "", Operation: ""}, // Task status update

        // NOTE: The following MCP methods are NOT included and will be DENIED by default:
        // - elicitation/create: User input prompting (requires new authorization feature)
        // - sampling/createMessage: LLM text generation (security-sensitive, requires new authorization feature)
        // - tasks/list, tasks/get, tasks/cancel, tasks/result: Task management (requires new authorization feature)
        //
        // To enable these methods, add appropriate authorization features/operations or add them
        // to the always-allowed list above after security review.
}

// shouldSkipInitialAuthorization checks if the request should skip authorization
// before inspecting parsed data. Every POST requires authorization regardless
// of Content-Type.
func shouldSkipInitialAuthorization(r *http.Request) bool {
        return r.Method != http.MethodPost
}

// shouldSkipSubsequentAuthorization checks if the request should skip authorization
// after parsing the JSON-RPC message.
func shouldSkipSubsequentAuthorization(method string) bool {
        // Skip authorization for methods that don't require it
        if method == "ping" || method == "initialize" {
                return true
        }

        return false
}

// missingTargetID reports whether a request that identifies its target by
// params.name or params.uri (tools/call, prompts/get, resources/read,
// skills/get) lacks a valid one. The parser yields an empty resource ID for an
// absent, empty or non-string value.
func missingTargetID(featureOp featureOperation, resourceID string) bool {
        _, targeted := targetedOperations[featureOp]
        return targeted && resourceID == ""
}

// targetedOperations are the feature/operations whose authorization target is
// params.name or params.uri.
var targetedOperations = map[featureOperation]struct{}{
        {Feature: authorizers.MCPFeatureTool, Operation: authorizers.MCPOperationCall}:     {},
        {Feature: authorizers.MCPFeaturePrompt, Operation: authorizers.MCPOperationGet}:    {},
        {Feature: authorizers.MCPFeatureResource, Operation: authorizers.MCPOperationRead}: {},
        {Feature: authorizers.MCPFeatureSkill, Operation: authorizers.MCPOperationGet}:     {},
}

// handleUnauthorized handles unauthorized requests. The client always sees the fixed
// "Unauthorized" message -- err (an authorizer failure) can carry policy detail that
// security.md forbids returning to callers, so it is logged server-side instead.
// Cedar's evaluation context can embed decoded JWT claim values (see the claim-keys-
// only rule in authorizers/cedar/core.go), so err must not be surfaced to the client,
// nor copied into additional log lines or fields beyond the single Warn below.
func handleUnauthorized(w http.ResponseWriter, msgID interface{}, err error) {
        if err != nil {
                slog.Warn("authorization denied", "error", err)
        }

        // Create a JSON-RPC error response
        id, convErr := mcp.ConvertToJSONRPC2ID(msgID)
        if convErr != nil {
                id = jsonrpc2.ID{} // Use empty ID if conversion fails
        }

        errorResponse := &jsonrpc2.Response{
                ID:    id,
                Error: jsonrpc2.NewError(mcp.JSONRPCCodeDenied, "Unauthorized"),
        }

        // The helper encodes before writing any header, so a marshal failure never
        // leaves a half-written response (e.g. a 403 header followed by a second
        // 500 write). Nothing to do on a write error for a denial body.
        _ = mcp.WriteJSONRPCError(w, http.StatusForbidden, errorResponse)
}

// rejectInvalidMCPRequest writes the 400 response for requests that arrive
// without a parsed MCP request, including client responses and POSTs that
// reached authorization without parsing middleware.
func rejectInvalidMCPRequest(w http.ResponseWriter) {
        http.Error(w, "Invalid or malformed MCP request", http.StatusBadRequest)
}

// Middleware creates an HTTP middleware that authorizes MCP requests.
// This middleware extracts the MCP message from the request, determines the feature,
// operation, and resource ID, and authorizes the request using the configured authorizer.
//
// For protected list operations (tools/list, prompts/list, resources/list,
// resources/templates/list, and skills/list), the middleware allows the request to
// proceed only when a response filter is registered, then filters out items that
// the user is not authorized to access based on the corresponding call/get/read
// policies. In particular, skills/list entries are filtered individually by
// get_skill authorization.
//
// An in-memory annotation cache is maintained per middleware instance. When a
// tools/list response passes through, tool annotations are captured. When a
// subsequent tools/call request arrives, the cached annotations are injected into
// the request context so that authorizers can use them for policy decisions.
//
// The authorizer parameter should implement the authorizers.Authorizer interface,
// which can be created using authz.CreateMiddlewareFromConfig() or directly
// from an authorizer package (e.g., cedar.NewCedarAuthorizer()).
func Middleware(a authorizers.Authorizer, next http.Handler, passThroughTools map[string]struct{}) http.Handler {
        // Cache is shared across requests for the same proxy.
        // Populated from tools/list responses, read during tools/call.
        annotationCache := NewAnnotationCache()

        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                // Check if we should skip authorization before checking parsed data
                if shouldSkipInitialAuthorization(r) {
                        next.ServeHTTP(w, r)
                        return
                }

                // Get parsed MCP request from context (set by parsing middleware)
                parsedRequest := mcp.GetParsedMCPRequest(r.Context())
                if parsedRequest == nil {
                        // Responses are not authorized here; missing parsing middleware also
                        // fails closed, regardless of Content-Type.
                        if marker, ok := mcp.AuthzDenialMarkerFromContext(r.Context()); ok {
                                marker.Denied = true
                        }
                        rejectInvalidMCPRequest(w)
                        return
                }

                // Check if we should skip authorization after parsing the message
                if shouldSkipSubsequentAuthorization(parsedRequest.Method) {
                        next.ServeHTTP(w, r)
                        return
                }

                // Methods whose authorization target is derived from the request body are
                // resolved first, ahead of the static map, so no always-allowed entry can
                // ever shadow their per-capability checks.
                if derive, ok := derivedAuthorizers[parsedRequest.Method]; ok {
                        authorizeDerivedAndServe(w, r, a, parsedRequest, derive, next)
                        return
                }

                // Get the feature and operation from the method
                featureOp, ok := MCPMethodToFeatureOperation[parsedRequest.Method]
                if !ok {
                        // Unknown method - deny by default for security. Methods must be
                        // explicitly added to MCPMethodToFeatureOperation to be allowed.
                        // This is expected traffic (e.g. a newer-spec client), not an
                        // operational failure, so it's Debug rather than the Warn
                        // handleUnauthorized logs for an actual authorizer error.
                        slog.Debug("MCP method denied by default", "method", parsedRequest.Method)
                        handleUnauthorized(w, parsedRequest.ID, nil)
                        return
                }

                // Methods with empty feature and operation are always allowed (protocol-level)
                if featureOp.Feature == "" && featureOp.Operation == "" {
                        next.ServeHTTP(w, r)
                        return
                }

                // These methods identify their target only by params.name or params.uri.
                // An absent, empty, or non-string value must never reach an authorizer,
                // whose policy might otherwise accidentally permit an empty identifier
                // (e.g. a permit that does not name the tool, beside a forbid that does).
                if missingTargetID(featureOp, parsedRequest.ResourceID) {
                        slog.Warn("denied MCP request without a valid target name or URI", "method", parsedRequest.Method)
                        handleUnauthorized(w, parsedRequest.ID, nil)
                        return
                }

                // Handle list operations differently: protected methods require a registered
                // response filter, while protocol-only methods pass through unchanged.
                if featureOp.Operation == authorizers.MCPOperationList {
                        authorizeListAndServe(w, r, a, parsedRequest, featureOp, annotationCache, passThroughTools, next)
                        return
                }

                // For tools/call, look up annotations and handle pass-through meta-tools.
                if featureOp.Feature == authorizers.MCPFeatureTool && featureOp.Operation == authorizers.MCPOperationCall {
                        handleToolsCall(w, r, a, parsedRequest, featureOp, annotationCache, passThroughTools, next)
                        return
                }

                // For non-list, non-tool operations, perform authorization using parsed data.
                authorizeAndServe(w, r, a, annotationCache,
                        featureOp.Feature, featureOp.Operation,
                        parsedRequest.ID, parsedRequest.ResourceID, parsedRequest.Arguments, next)
        })
}

// authorizeListAndServe intercepts a list response and applies its registered
// per-item authorization filter. Protected methods without a filter fail closed
// before backend dispatch; protocol-only list methods remain pass-through.
func authorizeListAndServe(
        w http.ResponseWriter,
        r *http.Request,
        a authorizers.Authorizer,
        parsedRequest *mcp.ParsedMCPRequest,
        featureOp featureOperation,
        annotationCache *AnnotationCache,
        passThroughTools map[string]struct{},
        next http.Handler,
) {
        // A protected list operation without a response filter would expose every
        // descriptor returned by the backend. Deny before dispatch so additions to
        // MCPMethodToFeatureOperation cannot silently create another filter bypass.
        if featureOp.Feature != "" && !requiresResponseFiltering(parsedRequest.Method) {
                slog.Error("protected MCP list method has no response filter; denying request",
                        "method", parsedRequest.Method)
                handleUnauthorized(w, parsedRequest.ID, nil)
                return
        }

        filteringWriter := NewResponseFilteringWriter(
                w, a, r, parsedRequest.Method, annotationCache, passThroughTools,
        )
        next.ServeHTTP(filteringWriter, r)

        if err := filteringWriter.FlushAndFilter(); err != nil {
                // The response may already be committed, so filtering errors can only be logged here.
                slog.Warn("error flushing filtered response", "error", err)
        }
}

// maxDerivedAuthzBudget bounds the total wall-clock time one request may spend
// making authorization decisions, across all of its derived checks.
//
// The count limit alone does not bound this. authorizers/http issues one
// outbound POST per decision with its own per-call timeout (30s by default), so
// a PDP that hangs turns N decisions into N times that timeout -- at the
// subscription limit, over twenty minutes of a request held open. The budget
// makes the worst case independent of N.
//
// It is set to the HTTP authorizer's own default per-call timeout on purpose:
// one request may spend no longer authorizing than a single decision was
// already allowed to take. That makes this a ceiling on fan-out rather than a
// new constraint on how slow an individual decision may be.
const maxDerivedAuthzBudget = 30 * time.Second

// authorizeChecks runs every check in order under one shared deadline, stopping
// at the first that does not pass. It reports whether all of them passed.
//
// The budget is applied only when there is more than one decision to make. A
// single decision has no fan-out to bound and is already limited by the
// authorizer's own per-call timeout, which an operator may have configured
// deliberately; wrapping it here would silently override that and make, say, a
// completion's get_prompt decision stricter than the identical decision for
// prompts/get.
//
// A budget overrun surfaces as an error from the authorizer call, which the
// caller maps to the same fail-closed denial any other authorizer error
// produces.
func authorizeChecks(
        ctx context.Context,
        a authorizers.Authorizer,
        checks []resourceCheck,
        budget time.Duration,
) (bool, error) {
        if len(checks) > 1 {
                var cancel context.CancelFunc
                ctx, cancel = context.WithTimeout(ctx, budget)
                defer cancel()
        }

        for i, check := range checks {
                // Arguments are deliberately nil, NOT the request's parsed arguments.
                //
                // Cedar prefixes every argument key with "arg_" and merges the result into
                // both the resource entity's attributes and the evaluation context
                // (authorizers/cedar/core.go preprocessArguments). The "arg_" namespace
                // therefore means "the arguments the authorized operation will run with".
                //
                // A derived method's params are NOT those arguments. pkg/mcp hands
                // completion/complete its ENTIRE params map as Arguments, so forwarding it
                // would let a caller put any key at the top level of a completion request
                // and have it evaluated as a prompt argument: a policy reading
                // context.arg_env would be satisfied by an attacker-chosen "env" that the
                // real prompts/get (which passes only params.arguments) could never forge.
                // For resources it is worse -- resources/read passes nil, so an
                // arg_-conditioned permit is unreachable there but would become reachable
                // through completion.
                //
                // nil is the fail-closed choice: an arg_-conditioned permit simply does not
                // match, so the request is denied rather than wrongly allowed. Passing real
                // completion context would need its own namespace, not this one.
                authorized, err := a.AuthorizeWithJWTClaims(ctx, check.Feature, check.Operation, check.resourceID, nil)
                if err != nil {
                        // Counts, never identifiers: how far the request got is the useful
                        // diagnostic for a budget overrun, and the URIs are caller-chosen.
                        return false, fmt.Errorf("authorizing decision %d of %d: %w", i+1, len(checks), err)
                }
                if !authorized {
                        return false, nil
                }
        }
        return true, nil
}

// authorizeDerivedAndServe authorizes a method whose targets are derived from the
// request body, dispatching only once every derived check passes.
//
// Checks are evaluated in order and the first denial ends the request, so a
// request naming several resources is admitted as a unit or not at all. A derive
// error is a denial too: it means no authorization target could be established,
// which must never be treated as "nothing to authorize".
func authorizeDerivedAndServe(
        w http.ResponseWriter,
        r *http.Request,
        a authorizers.Authorizer,
        parsedRequest *mcp.ParsedMCPRequest,
        derive deriveChecks,
        next http.Handler,
) {
        checks, err := derive(parsedRequest.Params)
        if err != nil {
                // WARN, not Debug: every derive failure is a request-SHAPE rejection (a
                // legacy bare-string ref, an unrecognised notifications member, a list over
                // the cap, undecodable params), not a policy decision. Each one means a
                // client or a policy needs updating, which is the repo's stated use for WARN
                // -- and it is the only signal an operator gets that an upgrade started
                // refusing traffic that previously passed. Policy denials stay silent, so
                // this does not log ordinary authorization outcomes.
                //
                // The reason is logged but never returned: the response body stays the fixed
                // "Unauthorized" so a prober cannot learn how its probe was malformed. Values
                // interpolated into these errors are truncated at the source, bounding what a
                // caller can write per line.
                slog.Warn("MCP request denied: no resolvable authorization target",
                        "method", parsedRequest.Method, "error", err)
                handleUnauthorized(w, parsedRequest.ID, nil)
                return
        }

        authorized, err := authorizeChecks(r.Context(), a, checks, maxDerivedAuthzBudget)
        if err != nil || !authorized {
                handleUnauthorized(w, parsedRequest.ID, err)
                return
        }

        next.ServeHTTP(w, r)
}

// authorizeAndServe injects tool annotations from the cache, authorizes the request,
// and calls next if authorized. It handles both the unauthorized response and the
// successful serve path, so callers do not need to do either after calling this.
func authorizeAndServe(
        w http.ResponseWriter,
        r *http.Request,
        a authorizers.Authorizer,
        annotationCache *AnnotationCache,
        feature authorizers.MCPFeature,
        operation authorizers.MCPOperation,
        msgID interface{},
        toolName string,
        args map[string]interface{},
        next http.Handler,
) {
        if ann := annotationCache.Get(toolName); ann != nil {
                r = r.WithContext(authorizers.WithToolAnnotations(r.Context(), ann))
        }
        authorized, err := a.AuthorizeWithJWTClaims(r.Context(), feature, operation, toolName, args)
        if err != nil || !authorized {
                handleUnauthorized(w, msgID, err)
                return
        }
        next.ServeHTTP(w, r)
}

// handleToolsCall handles tools/call authorization, including pass-through meta-tools.
// It always fully handles the request (authorization, unauthorized response, or serving).
//
// For pass-through meta-tools (find_tool, call_tool):
//   - call_tool: authorizes the real inner tool name, decoded from the request
//     arguments exactly as dispatch decodes them, so the two cannot disagree about
//     which tool a request names. Arguments that do not decode are denied.
//   - find_tool (and other pass-through tools without a tool_name): allowed through
//     as a discovery operation with no policy check.
//
// For normal tools: injects annotations from the cache and authorizes before serving.
func handleToolsCall(
        w http.ResponseWriter,
        r *http.Request,
        a authorizers.Authorizer,
        parsedRequest *mcp.ParsedMCPRequest,
        featureOp featureOperation,
        annotationCache *AnnotationCache,
        passThroughTools map[string]struct{},
        next http.Handler,
) {
        if _, isPassThrough := passThroughTools[parsedRequest.ResourceID]; isPassThrough {
                // Decode with the same call the two call_tool dispatch sites use rather than
                // indexing the arguments map. encoding/json matches struct fields
                // case-insensitively, so a map index on "tool_name" misses a request carrying
                // "Tool_Name" that dispatch resolves and runs. Going through CallToolInput
                // also applies the nested-tool_name hoist, so the name authorized here is the
                // one that will execute.
                input, err := schema.Translate[optimizer.CallToolInput](parsedRequest.Arguments)
                if err != nil {
                        // The arguments are not a decodable call_tool payload, so the tool they
                        // target cannot be established. Deny rather than pass through; dispatch
                        // decodes the same map with the same call and rejects it too, so no
                        // legitimate invocation is lost.
                        slog.Warn("denying pass-through tool call with undecodable arguments",
                                "tool", parsedRequest.ResourceID, "error", err)
                        handleUnauthorized(w, parsedRequest.ID, nil)
                        return
                }
                if input.ToolName != "" {
                        // call_tool: authorize the real backend tool name.
                        authorizeAndServe(w, r, a, annotationCache,
                                featureOp.Feature, featureOp.Operation,
                                parsedRequest.ID, input.ToolName, input.Parameters, next)
                        return
                }
                // find_tool: allow through but filter the tools list in the response so
                // callers cannot discover tools they are not authorized to call.
                if parsedRequest.ResourceID == optimizerdec.FindToolName {
                        filteringWriter := NewResponseFilteringWriter(w, a, r, optimizerdec.FindToolName, annotationCache, passThroughTools)
                        next.ServeHTTP(filteringWriter, r)
                        if err := filteringWriter.FlushAndFilter(); err != nil {
                                slog.Warn("error filtering find_tool response", "error", err)
                        }
                        return
                }
                // Other pass-through tools without a wrapped toolName: allow through.
                next.ServeHTTP(w, r)
                return
        }

        // Normal tool: inject annotations and authorize.
        authorizeAndServe(w, r, a, annotationCache,
                featureOp.Feature, featureOp.Operation,
                parsedRequest.ID, parsedRequest.ResourceID, parsedRequest.Arguments, next)
}

// Factory middleware type constant
const (
        MiddlewareType = "authorization"
)

// FactoryMiddlewareParams represents the parameters for authorization middleware
type FactoryMiddlewareParams struct {
        ConfigPath string  `json:"config_path,omitempty"` // Kept for backwards compatibility
        ConfigData *Config `json:"config_data,omitempty"` // New field for config contents
}

// FactoryMiddleware wraps authorization middleware functionality for factory pattern
type FactoryMiddleware struct {
        middleware types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *FactoryMiddleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (*FactoryMiddleware) Close() error {
        // Authorization middleware doesn't need cleanup
        return nil
}

// CreateMiddleware factory function for authorization middleware
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {

        var params FactoryMiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal authorization middleware parameters: %w", err)
        }

        var authzConfig *Config
        var err error

        if params.ConfigData != nil {
                // Use provided config data (preferred method)
                authzConfig = params.ConfigData
        } else if params.ConfigPath != "" {
                // Load config from file (backwards compatibility)
                authzConfig, err = LoadConfig(params.ConfigPath)
                if err != nil {
                        return fmt.Errorf("failed to load authorization configuration: %w", err)
                }
        } else {
                return fmt.Errorf("either config_data or config_path is required for authorization middleware")
        }

        middleware, err := CreateMiddlewareFromConfig(authzConfig, runner.GetConfig().GetName(), nil)
        if err != nil {
                return fmt.Errorf("failed to create authorization middleware: %w", err)
        }

        authzMw := &FactoryMiddleware{middleware: middleware}
        runner.AddMiddleware(config.Type, authzMw)
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package authz provides authorization utilities for MCP servers.
package authz

import (
        "bytes"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "reflect"
        "strings"

        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive/pkg/authz/authorizers"
        mcpparser "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer"
        "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec"
)

var errBug = errors.New("there's a bug")

// ResponseFilteringWriter wraps an http.ResponseWriter to intercept and filter responses
type ResponseFilteringWriter struct {
        http.ResponseWriter
        authorizer       authorizers.Authorizer
        request          *http.Request
        method           string
        buffer           *bytes.Buffer
        statusCode       int
        annotationCache  *AnnotationCache
        passThroughTools map[string]struct{}
}

// NewResponseFilteringWriter creates a new response filtering writer.
// The annotationCache parameter is optional; pass nil to disable annotation caching.
// The passThroughTools parameter is optional; tools whose names appear in this set
// bypass policy filtering because authorization is enforced elsewhere (e.g., inside
// the optimizer decorator for find_tool/call_tool).
func NewResponseFilteringWriter(
        w http.ResponseWriter, authorizer authorizers.Authorizer, r *http.Request, method string,
        annotationCache *AnnotationCache, passThroughTools map[string]struct{},
) *ResponseFilteringWriter {
        return &ResponseFilteringWriter{
                ResponseWriter:   w,
                authorizer:       authorizer,
                request:          r,
                method:           method,
                buffer:           &bytes.Buffer{},
                statusCode:       http.StatusOK,
                annotationCache:  annotationCache,
                passThroughTools: passThroughTools,
        }
}

// Write captures the response body for filtering
func (rfw *ResponseFilteringWriter) Write(data []byte) (int, error) {
        return rfw.buffer.Write(data)
}

// WriteHeader captures the status code
func (rfw *ResponseFilteringWriter) WriteHeader(statusCode int) {
        rfw.statusCode = statusCode
}

// FlushAndFilter processes the captured response and applies filtering if needed.
// Returns an error if filtering or writing fails.
func (rfw *ResponseFilteringWriter) FlushAndFilter() error {
        rfw.applyResponseCachePolicy()

        // Check if this response needs filtering
        if !requiresResponseFiltering(rfw.method) {
                rfw.ResponseWriter.WriteHeader(rfw.statusCode)
                _, err := rfw.ResponseWriter.Write(rfw.buffer.Bytes()) //nolint:gosec // G705 - JSON-RPC response, not rendered as HTML
                return err
        }

        rawResponse := rfw.buffer.Bytes()

        // A client strips a leading BOM per the WHATWG UTF-8 decode algorithm
        // before parsing, so strip it here too. Otherwise a BOM-prefixed list
        // response fails every decode and sniff below (Go's encoding/json treats
        // EF BB BF as a syntax error, not whitespace) and passes through
        // unfiltered. This mirrors the SSE-path fix in processSSEResponse, and
        // doing it once here covers the JSON decode path as well as both the
        // JSON and SSE smuggled-result sniffs in the default branch below.
        rawResponse = bytes.TrimPrefix(rawResponse, mcpparser.UTF8BOM)

        // Skip filtering for empty responses (common in SSE scenarios where actual data comes via SSE stream)
        if len(rawResponse) == 0 {
                rfw.ResponseWriter.WriteHeader(rfw.statusCode)
                _, err := rfw.ResponseWriter.Write(rawResponse) //nolint:gosec // G705 - JSON-RPC response, not rendered as HTML
                return err
        }

        // Media type names are case-insensitive and parameters may be preceded by
        // whitespace (RFC 9110 section 8.3.1), so normalize before matching.
        // Without this, "Application/JSON" or "application/json ; charset=utf-8"
        // -- both valid labels for a JSON body that clients happily parse -- would
        // fall through to the default branch and skip filtering entirely.
        contentType := rfw.ResponseWriter.Header().Get("Content-Type")
        mimeType := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0]))

        switch mimeType {
        case "application/json":
                // Remove the upstream Content-Length header. The reverse proxy copies it
                // from the backend response via Header() (which we don't override), but
                // filtering changes the body size. Without this, Go's HTTP server detects
                // the mismatch and tears down the connection.
                rfw.ResponseWriter.Header().Del("Content-Length")
                return rfw.processJSONResponse(rawResponse)
        case "text/event-stream":
                // Same issue: filtering changes the SSE payload size.
                rfw.ResponseWriter.Header().Del("Content-Length")
                return rfw.processSSEResponse(rawResponse)
        default:
                // A response to a method whose result must be filtered,
                // yet labeled with neither MCP-supported media type. That could be an
                // accident, or a backend deliberately mislabeling the response to
                // smuggle an unfiltered list past the filter -- the same
                // disguised-result concern as #5257, handled for the tool filter in
                // pkg/mcp/tool_filter.go's processUnrecognizedMimeType. Sniff the
                // body under both supported shapes and, when it carries a JSON-RPC
                // result, process it exactly as if it had been labeled correctly
                // (which filters clean frames and fails closed on dirty ones). Only a
                // body carrying no result under either shape passes through.
                if carriesResult(rawResponse) {
                        slog.Warn("response with unrecognized media type carries a JSON-RPC result; filtering as JSON",
                                "method", rfw.method, "contentType", contentType)
                        rfw.ResponseWriter.Header().Del("Content-Length")
                        return rfw.processJSONResponse(rawResponse)
                }
                if sseCarriesResult(rawResponse) {
                        slog.Warn("response with unrecognized media type carries a JSON-RPC result; filtering as SSE",
                                "method", rfw.method, "contentType", contentType)
                        rfw.ResponseWriter.Header().Del("Content-Length")
                        return rfw.processSSEResponse(rawResponse)
                }
                // The leading BOM was removed above, so the backend's Content-Length
                // is now three bytes too large even though this fallback doesn't filter
                // the body. Clear it before writing to avoid a truncated response.
                rfw.ResponseWriter.Header().Del("Content-Length")
                rfw.ResponseWriter.WriteHeader(rfw.statusCode)
                _, err := rfw.ResponseWriter.Write(rawResponse)
                return err
        }
}

// Flush implements http.Flusher if the underlying ResponseWriter supports it.
// This method is required for streaming support (SSE, streamable-http).
//
// We must delete the Content-Length header before flushing because
// httputil.ReverseProxy (with FlushInterval: -1) calls Flush() after copying
// the backend response. The first Flush() on the underlying writer triggers an
// implicit WriteHeader(200), sending headers to the wire. If the stale
// Content-Length is still present at that point, it's too late to remove it in
// FlushAndFilter().
//
// Commit the recorded status before the first flush. Without this, the
// implicit 200 would rewrite a non-2xx backend status (e.g. 500) on the wire.
// Committing here keeps the wire status identical to the recorded backend
// status; SSE (statusCode 200) is unaffected and later WriteHeader calls in
// FlushAndFilter become no-ops instead of corrupting the status.
func (rfw *ResponseFilteringWriter) Flush() {
        if flusher, ok := rfw.ResponseWriter.(http.Flusher); ok {
                rfw.applyResponseCachePolicy()
                rfw.ResponseWriter.Header().Del("Content-Length")
                rfw.ResponseWriter.WriteHeader(rfw.statusCode)
                flusher.Flush()
        }
}

// applyResponseCachePolicy prevents a caller-specific tools, prompts,
// resources, or resource-template view from being reused across authorization
// contexts by an HTTP cache. It applies to every filtered list, including a
// list result that carries no caching hints of its own. Flush calls this
// before an SSE response commits its headers; FlushAndFilter covers buffered
// JSON.
func (rfw *ResponseFilteringWriter) applyResponseCachePolicy() {
        switch responseFilterForMethod(rfw.method) {
        case responseFilterTools, responseFilterPrompts, responseFilterResources, responseFilterResourceTemplates:
                rfw.ResponseWriter.Header().Set("Cache-Control", "private, no-store")
        case responseFilterNone, responseFilterSkills, responseFilterFindTool:
                // These filters set no private caching hints in the result.
        }
}

func (rfw *ResponseFilteringWriter) processJSONResponse(rawResponse []byte) error {
        message, err := jsonrpc2.DecodeMessage(rawResponse)
        response, ok := message.(*jsonrpc2.Response)
        if err != nil || !ok {
                // Not a clean Response. The disguised-result bypass (#5257) is
                // transport-independent, so apply the same check here as on the SSE
                // path: if a non-Response, undecodable, or batch frame still carries a
                // result, fail closed rather than passing the smuggled list through.
                if carriesResult(rawResponse) {
                        slog.Warn("JSON response carried a result outside a clean Response frame; dropping as a protocol violation",
                                "method", rfw.method)
                        // A batch is replaced wholesale by this single error envelope
                        // even if only one of many elements carried the smuggled
                        // result: JSON-RPC does define batch responses, but we don't
                        // attempt a partial one here, and mcpparser.ParsingMiddleware
                        // already rejects batch requests outright (see #5745), so a
                        // batch response on this path is already off-spec traffic.
                        // Fail-closed is the right side to err on regardless.
                        return rfw.writeErrorResponse(rfw.requestID(),
                                errors.New("dropped a frame carrying a result outside a clean Response"))
                }
                rfw.ResponseWriter.WriteHeader(rfw.statusCode)
                _, werr := rfw.ResponseWriter.Write(rawResponse)
                return werr
        }

        filteredData, err := rfw.filterAndEncode(response)
        if err != nil {
                return rfw.writeErrorResponse(response.ID, err)
        }

        rfw.ResponseWriter.WriteHeader(rfw.statusCode)
        _, err = rfw.ResponseWriter.Write(filteredData)
        return err
}

// filterAndEncode filters a Response and encodes the result for the wire.
func (rfw *ResponseFilteringWriter) filterAndEncode(response *jsonrpc2.Response) ([]byte, error) {
        filteredResponse, err := rfw.filterListResponse(response)
        if err != nil {
                return nil, err
        }
        return jsonrpc2.EncodeMessage(filteredResponse)
}

// sseLine pairs one scanned SSE line with the terminator that followed it in
// the raw body: "\n", or nil for a final line with none. Lines are split on
// LF only (see processSSEResponse), so a line that was actually terminated by
// "\r\n" keeps its trailing "\r" as the last byte of text -- text+term always
// reproduces exactly the bytes that were scanned.
type sseLine struct {
        text []byte
        term []byte
}

// processSSEResponse rewrites a text/event-stream body per SSE framing rules
// rather than treating it as an arbitrary sequence of lines: within one
// event, all `data:` field values concatenate (joined by LF, per spec) into a
// single payload, and a blank line is structural — it dispatches the event
// rather than being emittable content in its own right.
//
// Lines are split on LF only, matching the client that actually consumes
// these streams (github.com/modelcontextprotocol/go-sdk's mcp/event.go,
// which ReadBytes('\n')s and never treats a lone CR as a line terminator) --
// not the WHATWG EventSource grammar, which additionally splits on a bare CR.
// The two disagree on a body containing an interior "\r" (e.g. inside a
// data: payload): WHATWG would cut the line there, but the client reads it as
// one line and parses through. Following the WHATWG grammar here decoded a
// SUPERSET of what the client would treat as one line, letting the shorter,
// wrongly-split pieces slip past filtering undetected -- a real regression,
// not a hypothetical one. Splitting on LF only means we parse a superset of
// what a stricter reader would call one line, so we decode and filter *more*
// than before, never less: worst case for a client that does honor a lone CR
// (e.g. a browser EventSource) is a payload it mis-frames identically to how
// it would mis-frame the backend's raw bytes directly, since output stays
// byte-preserving -- denial, never disclosure.
//
// Every line is re-emitted with the same terminator it was scanned with, so
// the output is structurally identical to the input (minus a leading UTF-8
// BOM, stripped below); only `data:` payloads are rewritten.
func (rfw *ResponseFilteringWriter) processSSEResponse(rawResponse []byte) error {
        // Note: this routine is adapted from the one in pkg/mcp/tool_filter.go.
        // I don't see an obvious way to factor out the commonalities, so I'm
        // duplicating it here, but we should refactor response parsing
        // respecting mime types to a common routine.
        // Commit the recorded status before writing the first event. On streaming
        // proxy paths Flush has already done this and WriteHeader is a no-op; on a
        // buffered non-2xx response this prevents the first body write from
        // implicitly changing the status to 200.
        rfw.ResponseWriter.WriteHeader(rfw.statusCode)

        // A client strips a leading BOM per the WHATWG UTF-8 decode algorithm
        // before parsing lines, so strip it here too: otherwise the first line's
        // "data:" prefix wouldn't match and the event would pass through
        // unfiltered.
        rawResponse = bytes.TrimPrefix(rawResponse, mcpparser.UTF8BOM)

        var outputLines []sseLine
        var event []sseLine
        for len(rawResponse) > 0 {
                idx := bytes.IndexByte(rawResponse, '\n')
                var line, term []byte
                if idx == -1 {
                        line, term, rawResponse = rawResponse, nil, nil
                } else {
                        line, term, rawResponse = rawResponse[:idx], rawResponse[idx:idx+1], rawResponse[idx+1:]
                }

                // A line is structurally blank once its trailing CR (left over from
                // a "\r\n" terminator the LF-only split doesn't consume) is ignored;
                // without this a CRLF-terminated blank line would scan as the
                // single byte "\r" and never be recognised as the event separator.
                // The blank line's own bytes (that "\r", if present) are preserved
                // in the output below rather than dropped, so structure stays
                // byte-identical to the input.
                if len(bytes.TrimRight(line, "\r")) == 0 {
                        outputLines = append(outputLines, rfw.resolveSSEEvent(event)...)
                        outputLines = append(outputLines, sseLine{text: line, term: term})
                        event = nil
                        continue
                }
                event = append(event, sseLine{text: line, term: term})
        }
        // A trailing event with no closing blank line: the go-sdk client (see
        // mcp/event.go's yieldEvent, called once more after its read loop hits
        // EOF) does still dispatch whatever is buffered at that point, so this
        // is filtered like any other event, not dropped. We write these bytes
        // (filtered) without fabricating a terminator the backend didn't send.
        outputLines = append(outputLines, rfw.resolveSSEEvent(event)...)

        for _, l := range outputLines {
                if _, err := rfw.ResponseWriter.Write(l.text); err != nil {
                        return fmt.Errorf("%w: %w", errBug, err)
                }
                if _, err := rfw.ResponseWriter.Write(l.term); err != nil {
                        return fmt.Errorf("%w: %w", errBug, err)
                }
        }

        return nil
}

// resolveSSEEvent assembles the `data:` payload of one SSE event (the lines
// between two blank lines, exclusive) and returns the lines to emit in its
// place. An event with no `data:` field (e.g. a bare "event: ping") or whose
// assembled payload passes filterSSEEventData unchanged is returned as-is.
func (rfw *ResponseFilteringWriter) resolveSSEEvent(event []sseLine) []sseLine {
        if len(event) == 0 {
                return nil
        }

        dataValues := eventDataValues(event)
        if len(dataValues) == 0 {
                return event
        }

        // No whole-buffer trim needed here: the per-value trim in eventDataValues
        // already leaves nothing but JSON-significant bytes at either edge of the join.
        assembled := bytes.Join(dataValues, []byte("\n"))
        if len(assembled) == 0 {
                // A client never dispatches an empty data buffer, so there's nothing
                // to classify or fail closed on.
                return event
        }

        replacement, failedClosed := rfw.filterSSEEventData(assembled)
        if replacement == nil {
                replacement, failedClosed = rfw.probeValuesForSmuggledResult(dataValues)
        }
        if replacement == nil {
                return event
        }
        return rebuildEventWithPayload(event, replacement, failedClosed)
}

// eventDataValues extracts one event's `data:` field values, in order, each
// trimmed exactly as the client trims it.
func eventDataValues(event []sseLine) [][]byte {
        var dataValues [][]byte
        for _, l := range event {
                // A line ending in "\r\n" keeps that "\r" as the trailing byte of
                // l.text (see sseLine), so trim it before matching the "data:"
                // prefix -- mirroring the client, which strips trailing "\r\n" from
                // every line before inspecting it (mcp/event.go).
                trimmed := bytes.TrimRight(l.text, "\r")
                // A bare "data" line with no colon, or one with a space before the
                // colon ("data :foo"), doesn't match the "data:" prefix below, so both
                // are excluded from assembly here, unlike strict WHATWG grammar (which
                // treats a bare "data" as a field with an empty value). That's safe
                // for two reasons: (1) the only effect on a spec-compliant client's
                // assembled buffer is one extra LF, which is JSON-insignificant
                // whitespace outside a string literal and can never turn JSON we'd
                // reject into JSON a client would accept; (2) for
                // "data :foo", bytes.Cut on the raw line yields before == "data " (not
                // "data"), so the go-sdk's own line parser ignores it too, exactly
                // like us and strict WHATWG.
                data, ok := bytes.CutPrefix(trimmed, []byte("data:"))
                if !ok {
                        continue
                }
                // Per the WHATWG EventSource grammar, exactly one leading space after
                // "data:" is part of the field delimiter, not the payload. Stripping
                // more (or a differently-positioned space) would corrupt a payload
                // split mid-string-literal across lines.
                data, _ = bytes.CutPrefix(data, []byte(" "))
                // TrimSpace each value independently, before joining, not on the
                // assembled buffer afterward. Go's JSON scanner treats only SP, TAB,
                // CR and LF as whitespace, but unicode.IsSpace also covers U+000B,
                // U+000C, U+0085, U+00A0, U+1680 and U+3000 -- and the go-sdk client
                // trims each data value with exactly this function before joining
                // (mcp/event.go), not the whole buffer once at the end. A
                // whole-buffer trim only catches one of those bytes sitting at the
                // very front or back of the assembled payload; one sitting at an
                // interior boundary between two data: lines of the SAME event
                // survived undetected, because it was never at either edge of the
                // joined buffer. Trimming here catches it regardless of position.
                data = bytes.TrimSpace(data)
                dataValues = append(dataValues, data)
        }
        return dataValues
}

// probeValuesForSmuggledResult is the fallback for an event whose assembled
// payload was undecodable as a whole. filterSSEEventData's carriesResult scan
// stops at the first undecodable JSON value in the joined payload, so a
// well-formed result-bearing value sitting right after a malformed one within
// the SAME event would otherwise ride through as "undecodable, nothing to
// filter" -- reopening the split-payload bypass this rewrite exists to close
// (#5257), just one line lower than before. Probing each value independently
// only ever turns a pass-through into a fail-closed error envelope, never the
// reverse, so it cannot itself reintroduce a bypass. A nil replacement means
// no value carried a result and the event genuinely needs no filtering.
func (rfw *ResponseFilteringWriter) probeValuesForSmuggledResult(dataValues [][]byte) (replacement []byte, failedClosed bool) {
        for _, v := range dataValues {
                if carriesResult(v) {
                        slog.Warn("SSE event's assembled payload was undecodable but an individual data value carries a result; failing closed",
                                "method", rfw.method)
                        return rfw.errorResponseBody(rfw.requestID(),
                                errors.New("dropped a frame carrying a result outside a clean Response")), true
                }
        }
        return nil, false
}

// rebuildEventWithPayload re-emits one event's lines with replacement standing
// in for the first `data:` line's value, dropping the later `data:` lines whose
// content is already folded into it.
func rebuildEventWithPayload(event []sseLine, replacement []byte, failedClosed bool) []sseLine {
        out := make([]sseLine, 0, len(event))
        replaced := false
        for _, l := range event {
                trimmed := bytes.TrimRight(l.text, "\r")
                // The go-sdk streamable client only dispatches unnamed ("message")
                // events, so a named event's own "event:" field would silently
                // swallow the fail-closed envelope substituted below, reproducing
                // the #6037 hang. Drop it: the substituted payload is either the
                // error envelope itself or a freshly filtered result, neither of
                // which the original event name describes any more.
                if failedClosed && bytes.HasPrefix(trimmed, []byte("event:")) {
                        continue
                }
                if !bytes.HasPrefix(trimmed, []byte("data:")) {
                        out = append(out, l)
                        continue
                }
                if replaced {
                        // Subsequent data: lines are already folded into the assembled
                        // payload; only the first line carries it.
                        continue
                }
                // A "\r\n"-terminated line keeps that "\r" as l.text's trailing
                // byte (see sseLine), not as part of l.term. Since the replacement
                // text discards l.text entirely, re-attach any such trailing "\r" to
                // the terminator here -- otherwise a CRLF-terminated data line would
                // silently downgrade to bare LF once filtered, while an untouched
                // (pass-through) line keeps its CRLF because it reuses l.text verbatim.
                term := l.term
                if cr := l.text[len(trimmed):]; len(cr) > 0 {
                        term = append(append([]byte{}, cr...), term...)
                }
                // jsonrpc2.EncodeMessage and json.Marshal both escape newlines, so
                // replacement can never contain a raw line separator and is always
                // safe to emit as a single data: line.
                out = append(out, sseLine{text: append([]byte("data: "), replacement...), term: term})
                replaced = true
        }
        return out
}

// filterSSEEventData classifies an event's assembled data payload and returns
// the payload to emit in its place, plus whether that payload is a fail-closed
// error envelope (as opposed to a successfully filtered result). A nil
// replacement means the event's lines pass through unchanged, and failedClosed
// is meaningless in that case.
func (rfw *ResponseFilteringWriter) filterSSEEventData(data []byte) (replacement []byte, failedClosed bool) {
        message, decodeErr := jsonrpc2.DecodeMessage(data)
        response, isResponse := message.(*jsonrpc2.Response)
        switch {
        case isResponse:
                filteredData, ferr := rfw.filterAndEncode(response)
                if ferr != nil {
                        // errorResponseBody logs ferr itself, so don't repeat it here.
                        slog.Warn("emitting a JSON-RPC error in place of an SSE frame that failed to filter",
                                "method", rfw.method)
                        return rfw.errorResponseBody(response.ID, ferr), true
                }
                return filteredData, false
        case carriesResult(data):
                // The frame is not a clean Response but still carries a result. This
                // covers a non-Response type (a request/notification frame smuggling a
                // result), a decode error (missing or invalid jsonrpc tag, a response
                // frame with no or a non-scalar id), and a batch array. All are
                // upstream-controlled shapes that would otherwise leak the very list
                // this filter scrubs (#5257). Fail closed with an explicit error
                // envelope rather than dropping the event silently: an event
                // dispatched with an empty data buffer isn't delivered to the
                // client's handler at all, leaving it hanging on its own timeout
                // (#6037).
                slog.Warn("SSE event carried a result outside a clean Response frame; failing closed as a protocol violation",
                        "method", rfw.method)
                return rfw.errorResponseBody(rfw.requestID(),
                        errors.New("dropped a frame carrying a result outside a clean Response")), true
        case decodeErr != nil:
                // Genuinely undecodable and no smuggled result. Pass through unfiltered.
                slog.Warn("SSE event data could not be decoded as JSON-RPC; passing through unfiltered",
                        "method", rfw.method, "error", decodeErr)
                return nil, false
        default:
                // Genuine non-Response frame (e.g. an interleaved notifications/*
                // message) with no result payload. Routine SSE traffic, so log at
                // Debug to keep the suspicious branches above from being buried.
                slog.Debug("SSE event data was not a JSON-RPC Response; passing through unfiltered",
                        "method", rfw.method)
                return nil, false
        }
}

type responseFilterKind uint8

const (
        responseFilterNone responseFilterKind = iota
        responseFilterTools
        responseFilterPrompts
        responseFilterResources
        responseFilterResourceTemplates
        responseFilterSkills
        responseFilterFindTool
)

// responseFilterForMethod is the authoritative mapping from an MCP method to
// its response filter. Keeping eligibility and dispatch behind this one mapping
// prevents a protected list method from being intercepted but passed through
// because it was added to only one of two method lists.
func responseFilterForMethod(method string) responseFilterKind {
        switch method {
        case string(mcp.MethodToolsList):
                return responseFilterTools
        case string(mcp.MethodPromptsList):
                return responseFilterPrompts
        case string(mcp.MethodResourcesList):
                return responseFilterResources
        case string(mcp.MethodResourcesTemplatesList):
                return responseFilterResourceTemplates
        case "skills/list":
                return responseFilterSkills
        case optimizerdec.FindToolName:
                return responseFilterFindTool
        default:
                return responseFilterNone
        }
}

// requiresResponseFiltering reports whether the method needs response filtering.
func requiresResponseFiltering(method string) bool {
        return responseFilterForMethod(method) != responseFilterNone
}

// carriesResult reports whether a data payload contains a JSON-RPC "result"
// field, either directly on an object or in any element of a batch array. A
// frame that is not a clean Response (a request/notification smuggling a
// result, a shape DecodeMessage rejects, or a batch array) but still carries a
// result must not pass through: it would leak the list the filter exists to
// scrub. See issue #5257.
//
// The payload may hold more than one top-level JSON value: an SSE event's
// assembled data can concatenate several messages (e.g. a notification and a
// response, one per data: line) into something DecodeMessage rejects as a
// single value. json.Decoder reads one JSON value at a time and stops at the
// first one it can't decode, rather than Unmarshal-ing the whole payload as
// one document, so a result-bearing value after the first is still caught
// instead of making the whole payload look undecodable. Stopping there
// (rather than trying to resync past the bad value) is safe: a payload the
// decoder can't parse value-by-value isn't a single parseable JSON document
// either, so no strict client can read past that point.
func carriesResult(data []byte) bool {
        dec := json.NewDecoder(bytes.NewReader(data))
        for {
                var value json.RawMessage
                if dec.Decode(&value) != nil {
                        return false // EOF, or malformed JSON with nothing left to check
                }
                if valueCarriesResult(value) {
                        return true
                }
        }
}

// valueCarriesResult applies carriesResult's check to a single JSON value.
func valueCarriesResult(value json.RawMessage) bool {
        trimmed := bytes.TrimSpace(value)
        if len(trimmed) == 0 {
                return false
        }
        switch trimmed[0] {
        case '{':
                return objectCarriesResult(trimmed)
        case '[':
                // A single level, not recursive: a JSON-RPC batch is a flat array of
                // message objects, so that's the whole legitimate surface. Recursing
                // through nested arrays is O(d^2) on encoding/json's own nesting cap
                // (10000), an easy amplification handle for something no client
                // actually flattens.
                var batch []json.RawMessage
                if json.Unmarshal(trimmed, &batch) != nil {
                        return false
                }
                for _, el := range batch {
                        elTrimmed := bytes.TrimSpace(el)
                        if len(elTrimmed) == 0 || elTrimmed[0] != '{' {
                                continue
                        }
                        if objectCarriesResult(elTrimmed) {
                                return true
                        }
                }
        }
        return false
}

// sseCarriesResult reports whether rawResponse contains an SSE "data:" line
// whose payload carries a JSON-RPC result. It is a lightweight detector, used
// only to decide whether a body with an unrecognized media type needs the
// full SSE processing path (which applies its own event-based filtering and
// fail-closed rules); it mirrors sniffSSEToolsList in pkg/mcp/tool_filter.go.
func sseCarriesResult(rawResponse []byte) bool {
        normalized := bytes.ReplaceAll(rawResponse, []byte("\r\n"), []byte("\n"))
        normalized = bytes.ReplaceAll(normalized, []byte("\r"), []byte("\n"))
        for _, line := range bytes.Split(normalized, []byte("\n")) {
                if data, ok := bytes.CutPrefix(line, []byte("data:")); ok && carriesResult(data) {
                        return true
                }
        }
        return false
}

// objectCarriesResult reports whether a trimmed JSON object value has a
// "result" key present, including an explicit `"result":null` — RawMessage's
// UnmarshalJSON runs even for a JSON null, so probe.Result is non-nil (the
// 4 bytes "null") whenever the key exists at all. Treating a null result as
// carrying one is deliberate: it's the fail-closed direction.
func objectCarriesResult(trimmedObject json.RawMessage) bool {
        var probe struct {
                Result json.RawMessage `json:"result"`
        }
        return json.Unmarshal(trimmedObject, &probe) == nil && probe.Result != nil
}

// filterListResponse filters the list response based on authorization policies
func (rfw *ResponseFilteringWriter) filterListResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) {
        if response.Error != nil {
                // A Response carrying both error and result is a protocol violation:
                // DecodeMessage populates both fields as given, and some clients
                // (e.g. the reference TS SDK's zod schema) silently strip the
                // unexpected "error" key and treat it as a successful result. Fail
                // closed rather than passing the smuggled list through under cover
                // of the error field. See #5257.
                // A literal `"result":null` is exempted: DecodeMessage sets Result to
                // the non-nil 4 bytes "null" whenever the key is present at all (see
                // objectCarriesResult), so without this a legitimate upstream error
                // that happens to carry an explicit null result would lose its own
                // code (e.g. -32601) to our generic internal-error envelope for no
                // security benefit -- a null result can never carry a list, so
                // failing closed on it buys nothing.
                if response.Result != nil && !bytes.Equal(bytes.TrimSpace(response.Result), []byte("null")) {
                        return nil, errors.New("response carried both error and result")
                }
                // If there's an error and no result (or an explicit null result), don't filter
                return response, nil
        }

        if response.Result == nil {
                // If there's no result, don't filter
                return response, nil
        }

        // Filter based on the method. responseFilterForMethod is shared with the
        // eligibility check in FlushAndFilter so these cases cannot drift apart.
        switch responseFilterForMethod(rfw.method) {
        case responseFilterTools:
                return rfw.filterToolsResponse(response)
        case responseFilterPrompts:
                return rfw.filterPromptsResponse(response)
        case responseFilterResources:
                return rfw.filterResourcesResponse(response)
        case responseFilterResourceTemplates:
                return rfw.filterResourceTemplatesResponse(response)
        case responseFilterSkills:
                return rfw.filterSkillsResponse(response)
        case responseFilterFindTool:
                return rfw.filterFindToolResponse(response)
        case responseFilterNone:
                return nil, fmt.Errorf("no response filter for method %q", rfw.method)
        }
        return nil, fmt.Errorf("unknown response filter for method %q", rfw.method)
}

// filterToolsResponse filters tools based on call_tool authorization
func (rfw *ResponseFilteringWriter) filterToolsResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) {
        members, rawTools, err := decodeFilterableListResult(response.Result, "tools", "name")
        if err != nil {
                return nil, fmt.Errorf("validating tools list response: %w", err)
        }
        if err := validateToolAnnotations(rawTools); err != nil {
                return nil, fmt.Errorf("validating tools list response: %w", err)
        }

        // The typed descriptors feed policy evaluation and the annotation cache;
        // the filtered result carries the backend's raw descriptors.
        var listResult mcp.ListToolsResult
        if err := json.Unmarshal(response.Result, &listResult); err != nil {
                return nil, fmt.Errorf("decoding tools list response: %w", err)
        }
        if len(listResult.Tools) != len(rawTools) {
                return nil, errors.New("decoded tools do not match the raw tools list")
        }

        // Populate annotation cache from tools/list response so that
        // subsequent tools/call requests can look up annotations.
        rfw.annotationCache.SetFromToolsList(listResult.Tools)

        // When the optimizer is enabled, its meta-tools (find_tool, call_tool) appear
        // in tools/list instead of real backend tools. These meta-tools won't match
        // any operator-written Cedar policy (which references real tool names), so
        // default-deny would filter them out — leaving the client with zero tools.
        // Authorization for the underlying backend tools is enforced by the authz
        // middleware: call_tool requests are intercepted and the inner tool_name
        // argument is authorized against Cedar policy before the request is served.
        // See: https://github.com/stacklok/toolhive/issues/4373
        permitted := make([]json.RawMessage, 0, len(rawTools))
        regular := []mcp.Tool{}
        regularIndexes := []int{}
        for i, t := range listResult.Tools {
                if _, ok := rfw.passThroughTools[t.Name]; ok {
                        permitted = append(permitted, rawTools[i])
                } else {
                        regular = append(regular, t)
                        regularIndexes = append(regularIndexes, i)
                }
        }

        // authorizedToolIndexes checks each tool against the caller's Cedar policies
        // (injecting annotations into context for when-clause evaluation) and returns
        // the positions of the tools the caller is authorized to call.
        for _, index := range authorizedToolIndexes(rfw.request.Context(), rfw.authorizer, regular) {
                permitted = append(permitted, rawTools[regularIndexes[index]])
        }

        return filteredListResponse(response.ID, members, "tools", permitted)
}

// filterPromptsResponse filters prompts based on get_prompt authorization
func (rfw *ResponseFilteringWriter) filterPromptsResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) {
        members, rawPrompts, err := decodeFilterableListResult(response.Result, "prompts", "name")
        if err != nil {
                return nil, fmt.Errorf("validating prompts list response: %w", err)
        }

        // The typed descriptors are validated and authorized; the filtered result
        // carries the backend's raw descriptors.
        var listResult mcp.ListPromptsResult
        if err := json.Unmarshal(response.Result, &listResult); err != nil {
                return nil, fmt.Errorf("decoding prompts list response: %w", err)
        }
        if len(listResult.Prompts) != len(rawPrompts) {
                return nil, errors.New("decoded prompts do not match the raw prompts list")
        }

        permitted := make([]json.RawMessage, 0, len(rawPrompts))
        for i, prompt := range listResult.Prompts {
                // Check if the user is authorized to get this prompt
                authorized, err := rfw.authorizer.AuthorizeWithJWTClaims(
                        rfw.request.Context(),
                        authorizers.MCPFeaturePrompt,
                        authorizers.MCPOperationGet,
                        prompt.Name,
                        nil, // No arguments for the authorization check
                )
                if err != nil {
                        slog.Warn("Authorization check failed for prompt, skipping",
                                "prompt", prompt.Name, "error", err)
                        continue
                }

                if authorized {
                        permitted = append(permitted, rawPrompts[i])
                } else {
                        slog.Debug("Prompt denied by authorization policy",
                                "prompt", prompt.Name)
                }
        }

        if denied := len(listResult.Prompts) - len(permitted); denied > 0 {
                slog.Debug("Authorization policy filtered prompts",
                        "total", len(listResult.Prompts), "allowed", len(permitted), "denied", denied)
        }

        return filteredListResponse(response.ID, members, "prompts", permitted)
}

// filterResourcesResponse filters resources based on read_resource authorization
func (rfw *ResponseFilteringWriter) filterResourcesResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) {
        members, rawResources, err := decodeFilterableListResult(response.Result, "resources", "uri")
        if err != nil {
                return nil, fmt.Errorf("validating resources list response: %w", err)
        }

        // The typed descriptors are validated and authorized; the filtered result
        // carries the backend's raw descriptors.
        var listResult mcp.ListResourcesResult
        if err := json.Unmarshal(response.Result, &listResult); err != nil {
                return nil, fmt.Errorf("decoding resources list response: %w", err)
        }
        if len(listResult.Resources) != len(rawResources) {
                return nil, errors.New("decoded resources do not match the raw resources list")
        }

        permitted := make([]json.RawMessage, 0, len(rawResources))
        for i, resource := range listResult.Resources {
                // Check if the user is authorized to read this resource
                authorized, err := rfw.authorizer.AuthorizeWithJWTClaims(
                        rfw.request.Context(),
                        authorizers.MCPFeatureResource,
                        authorizers.MCPOperationRead,
                        resource.URI,
                        nil, // No arguments for the authorization check
                )
                if err != nil {
                        slog.Warn("Authorization check failed for resource, skipping",
                                "resource", resource.URI, "error", err)
                        continue
                }

                if authorized {
                        permitted = append(permitted, rawResources[i])
                } else {
                        slog.Debug("Resource denied by authorization policy",
                                "resource", resource.URI)
                }
        }

        if denied := len(listResult.Resources) - len(permitted); denied > 0 {
                slog.Debug("Authorization policy filtered resources",
                        "total", len(listResult.Resources), "allowed", len(permitted), "denied", denied)
        }

        return filteredListResponse(response.ID, members, "resources", permitted)
}

// cachingHintMembers are the MCP result members that tell a client or a shared
// cache whether, and for how long, a list result may be reused.
var cachingHintMembers = []string{"cacheScope", "ttlMs"}

// decodeFilterableListResult decodes a tools, prompts, or resources list
// result for filtering. On top of decodeListResult's checks it rejects
// ambiguous spellings of the caching hints, because filteredListResponse
// overrides them and a client must not be able to read the backend's value
// under an alias instead.
func decodeFilterableListResult(
        data json.RawMessage,
        listField, identifierField string,
) ([]jsonObjectMember, []json.RawMessage, error) {
        members, items, err := decodeListResult(data, listField, identifierField)
        if err != nil {
                return nil, nil, err
        }
        for _, hint := range cachingHintMembers {
                if _, _, err := uniqueCanonicalMember(members, hint); err != nil {
                        return nil, nil, err
                }
        }
        return members, items, nil
}

// toolAnnotationMembers are the annotation members that encoding/json binds
// to the typed tool annotations used for policy evaluation and the annotation
// cache.
var toolAnnotationMembers = jsonMemberNames(reflect.TypeFor[mcp.ToolAnnotation]())

// validateToolAnnotations rejects ambiguous spellings of a tool's annotations
// member and of every member the typed annotations decode. encoding/json
// resolves case-folded aliases and duplicate members, and merges a repeated
// annotations object, so policies could otherwise evaluate hints that differ
// from those in the raw descriptor a filter forwards to the client.
func validateToolAnnotations(rawTools []json.RawMessage) error {
        for i, rawTool := range rawTools {
                toolMembers, err := decodeJSONObjectMembers(rawTool)
                if err != nil {
                        return fmt.Errorf("decoding tools item at index %d: %w", i, err)
                }
                rawAnnotations, ok, err := uniqueCanonicalMember(toolMembers, "annotations")
                if err != nil {
                        return fmt.Errorf("tools item at index %d: %w", i, err)
                }
                if !ok || isJSONNull(rawAnnotations) {
                        continue
                }
                annotationMembers, err := decodeJSONObjectMembers(rawAnnotations)
                if err != nil {
                        return fmt.Errorf("decoding annotations of tools item at index %d: %w", i, err)
                }
                for _, name := range toolAnnotationMembers {
                        if _, _, err := uniqueCanonicalMember(annotationMembers, name); err != nil {
                                return fmt.Errorf("annotations of tools item at index %d: %w", i, err)
                        }
                }
        }
        return nil
}

// filteredListResponse builds the response for a filtered tools, prompts, or
// resources list. The list member is rewritten to the permitted raw
// descriptors. The permitted list depends on the caller's authorization, so a
// cacheScope or ttlMs member the backend sent is replaced with "private" or 0,
// marking the result private and immediately stale in place of the backend's
// caching hints for its unfiltered list. A caching hint the backend omitted
// stays absent, so a result from a backend that predates these members keeps
// its shape. Every other result member (resultType, _meta, nextCursor, and any
// extension) is kept exactly as the backend sent it.
func filteredListResponse(
        id jsonrpc2.ID,
        members []jsonObjectMember,
        listField string,
        permitted []json.RawMessage,
) (*jsonrpc2.Response, error) {
        replacements := map[string]json.RawMessage{
                listField:    encodeJSONArrayValues(permitted),
                "cacheScope": json.RawMessage(`"private"`),
                "ttlMs":      json.RawMessage(`0`),
        }

        filteredResult, err := encodeJSONObjectMembers(members, replacements)
        if err != nil {
                return nil, err
        }
        return &jsonrpc2.Response{
                ID:     id,
                Result: json.RawMessage(filteredResult),
        }, nil
}

// filterResourceTemplatesResponse filters resource templates based on
// read_resource authorization. A template's RFC 6570 URI template is treated
// as the resource identifier, matching the admission semantics used by vMCP.
func (rfw *ResponseFilteringWriter) filterResourceTemplatesResponse(
        response *jsonrpc2.Response,
) (*jsonrpc2.Response, error) {
        result, resourceTemplates, uriTemplates, err := decodeResourceTemplatesListResult(response.Result)
        if err != nil {
                return nil, err
        }

        filteredTemplates := make([]json.RawMessage, 0, len(resourceTemplates))
        for i, resourceTemplate := range resourceTemplates {
                authorized, err := rfw.authorizer.AuthorizeWithJWTClaims(
                        rfw.request.Context(),
                        authorizers.MCPFeatureResource,
                        authorizers.MCPOperationRead,
                        uriTemplates[i],
                        nil,
                )
                if err != nil {
                        slog.Warn("Authorization check failed for resource template, skipping",
                                "resourceTemplate", uriTemplates[i], "error", err)
                        continue
                }

                if authorized {
                        filteredTemplates = append(filteredTemplates, resourceTemplate)
                } else {
                        slog.Debug("Resource template denied by authorization policy",
                                "resourceTemplate", uriTemplates[i])
                }
        }

        if denied := len(resourceTemplates) - len(filteredTemplates); denied > 0 {
                slog.Debug("Authorization policy filtered resource templates",
                        "total", len(resourceTemplates), "allowed", len(filteredTemplates), "denied", denied)
        }

        filteredTemplatesData, err := json.Marshal(filteredTemplates)
        if err != nil {
                return nil, err
        }
        result["resourceTemplates"] = filteredTemplatesData
        result["cacheScope"] = json.RawMessage(`"private"`)
        result["ttlMs"] = json.RawMessage(`0`)

        filteredResultData, err := json.Marshal(result)
        if err != nil {
                return nil, err
        }

        return &jsonrpc2.Response{
                ID:     response.ID,
                Result: json.RawMessage(filteredResultData),
        }, nil
}

// decodeResourceTemplatesListResult validates security-sensitive list and
// descriptor members before authorization while retaining raw descriptors and
// result extensions for the filtered response.
func decodeResourceTemplatesListResult(
        data json.RawMessage,
) (map[string]json.RawMessage, []json.RawMessage, []string, error) {
        members, err := decodeJSONObjectMembers(data)
        if err != nil {
                return nil, nil, nil, fmt.Errorf("decoding resource templates list response: %w", err)
        }
        result := make(map[string]json.RawMessage, len(members)+2)
        for _, member := range members {
                result[member.name] = member.value
        }

        rawTemplates, ok, err := uniqueCanonicalMember(members, "resourceTemplates")
        if err != nil {
                return nil, nil, nil, err
        }
        if !ok {
                return nil, nil, nil, errors.New("resource templates list result is missing resourceTemplates")
        }
        if _, _, err := uniqueCanonicalMember(members, "cacheScope"); err != nil {
                return nil, nil, nil, err
        }
        if _, _, err := uniqueCanonicalMember(members, "ttlMs"); err != nil {
                return nil, nil, nil, err
        }

        var resourceTemplates []json.RawMessage
        if err := json.Unmarshal(rawTemplates, &resourceTemplates); err != nil {
                return nil, nil, nil, fmt.Errorf("decoding resourceTemplates: %w", err)
        }
        if resourceTemplates == nil {
                return nil, nil, nil, errors.New("resourceTemplates must be an array")
        }
        uriTemplates, err := decodeResourceTemplateURIs(resourceTemplates)
        if err != nil {
                return nil, nil, nil, err
        }
        return result, resourceTemplates, uriTemplates, nil
}

// validateListResult rejects malformed list results and ambiguous spellings
// of the fields used to decide authorization. encoding/json otherwise accepts
// case-folded aliases and resolves duplicate members using the last value,
// which could make authorization inspect a different list or identifier than
// a client consuming the response.
func validateListResult(data json.RawMessage, listField, identifierField string) error {
        _, _, err := decodeListResult(data, listField, identifierField)
        return err
}

// decodeListResult applies validateListResult's checks and returns the
// result's members in their original order along with the raw list items, so
// a filter can drop items without re-encoding anything else in the result.
func decodeListResult(
        data json.RawMessage,
        listField, identifierField string,
) ([]jsonObjectMember, []json.RawMessage, error) {
        members, err := decodeJSONObjectMembers(data)
        if err != nil {
                return nil, nil, err
        }
        rawItems, ok, err := uniqueCanonicalMember(members, listField)
        if err != nil {
                return nil, nil, err
        }
        if !ok {
                return nil, nil, fmt.Errorf("result is missing %q", listField)
        }

        var items []json.RawMessage
        if err := json.Unmarshal(rawItems, &items); err != nil {
                return nil, nil, fmt.Errorf("decoding %s: %w", listField, err)
        }
        if items == nil {
                return nil, nil, fmt.Errorf("field %q must be an array", listField)
        }
        for i, item := range items {
                itemMembers, err := decodeJSONObjectMembers(item)
                if err != nil {
                        return nil, nil, fmt.Errorf("decoding %s item at index %d: %w", listField, i, err)
                }
                rawIdentifier, ok, err := uniqueCanonicalMember(itemMembers, identifierField)
                if err != nil {
                        return nil, nil, fmt.Errorf("%s item at index %d: %w", listField, i, err)
                }
                if !ok {
                        return nil, nil, fmt.Errorf("%s item at index %d is missing %q", listField, i, identifierField)
                }
                var identifier *string
                if err := json.Unmarshal(rawIdentifier, &identifier); err != nil {
                        return nil, nil, fmt.Errorf("%s item at index %d has an invalid %s: %w", listField, i, identifierField, err)
                }
                if identifier == nil {
                        return nil, nil, fmt.Errorf("%s item at index %d is missing a string %s", listField, i, identifierField)
                }
        }
        return members, items, nil
}

// decodeResourceTemplateURIs validates every descriptor before any
// authorization calls. The caller separately retains the original RawMessages
// so standard fields and backend extensions survive filtering unchanged.
func decodeResourceTemplateURIs(resourceTemplates []json.RawMessage) ([]string, error) {
        uriTemplates := make([]string, len(resourceTemplates))
        for i, rawTemplate := range resourceTemplates {
                descriptorMembers, err := decodeJSONObjectMembers(rawTemplate)
                if err != nil {
                        return nil, fmt.Errorf("decoding resource template at index %d: %w", i, err)
                }
                rawURITemplate, ok, err := uniqueCanonicalMember(descriptorMembers, "uriTemplate")
                if err != nil {
                        return nil, fmt.Errorf("resource template at index %d: %w", i, err)
                }
                if !ok {
                        return nil, fmt.Errorf("resource template at index %d is missing a string uriTemplate", i)
                }
                var uriTemplate *string
                if err := json.Unmarshal(rawURITemplate, &uriTemplate); err != nil {
                        return nil, fmt.Errorf("resource template at index %d has an invalid uriTemplate: %w", i, err)
                }
                if uriTemplate == nil {
                        return nil, fmt.Errorf("resource template at index %d is missing a string uriTemplate", i)
                }
                uriTemplates[i] = *uriTemplate
        }
        return uriTemplates, nil
}

type jsonObjectMember struct {
        name  string
        value json.RawMessage
}

// decodeJSONObjectMembers retains object member order and duplicates so
// security-sensitive keys can be validated before encoding/json's usual
// case-insensitive matching or last-value-wins behavior can hide them.
func decodeJSONObjectMembers(data []byte) ([]jsonObjectMember, error) {
        decoder := json.NewDecoder(bytes.NewReader(data))
        opening, err := decoder.Token()
        if err != nil {
                return nil, err
        }
        if delimiter, ok := opening.(json.Delim); !ok || delimiter != '{' {
                return nil, errors.New("value must be an object")
        }

        members := make([]jsonObjectMember, 0)
        for decoder.More() {
                nameToken, err := decoder.Token()
                if err != nil {
                        return nil, err
                }
                name, ok := nameToken.(string)
                if !ok {
                        return nil, errors.New("object member name must be a string")
                }
                var value json.RawMessage
                if err := decoder.Decode(&value); err != nil {
                        return nil, err
                }
                members = append(members, jsonObjectMember{name: name, value: value})
        }

        closing, err := decoder.Token()
        if err != nil {
                return nil, err
        }
        if delimiter, ok := closing.(json.Delim); !ok || delimiter != '}' {
                return nil, errors.New("unterminated object")
        }
        var trailing json.RawMessage
        if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
                if err == nil {
                        return nil, errors.New("unexpected data after object")
                }
                return nil, err
        }
        return members, nil
}

// uniqueCanonicalMember returns the value of canonical when it occurs exactly
// once. Case-folded aliases are rejected because common Go JSON decoders may
// treat them as the canonical field and select a different value than the one
// used for authorization or filtering.
func uniqueCanonicalMember(
        members []jsonObjectMember,
        canonical string,
) (json.RawMessage, bool, error) {
        var value json.RawMessage
        found := false
        for _, member := range members {
                if !strings.EqualFold(member.name, canonical) {
                        continue
                }
                if member.name != canonical {
                        return nil, false, fmt.Errorf("field %q has non-canonical alias %q", canonical, member.name)
                }
                if found {
                        return nil, false, fmt.Errorf("field %q occurs more than once", canonical)
                }
                value = member.value
                found = true
        }
        return value, found, nil
}

// jsonMemberNames returns the object member names that encoding/json matches
// to the fields of the struct type t: each field's json tag name, else its Go
// field name, skipping unexported fields and fields tagged "-", with the fields
// of untagged embedded structs promoted. It does not apply encoding/json's
// fallback for invalid tag names or its precedence rules for conflicting
// names, so it suits only struct types that use neither, such as
// mcp.ToolAnnotation.
func jsonMemberNames(t reflect.Type) []string {
        var names []string
        for i := range t.NumField() {
                field := t.Field(i)
                tag := field.Tag.Get("json")
                if tag == "-" {
                        continue
                }
                name, _, _ := strings.Cut(tag, ",")
                fieldType := field.Type
                if fieldType.Kind() == reflect.Pointer {
                        fieldType = fieldType.Elem()
                }
                if field.Anonymous && name == "" && fieldType.Kind() == reflect.Struct {
                        names = append(names, jsonMemberNames(fieldType)...)
                        continue
                }
                if !field.IsExported() {
                        continue
                }
                if name == "" {
                        name = field.Name
                }
                names = append(names, name)
        }
        return names
}

// filterSkillsResponse filters skills/list entries by get_skill authorization.
// It retains each permitted entry as its original JSON value and preserves all
// result-level fields, including SEP extension fields the proxy does not own.
func (rfw *ResponseFilteringWriter) filterSkillsResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) {
        if err := validateListResult(response.Result, "skills", "uri"); err != nil {
                return nil, fmt.Errorf("validating skills list response: %w", err)
        }

        var result map[string]json.RawMessage
        if err := json.Unmarshal(response.Result, &result); err != nil {
                return nil, fmt.Errorf("decoding skills list response: %w", err)
        }
        var entries []json.RawMessage
        if err := json.Unmarshal(result["skills"], &entries); err != nil {
                return nil, fmt.Errorf("decoding skills: %w", err)
        }

        permitted := make([]json.RawMessage, 0, len(entries))
        for _, entry := range entries {
                var skill struct {
                        URI string `json:"uri"`
                }
                if err := json.Unmarshal(entry, &skill); err != nil {
                        return nil, fmt.Errorf("decoding skill entry: %w", err)
                }
                if skill.URI == "" {
                        return nil, errors.New("skill entry has an empty uri")
                }
                authorized, err := rfw.authorizer.AuthorizeWithJWTClaims(
                        rfw.request.Context(), authorizers.MCPFeatureSkill, authorizers.MCPOperationGet, skill.URI, nil,
                )
                if err != nil {
                        slog.Warn("authorization check failed for skill, skipping", "uri", skill.URI, "error", err)
                        continue
                }
                if authorized {
                        permitted = append(permitted, entry)
                }
        }

        filteredSkills, err := json.Marshal(permitted)
        if err != nil {
                return nil, fmt.Errorf("encoding filtered skills: %w", err)
        }
        result["skills"] = filteredSkills
        filteredResult, err := json.Marshal(result)
        if err != nil {
                return nil, fmt.Errorf("encoding filtered skills list result: %w", err)
        }
        return &jsonrpc2.Response{ID: response.ID, Result: json.RawMessage(filteredResult)}, nil
}

// errorResponseBody logs the full filtering error server-side and encodes a
// JSON-RPC error response carrying a deliberately generic client-visible
// message: err can originate in policy evaluation and name tools or resources,
// so security.md forbids returning it verbatim (#6066). Both transports build
// their body here, so neither can leak it. id is echoed so the client can
// correlate the error with its request.
func (rfw *ResponseFilteringWriter) errorResponseBody(id jsonrpc2.ID, err error) []byte {
        slog.Error("error filtering response", "method", rfw.method, "error", err)
        errorResponse := &jsonrpc2.Response{
                ID:    id,
                Error: jsonrpc2.NewError(mcpparser.CodeInternalError, "internal error"),
        }

        body, encErr := jsonrpc2.EncodeMessage(errorResponse)
        if encErr != nil {
                // Unreachable in practice: errorResponse is always a well-formed
                // jsonrpc2.Response built from a valid ID and message above. Fall back
                // to a hardcoded valid JSON-RPC error body rather than writing nothing.
                slog.Error("failed to encode JSON-RPC error response", "error", encErr)
                return fmt.Appendf(nil, `{"jsonrpc":"2.0","error":{"code":%d,"message":"internal error"}}`,
                        mcpparser.CodeInternalError)
        }
        return body
}

// writeErrorResponse writes an error response for the application/json path.
// The SSE path must NOT use this: WriteHeader is a no-op once Flush() has
// committed the headers, which is the case for any real SSE response. On the
// SSE path a filtering failure instead becomes the replacement payload of the
// event's data: line (see filterSSEEventData), written through
// processSSEResponse's normal line writer alongside every other line — no
// separate write path or WriteHeader call is needed there.
func (rfw *ResponseFilteringWriter) writeErrorResponse(id jsonrpc2.ID, err error) error {
        rfw.ResponseWriter.WriteHeader(http.StatusInternalServerError)
        _, writeErr := rfw.ResponseWriter.Write(rfw.errorResponseBody(id, err))
        return writeErr
}

// requestID recovers the JSON-RPC id of the original request so an error
// response written mid-filtering can still be correlated by the client.
// Falls back to an empty jsonrpc2.ID (which EncodeMessage omits from the
// wire entirely) if the request was never parsed or its id doesn't convert.
func (rfw *ResponseFilteringWriter) requestID() jsonrpc2.ID {
        parsed := mcpparser.GetParsedMCPRequest(rfw.request.Context())
        if parsed == nil {
                return jsonrpc2.ID{}
        }
        id, err := mcpparser.ConvertToJSONRPC2ID(parsed.ID)
        if err != nil {
                return jsonrpc2.ID{}
        }
        return id
}

// filterFindToolResponse filters the tools list embedded in a find_tool tools/call
// response. Successful output must have exactly one text content item containing a
// JSON-encoded optimizer.FindToolOutput. Modern responses may repeat that output in
// structuredContent; when present, both representations must agree.
func (rfw *ResponseFilteringWriter) filterFindToolResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) {
        callResult, err := decodeFindToolCallResult(response.Result)
        if err != nil {
                return nil, fmt.Errorf("validating find_tool response: %w", err)
        }

        if callResult.isError {
                if len(callResult.outputs) != 0 || callResult.structuredContent != nil {
                        return nil, errors.New("find_tool error result carried output")
                }
                return response, nil
        }
        if len(callResult.content) != 1 {
                return nil, fmt.Errorf("find_tool result must contain exactly one content item, got %d", len(callResult.content))
        }
        if len(callResult.outputs) != 1 {
                return nil, fmt.Errorf("find_tool result must contain exactly one text output, got %d", len(callResult.outputs))
        }

        textOutput := callResult.outputs[0]
        if callResult.structuredContent != nil {
                equivalent, err := equivalentJSONValues(textOutput.rawOutput, callResult.structuredContent.rawOutput)
                if err != nil {
                        return nil, fmt.Errorf("comparing find_tool output representations: %w", err)
                }
                if !equivalent {
                        return nil, errors.New("find_tool text and structured outputs differ")
                }
        }

        // Populate annotation cache before filtering, mirroring filterToolsResponse.
        // Subsequent call_tool requests use these annotations for Cedar when-clause evaluation
        // (e.g. resource.readOnlyHint). The cache is populated from the unfiltered list so
        // that annotations are available even for tools that Cedar will deny.
        rfw.annotationCache.SetFromToolsList(textOutput.output.Tools)

        authorizedIndexes := authorizedToolIndexes(rfw.request.Context(), rfw.authorizer, textOutput.output.Tools)

        filteredResult, err := callResult.withFilteredToolIndexes(authorizedIndexes)
        if err != nil {
                return nil, fmt.Errorf("re-encoding call result: %w", err)
        }

        return &jsonrpc2.Response{
                ID:     response.ID,
                Result: json.RawMessage(filteredResult),
        }, nil
}

type findToolOutputCarrier struct {
        contentIndex   int
        contentMembers []jsonObjectMember
        rawOutput      json.RawMessage
        outputMembers  []jsonObjectMember
        rawTools       []json.RawMessage
        output         optimizer.FindToolOutput
}

type decodedFindToolCallResult struct {
        members           []jsonObjectMember
        content           []json.RawMessage
        isError           bool
        outputs           []findToolOutputCarrier
        structuredContent *findToolOutputCarrier
}

type findToolCallResultEnvelope struct {
        members       []jsonObjectMember
        content       []json.RawMessage
        isError       bool
        rawStructured json.RawMessage
        hasStructured bool
}

// decodeFindToolCallResult validates every output representation before any
// authorization or cache mutation while retaining raw metadata and extensions.
func decodeFindToolCallResult(data json.RawMessage) (*decodedFindToolCallResult, error) {
        envelope, err := decodeFindToolCallResultEnvelope(data)
        if err != nil {
                return nil, err
        }

        decoded := &decodedFindToolCallResult{
                members: envelope.members,
                content: envelope.content,
                isError: envelope.isError,
                outputs: make([]findToolOutputCarrier, 0, 1),
        }
        for i, item := range envelope.content {
                carrier, ok, err := decodeFindToolTextContent(i, item)
                if err != nil {
                        return nil, err
                }
                if ok {
                        decoded.outputs = append(decoded.outputs, *carrier)
                }
        }

        if envelope.hasStructured {
                carrier, ok, err := decodeFindToolOutput(envelope.rawStructured)
                if err != nil {
                        return nil, fmt.Errorf("decoding find_tool structuredContent: %w", err)
                }
                if !ok {
                        if !envelope.isError {
                                return nil, errors.New("find_tool structuredContent is missing tools")
                        }
                        return decoded, nil
                }
                decoded.structuredContent = carrier
        }

        return decoded, nil
}

func decodeFindToolCallResultEnvelope(data json.RawMessage) (*findToolCallResultEnvelope, error) {
        members, err := decodeJSONObjectMembers(data)
        if err != nil {
                return nil, err
        }

        rawContent, ok, err := uniqueCanonicalMember(members, "content")
        if err != nil {
                return nil, err
        }
        if !ok {
                return nil, errors.New("find_tool result is missing content")
        }
        rawIsError, hasIsError, err := uniqueCanonicalMember(members, "isError")
        if err != nil {
                return nil, err
        }
        rawStructured, hasStructured, err := uniqueCanonicalMember(members, "structuredContent")
        if err != nil {
                return nil, err
        }

        var content []json.RawMessage
        if err := json.Unmarshal(rawContent, &content); err != nil {
                return nil, fmt.Errorf("decoding find_tool content: %w", err)
        }
        if content == nil {
                return nil, errors.New("find_tool content must be an array")
        }

        isError := false
        if hasIsError {
                var value *bool
                if err := json.Unmarshal(rawIsError, &value); err != nil {
                        return nil, fmt.Errorf("decoding find_tool isError: %w", err)
                }
                if value == nil {
                        return nil, errors.New("find_tool isError must be a boolean")
                }
                isError = *value
        }

        return &findToolCallResultEnvelope{
                members:       members,
                content:       content,
                isError:       isError,
                rawStructured: rawStructured,
                hasStructured: hasStructured,
        }, nil
}

func decodeFindToolTextContent(index int, data json.RawMessage) (*findToolOutputCarrier, bool, error) {
        members, err := decodeJSONObjectMembers(data)
        if err != nil {
                return nil, false, fmt.Errorf("decoding find_tool content at index %d: %w", index, err)
        }
        rawType, ok, err := uniqueCanonicalMember(members, "type")
        if err != nil {
                return nil, false, fmt.Errorf("find_tool content at index %d: %w", index, err)
        }
        if !ok {
                return nil, false, fmt.Errorf("find_tool content at index %d is missing type", index)
        }
        var contentType *string
        if err := json.Unmarshal(rawType, &contentType); err != nil || contentType == nil {
                return nil, false, fmt.Errorf("find_tool content at index %d has an invalid type", index)
        }
        if *contentType != "text" {
                return nil, false, nil
        }

        rawText, ok, err := uniqueCanonicalMember(members, "text")
        if err != nil {
                return nil, false, fmt.Errorf("find_tool text content at index %d: %w", index, err)
        }
        if !ok {
                return nil, false, fmt.Errorf("find_tool text content at index %d is missing text", index)
        }
        var text *string
        if err := json.Unmarshal(rawText, &text); err != nil || text == nil {
                return nil, false, fmt.Errorf("find_tool text content at index %d has invalid text", index)
        }

        carrier, isOutput, err := decodeFindToolOutput(json.RawMessage(*text))
        if err != nil {
                return nil, false, fmt.Errorf("decoding find_tool text content at index %d: %w", index, err)
        }
        if !isOutput {
                return nil, false, nil
        }
        carrier.contentIndex = index
        carrier.contentMembers = members
        return carrier, true, nil
}

// decodeFindToolOutput returns ok=false for ancillary text that is not a JSON
// object or does not contain a tools member. A tools-bearing object is always
// validated strictly and decoded into the concrete output type.
func decodeFindToolOutput(data json.RawMessage) (*findToolOutputCarrier, bool, error) {
        members, err := decodeJSONObjectMembers(data)
        if err != nil {
                if trimmed := bytes.TrimSpace(data); len(trimmed) != 0 && trimmed[0] == '{' {
                        return nil, false, fmt.Errorf("decoding tools-bearing object: %w", err)
                }
                return nil, false, nil
        }
        rawToolsValue, ok, err := uniqueCanonicalMember(members, "tools")
        if err != nil {
                return nil, false, err
        }
        if !ok {
                return nil, false, nil
        }
        if err := validateListResult(data, "tools", "name"); err != nil {
                return nil, false, err
        }
        var rawTools []json.RawMessage
        if err := json.Unmarshal(rawToolsValue, &rawTools); err != nil {
                return nil, false, fmt.Errorf("decoding raw find_tool tools: %w", err)
        }
        if err := validateToolAnnotations(rawTools); err != nil {
                return nil, false, err
        }

        var output optimizer.FindToolOutput
        if err := json.Unmarshal(data, &output); err != nil {
                return nil, false, fmt.Errorf("decoding typed find_tool output: %w", err)
        }
        return &findToolOutputCarrier{
                rawOutput:     data,
                outputMembers: members,
                rawTools:      rawTools,
                output:        output,
        }, true, nil
}

func equivalentJSONValues(left, right json.RawMessage) (bool, error) {
        var leftValue any
        leftDecoder := json.NewDecoder(bytes.NewReader(left))
        leftDecoder.UseNumber()
        if err := leftDecoder.Decode(&leftValue); err != nil {
                return false, err
        }
        var rightValue any
        rightDecoder := json.NewDecoder(bytes.NewReader(right))
        rightDecoder.UseNumber()
        if err := rightDecoder.Decode(&rightValue); err != nil {
                return false, err
        }
        return reflect.DeepEqual(leftValue, rightValue), nil
}

func (result *decodedFindToolCallResult) withFilteredToolIndexes(indexes []int) ([]byte, error) {
        textCarrier := result.outputs[0]
        filteredTextOutput, err := textCarrier.withFilteredToolIndexes(indexes)
        if err != nil {
                return nil, err
        }
        textValue, err := json.Marshal(string(filteredTextOutput))
        if err != nil {
                return nil, err
        }
        filteredTextContent, err := encodeJSONObjectMembers(
                textCarrier.contentMembers,
                map[string]json.RawMessage{"text": textValue},
        )
        if err != nil {
                return nil, err
        }
        result.content[textCarrier.contentIndex] = filteredTextContent

        filteredContent, err := json.Marshal(result.content)
        if err != nil {
                return nil, err
        }
        replacements := map[string]json.RawMessage{"content": filteredContent}
        if result.structuredContent != nil {
                filteredStructuredOutput, err := result.structuredContent.withFilteredToolIndexes(indexes)
                if err != nil {
                        return nil, err
                }
                replacements["structuredContent"] = filteredStructuredOutput
        }
        return encodeJSONObjectMembers(result.members, replacements)
}

func (carrier *findToolOutputCarrier) withFilteredToolIndexes(indexes []int) ([]byte, error) {
        filteredTools := make([]json.RawMessage, 0, len(indexes))
        for _, index := range indexes {
                if index < 0 || index >= len(carrier.rawTools) {
                        return nil, fmt.Errorf("authorized tool index %d is out of range", index)
                }
                filteredTools = append(filteredTools, carrier.rawTools[index])
        }
        return encodeJSONObjectMembers(
                carrier.outputMembers,
                map[string]json.RawMessage{"tools": encodeJSONArrayValues(filteredTools)},
        )
}

func encodeJSONArrayValues(values []json.RawMessage) json.RawMessage {
        var encoded bytes.Buffer
        encoded.WriteByte('[')
        for i, value := range values {
                if i != 0 {
                        encoded.WriteByte(',')
                }
                encoded.Write(value)
        }
        encoded.WriteByte(']')
        return encoded.Bytes()
}

// encodeJSONObjectMembers replaces canonical values while retaining raw extensions.
func encodeJSONObjectMembers(
        members []jsonObjectMember,
        replacements map[string]json.RawMessage,
) ([]byte, error) {
        var encoded bytes.Buffer
        encoded.WriteByte('{')
        for i, member := range members {
                if i != 0 {
                        encoded.WriteByte(',')
                }
                name, err := json.Marshal(member.name)
                if err != nil {
                        return nil, err
                }
                encoded.Write(name)
                encoded.WriteByte(':')
                value := member.value
                if replacement, ok := replacements[member.name]; ok {
                        value = replacement
                }
                encoded.Write(value)
        }
        encoded.WriteByte('}')
        return encoded.Bytes(), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authz

import (
        "context"
        "log/slog"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive/pkg/authz/authorizers"
)

// authorizedToolIndexes returns the positions, in tools, of the tools the caller
// (identified by JWT claims in ctx) is authorized to call, so a filter can keep
// the matching raw descriptors. If authorizer is nil, every position is returned.
func authorizedToolIndexes(ctx context.Context, a authorizers.Authorizer, tools []mcp.Tool) []int {
        if a == nil {
                indexes := make([]int, len(tools))
                for i := range tools {
                        indexes[i] = i
                }
                return indexes
        }

        authorizedIndexes := make([]int, 0, len(tools))
        for i, tool := range tools {
                // Inject this tool's annotations into the context so Cedar policies
                // that use when clauses on resource attributes (e.g. resource.readOnlyHint)
                // can evaluate correctly. Without this, the authorization check runs
                // against a context with no annotations and all when clauses fail.
                toolCtx := ctx
                ann := &tools[i].Annotations
                if hasAnyHint(ann) {
                        toolCtx = authorizers.WithToolAnnotations(toolCtx, convertMCPAnnotation(ann))
                }

                authorized, err := authorizeToolCall(toolCtx, a, tool.Name, nil)
                if err != nil {
                        slog.Warn("Authorization check failed for tool, skipping",
                                "tool", tool.Name, "error", err)
                        continue
                }

                if authorized {
                        authorizedIndexes = append(authorizedIndexes, i)
                } else {
                        slog.Debug("Tool denied by authorization policy",
                                "tool", tool.Name)
                }
        }

        if denied := len(tools) - len(authorizedIndexes); denied > 0 {
                slog.Debug("Authorization policy filtered tools",
                        "total", len(tools), "allowed", len(authorizedIndexes), "denied", denied)
        }

        return authorizedIndexes
}

// authorizeToolCall checks whether the caller is authorized to call a specific tool
// with the given arguments. Returns true if authorized, false if denied.
// If authorizer is nil, returns true (no-op).
func authorizeToolCall(
        ctx context.Context, a authorizers.Authorizer, toolName string, arguments map[string]interface{},
) (bool, error) {
        if a == nil {
                return true, nil
        }

        return a.AuthorizeWithJWTClaims(
                ctx,
                authorizers.MCPFeatureTool,
                authorizers.MCPOperationCall,
                toolName,
                arguments,
        )
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package bodylimit provides HTTP middleware that caps the size of request
// bodies, rejecting oversized requests with 413 Request Entity Too Large.
//
// It is used both directly as a net/http middleware (the management API and the
// vMCP server) and via the runner middleware registry (the MCP proxies), so
// every inbound listener can bound the memory a single request may consume
// before it is buffered by handlers that call io.ReadAll.
package bodylimit

import (
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net/http"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // MiddlewareType is the type constant for the body limit middleware in the
        // runner middleware registry.
        MiddlewareType = "bodylimit"

        // DefaultMaxRequestBodySize is the fail-safe cap applied to inbound request
        // bodies when no limit is configured. It is sized to accommodate legitimate
        // MCP tools/call payloads with large inline content (e.g. base64 images or
        // documents) while still bounding per-request memory; operators can override
        // it. Note this caps requests only — server-produced response content is not
        // limited. Zero must never mean "unlimited" (see go-style rules), so callers
        // passing a non-positive limit fall back to this value.
        DefaultMaxRequestBodySize int64 = 8 << 20 // 8 MB
)

// MiddlewareParams holds the parameters for the body limit middleware factory.
type MiddlewareParams struct {
        // MaxBytes is the maximum request body size in bytes. Values <= 0 are
        // treated as DefaultMaxRequestBodySize (zero never means "unlimited").
        MaxBytes int64 `json:"max_bytes"`
}

// bodyLimitMiddleware adapts the body-limit handler to the types.Middleware
// interface expected by the runner middleware registry.
type bodyLimitMiddleware struct {
        handler types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *bodyLimitMiddleware) Handler() types.MiddlewareFunction {
        return m.handler
}

// Close releases resources held by the middleware. The body-limit middleware
// holds none, so this is a no-op.
func (*bodyLimitMiddleware) Close() error {
        return nil
}

// Middleware returns a net/http middleware that rejects requests whose body
// exceeds maxBytes with 413 Request Entity Too Large. It rejects early based on
// the Content-Length header, and wraps the body in http.MaxBytesReader as a
// safety net for requests that omit (or understate) Content-Length. A
// non-positive maxBytes falls back to DefaultMaxRequestBodySize.
func Middleware(maxBytes int64) func(http.Handler) http.Handler {
        if maxBytes <= 0 {
                maxBytes = DefaultMaxRequestBodySize
        }
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // Check Content-Length header first for early rejection.
                        if r.ContentLength > maxBytes {
                                slog.Warn("request body size exceeds limit", //nolint:gosec // G706: request metadata for diagnostics
                                        "content_length", r.ContentLength, "limit", maxBytes, "method", r.Method, "path", r.URL.Path)
                                http.Error(w, "Request Entity Too Large", http.StatusRequestEntityTooLarge)
                                return
                        }

                        // Track if MaxBytesReader's limit is exceeded.
                        limitExceeded := false

                        // Wrap ResponseWriter to intercept only MaxBytesReader errors.
                        wrappedWriter := &bodySizeResponseWriter{
                                ResponseWriter: w,
                                limitExceeded:  &limitExceeded,
                                written:        false,
                        }

                        // Set MaxBytesReader as a safety net for requests without Content-Length.
                        limitedBody := http.MaxBytesReader(wrappedWriter, r.Body, maxBytes)

                        // Wrap the limited body to detect when the size limit is actually
                        // exceeded (signalled by an *http.MaxBytesError on Read).
                        tracker := &maxBytesTracker{
                                ReadCloser:    limitedBody,
                                limitExceeded: &limitExceeded,
                        }
                        r.Body = tracker

                        next.ServeHTTP(wrappedWriter, r)
                })
        }
}

// IsRequestTooLarge reports whether err is the *http.MaxBytesError that
// http.MaxBytesReader returns when a request body exceeds the configured limit.
// Handlers that read the body directly use this to translate the read error
// into a 413 instead of a generic 500.
func IsRequestTooLarge(err error) bool {
        var maxBytesErr *http.MaxBytesError
        return errors.As(err, &maxBytesErr)
}

// CreateMiddleware is the types.MiddlewareFactory registered in the runner's
// GetSupportedMiddlewareFactories. It unmarshals MiddlewareParams, builds the
// handler via Middleware, and registers it with the runner.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal body limit middleware parameters: %w", err)
        }

        mw := &bodyLimitMiddleware{
                handler: Middleware(params.MaxBytes),
        }
        runner.AddMiddleware(MiddlewareType, mw)
        return nil
}

// maxBytesTracker wraps an io.ReadCloser to detect when http.MaxBytesReader
// rejects an oversized body. It flags only genuine overflows so that
// bodySizeResponseWriter rewrites a handler 400 to 413 exclusively when the
// body actually exceeded the limit — a body of exactly the limit is within
// bounds and never sets the flag.
type maxBytesTracker struct {
        io.ReadCloser
        limitExceeded *bool
}

func (t *maxBytesTracker) Read(p []byte) (n int, err error) {
        n, err = t.ReadCloser.Read(p)

        // http.MaxBytesReader returns an *http.MaxBytesError only once the body
        // has grown past the limit; reading exactly the limit returns no error.
        var maxBytesErr *http.MaxBytesError
        if errors.As(err, &maxBytesErr) {
                *t.limitExceeded = true
        }

        return n, err
}

// bodySizeResponseWriter wraps http.ResponseWriter to convert 400 to 413 only
// when MaxBytesReader's limit was exceeded (not for validation errors).
type bodySizeResponseWriter struct {
        http.ResponseWriter
        limitExceeded *bool
        written       bool
}

func (w *bodySizeResponseWriter) WriteHeader(statusCode int) {
        // Only convert 400 to 413 if MaxBytesReader's limit was actually exceeded.
        if statusCode == http.StatusBadRequest && !w.written && *w.limitExceeded {
                statusCode = http.StatusRequestEntityTooLarge
        }
        w.written = true
        w.ResponseWriter.WriteHeader(statusCode)
}

func (w *bodySizeResponseWriter) Write(b []byte) (int, error) {
        if !w.written {
                w.WriteHeader(http.StatusOK)
        }
        return w.ResponseWriter.Write(b)
}

// Unwrap exposes the underlying ResponseWriter so http.ResponseController (and
// other unwrap-aware callers) can reach optional capabilities such as
// SetWriteDeadline. Without this, wrapping the writer would break long-lived
// SSE streams whose transport clears the server write deadline.
func (w *bodySizeResponseWriter) Unwrap() http.ResponseWriter {
        return w.ResponseWriter
}

// Flush forwards to the underlying ResponseWriter when it supports flushing,
// preserving streaming (SSE) responses for callers that assert http.Flusher
// directly rather than going through http.ResponseController. http.Hijacker is
// intentionally not forwarded: the transports this middleware guards stream over
// SSE and do not hijack the connection.
func (w *bodySizeResponseWriter) Flush() {
        if f, ok := w.ResponseWriter.(http.Flusher); ok {
                f.Flush()
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package certs provides utilities for certificate validation and handling.
package certs

import (
        "crypto/x509"
        "encoding/pem"
        "fmt"
        "log/slog"
)

// ValidateCACertificate validates that the provided data contains a valid PEM-encoded certificate
func ValidateCACertificate(certData []byte) error {
        // Check if the data contains PEM blocks
        block, _ := pem.Decode(certData)
        if block == nil {
                return fmt.Errorf("no PEM data found in certificate file")
        }

        // Check if it's a certificate block
        if block.Type != "CERTIFICATE" {
                return fmt.Errorf("PEM block is not a certificate (found: %s)", block.Type)
        }

        // Parse the certificate to ensure it's valid
        cert, err := x509.ParseCertificate(block.Bytes)
        if err != nil {
                return fmt.Errorf("failed to parse certificate: %w", err)
        }

        // Basic validation - check if it's a CA certificate
        if !cert.IsCA {
                // Log a warning but don't fail - some corporate proxies use non-CA certificates
                slog.Warn("Certificate is not marked as a CA certificate, but proceeding anyway")
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "net/url"
        "os"
        "path/filepath"
        "strings"

        "github.com/stacklok/toolhive/pkg/fileutils"
)

// Claude Desktop in third-party mode takes remote MCP servers ("connectors")
// only from the managedMcpServers key of its applied configuration. ToolHive
// writes them into the same configLibrary document "thv llm setup --client
// claude-desktop" owns: Desktop applies exactly one document at a time, so a
// separate document would displace the LLM gateway configuration. Desktop
// connects to each server directly and runs the server's OAuth sign-in itself,
// so no local proxy is involved.
//
// Re-running "thv llm setup" keeps these entries when it rewrites the document.
// "thv llm teardown" deletes the document along with any entries still in it
// and logs a warning naming them, since Desktop ignores them without the
// inference settings.

const claudeDesktopManagedMCPServersKey = "managedMcpServers"

var (
        // ErrConnectorsUnsupported is returned for a client that cannot be
        // configured with remote MCP connectors.
        ErrConnectorsUnsupported = errors.New("client does not support connectors")
        // ErrConnectorsNotSetUp is returned when Claude Desktop has no
        // ToolHive-owned configuration document to add connectors to.
        ErrConnectorsNotSetUp = errors.New(
                "no ToolHive configuration found for Claude Desktop; run \"thv llm setup --client claude-desktop\" first")
        // ErrConnectorNotFound is returned when removing a connector that does not exist.
        ErrConnectorNotFound = errors.New("connector not found")
)

// Connector is a remote MCP server configured directly in a client.
type Connector struct {
        Name string `json:"name"`
        URL  string `json:"url"`
}

// AddConnector adds a remote MCP server to the client, or updates the URL of
// an existing connector with the same name, keeping its other fields. The
// client signs in to the server itself (OAuth), so nothing is proxied through
// ToolHive. A same-name entry that is not a remote server is an error.
//
// The write succeeds but logs a warning when Claude Desktop will not load the
// connector: when a managed-preferences (MDM) profile overrides the local
// configuration, or when the user has applied a different saved configuration.
func (cm *ClientManager) AddConnector(clientType ClientApp, name, serverURL string) error {
        if err := validateConnector(name, serverURL); err != nil {
                return err
        }
        applied, err := cm.updateConnectors(clientType, func(servers []any) ([]any, error) {
                for _, s := range servers {
                        m, ok := s.(map[string]any)
                        if !ok || m["name"] != name {
                                continue
                        }
                        if !isRemoteServer(m) {
                                return nil, fmt.Errorf("a server named %q that is not a remote connector already exists", name)
                        }
                        m["url"] = serverURL
                        return servers, nil
                }
                return append(servers, map[string]any{
                        "name":      name,
                        "transport": "http",
                        "url":       serverURL,
                        "oauth":     true,
                }), nil
        })
        if err != nil {
                return err
        }
        if cm.IsManaged(clientType) {
                slog.Warn("a managed configuration overrides Claude Desktop's local configuration, "+
                        "so it will not load this connector; add it to the managed configuration instead, "+
                        "or fix the managed configuration if Desktop rejected it",
                        "connector", name)
        }
        if !applied {
                slog.Warn("Claude Desktop is using a different saved configuration, so it will not load this connector "+
                        "until \""+claudeDesktopManagedEntryName+"\" is applied again", "connector", name)
        }
        return nil
}

// RemoveConnector removes the connector with the given name from the client.
// Entries that are not remote servers are never removed, matching
// ListConnectors.
func (cm *ClientManager) RemoveConnector(clientType ClientApp, name string) error {
        _, err := cm.updateConnectors(clientType, func(servers []any) ([]any, error) {
                kept := make([]any, 0, len(servers))
                for _, s := range servers {
                        if m, ok := s.(map[string]any); ok && m["name"] == name && isRemoteServer(m) {
                                continue
                        }
                        kept = append(kept, s)
                }
                if len(kept) == len(servers) {
                        return nil, fmt.Errorf("%w: %q", ErrConnectorNotFound, name)
                }
                return kept, nil
        })
        return err
}

// ListConnectors returns the remote MCP servers configured in the client.
// Entries without a name or URL (for example local stdio servers a user added
// by hand) are skipped.
func (cm *ClientManager) ListConnectors(clientType ClientApp) ([]Connector, error) {
        docPath, err := cm.connectorDocPath(clientType)
        if err != nil {
                return nil, err
        }
        doc, err := readClaudeDesktopDoc(docPath)
        if err != nil {
                return nil, err
        }
        servers, err := managedMCPServers(doc, docPath)
        if err != nil {
                return nil, err
        }
        return connectorsFrom(servers), nil
}

// updateConnectors applies fn to the managedMcpServers array of ToolHive's
// Claude Desktop configuration document, under the same lock "thv llm setup"
// and "thv llm teardown" hold, preserving every other key in the document. It
// reports whether that document is the one Claude Desktop currently applies.
func (cm *ClientManager) updateConnectors(clientType ClientApp, fn func([]any) ([]any, error)) (bool, error) {
        appCfg, err := cm.connectorAppConfig(clientType)
        if err != nil {
                return false, err
        }
        metaPath := cm.buildLLMSettingsPath(appCfg)
        // The lock file lives in configLibrary, which only exists once setup has run.
        if !fileExistsAt(filepath.Dir(metaPath)) {
                return false, ErrConnectorsNotSetUp
        }
        var applied bool
        err = fileutils.WithFileLock(metaPath, func() error {
                docPath, isApplied, err := connectorDocPathLocked(metaPath)
                if err != nil {
                        return err
                }
                applied = isApplied
                doc, err := readClaudeDesktopDoc(docPath)
                if err != nil {
                        return err
                }
                servers, err := managedMCPServers(doc, docPath)
                if err != nil {
                        return err
                }
                servers, err = fn(servers)
                if err != nil {
                        return err
                }
                if len(servers) == 0 {
                        delete(doc, claudeDesktopManagedMCPServersKey)
                } else {
                        doc[claudeDesktopManagedMCPServersKey] = servers
                }
                data, err := json.MarshalIndent(doc, "", "  ")
                if err != nil {
                        return fmt.Errorf("encoding %s: %w", docPath, err)
                }
                if err := fileutils.AtomicWriteFile(docPath, data, 0o600); err != nil {
                        return fmt.Errorf("writing %s: %w", docPath, err)
                }
                return nil
        })
        return applied, err
}

// connectorAppConfig returns the integration config for a client that
// supports connectors. Only Claude Desktop does today.
func (cm *ClientManager) connectorAppConfig(clientType ClientApp) (*clientAppConfig, error) {
        if clientType != ClientApp(ClaudeDesktop) {
                return nil, fmt.Errorf("%w: %q (supported: %s)", ErrConnectorsUnsupported, clientType, ClaudeDesktop)
        }
        appCfg := cm.lookupClientAppConfig(clientType)
        if appCfg == nil {
                return nil, fmt.Errorf("%w: %q", ErrConnectorsUnsupported, clientType)
        }
        return appCfg, nil
}

// connectorDocPath resolves ToolHive's Claude Desktop configuration document
// without taking the lock; used by read-only callers.
func (cm *ClientManager) connectorDocPath(clientType ClientApp) (string, error) {
        appCfg, err := cm.connectorAppConfig(clientType)
        if err != nil {
                return "", err
        }
        docPath, _, err := connectorDocPathLocked(cm.buildLLMSettingsPath(appCfg))
        return docPath, err
}

// connectorDocPathLocked finds the document named by ToolHive's _meta.json
// entry, and reports whether _meta.json's appliedId still selects it. Callers
// that write must hold the metaPath lock.
func connectorDocPathLocked(metaPath string) (string, bool, error) {
        meta, err := readClaudeDesktopMeta(metaPath)
        if err != nil {
                return "", false, err
        }
        id := metaEntryID(meta, claudeDesktopManagedEntryName)
        if id == "" {
                return "", false, ErrConnectorsNotSetUp
        }
        docPath := filepath.Join(filepath.Dir(metaPath), id+".json")
        if !fileExistsAt(docPath) {
                return "", false, ErrConnectorsNotSetUp
        }
        applied, _ := meta["appliedId"].(string)
        return docPath, applied == id, nil
}

// readClaudeDesktopDoc reads a configuration document into a generic map so
// keys ToolHive does not own are preserved on write-back.
func readClaudeDesktopDoc(path string) (map[string]any, error) {
        data, err := os.ReadFile(path) // #nosec G304 -- path is derived from the configLibrary dir
        if err != nil {
                return nil, fmt.Errorf("reading %s: %w", path, err)
        }
        var doc map[string]any
        if err := json.Unmarshal(data, &doc); err != nil {
                return nil, fmt.Errorf("parsing %s: %w", path, err)
        }
        if doc == nil {
                doc = map[string]any{}
        }
        return doc, nil
}

// managedMCPServers returns the document's managedMcpServers array. A value of
// the wrong shape is an error rather than treated as empty, so a write never
// overwrites data ToolHive does not understand.
func managedMCPServers(doc map[string]any, path string) ([]any, error) {
        v, ok := doc[claudeDesktopManagedMCPServersKey]
        if !ok || v == nil {
                return nil, nil
        }
        servers, ok := v.([]any)
        if !ok {
                return nil, fmt.Errorf("parsing %s: %q must be a JSON array", path, claudeDesktopManagedMCPServersKey)
        }
        return servers, nil
}

// connectorsFrom converts managedMcpServers entries to Connectors, skipping
// entries that are not remote servers.
func connectorsFrom(servers []any) []Connector {
        connectors := make([]Connector, 0, len(servers))
        for _, s := range servers {
                m, ok := s.(map[string]any)
                if !ok || !isRemoteServer(m) {
                        continue
                }
                connectors = append(connectors, Connector{Name: m["name"].(string), URL: m["url"].(string)})
        }
        return connectors
}

// isRemoteServer reports whether a managedMcpServers entry is a remote server
// with a name, the only kind of entry the connector commands read or change.
func isRemoteServer(m map[string]any) bool {
        name, _ := m["name"].(string)
        serverURL, _ := m["url"].(string)
        return name != "" && serverURL != ""
}

// validateConnector checks the name and URL before anything is written.
// Claude Desktop only connects to remote servers over HTTPS.
func validateConnector(name, serverURL string) error {
        if strings.TrimSpace(name) == "" || strings.TrimSpace(name) != name {
                return fmt.Errorf("invalid connector name %q: must be non-empty with no leading or trailing spaces", name)
        }
        u, err := url.Parse(serverURL)
        if err != nil || u.Scheme != "https" || u.Host == "" {
                return fmt.Errorf("invalid connector URL %q: must be an absolute https URL", serverURL)
        }
        return nil
}

// previousManagedMCPServers returns the managedMcpServers entries of the
// document "thv llm setup" is about to rewrite, so re-running setup keeps them.
// A missing document has none. An unreadable one is still replaced, as before
// connectors existed, with a warning that anything in it is dropped.
func previousManagedMCPServers(path string) []any {
        if !fileExistsAt(path) {
                return nil
        }
        servers, err := func() ([]any, error) {
                doc, err := readClaudeDesktopDoc(path)
                if err != nil {
                        return nil, err
                }
                return managedMCPServers(doc, path)
        }()
        if err != nil {
                slog.Warn("could not read connectors from the previous Claude Desktop configuration; "+
                        "they will not be carried over", "path", path, "error", err)
                return nil
        }
        return servers
}

// connectorNamesInDoc returns the names of the connectors in a configuration
// document, for the teardown warning. Errors are ignored: the warning is
// best-effort and must not block teardown.
func connectorNamesInDoc(path string) []string {
        doc, err := readClaudeDesktopDoc(path)
        if err != nil {
                return nil
        }
        servers, err := managedMCPServers(doc, path)
        if err != nil {
                return nil
        }
        names := make([]string, 0, len(servers))
        for _, c := range connectorsFrom(servers) {
                names = append(names, c.Name)
        }
        return names
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "bytes"
        "fmt"
        "log/slog"
        "os"
        "os/exec"
        "path/filepath"
)

const (
        codexBundleIdentifier     = "com.openai.codex"
        codexWindowsPackageFamily = "OpenAI.Codex_2p2nqsd0c76g0"
        codexWindowsPackageName   = "OpenAI.Codex"
        codexWindowsInstalledMark = "installed"
        // plutilPath is the absolute path to plutil on macOS. Using the absolute
        // path avoids PATH-resolution failures in stripped-environment contexts
        // (launchd jobs, sandboxes, CI) where /usr/bin is not on $PATH.
        plutilPath = "/usr/bin/plutil"
)

type codexDesktopDetector struct {
        homeDir string
        goos    string
        stat    func(string) (os.FileInfo, error)
        command func(string, ...string) ([]byte, error)
}

func newCodexDesktopDetector(homeDir, goos string) codexDesktopDetector {
        return codexDesktopDetector{
                homeDir: homeDir,
                goos:    goos,
                stat:    os.Stat,
                command: func(name string, args ...string) ([]byte, error) {
                        // #nosec G204 -- the command name is a fixed compile-time literal
                        // ("plutil" or "powershell.exe"); args include derived file paths
                        // from os.UserHomeDir, not user input.
                        return exec.Command(name, args...).Output()
                },
        }
}

func (d codexDesktopDetector) installed() (bool, error) {
        switch d.goos {
        case "darwin":
                return d.installedOnDarwin()
        case "windows":
                return d.installedOnWindows()
        default:
                return false, nil
        }
}

func (d codexDesktopDetector) installedOnDarwin() (bool, error) {
        applicationDirs := []string{"/Applications", filepath.Join(d.homeDir, "Applications")}
        for _, dir := range applicationDirs {
                for _, appName := range []string{"Codex.app", "ChatGPT.app"} {
                        infoPlist := filepath.Join(dir, appName, "Contents", "Info.plist")
                        if _, err := d.stat(infoPlist); err != nil {
                                continue
                        }
                        output, err := d.command(plutilPath, "-extract", "CFBundleIdentifier", "raw", "-o", "-", infoPlist)
                        if err != nil {
                                // Log the plutil failure so it's distinguishable from "app absent."
                                // A failure here means the plist exists but its bundle identifier
                                // could not be extracted (e.g. plist corruption or plutil error).
                                slog.Warn("failed to extract CFBundleIdentifier from Codex app plist",
                                        "app", appName, "plist", infoPlist, "error", err)
                                continue
                        }
                        if string(bytes.TrimSpace(output)) == codexBundleIdentifier {
                                return true, nil
                        }
                }
        }
        return false, nil
}

func (d codexDesktopDetector) installedOnWindows() (bool, error) {
        const query = "$p = Get-AppxPackage -Name '" + codexWindowsPackageName + "' -ErrorAction SilentlyContinue | " +
                "Where-Object { $_.PackageFamilyName -eq '" + codexWindowsPackageFamily + "' -and $_.SignatureKind -eq 'Store' } | " +
                "Select-Object -First 1; if ($p) { Write-Output '" + codexWindowsInstalledMark + "' }"
        output, err := d.command("powershell.exe", "-NoProfile", "-NonInteractive", "-Command", query)
        if err != nil {
                return false, fmt.Errorf("querying Codex Microsoft Store package: %w", err)
        }
        return string(bytes.TrimSpace(output)) == codexWindowsInstalledMark, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package client provides utilities for managing client configurations
// and interacting with MCP servers.
package client

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "runtime"
        "sort"
        "strings"

        "github.com/pelletier/go-toml/v2"
        "github.com/tailscale/hujson"
        "go.yaml.in/yaml/v3"

        "github.com/stacklok/toolhive/pkg/llmgateway"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// defaultURLFieldName is the default URL field name used when no specific mapping exists
const defaultURLFieldName = "url"

const httpTransportLabel = "http"
const skillsDirName = "skills"

// ClientApp is an enum of MCP-capable AI clients (IDEs, editors, and coding tools).
// Only clients that support MCP registration appear here; LLM-gateway-only
// tools (e.g. Xcode) are represented by the separate LLMClientApp type so
// that code generators (swag) do not include them in the MCP API enum.
//
//nolint:revive // ClientApp is intentionally named for clarity across packages
type ClientApp string

// LLMClientApp identifies tools that support the LLM gateway but do not
// support MCP client registration (e.g. GitHub Copilot for Xcode).
// Keeping this type separate from ClientApp prevents swag from including
// these tools in the MCP API ClientApp enum.
type LLMClientApp string

const (
        // RooCode represents the Roo Code extension for VS Code.
        RooCode ClientApp = "roo-code"
        // Cline represents the Cline extension for VS Code.
        Cline ClientApp = "cline"
        // Cursor represents the Cursor editor.
        Cursor ClientApp = "cursor"
        // VSCodeInsider represents the VS Code Insiders editor.
        VSCodeInsider ClientApp = "vscode-insider"
        // VSCode represents the standard VS Code editor.
        VSCode ClientApp = "vscode"
        // ClaudeCode represents the Claude Code CLI.
        ClaudeCode ClientApp = "claude-code"
        // Windsurf represents the Windsurf IDE.
        Windsurf ClientApp = "windsurf"
        // WindsurfJetBrains represents the Windsurf plugin for JetBrains.
        WindsurfJetBrains ClientApp = "windsurf-jetbrains"
        // AmpCli represents the Sourcegraph Amp CLI.
        AmpCli ClientApp = "amp-cli"
        // LMStudio represents the LM Studio application.
        LMStudio ClientApp = "lm-studio"
        // Goose represents the Goose AI agent.
        Goose ClientApp = "goose"
        // Trae represents the Trae IDE.
        Trae ClientApp = "trae"
        // Continue represents the Continue.dev IDE plugins.
        Continue ClientApp = "continue"
        // OpenCode represents the OpenCode editor.
        OpenCode ClientApp = "opencode"
        // Kiro represents the Kiro AI IDE.
        Kiro ClientApp = "kiro"
        // Antigravity represents the Google Antigravity IDE.
        Antigravity ClientApp = "antigravity"
        // Zed represents the Zed editor.
        Zed ClientApp = "zed"
        // GeminiCli represents the Google Gemini CLI.
        GeminiCli ClientApp = "gemini-cli"
        // VSCodeServer represents Microsoft's VS Code Server (remote development).
        VSCodeServer ClientApp = "vscode-server"
        // MistralVibe represents the Mistral Vibe IDE.
        MistralVibe ClientApp = "mistral-vibe"
        // Codex represents the OpenAI Codex CLI and desktop app.
        Codex ClientApp = "codex"
        // KimiCli represents the Kimi Code CLI.
        KimiCli ClientApp = "kimi-cli"
        // Factory represents the Factory.ai Droid CLI.
        Factory ClientApp = "factory"
        // CopilotCli represents the GitHub Copilot CLI.
        CopilotCli ClientApp = "copilot-cli"
        // Qoder represents the Qoder IDE.
        Qoder ClientApp = "qoder"
        // Pi represents the Pi coding agent.
        Pi ClientApp = "pi"
)

const (
        // Xcode represents GitHub Copilot for Xcode.
        // Xcode does not support MCP; it is an LLM-gateway-only tool.
        // It is declared as LLMClientApp (not ClientApp) so that code generators
        // such as swag do not include "xcode" in the MCP API ClientApp enum.
        Xcode LLMClientApp = "xcode"

        // ClaudeDesktop represents the Claude Desktop app configured for LLM gateway
        // routing via its "third-party inference" surface. ToolHive workloads are not
        // registered with it (remote MCP servers are added with "thv client
        // connector" instead), so it is declared as LLMClientApp — same rationale as
        // Xcode above.
        ClaudeDesktop LLMClientApp = "claude-desktop"
)

// Extension is extension of the client config file.
type Extension string

const (
        // JSON represents a JSON extension.
        JSON Extension = "json"
        // YAML represents a YAML extension.
        YAML Extension = "yaml"
        // TOML represents a TOML extension.
        TOML Extension = "toml"
)

// YAMLStorageType represents how servers are stored in YAML configuration files.
type YAMLStorageType string

const (
        // YAMLStorageTypeMap represents servers stored as a map with server names as keys.
        YAMLStorageTypeMap YAMLStorageType = "map"
        // YAMLStorageTypeArray represents servers stored as an array of objects.
        YAMLStorageTypeArray YAMLStorageType = "array"
)

// TOMLStorageType represents how servers are stored in TOML configuration files.
type TOMLStorageType string

const (
        // TOMLStorageTypeMap represents servers stored as nested tables [section.servername].
        // Example: [mcp_servers.myserver]
        TOMLStorageTypeMap TOMLStorageType = "map"
        // TOMLStorageTypeArray represents servers stored as array of tables [[section]].
        // Example: [[mcp_servers]]
        TOMLStorageTypeArray TOMLStorageType = "array"
)

// Platform represents a runtime.GOOS value used as a key in platform-specific path maps.
type Platform string

const (
        // PlatformLinux is the Linux platform.
        PlatformLinux Platform = "linux"
        // PlatformDarwin is the macOS platform.
        PlatformDarwin Platform = "darwin"
        // PlatformWindows is the Windows platform.
        PlatformWindows Platform = "windows"
)

// LLMGatewayKeySpec describes a single JSON key to patch when configuring
// (or reverting) LLM gateway access for a tool. JSONPointer is an RFC 6901
// path (e.g. "/apiKeyHelper" or "/env/ANTHROPIC_BASE_URL"). Dots in flat
// top-level key names (e.g. "/cursor.general.openAIBaseURL") are treated as
// literals by hujson.Patch.
//
// Exactly one of ValueField or Literal must be set:
//   - ValueField names which ApplyConfig field to write. Valid values:
//     "GatewayURL", "AnthropicBaseURL", "ProxyBaseURL", "ProxyOrigin",
//     "TokenHelperCommand", "PlaceholderAPIKey", "ClaudeCodeHelperTTLMillis",
//     "NodeTLSRejectUnauthorized", "PromptCacheTTL",
//     "PromptCache1HLegacy",
//     "BedrockDisableExperimentalBetas", "BedrockHaikuModel",
//     "BedrockOpusModel", "BedrockSonnetModel", "BedrockFableModel". An unrecognised ValueField is
//     a programming error and causes ConfigureLLMGateway to return an error.
//   - Literal is written verbatim into the settings key (e.g. a fixed auth
//     type string). Use Literal instead of ValueField for constant values so
//     that typos in ValueField are caught as errors rather than silently
//     written as unexpected strings.
//
// ClearWhenEmpty: when true and the resolved value is empty, the key is
// removed from the settings file rather than skipped. Use for conditional
// keys like NODE_TLS_REJECT_UNAUTHORIZED that must be cleaned up when the
// flag is cleared. Ignored when Literal is set (literals are never empty).
type LLMGatewayKeySpec struct {
        JSONPointer string // RFC 6901 path
        // ValueField: "GatewayURL" | "AnthropicBaseURL" | "ProxyBaseURL" | "ProxyOrigin" |
        // "TokenHelperCommand" | "PlaceholderAPIKey" | "ClaudeCodeHelperTTLMillis" |
        // "NodeTLSRejectUnauthorized" | "PromptCacheTTL" | "PromptCache1HLegacy" |
        // "BedrockDisableExperimentalBetas" | "BedrockHaikuModel" |
        // "BedrockOpusModel" | "BedrockSonnetModel" | "BedrockFableModel"
        ValueField     string
        Literal        string // constant value written verbatim; mutually exclusive with ValueField
        ClearWhenEmpty bool   // remove the key when the resolved value is empty (ignored for Literal)
}

// LLMEnvFileKeySpec describes a single KEY=value entry to write to a dotenv
// file when configuring (or reverting) LLM gateway access for a tool.
//
// Exactly one of ValueField or Literal must be set. ValueField semantics are
// identical to LLMGatewayKeySpec.ValueField.
type LLMEnvFileKeySpec struct {
        // Name is the environment variable name (e.g. "GEMINI_API_KEY").
        Name string
        // ValueField names which ApplyConfig field to resolve. Valid values are the
        // same as LLMGatewayKeySpec.ValueField. Mutually exclusive with Literal.
        ValueField string
        // Literal is written verbatim as the variable value. Mutually exclusive with ValueField.
        Literal string
}

// clientAppConfig represents a configuration path for a supported MCP client.
type clientAppConfig struct {
        ClientType  ClientApp
        Description string
        // Deprecated marks a client integration whose upstream tool is no longer
        // maintained. Deprecated clients still function but are flagged in the CLI
        // client list and trigger a warning when registered or removed.
        Deprecated bool
        // DeprecationMessage is the full warning text shown to users (on stderr)
        // when they touch a deprecated client. Only set when Deprecated is true.
        DeprecationMessage            string
        RelPath                       []string
        SettingsFile                  string
        PlatformPrefix                map[Platform][]string
        MCPServersPathPrefix          string
        Extension                     Extension
        SupportedTransportTypesMap    map[types.TransportType]string // stdio mapped to streamable-http (SSE deprecated)
        IsTransportTypeFieldSupported bool
        // MCPServersUrlLabelMap maps transport type to URL field name (e.g., "url", "serverUrl", "httpUrl")
        MCPServersUrlLabelMap map[types.TransportType]string
        // YAML-specific configuration (only used when Extension == YAML)
        YAMLStorageType     YAMLStorageType // How servers are stored in YAML (map or array)
        YAMLIdentifierField string          // For array type: field name that identifies the server
        YAMLDefaults        map[string]any  // Default values to add to entries
        // TOML-specific configuration (only used when Extension == TOML)
        TOMLStorageType TOMLStorageType // How servers are stored in TOML (map or array)
        // ConfigDirEnvVar names an environment variable that, when set, relocates
        // the client's config directory. It replaces the home-relative directory
        // holding SettingsFile and LLMSettingsFile, and the global skills directory
        // becomes <dir>/skills. A leading "~" expands to the home directory and a
        // relative value resolves against the working directory.
        ConfigDirEnvVar string
        // Skill-specific configuration
        SupportsSkills    bool     // Whether this client supports skills
        SkillsGlobalPath  []string // Path segments for global skills dir (from home dir)
        SkillsProjectPath []string // Path segments for project-local skills dir (from project root)
        // SkillsPlatformPrefix maps Platform values to path segments inserted between
        // home dir and SkillsGlobalPath. Needed for clients following platform conventions
        // (e.g., XDG ~/.config/ on Linux/macOS).
        // If nil or missing an entry for the current OS, no prefix is added.
        SkillsPlatformPrefix map[Platform][]string
        // Plugin-specific configuration
        SupportsPlugins    bool     // Whether this client supports plugins
        PluginsGlobalPath  []string // Path segments for global plugins dir (from home dir)
        PluginsProjectPath []string // Path segments for project-local plugins dir (from project root)
        // PluginsPlatformPrefix maps Platform values to path segments inserted between
        // home dir and PluginsGlobalPath. Same semantics as SkillsPlatformPrefix.
        // If nil or missing an entry for the current OS, no prefix is added.
        PluginsPlatformPrefix map[Platform][]string
        // LLM gateway configuration ─────────────────────────────────────────────
        // LLMGatewayMode identifies the gateway integration strategy (direct token
        // helper, proxy, credential helper, or Codex auth), or "" when unsupported.
        LLMGatewayMode string
        // LLMBinaryName is the executable name looked up via exec.LookPath to
        // confirm the tool is actually installed (not just a leftover config
        // directory). Leave empty for tools that are not on $PATH (e.g. macOS
        // GUI apps).
        LLMBinaryName string
        // LLMGatewayOnly marks tools that support LLM gateway but not MCP (e.g. Xcode).
        // Entries with this flag are excluded from the MCP client list.
        LLMGatewayOnly bool
        // LLMSettingsFile is the filename of the settings file to patch for LLM
        // gateway (may differ from SettingsFile used for MCP).
        LLMSettingsFile string
        // LLMSettingsRelPath is the path segments from home dir + platform prefix
        // to the directory containing LLMSettingsFile.
        LLMSettingsRelPath []string
        // LLMSettingsPlatformPrefix maps Platform to path segments prepended before
        // LLMSettingsRelPath (same semantics as PlatformPrefix).
        LLMSettingsPlatformPrefix map[Platform][]string
        // LLMGatewayKeys lists the JSON Pointer paths and value-field mappings to
        // apply when setting up (or reverting) LLM gateway access.
        LLMGatewayKeys []LLMGatewayKeySpec
        // LLMEnvFileRelPath is the path segments from home dir to the directory
        // containing the .env file to manage for LLM gateway (e.g. []string{".gemini"}).
        // Empty means this client has no .env file to manage.
        LLMEnvFileRelPath []string
        // LLMEnvFileName is the filename of the .env file (e.g. ".env").
        LLMEnvFileName string
        // LLMEnvFileKeys lists the key=value entries to write to the .env file when
        // setting up (or reverting) LLM gateway access.
        LLMEnvFileKeys []LLMEnvFileKeySpec
        // LLMDetectRelPath and LLMDetectPlatformPrefix locate a directory whose
        // existence indicates the tool is installed. Used for GUI apps that are not
        // on $PATH (LLMBinaryName is empty) and whose LLM settings directory does not
        // exist until first configured — so the settings dir cannot be the detection
        // signal. When set, DetectedLLMGatewayClients checks this directory instead of
        // the settings directory. Resolved with the same semantics as
        // LLMSettingsRelPath / LLMSettingsPlatformPrefix.
        LLMDetectRelPath        []string
        LLMDetectPlatformPrefix map[Platform][]string
        // LLMInstalledDetector is an optional per-client detection hook that runs
        // in addition to the shared settings-dir + binary-on-PATH check. Used by
        // clients that have installation evidence beyond the CLI (e.g. Codex's
        // desktop app). When set, the client is considered installed if either
        // the shared check or this hook returns true. The hook must be set at
        // construction time (NewClientManager), not in static config.
        LLMInstalledDetector func() (bool, error)
        // LLMManagedProfileDomain is the macOS managed-preferences plist domain
        // (e.g. "com.anthropic.claudefordesktop.plist") that, when present, overrides
        // the client's local config. Setup warns when detected. Empty when the client
        // has no managed-profile surface.
        LLMManagedProfileDomain string
        // LLMManagedSettingsRelPath is the path segments, under /etc, of the Linux
        // managed settings file (e.g. claude-desktop/managed-settings.json) that,
        // when present, overrides the client's local config. Setup warns when
        // detected. Empty when the client has no Linux managed-settings surface.
        LLMManagedSettingsRelPath []string
}

// extractServersKeyFromConfig extracts the servers key from MCPServersPathPrefix
// by removing the leading "/" (e.g., "/mcpServers" -> "mcpServers").
func extractServersKeyFromConfig(cfg *clientAppConfig) string {
        return strings.TrimPrefix(cfg.MCPServersPathPrefix, "/")
}

// extractURLLabelFromConfig extracts the URL field label from MCPServersUrlLabelMap.
// It checks transport types in priority order: StreamableHTTP, then Stdio.
// Returns defaultURLFieldName if no mapping is found.
func extractURLLabelFromConfig(cfg *clientAppConfig) string {
        if cfg.MCPServersUrlLabelMap != nil {
                if label, ok := cfg.MCPServersUrlLabelMap[types.TransportTypeStreamableHTTP]; ok {
                        return label
                }
                if label, ok := cfg.MCPServersUrlLabelMap[types.TransportTypeStdio]; ok {
                        return label
                }
        }
        return defaultURLFieldName
}

var (
        // ErrConfigFileNotFound is returned when a client configuration file is not found
        ErrConfigFileNotFound = fmt.Errorf("client config file not found")
        // ErrUnsupportedClientType is returned when an unsupported client type is provided
        ErrUnsupportedClientType = fmt.Errorf("unsupported client type")
)

var supportedClientIntegrations = []clientAppConfig{
        {
                ClientType:  RooCode,
                Description: "VS Code Roo Code extension",
                Deprecated:  true,
                DeprecationMessage: "The Roo Code VS Code extension has been discontinued (last release May 15, 2026)\n" +
                        "and its repository has been archived. Support for this client will be removed in a\n" +
                        "future ToolHive release. The Roo Code team recommends migrating to Cline:\n" +
                        "  thv client register cline",
                SettingsFile: "mcp_settings.json",
                RelPath: []string{
                        "Code", "User", "globalStorage", "rooveterinaryinc.roo-cline", "settings",
                },
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                MCPServersPathPrefix: "/mcpServers",
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "streamable-http",
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: "streamable-http",
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".roo", skillsDirName},
                SkillsProjectPath: []string{".roo", skillsDirName},
        },
        {
                ClientType:   Cline,
                Description:  "VS Code Cline extension",
                SettingsFile: "cline_mcp_settings.json",
                RelPath: []string{
                        "Code", "User", "globalStorage", "saoudrizwan.claude-dev", "settings",
                },
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                MCPServersPathPrefix: "/mcpServers",
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "streamableHttp",
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: "streamableHttp",
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".cline", skillsDirName},
                SkillsProjectPath: []string{".cline", skillsDirName},
        },
        {
                ClientType:   VSCodeInsider,
                Description:  "Visual Studio Code Insiders",
                SettingsFile: "mcp.json",
                RelPath: []string{
                        "Code - Insiders", "User",
                },
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                MCPServersPathPrefix: "/servers",
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".copilot", skillsDirName},
                SkillsProjectPath: []string{".github", skillsDirName},
                // VS Code reads custom LLM provider groups from chatLanguageModels.json.
                LLMGatewayMode:     llmgateway.ModeVSCode,
                LLMSettingsFile:    "chatLanguageModels.json",
                LLMSettingsRelPath: []string{"Code - Insiders", "User"},
                LLMSettingsPlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
        },
        {
                ClientType:   VSCode,
                Description:  "Visual Studio Code",
                SettingsFile: "mcp.json",
                RelPath: []string{
                        "Code", "User",
                },
                MCPServersPathPrefix: "/servers",
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                Extension: JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".copilot", skillsDirName},
                SkillsProjectPath: []string{".github", skillsDirName},
                // VS Code reads custom LLM provider groups from chatLanguageModels.json.
                LLMGatewayMode:     llmgateway.ModeVSCode,
                LLMSettingsFile:    "chatLanguageModels.json",
                LLMSettingsRelPath: []string{"Code", "User"},
                LLMSettingsPlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
        },
        {
                ClientType:           Cursor,
                Description:          "Cursor editor",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".cursor"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                // Adding type field is not explicitly required though, Cursor auto-detects and is able to
                // connect to both sse and streamable-http types
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".cursor", skillsDirName},
                SkillsProjectPath: []string{".cursor", skillsDirName},
                // LLM gateway: patches the editor settings.json (different from the MCP mcp.json)
                LLMGatewayMode:     llmgateway.ModeProxy,
                LLMSettingsFile:    "settings.json",
                LLMSettingsRelPath: []string{"Cursor", "User"},
                LLMSettingsPlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                LLMGatewayKeys: []LLMGatewayKeySpec{
                        {JSONPointer: "/cursor.general.openAIBaseURL", ValueField: "ProxyBaseURL"},
                        {JSONPointer: "/cursor.general.openAIAPIKey", ValueField: "PlaceholderAPIKey"},
                },
        },
        {
                ClientType:           ClaudeCode,
                Description:          "Claude Code CLI",
                SettingsFile:         ".claude.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:     true,
                SkillsGlobalPath:   []string{".claude", skillsDirName},
                SkillsProjectPath:  []string{".claude", skillsDirName},
                SupportsPlugins:    true,
                PluginsGlobalPath:  []string{".claude", "plugins"},
                PluginsProjectPath: []string{".claude", "plugins"},
                // LLM gateway: patches ~/.claude/settings.json (different from the MCP .claude.json)
                LLMGatewayMode:     llmgateway.ModeDirect,
                LLMBinaryName:      "claude",
                LLMSettingsFile:    "settings.json",
                LLMSettingsRelPath: []string{".claude"},
                LLMGatewayKeys: []LLMGatewayKeySpec{
                        {JSONPointer: "/apiKeyHelper", ValueField: "TokenHelperCommand"},
                        {JSONPointer: "/env/ANTHROPIC_BASE_URL", ValueField: "AnthropicBaseURL"},
                        // Pin the apiKeyHelper re-invocation cadence so it stays below the token
                        // source's preemptive refresh window — together they guarantee the helper
                        // never returns an about-to-expire token.
                        {JSONPointer: "/env/CLAUDE_CODE_API_KEY_HELPER_TTL_MS", ValueField: "ClaudeCodeHelperTTLMillis"},
                        // Populate the /model picker from the gateway's GET /v1/models on each
                        // startup, so models the gateway adds appear without re-running setup.
                        {JSONPointer: "/env/CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", Literal: "1"},
                        // NODE_TLS_REJECT_UNAUTHORIZED is only written when --tls-skip-verify is set.
                        // ClearWhenEmpty ensures it is removed when the flag is later cleared.
                        {JSONPointer: "/env/NODE_TLS_REJECT_UNAUTHORIZED", ValueField: "NodeTLSRejectUnauthorized", ClearWhenEmpty: true},
                        // Current Claude Code versions expose separate controls for the main
                        // conversation and auxiliary requests. ENABLE_PROMPT_CACHING_1H is the
                        // fallback for versions that predate those per-bucket settings.
                        {JSONPointer: "/promptCacheTtl", ValueField: "PromptCacheTTL", ClearWhenEmpty: true},
                        {JSONPointer: "/subagentPromptCacheTtl", ValueField: "PromptCacheTTL", ClearWhenEmpty: true},
                        {JSONPointer: "/env/ENABLE_PROMPT_CACHING_1H", ValueField: "PromptCache1HLegacy", ClearWhenEmpty: true},
                        // Bedrock-compat keys (written only with --bedrock-compat). Bedrock rejects
                        // Claude Code's experimental anthropic-beta headers, so betas are disabled;
                        // the per-tier model IDs pin Bedrock inference-profile IDs. All use
                        // ClearWhenEmpty so a non-Bedrock setup removes any stale keys.
                        {
                                JSONPointer:    "/env/CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",
                                ValueField:     "BedrockDisableExperimentalBetas",
                                ClearWhenEmpty: true,
                        },
                        {JSONPointer: "/env/ANTHROPIC_DEFAULT_HAIKU_MODEL", ValueField: "BedrockHaikuModel", ClearWhenEmpty: true},
                        {JSONPointer: "/env/ANTHROPIC_DEFAULT_OPUS_MODEL", ValueField: "BedrockOpusModel", ClearWhenEmpty: true},
                        {JSONPointer: "/env/ANTHROPIC_DEFAULT_SONNET_MODEL", ValueField: "BedrockSonnetModel", ClearWhenEmpty: true},
                        {JSONPointer: "/env/ANTHROPIC_DEFAULT_FABLE_MODEL", ValueField: "BedrockFableModel", ClearWhenEmpty: true},
                },
        },
        {
                ClientType:           Windsurf,
                Description:          "Windsurf IDE",
                SettingsFile:         "mcp_config.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".codeium", "windsurf"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "serverUrl",
                        types.TransportTypeSSE:            "serverUrl",
                        types.TransportTypeStreamableHTTP: "serverUrl",
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".codeium", "windsurf", skillsDirName},
                SkillsProjectPath: []string{".windsurf", skillsDirName},
        },
        {
                ClientType:           WindsurfJetBrains,
                Description:          "Windsurf plugin for JetBrains IDEs",
                SettingsFile:         "mcp_config.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".codeium"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "serverUrl",
                        types.TransportTypeSSE:            "serverUrl",
                        types.TransportTypeStreamableHTTP: "serverUrl",
                },
        },
        {
                ClientType:           AmpCli,
                Description:          "Sourcegraph Amp CLI",
                SettingsFile:         "settings.json",
                MCPServersPathPrefix: "/amp.mcpServers",
                RelPath:              []string{"amp"},
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {".config"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                Extension: JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".agents", skillsDirName},
                SkillsProjectPath: []string{".agents", skillsDirName},
        },
        {
                ClientType:           LMStudio,
                Description:          "LM Studio application",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".lmstudio"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
        },
        {
                ClientType:           Goose,
                Description:          "Goose AI agent",
                SettingsFile:         "config.yaml",
                MCPServersPathPrefix: "/extensions",
                RelPath:              []string{"goose"},
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {".config"},
                        PlatformWindows: {"AppData", "Block"},
                },
                Extension: YAML,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "streamable_http",
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: "streamable_http",
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "uri",
                        types.TransportTypeSSE:            "uri",
                        types.TransportTypeStreamableHTTP: "uri",
                },
                // YAML configuration
                YAMLStorageType: YAMLStorageTypeMap,
                YAMLDefaults: map[string]interface{}{
                        "enabled":     true,
                        "timeout":     60,
                        "description": "",
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".agents", skillsDirName},
                SkillsProjectPath: []string{".agents", skillsDirName},
        },
        {
                ClientType:           Trae,
                Description:          "Trae IDE",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{"Trae", "User"},
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                Extension: JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: false,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".agents", skillsDirName},
                SkillsProjectPath: []string{".agents", skillsDirName},
        },
        {
                ClientType:           Continue,
                Description:          "Continue.dev IDE plugins",
                SettingsFile:         "config.yaml",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".continue"},
                Extension:            YAML,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "streamable-http",
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: "streamable-http",
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                // YAML configuration
                YAMLStorageType:     YAMLStorageTypeArray,
                YAMLIdentifierField: "name",
        },
        {
                ClientType:           OpenCode,
                Description:          "OpenCode editor",
                SettingsFile:         "opencode.json",
                MCPServersPathPrefix: "/mcp",
                RelPath:              []string{".config", "opencode"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "remote", // OpenCode requires "type": "remote" for URL-based servers
                        types.TransportTypeSSE:            "remote",
                        types.TransportTypeStreamableHTTP: "remote",
                },
                IsTransportTypeFieldSupported: true, // OpenCode requires "type": "remote" for URL-based servers
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{"opencode", skillsDirName},
                SkillsProjectPath: []string{".opencode", skillsDirName},
                SkillsPlatformPrefix: map[Platform][]string{
                        PlatformLinux:  {".config"},
                        PlatformDarwin: {".config"},
                },
        },
        {
                ClientType:           Kiro,
                Description:          "Kiro AI IDE",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".kiro", "settings"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: false,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".kiro", skillsDirName},
                SkillsProjectPath: []string{".kiro", skillsDirName},
        },
        {
                ClientType:                    Antigravity,
                Description:                   "Google Antigravity IDE",
                SettingsFile:                  "mcp_config.json",
                MCPServersPathPrefix:          "/mcpServers",
                RelPath:                       []string{".gemini", "antigravity"},
                Extension:                     JSON,
                IsTransportTypeFieldSupported: false,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "serverUrl",
                        types.TransportTypeSSE:            "serverUrl",
                        types.TransportTypeStreamableHTTP: "serverUrl",
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".agents", skillsDirName},
                SkillsProjectPath: []string{".agents", skillsDirName},
        },
        {
                ClientType:           Zed,
                Description:          "Zed editor",
                SettingsFile:         "settings.json",
                MCPServersPathPrefix: "/context_servers",
                RelPath:              []string{"zed"},
                PlatformPrefix: map[Platform][]string{
                        PlatformLinux:   {".config"},
                        PlatformDarwin:  {".config"},
                        PlatformWindows: {"AppData", "Roaming"},
                },
                Extension:                     JSON,
                IsTransportTypeFieldSupported: false,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
        },
        {
                ClientType:           GeminiCli,
                Description:          "Google Gemini CLI",
                SettingsFile:         "settings.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".gemini"},
                Extension:            JSON,
                // Gemini CLI uses different URL fields based on transport type:
                // - SSE transport uses "url" field
                // - Streamable HTTP transport uses "httpUrl" field
                IsTransportTypeFieldSupported: false,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "httpUrl",
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: "httpUrl",
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".agents", skillsDirName},
                SkillsProjectPath: []string{".agents", skillsDirName},
                // LLM gateway: patches the same settings.json used for MCP.
                // Gemini CLI has no dynamic token-command equivalent, so it uses the
                // proxy path. GOOGLE_GEMINI_BASE_URL is only honoured when
                // security.auth.selectedType is "gemini-api-key" (fixed in
                // gemini-cli v0.40.0, PR #25357); OAuth auth ignores the override.
                //
                // NODE_TLS_REJECT_UNAUTHORIZED is intentionally omitted: in proxy mode
                // the tool connects to the local proxy over plain HTTP, so there is no
                // TLS handshake on that leg. Setting the env var would globally disable
                // TLS verification for all other HTTPS requests the Gemini CLI process
                // makes, which is an unacceptable side-effect.
                LLMGatewayMode:     llmgateway.ModeProxy,
                LLMBinaryName:      "gemini",
                LLMSettingsFile:    "settings.json",
                LLMSettingsRelPath: []string{".gemini"},
                LLMGatewayKeys: []LLMGatewayKeySpec{
                        // Force API-key auth so GOOGLE_GEMINI_BASE_URL is respected.
                        {JSONPointer: "/security/auth/selectedType", Literal: "gemini-api-key"},
                },
                // Gemini CLI reads GEMINI_API_KEY and GOOGLE_GEMINI_BASE_URL from
                // process.env (not from settings.json), so they are injected via .env.
                LLMEnvFileRelPath: []string{".gemini"},
                LLMEnvFileName:    ".env",
                LLMEnvFileKeys: []LLMEnvFileKeySpec{
                        {Name: "GEMINI_API_KEY", Literal: llmPlaceholderAPIKey},
                        {Name: "GOOGLE_GEMINI_BASE_URL", ValueField: "ProxyOrigin"},
                },
        },
        {
                ClientType:   VSCodeServer,
                Description:  "Microsoft's VS Code Server (remote development)",
                SettingsFile: "mcp.json",
                RelPath: []string{
                        ".vscode-server", "data", "User",
                },
                MCPServersPathPrefix: "/servers",
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
        },
        {
                ClientType:           MistralVibe,
                Description:          "Mistral Vibe IDE",
                SettingsFile:         "config.toml",
                MCPServersPathPrefix: "/mcp_servers",
                RelPath:              []string{".vibe"},
                Extension:            TOML,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          "streamable-http",
                        types.TransportTypeSSE:            httpTransportLabel,
                        types.TransportTypeStreamableHTTP: "streamable-http",
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                // TOML configuration: uses array-of-tables format [[mcp_servers]]
                TOMLStorageType:   TOMLStorageTypeArray,
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".vibe", skillsDirName},
                SkillsProjectPath: []string{".vibe", skillsDirName},
        },
        {
                ClientType:           Codex,
                Description:          "OpenAI Codex CLI",
                SettingsFile:         "config.toml",
                MCPServersPathPrefix: "/mcp_servers",
                RelPath:              []string{".codex"},
                Extension:            TOML,
                // Codex doesn't support a transport type field - it auto-detects
                IsTransportTypeFieldSupported: false,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                // TOML configuration: uses nested tables format [mcp_servers.servername]
                TOMLStorageType:   TOMLStorageTypeMap,
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".agents", skillsDirName},
                SkillsProjectPath: []string{".agents", skillsDirName},
                SupportsPlugins:   true,
                // Codex discovers marketplaces at ~/.agents/plugins/marketplace.json
                // (personal) and $REPO_ROOT/.agents/plugins/marketplace.json (project).
                // ToolHive lays each plugin's source under that marketplace root,
                // namespaced by the "toolhive" marketplace name, so the manifest can
                // reference it with a relative "./toolhive/<name>" source.
                PluginsGlobalPath:  []string{".agents", "plugins", "toolhive"},
                PluginsProjectPath: []string{".agents", "plugins", "toolhive"},
                // LLM gateway: patches the same config.toml its MCP config uses, via a
                // dedicated TOML writer (llm_gateway_codex.go) rather than LLMGatewayKeys.
                LLMGatewayMode:     llmgateway.ModeCodexAuth,
                LLMBinaryName:      "codex",
                LLMSettingsFile:    "config.toml",
                LLMSettingsRelPath: []string{".codex"},
        },
        {
                ClientType:           KimiCli,
                Description:          "Kimi Code CLI",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".kimi"},
                Extension:            JSON,
                // Kimi CLI does not use a transport type field in the config file
                IsTransportTypeFieldSupported: false,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".kimi", skillsDirName},
                SkillsProjectPath: []string{".kimi", skillsDirName},
        },
        {
                ClientType:           Factory,
                Description:          "Factory.ai Droid CLI",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".factory"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".factory", skillsDirName},
                SkillsProjectPath: []string{".factory", skillsDirName},
        },
        {
                ClientType:           CopilotCli,
                Description:          "GitHub Copilot CLI",
                SettingsFile:         "mcp-config.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".copilot"},
                Extension:            JSON,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
        },
        {
                ClientType:           Qoder,
                Description:          "Qoder IDE",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{},
                PlatformPrefix: map[Platform][]string{
                        PlatformDarwin:  {".qoder"},
                        PlatformLinux:   {".qoder"},
                        PlatformWindows: {"AppData", "Roaming", "Qoder", "SharedClientCache"},
                },
                Extension:                     JSON,
                IsTransportTypeFieldSupported: true,
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeSSE:            "sse",
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                SupportsSkills:    true,
                SkillsGlobalPath:  []string{".qoder", skillsDirName},
                SkillsProjectPath: []string{".qoder", skillsDirName},
        },
        {
                ClientType:           Pi,
                Description:          "Pi coding agent",
                SettingsFile:         "mcp.json",
                MCPServersPathPrefix: "/mcpServers",
                RelPath:              []string{".pi", "agent"},
                Extension:            JSON,
                // Pi rejects the legacy "sse" type, so SSE is intentionally unmapped
                // and no type field is written for SSE servers.
                SupportedTransportTypesMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          httpTransportLabel,
                        types.TransportTypeStreamableHTTP: httpTransportLabel,
                },
                IsTransportTypeFieldSupported: true,
                MCPServersUrlLabelMap: map[types.TransportType]string{
                        types.TransportTypeStdio:          defaultURLFieldName,
                        types.TransportTypeSSE:            defaultURLFieldName,
                        types.TransportTypeStreamableHTTP: defaultURLFieldName,
                },
                ConfigDirEnvVar:    "PI_CODING_AGENT_DIR",
                SupportsSkills:     true,
                SkillsGlobalPath:   []string{".pi", "agent", skillsDirName},
                SkillsProjectPath:  []string{".pi", skillsDirName},
                LLMGatewayMode:     llmgateway.ModePi,
                LLMBinaryName:      "pi",
                LLMSettingsFile:    "models.json",
                LLMSettingsRelPath: []string{".pi", "agent"},
        },
        {
                // Xcode does not support MCP; it is an LLM-gateway-only entry.
                // Cast LLMClientApp → ClientApp for internal config storage; the type
                // distinction matters only for swag enum generation (see LLMClientApp).
                ClientType:     ClientApp(Xcode),
                Description:    "GitHub Copilot for Xcode",
                LLMGatewayOnly: true,
                LLMGatewayMode: llmgateway.ModeProxy,
                // Full path is macOS-specific; on Linux/Windows this directory will not
                // exist, so DetectedLLMGatewayClients() naturally returns false there.
                LLMSettingsFile:    "editorSettings.json",
                LLMSettingsRelPath: []string{"Library", "Application Support", "GitHub Copilot for Xcode"},
                LLMGatewayKeys: []LLMGatewayKeySpec{
                        {JSONPointer: "/openAIBaseURL", ValueField: "ProxyBaseURL"},
                        {JSONPointer: "/apiKey", ValueField: "PlaceholderAPIKey"},
                },
        },
        {
                // Claude Desktop routes LLM traffic through the gateway via its
                // "third-party inference" surface. Unlike Claude Code (a single JSON
                // settings file), Desktop uses a configLibrary directory: one config
                // document per saved config plus a _meta.json selector naming the active
                // one. That document model does not fit JSON-key patching, so this entry
                // uses LLMGatewayMode "credential-helper" and carries no LLMGatewayKeys —
                // the dedicated credential-helper writer handles it. LLM-gateway-only for
                // workload registration; remote MCP servers are added to the same document
                // by "thv client connector" instead. Cast LLMClientApp → ClientApp for
                // internal storage.
                ClientType:     ClientApp(ClaudeDesktop),
                Description:    "Claude Desktop",
                LLMGatewayOnly: true,
                LLMGatewayMode: llmgateway.ModeCredentialHelper,
                // Settings file is the _meta.json selector; the writer derives the
                // containing configLibrary directory from it.
                LLMSettingsFile:    "_meta.json",
                LLMSettingsRelPath: []string{"Claude-3p", "configLibrary"},
                LLMSettingsPlatformPrefix: map[Platform][]string{
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Local"},
                        PlatformLinux:   {".config"},
                },
                // Detect via the app's user-data directory, which exists once Claude
                // Desktop has run — the configLibrary directory above does not exist
                // until first configured, so it cannot be the detection signal.
                LLMDetectRelPath: []string{"Claude"},
                LLMDetectPlatformPrefix: map[Platform][]string{
                        PlatformDarwin:  {"Library", "Application Support"},
                        PlatformWindows: {"AppData", "Roaming"},
                        PlatformLinux:   {".config"},
                },
                LLMManagedProfileDomain:   "com.anthropic.claudefordesktop.plist",
                LLMManagedSettingsRelPath: []string{"claude-desktop", "managed-settings.json"},
        },
}

// GetAllClients returns a slice of all supported MCP client types, sorted alphabetically.
// This is the single source of truth for valid client types.
func GetAllClients() []ClientApp {
        clients := make([]ClientApp, 0, len(supportedClientIntegrations))
        for _, config := range supportedClientIntegrations {
                if !config.LLMGatewayOnly {
                        clients = append(clients, config.ClientType)
                }
        }
        // Sort alphabetically
        sort.Slice(clients, func(i, j int) bool {
                return clients[i] < clients[j]
        })
        return clients
}

// IsValidClient checks if the provided client type is supported.
func IsValidClient(clientType string) bool {
        for _, config := range supportedClientIntegrations {
                if string(config.ClientType) == clientType && !config.LLMGatewayOnly {
                        return true
                }
        }
        return false
}

// GetClientDescription returns the description for a given client type.
// Returns an empty string if the client type is not found.
func GetClientDescription(clientType ClientApp) string {
        for _, config := range supportedClientIntegrations {
                if config.ClientType == clientType {
                        return config.Description
                }
        }
        return ""
}

// GetClientDeprecation returns the deprecation warning message for a client type
// and whether the client is deprecated. The message is empty when the client is
// not deprecated.
func GetClientDeprecation(clientType ClientApp) (string, bool) {
        for _, config := range supportedClientIntegrations {
                if config.ClientType == clientType {
                        return config.DeprecationMessage, config.Deprecated
                }
        }
        return "", false
}

// GetClientListFormatted returns a formatted multi-line string listing all supported clients
// with their descriptions, sorted alphabetically. This is suitable for use in CLI help text.
func GetClientListFormatted() string {
        // Create a sorted copy of MCP-capable configurations (exclude LLM-gateway-only entries).
        var configs []clientAppConfig
        for _, cfg := range supportedClientIntegrations {
                if !cfg.LLMGatewayOnly {
                        configs = append(configs, cfg)
                }
        }
        sort.Slice(configs, func(i, j int) bool {
                return configs[i].ClientType < configs[j].ClientType
        })

        var sb strings.Builder
        for _, config := range configs {
                description := config.Description
                if config.Deprecated {
                        description += " (deprecated)"
                }
                fmt.Fprintf(&sb, "  - %s: %s\n", config.ClientType, description)
        }
        return strings.TrimSuffix(sb.String(), "\n")
}

// GetClientListCSV returns a comma-separated list of all supported client types, sorted alphabetically.
// This is suitable for use in error messages.
func GetClientListCSV() string {
        clients := GetAllClients() // GetAllClients already returns sorted list
        clientStrs := make([]string, len(clients))
        for i, client := range clients {
                clientStrs[i] = string(client)
        }
        return strings.Join(clientStrs, ", ")
}

// ConfigFile represents a client configuration file
type ConfigFile struct {
        Path          string
        ClientType    ClientApp
        ConfigUpdater ConfigUpdater
        Extension     Extension
}

// FindClientConfig returns the client configuration file for a given client type.
func FindClientConfig(clientType ClientApp) (*ConfigFile, error) {
        manager, err := NewClientManager()
        if err != nil {
                return nil, err
        }
        return manager.FindClientConfig(clientType)
}

// FindClientConfig returns the client configuration file for a given client type using this manager's dependencies.
func (cm *ClientManager) FindClientConfig(clientType ClientApp) (*ConfigFile, error) {
        // retrieve the metadata of the config files
        configFile, err := cm.retrieveConfigFileMetadata(clientType)
        if err != nil {
                if errors.Is(err, ErrConfigFileNotFound) {
                        // Propagate the error if the file is not found
                        return nil, fmt.Errorf("%w for client %s", ErrConfigFileNotFound, clientType)
                }
                return nil, err
        }

        // validate the format of the config files
        err = validateConfigFileFormat(configFile)
        if err != nil {
                return nil, fmt.Errorf("failed to validate config file format: %w", err)
        }
        return configFile, nil
}

// FindRegisteredClientConfigs finds all registered client configs using this manager's dependencies
func (cm *ClientManager) FindRegisteredClientConfigs(ctx context.Context) ([]ConfigFile, error) {
        clientStatuses, err := cm.GetClientStatus(ctx)
        if err != nil {
                return nil, fmt.Errorf("failed to get client status: %w", err)
        }

        var configFiles []ConfigFile
        for _, clientStatus := range clientStatuses {
                if !clientStatus.Installed || !clientStatus.Registered {
                        continue
                }
                cf, err := cm.FindClientConfig(clientStatus.ClientType)
                if err != nil {
                        if errors.Is(err, ErrConfigFileNotFound) {
                                slog.Debug("client config file not found, creating", "client", clientStatus.ClientType)
                                cf, err = cm.CreateClientConfig(clientStatus.ClientType)
                                if err != nil {
                                        slog.Warn("unable to create client config", "client", clientStatus.ClientType, "error", err)
                                        continue
                                }
                                slog.Debug("successfully created client config file", "client", clientStatus.ClientType)
                        } else {
                                slog.Warn("unable to process client config", "client", clientStatus.ClientType, "error", err)
                                continue
                        }
                }
                configFiles = append(configFiles, *cf)
        }

        return configFiles, nil
}

// CreateClientConfig creates a new client configuration file for a given client type.
func CreateClientConfig(clientType ClientApp) (*ConfigFile, error) {
        manager, err := NewClientManager()
        if err != nil {
                return nil, err
        }
        return manager.CreateClientConfig(clientType)
}

// CreateClientConfig creates a new client configuration file for a given client type using this manager's dependencies.
func (cm *ClientManager) CreateClientConfig(clientType ClientApp) (*ConfigFile, error) {
        // Find the configuration for the requested client type
        clientCfg := cm.lookupClientAppConfig(clientType)
        if clientCfg == nil {
                return nil, fmt.Errorf("%w: %s", ErrUnsupportedClientType, clientType)
        }
        if clientCfg.LLMGatewayOnly {
                return nil, fmt.Errorf("%w: %s does not support MCP configuration", ErrUnsupportedClientType, clientType)
        }

        // Build the path to the configuration file
        path := cm.buildMCPConfigPath(clientCfg)

        // Validate that the file does not already exist
        if _, err := os.Stat(path); !os.IsNotExist(err) {
                return nil, fmt.Errorf("client config file already exists at %s", path)
        }

        // Create the file if it does not exist
        slog.Debug("creating new client config file", "path", path)

        // Create parent directories if they don't exist
        parentDir := filepath.Dir(path)
        if err := os.MkdirAll(parentDir, 0700); err != nil {
                return nil, fmt.Errorf("failed to create parent directories for %s: %w", path, err)
        }

        var initialContent []byte
        switch clientCfg.Extension {
        case YAML, TOML:
                // For YAML and TOML files, create an empty file - the updater will initialize structure as needed
                initialContent = []byte("")
        case JSON:
                // JSON files get empty object
                initialContent = []byte("{}")
        }

        err := os.WriteFile(path, initialContent, 0600)
        if err != nil {
                return nil, fmt.Errorf("failed to create client config file: %w", err)
        }

        return cm.FindClientConfig(clientType)
}

// Upsert updates/inserts an MCP server in a client configuration file
// It is a wrapper around the ConfigUpdater.Upsert method. Because the
// ConfigUpdater is different for each client type, we need to handle
// the different types of McpServer objects. For example, VSCode and ClaudeCode allows
// for a `type` field, but Cursor and others do not. This allows us to
// build up more complex MCP server configurations for different clients
// without leaking them into the CMD layer.
func Upsert(cf ConfigFile, name string, url string, transportType string) error {
        manager, err := NewClientManager()
        if err != nil {
                return err
        }
        return manager.Upsert(cf, name, url, transportType)
}

// Upsert updates/inserts an MCP server in a client configuration file using this manager's dependencies.
func (cm *ClientManager) Upsert(cf ConfigFile, name string, url string, transportType string) error {
        cfg := cm.lookupClientAppConfig(cf.ClientType)
        if cfg == nil {
                return nil
        }
        server := buildMCPServer(url, transportType, cfg)
        return cf.ConfigUpdater.Upsert(name, server)
}

// buildMCPServer constructs an MCPServer struct with the appropriate URL field and optional type field.
// The URL field name is determined by looking up the transport type in MCPServersUrlLabelMap.
// If the map is nil or the transport type is not found, it falls back to "url" as the default.
// For most clients, all transport types map to the same URL field (e.g., "url"), but some clients
// like Gemini CLI use different URL fields per transport type (e.g., "url" for SSE, "httpUrl" for streamable HTTP).
func buildMCPServer(url, transportType string, clientCfg *clientAppConfig) MCPServer {
        server := MCPServer{}

        // Determine the URL field name from the transport type using MCPServersUrlLabelMap
        urlFieldName := defaultURLFieldName // default fallback
        if clientCfg.MCPServersUrlLabelMap != nil {
                if mappedUrlField, ok := clientCfg.MCPServersUrlLabelMap[types.TransportType(transportType)]; ok {
                        urlFieldName = mappedUrlField
                }
        }

        // Set the URL in the appropriate field
        switch urlFieldName {
        case "serverUrl":
                server.ServerUrl = url
        case "httpUrl":
                server.HttpUrl = url
        case "uri":
                server.Uri = url
        default:
                server.Url = url
        }

        // Add transport type field if supported by the client
        if clientCfg.IsTransportTypeFieldSupported {
                if mappedType, ok := clientCfg.SupportedTransportTypesMap[types.TransportType(transportType)]; ok {
                        server.Type = mappedType
                }
        }

        return server
}

// retrieveConfigFileMetadata retrieves the metadata for client configuration files using this manager's dependencies.
func (cm *ClientManager) retrieveConfigFileMetadata(clientType ClientApp) (*ConfigFile, error) {
        // Find the configuration for the requested client type
        clientCfg := cm.lookupClientAppConfig(clientType)
        if clientCfg == nil {
                return nil, fmt.Errorf("%w: %s", ErrUnsupportedClientType, clientType)
        }

        // Build the path to the configuration file
        path := cm.buildMCPConfigPath(clientCfg)

        // Validate that the file exists
        if err := validateConfigFileExists(path); err != nil {
                return nil, err
        }

        // Create a config updater for this file based on the extension
        var configUpdater ConfigUpdater
        switch clientCfg.Extension {
        case YAML:
                // Use the generic YAML converter with configuration from clientAppConfig
                converter := NewGenericYAMLConverter(clientCfg)
                configUpdater = &YAMLConfigUpdater{
                        Path:      path,
                        Converter: converter,
                }
        case TOML:
                serversKey := extractServersKeyFromConfig(clientCfg)
                urlLabel := extractURLLabelFromConfig(clientCfg)

                // Choose TOML updater based on storage type
                if clientCfg.TOMLStorageType == TOMLStorageTypeMap {
                        // Use map-based format [section.servername] (e.g., Codex)
                        configUpdater = &TOMLMapConfigUpdater{
                                Path:       path,
                                ServersKey: serversKey,
                                URLField:   urlLabel,
                        }
                } else {
                        // Default to array-of-tables format [[section]] (e.g., Mistral Vibe)
                        configUpdater = &TOMLConfigUpdater{
                                Path:            path,
                                ServersKey:      serversKey,
                                IdentifierField: "name", // TOML configs use "name" as the identifier
                                URLField:        urlLabel,
                        }
                }
        case JSON:
                configUpdater = &JSONConfigUpdater{
                        Path:                 path,
                        MCPServersPathPrefix: clientCfg.MCPServersPathPrefix,
                }
        }

        // Return the configuration file metadata
        return &ConfigFile{
                Path:          path,
                ConfigUpdater: configUpdater,
                ClientType:    clientCfg.ClientType,
                Extension:     clientCfg.Extension,
        }, nil
}

// buildMCPConfigPath resolves the absolute path to the client's MCP settings
// file, honoring cfg.ConfigDirEnvVar when it is set.
func (cm *ClientManager) buildMCPConfigPath(cfg *clientAppConfig) string {
        if dir := cm.configDirOverride(cfg); dir != "" {
                return filepath.Join(dir, cfg.SettingsFile)
        }
        return buildConfigFilePath(cfg.SettingsFile, cfg.RelPath, cfg.PlatformPrefix, []string{cm.homeDir})
}

// configDirOverride returns the absolute directory named by cfg.ConfigDirEnvVar,
// or "" when the client has no override variable or it is unset.
func (cm *ClientManager) configDirOverride(cfg *clientAppConfig) string {
        if cfg.ConfigDirEnvVar == "" {
                return ""
        }
        dir := os.Getenv(cfg.ConfigDirEnvVar)
        switch {
        case dir == "":
                return ""
        case dir == "~":
                dir = cm.homeDir
        case strings.HasPrefix(dir, "~/"):
                dir = filepath.Join(cm.homeDir, dir[2:])
        }
        if !filepath.IsAbs(dir) {
                if absolute, err := filepath.Abs(dir); err == nil {
                        dir = absolute
                }
        }
        return dir
}

func buildConfigFilePath(settingsFile string, relPath []string, platformPrefix map[Platform][]string, path []string) string {
        if prefix, ok := platformPrefix[Platform(runtime.GOOS)]; ok {
                path = append(path, prefix...)
        }
        path = append(path, relPath...)
        path = append(path, settingsFile)
        return filepath.Clean(filepath.Join(path...))
}

// validateConfigFileExists validates that a client configuration file exists.
func validateConfigFileExists(path string) error {
        if _, err := os.Stat(path); os.IsNotExist(err) {
                return ErrConfigFileNotFound
        }
        return nil
}

func validateConfigFileFormat(cf *ConfigFile) error {
        data, err := os.ReadFile(cf.Path)
        if err != nil {
                return fmt.Errorf("failed to read file %s: %w", cf.Path, err)
        }

        // For YAML and TOML files, empty content is valid
        // For JSON files, default to empty object if the file is empty
        if len(data) == 0 {
                switch cf.Extension {
                case YAML, TOML:
                        return nil // Empty YAML/TOML files are valid
                case JSON:
                        data = []byte("{}") // Default to an empty JSON object
                }
        }

        switch cf.Extension {
        case YAML:
                var temp any
                err = yaml.Unmarshal(data, &temp)
                if err != nil {
                        return fmt.Errorf("failed to parse YAML for file %s: %w", cf.Path, err)
                }
        case TOML:
                var temp any
                err = toml.Unmarshal(data, &temp)
                if err != nil {
                        return fmt.Errorf("failed to parse TOML for file %s: %w", cf.Path, err)
                }
        case JSON:
                _, err = hujson.Parse(data)
                if err != nil {
                        return fmt.Errorf("failed to parse JSON for file %s: %w", cf.Path, err)
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// config_editor.go provides ConfigUpdater implementations for editing MCP client
// configuration files in JSON, YAML, and TOML formats.
//
// # Error Handling
//
// All ConfigUpdater methods (Upsert/Remove) return errors to their callers rather
// than handling them internally. This design allows callers to decide the appropriate
// action based on context:
//
//   - CLI commands (e.g., "thv client register"): Errors propagate up to Cobra's
//     RunE function, which prints the error to stderr and exits with code 1.
//     This is the correct behavior for explicit user commands.
//
//   - Background operations (e.g., RemoveServerFromClients during workload cleanup):
//     Callers log errors as warnings and continue processing other clients.
//     This allows partial success when some clients fail.
//
//   - Migrations: Errors are logged as warnings and the migration continues,
//     allowing best-effort migration of client configurations.
//
// Write failures are logged at WARN level (not ERROR) because:
//  1. The error is also returned to the caller who decides the severity
//  2. Many callers (RemoveServerFromClients, migrations) treat these as non-fatal
//  3. This avoids misleading ERROR logs for expected failure scenarios
//
// # File Locking
//
// All operations use file-based locking via fileutils.WithFileLock() to ensure safe concurrent
// access. Each config file has a corresponding ".lock" file that is acquired before
// any read-modify-write operation.

package client

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "os"
        "strings"

        "github.com/pelletier/go-toml/v2"
        "github.com/tailscale/hujson"
        "github.com/tidwall/gjson"
        "go.yaml.in/yaml/v3"

        "github.com/stacklok/toolhive/pkg/fileutils"
)

// ConfigUpdater defines the interface for types which can edit MCP client config files.
// All methods return errors rather than handling them internally, allowing callers to
// determine the appropriate response (fatal error, warning, or ignore) based on context.
// See the package-level documentation for details on error handling patterns.
type ConfigUpdater interface {
        // Upsert inserts or updates an MCP server configuration.
        // Returns an error if the operation fails (file read/write, parsing, marshaling).
        Upsert(serverName string, data MCPServer) error

        // Remove removes an MCP server configuration.
        // Returns nil if the server doesn't exist (idempotent).
        // Returns an error only for actual failures (file read/write, parsing).
        Remove(serverName string) error
}

// MCPServer represents an MCP server in a MCP client config file
type MCPServer struct {
        Url       string `json:"url,omitempty"`
        ServerUrl string `json:"serverUrl,omitempty"`
        HttpUrl   string `json:"httpUrl,omitempty"`
        Uri       string `json:"uri,omitempty"`
        Type      string `json:"type,omitempty"`
}

// --- Shared helper functions ---

// JSONConfigUpdater is a ConfigUpdater that is responsible for updating
// JSON config files.
type JSONConfigUpdater struct {
        Path                 string
        MCPServersPathPrefix string
}

// Upsert inserts or updates an MCP server in the MCP client config file
func (jcu *JSONConfigUpdater) Upsert(serverName string, data MCPServer) error {
        return fileutils.WithFileLock(jcu.Path, func() error {
                content, err := os.ReadFile(jcu.Path)
                if err != nil && !os.IsNotExist(err) {
                        return fmt.Errorf("failed to read file: %w", err)
                }

                if len(content) == 0 {
                        // If the file is empty, we need to initialize it with an empty JSON object
                        content = []byte("{}")
                }

                content = ensurePathExists(content, jcu.MCPServersPathPrefix)

                v, err := hujson.Parse(content)
                if err != nil {
                        return fmt.Errorf("failed to parse JSON: %w", err)
                }

                dataJSON, err := json.Marshal(data)
                if err != nil {
                        return fmt.Errorf("failed to marshal MCPServer to JSON: %w", err)
                }

                patch := fmt.Sprintf(`[{ "op": "add", "path": "%s/%s", "value": %s } ]`, jcu.MCPServersPathPrefix, serverName, dataJSON)
                if err := v.Patch([]byte(patch)); err != nil {
                        return fmt.Errorf("failed to patch JSON: %w", err)
                }

                formatted, err := hujson.Format(v.Pack())
                if err != nil {
                        return fmt.Errorf("failed to format JSON: %w", err)
                }

                // Write back to the file atomically
                if err := fileutils.AtomicWriteFile(jcu.Path, formatted, 0600); err != nil {
                        slog.Warn("failed to write JSON config file", "error", err)
                        return fmt.Errorf("failed to write file: %w", err)
                }

                slog.Debug("successfully updated client config file", "server", serverName)
                return nil
        })
}

// Remove removes an MCP server from the MCP client config file
func (jcu *JSONConfigUpdater) Remove(serverName string) error {
        return fileutils.WithFileLock(jcu.Path, func() error {
                content, err := os.ReadFile(jcu.Path)
                if err != nil {
                        if os.IsNotExist(err) {
                                // File doesn't exist, nothing to remove
                                return nil
                        }
                        return fmt.Errorf("failed to read file: %w", err)
                }

                if len(content) == 0 {
                        // If the file is empty, there is nothing to remove.
                        return nil
                }

                v, err := hujson.Parse(content)
                if err != nil {
                        return fmt.Errorf("failed to parse JSON: %w", err)
                }

                // Check if the server exists by attempting the patch and handling the error gracefully
                patch := fmt.Sprintf(`[{ "op": "remove", "path": "%s/%s" } ]`, jcu.MCPServersPathPrefix, serverName)
                if err := v.Patch([]byte(patch)); err != nil {
                        // If the patch fails because the path doesn't exist, that's fine - nothing to remove
                        if strings.Contains(err.Error(), "value not found") || strings.Contains(err.Error(), "path not found") {
                                slog.Debug("mcpserver not found in client config file, nothing to remove", "server", serverName)
                                return nil
                        }
                        // For other errors, return the error
                        return fmt.Errorf("failed to patch JSON: %w", err)
                }

                formatted, err := hujson.Format(v.Pack())
                if err != nil {
                        return fmt.Errorf("failed to format JSON: %w", err)
                }

                // Write back to the file atomically
                if err := fileutils.AtomicWriteFile(jcu.Path, formatted, 0600); err != nil {
                        slog.Warn("failed to write JSON config file", "error", err)
                        return fmt.Errorf("failed to write file: %w", err)
                }

                slog.Debug("successfully removed mcpserver from client config file", "server", serverName)
                return nil
        })
}

// YAMLConfigUpdater is a ConfigUpdater that is responsible for updating
// YAML config files using a converter interface for flexibility.
type YAMLConfigUpdater struct {
        Path      string
        Converter YAMLConverter
}

// Upsert inserts or updates an MCP server in the config.yaml file using the converter
func (ycu *YAMLConfigUpdater) Upsert(serverName string, data MCPServer) error {
        return fileutils.WithFileLock(ycu.Path, func() error {
                content, err := os.ReadFile(ycu.Path)
                if err != nil && !os.IsNotExist(err) {
                        return fmt.Errorf("failed to read file: %w", err)
                }

                // Use a generic map to preserve all existing fields, not just extensions
                var config map[string]any

                // If file exists and is not empty, unmarshal existing config into generic map
                if len(content) > 0 {
                        if err := yaml.Unmarshal(content, &config); err != nil {
                                return fmt.Errorf("failed to parse existing YAML config: %w", err)
                        }
                } else {
                        // Initialize empty map if file doesn't exist or is empty
                        config = make(map[string]any)
                }

                // Convert MCPServer using the converter
                entry, err := ycu.Converter.ConvertFromMCPServer(serverName, data)
                if err != nil {
                        return fmt.Errorf("failed to convert MCPServer: %w", err)
                }

                // Upsert the entry using the converter
                if err := ycu.Converter.UpsertEntry(config, serverName, entry); err != nil {
                        return fmt.Errorf("failed to upsert entry: %w", err)
                }

                // Marshal back to YAML
                updatedContent, err := yaml.Marshal(config)
                if err != nil {
                        return fmt.Errorf("failed to marshal YAML: %w", err)
                }

                // Write back to file atomically
                if err := fileutils.AtomicWriteFile(ycu.Path, updatedContent, 0600); err != nil {
                        slog.Warn("failed to write YAML config file", "error", err)
                        return fmt.Errorf("failed to write file: %w", err)
                }

                slog.Debug("successfully updated YAML client config file", "server", serverName)
                return nil
        })
}

// Remove removes an entry from the config.yaml file using the converter
func (ycu *YAMLConfigUpdater) Remove(serverName string) error {
        return fileutils.WithFileLock(ycu.Path, func() error {
                // Read existing config
                content, err := os.ReadFile(ycu.Path)
                if err != nil {
                        if os.IsNotExist(err) {
                                // File doesn't exist, nothing to remove
                                return nil
                        }
                        return fmt.Errorf("failed to read file: %w", err)
                }

                if len(content) == 0 {
                        // File is empty, nothing to remove
                        return nil
                }

                // Use a generic map to preserve all existing fields, not just extensions
                var config map[string]any
                if err := yaml.Unmarshal(content, &config); err != nil {
                        return fmt.Errorf("failed to parse YAML: %w", err)
                }

                if err := ycu.Converter.RemoveEntry(config, serverName); err != nil {
                        return fmt.Errorf("failed to remove entry: %w", err)
                }

                updatedContent, err := yaml.Marshal(config)
                if err != nil {
                        return fmt.Errorf("failed to marshal YAML: %w", err)
                }

                // Write back to file atomically
                if err := fileutils.AtomicWriteFile(ycu.Path, updatedContent, 0600); err != nil {
                        slog.Warn("failed to write YAML config file", "error", err)
                        return fmt.Errorf("failed to write file: %w", err)
                }

                slog.Debug("successfully removed server from YAML config file", "server", serverName)
                return nil
        })
}

// --- Shared TOML helper functions ---

// readTOMLConfig reads and parses a TOML config file from the specified path.
func readTOMLConfig(path string) (map[string]any, error) {
        // #nosec G304 -- path is controlled by internal code (TOMLConfigUpdater/TOMLMapConfigUpdater structs)
        content, err := os.ReadFile(path)
        if err != nil && !os.IsNotExist(err) {
                return nil, fmt.Errorf("failed to read file: %w", err)
        }

        if len(content) == 0 {
                return make(map[string]any), nil
        }

        var config map[string]any
        if err := toml.Unmarshal(content, &config); err != nil {
                return nil, fmt.Errorf("failed to parse existing TOML config: %w", err)
        }
        return config, nil
}

// writeTOMLConfig marshals and writes the config to the specified TOML file path atomically.
func writeTOMLConfig(path string, config map[string]any) error {
        updatedContent, err := toml.Marshal(config)
        if err != nil {
                return fmt.Errorf("failed to marshal TOML: %w", err)
        }
        if err := fileutils.AtomicWriteFile(path, updatedContent, 0600); err != nil {
                slog.Warn("failed to write TOML config file", "error", err)
                return fmt.Errorf("failed to write file: %w", err)
        }
        return nil
}

// extractURLFromMCPServer extracts the URL value from an MCPServer struct,
// checking fields in priority order based on client specificity:
// Uri (Goose) → ServerUrl (Windsurf) → HttpUrl (Gemini) → Url (default/most common).
func extractURLFromMCPServer(data MCPServer) string {
        switch {
        case data.Uri != "":
                return data.Uri
        case data.ServerUrl != "":
                return data.ServerUrl
        case data.HttpUrl != "":
                return data.HttpUrl
        case data.Url != "":
                return data.Url
        default:
                return ""
        }
}

// convertToAnySlice converts various slice types to []any.
func convertToAnySlice(v any) []any {
        switch s := v.(type) {
        case []any:
                return s
        case []map[string]any:
                result := make([]any, len(s))
                for i, item := range s {
                        result[i] = item
                }
                return result
        default:
                return nil
        }
}

// --- TOMLConfigUpdater (array-of-tables format) ---

// TOMLConfigUpdater is a ConfigUpdater that is responsible for updating
// TOML config files with array-of-tables format (used by Mistral Vibe).
type TOMLConfigUpdater struct {
        Path            string
        ServersKey      string // The TOML array key (e.g., "mcp_servers")
        IdentifierField string // The field name used to identify servers (e.g., "name")
        URLField        string // The field name for URL (e.g., "url")
}

// Upsert inserts or updates an MCP server in the TOML config file
func (tcu *TOMLConfigUpdater) Upsert(serverName string, data MCPServer) error {
        return fileutils.WithFileLock(tcu.Path, func() error {
                config, err := readTOMLConfig(tcu.Path)
                if err != nil {
                        return err
                }

                servers := tcu.getServersArray(config)
                newEntry := tcu.buildServerEntry(serverName, data)
                servers = tcu.upsertServerEntry(servers, serverName, newEntry)
                config[tcu.ServersKey] = servers

                if err := writeTOMLConfig(tcu.Path, config); err != nil {
                        return err
                }

                slog.Debug("successfully updated TOML client config file", "server", serverName)
                return nil
        })
}

// Remove removes an MCP server from the TOML config file
func (tcu *TOMLConfigUpdater) Remove(serverName string) error {
        return fileutils.WithFileLock(tcu.Path, func() error {
                config, err := readTOMLConfig(tcu.Path)
                if err != nil {
                        return err
                }

                // If config is empty (file didn't exist or was empty), nothing to remove
                if len(config) == 0 {
                        return nil
                }

                existingServers, ok := config[tcu.ServersKey]
                if !ok {
                        return nil // No servers section, nothing to remove
                }

                servers := convertToAnySlice(existingServers)
                if servers == nil {
                        return nil // Unknown format, nothing to remove
                }

                config[tcu.ServersKey] = tcu.filterOutServer(servers, serverName)

                if err := writeTOMLConfig(tcu.Path, config); err != nil {
                        return err
                }

                slog.Debug("successfully removed server from TOML config file", "server", serverName)
                return nil
        })
}

// getServersArray extracts or initializes the servers array from config
func (tcu *TOMLConfigUpdater) getServersArray(config map[string]any) []any {
        existingServers, ok := config[tcu.ServersKey]
        if !ok {
                return []any{}
        }
        servers := convertToAnySlice(existingServers)
        if servers == nil {
                return []any{}
        }
        return servers
}

// upsertServerEntry updates an existing server or appends a new one
func (tcu *TOMLConfigUpdater) upsertServerEntry(servers []any, serverName string, newEntry map[string]any) []any {
        for i, s := range servers {
                if serverEntry, ok := s.(map[string]any); ok {
                        if name, exists := serverEntry[tcu.IdentifierField]; exists && name == serverName {
                                servers[i] = newEntry
                                return servers
                        }
                }
        }
        return append(servers, newEntry)
}

// filterOutServer removes the server with the given name from the slice
func (tcu *TOMLConfigUpdater) filterOutServer(servers []any, serverName string) []any {
        filtered := make([]any, 0, len(servers))
        for _, s := range servers {
                serverEntry, ok := s.(map[string]any)
                if !ok {
                        filtered = append(filtered, s)
                        continue
                }
                name, exists := serverEntry[tcu.IdentifierField]
                if !exists || name != serverName {
                        filtered = append(filtered, s)
                }
        }
        return filtered
}

// buildServerEntry creates a server entry map from MCPServer data
func (tcu *TOMLConfigUpdater) buildServerEntry(serverName string, data MCPServer) map[string]any {
        entry := map[string]any{
                tcu.IdentifierField: serverName,
        }

        if url := extractURLFromMCPServer(data); url != "" {
                entry[tcu.URLField] = url
        }

        // Add transport type if specified
        if data.Type != "" {
                entry["transport"] = data.Type
        }

        return entry
}

// --- TOMLMapConfigUpdater (nested tables format) ---

// TOMLMapConfigUpdater is a ConfigUpdater that is responsible for updating
// TOML config files with nested tables format [section.servername] (used by Codex).
type TOMLMapConfigUpdater struct {
        Path       string
        ServersKey string // The TOML section key (e.g., "mcp_servers")
        URLField   string // The field name for URL (e.g., "url")
}

// Upsert inserts or updates an MCP server in the TOML config file using map format
func (tmu *TOMLMapConfigUpdater) Upsert(serverName string, data MCPServer) error {
        return fileutils.WithFileLock(tmu.Path, func() error {
                config, err := readTOMLConfig(tmu.Path)
                if err != nil {
                        return err
                }

                // Get or create the servers map
                serversMap := tmu.getServersMap(config)

                // Build the server entry (without the name field since it's the key)
                serverEntry := tmu.buildServerEntry(data)

                // Set the server entry
                serversMap[serverName] = serverEntry
                config[tmu.ServersKey] = serversMap

                if err := writeTOMLConfig(tmu.Path, config); err != nil {
                        return err
                }

                slog.Debug("successfully updated TOML client config file", "server", serverName)
                return nil
        })
}

// Remove removes an MCP server from the TOML config file
func (tmu *TOMLMapConfigUpdater) Remove(serverName string) error {
        return fileutils.WithFileLock(tmu.Path, func() error {
                config, err := readTOMLConfig(tmu.Path)
                if err != nil {
                        return err
                }

                // If config is empty (file didn't exist or was empty), nothing to remove
                if len(config) == 0 {
                        return nil
                }

                serversSection, ok := config[tmu.ServersKey]
                if !ok {
                        return nil // No servers section, nothing to remove
                }

                serversMap, ok := serversSection.(map[string]any)
                if !ok {
                        return nil // Unknown format, nothing to remove
                }

                // Remove the server if it exists
                delete(serversMap, serverName)
                config[tmu.ServersKey] = serversMap

                if err := writeTOMLConfig(tmu.Path, config); err != nil {
                        return err
                }

                slog.Debug("successfully removed server from TOML config file", "server", serverName)
                return nil
        })
}

// getServersMap extracts or initializes the servers map from config
func (tmu *TOMLMapConfigUpdater) getServersMap(config map[string]any) map[string]any {
        existingServers, ok := config[tmu.ServersKey]
        if !ok {
                return make(map[string]any)
        }
        serversMap, ok := existingServers.(map[string]any)
        if !ok {
                return make(map[string]any)
        }
        return serversMap
}

// buildServerEntry creates a server entry map from MCPServer data
func (tmu *TOMLMapConfigUpdater) buildServerEntry(data MCPServer) map[string]any {
        entry := make(map[string]any)

        if url := extractURLFromMCPServer(data); url != "" {
                entry[tmu.URLField] = url
        }

        // Add transport type if specified
        if data.Type != "" {
                entry["transport"] = data.Type
        }

        return entry
}

// ensurePathExists ensures that the path exists in the JSON content
// and returns the updated content.
// For example:
//   - if the path is "/mcp/servers",
//     the function will ensure that the path "/mcp/servers" exists
//     and returns the updated content.
//   - if the path is "/mcpServers",
//     the function will ensure that the path "/mcpServers" exists
//     and returns the updated content.
//
// This is necessary because the MCP client config file is a JSON object,
// and we need to ensure that the path exists before we can add a new key to it.
func ensurePathExists(content []byte, path string) []byte {
        // Special case: if path is root ("/"), just return everything (formatted)
        if path == "/" {
                v, _ := hujson.Parse(content)
                formatted, _ := hujson.Format(v.Pack())
                return formatted
        }

        segments := strings.Split(path, "/")

        // Navigate through the JSON structure
        var pathSoFarForPatch string
        var pathSoFarForRetrieval string
        for i, segment := range segments[:] {
                // we want to skip the first segments because it is the root
                if path[0] == '/' && (i == 0) {
                        continue
                }

                // We build the path up to this segment so that we can check if it exists
                // and if it doesn't, we can create it as an empty object.
                // The "/" is added to the path for the patch operation because the path
                // is a JSON pointer, and JSON pointers are prefixed with "/".
                // The "." is added to the path for the retrieval operation.
                // - gjson (used for retrieval) treats `.` as a special (traversal) character,
                // so any json keys which contain `.` must have the `.` "escaped" with a single
                // '\'. In it, key `a.b` would be matched by `a\.b` but not `a.b`.
                // - hujson (used for the patch) treats "." and "\" as ordinary characters in a
                // json key. In it, key `a.b` would be matched by `a.b` but not `a\.b`.
                // So we need to "escape" json keys this way for retrieval, but not for patch.
                if len(pathSoFarForPatch) == 0 {
                        pathSoFarForPatch = "/" + segment
                        pathSoFarForRetrieval = strings.ReplaceAll(segment, ".", `\.`)
                } else {
                        pathSoFarForPatch = pathSoFarForPatch + "/" + segment
                        pathSoFarForRetrieval = pathSoFarForRetrieval + "." + strings.ReplaceAll(segment, ".", `\.`)
                }

                // We retrieve the segment from the content so that we can check if it exists
                // and if it doesn't, we can create it as an empty object. If it does exist,
                // we can skip the patch operation onto the next segment.
                segmentPath := gjson.GetBytes(content, pathSoFarForRetrieval).Raw
                if segmentPath != "" {
                        continue
                }

                // Create a JSON patch to add an empty object at this path
                patch := fmt.Sprintf(`[{ "op": "add", "path": "%s", "value": {} }]`, pathSoFarForPatch)

                // Parse the current content and apply the patch
                v, _ := hujson.Parse(content)
                err := v.Patch([]byte(patch))
                if err != nil {
                        slog.Error("failed to patch file", "error", err)
                }

                // Update the content with the patched version
                content = v.Pack()
        }
        // Parse the updated content with hujson to maintain formatting
        v, _ := hujson.Parse(content)
        formatted, _ := hujson.Format(v.Pack())
        return formatted
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "fmt"
)

// YAMLConverter defines an interface for converting MCPServer data to different YAML config formats
type YAMLConverter interface {
        ConvertFromMCPServer(serverName string, server MCPServer) (interface{}, error)
        UpsertEntry(config interface{}, serverName string, entry interface{}) error
        RemoveEntry(config interface{}, serverName string) error
}

// GenericYAMLConverter implements YAMLConverter using configuration from clientAppConfig
type GenericYAMLConverter struct {
        storageType     YAMLStorageType        // How servers are stored in YAML (map or array)
        serversPath     string                 // path to servers section (e.g., "extensions" or "mcpServers")
        identifierField string                 // for array type: field that identifies the server
        defaults        map[string]interface{} // default values for fields
        urlLabel        string                 // label for URL field (e.g., "url", "uri", "serverUrl")
}

// NewGenericYAMLConverter creates a converter from clientAppConfig
func NewGenericYAMLConverter(cfg *clientAppConfig) *GenericYAMLConverter {
        return &GenericYAMLConverter{
                storageType:     cfg.YAMLStorageType,
                serversPath:     extractServersKeyFromConfig(cfg),
                identifierField: cfg.YAMLIdentifierField,
                defaults:        cfg.YAMLDefaults,
                urlLabel:        extractURLLabelFromConfig(cfg),
        }
}

// ConvertFromMCPServer converts an MCPServer to the appropriate format based on configuration
func (g *GenericYAMLConverter) ConvertFromMCPServer(serverName string, server MCPServer) (interface{}, error) {
        result := make(map[string]interface{})

        // Add name field
        result["name"] = serverName

        // Handle URL field - extract from whichever MCPServer field has a value
        // and use the configured URL label for the output key.
        if url := extractURLFromMCPServer(server); url != "" {
                // Use the configured URL label (e.g., "uri" for Goose, "url" for Continue)
                if g.urlLabel != "" {
                        result[g.urlLabel] = url
                } else {
                        result[defaultURLFieldName] = url // Default fallback
                }
        }

        // Add type field
        if server.Type != "" {
                result["type"] = server.Type
        }

        // Apply defaults (e.g., enabled, timeout for Goose)
        for key, value := range g.defaults {
                if _, exists := result[key]; !exists {
                        result[key] = value
                }
        }

        return result, nil
}

// UpsertEntry adds or updates an entry based on storage type (map or array)
func (g *GenericYAMLConverter) UpsertEntry(config interface{}, serverName string, entry interface{}) error {
        configMap, ok := config.(map[string]interface{})
        if !ok {
                return fmt.Errorf("invalid config format")
        }

        // Initialize servers section if it doesn't exist
        if configMap[g.serversPath] == nil {
                if g.storageType == YAMLStorageTypeMap {
                        configMap[g.serversPath] = make(map[string]interface{})
                } else {
                        configMap[g.serversPath] = []interface{}{}
                }
        }

        // Convert entry to map for YAML marshaling
        entryMap, ok := entry.(map[string]interface{})
        if !ok {
                return fmt.Errorf("entry must be a map[string]interface{}")
        }

        if g.storageType == YAMLStorageTypeMap {
                return g.upsertMapEntry(configMap, serverName, entryMap)
        }
        return g.upsertArrayEntry(configMap, serverName, entryMap)
}

// upsertMapEntry handles map-based storage (like Goose)
func (g *GenericYAMLConverter) upsertMapEntry(
        configMap map[string]interface{}, serverName string, entryMap map[string]interface{},
) error {
        servers, ok := configMap[g.serversPath].(map[string]interface{})
        if !ok {
                servers = make(map[string]interface{})
                configMap[g.serversPath] = servers
        }

        servers[serverName] = entryMap
        return nil
}

// upsertArrayEntry handles array-based storage (like Continue)
func (g *GenericYAMLConverter) upsertArrayEntry(
        configMap map[string]interface{}, serverName string, entryMap map[string]interface{},
) error {
        var servers []interface{}

        // Get the servers array, handling different types
        switch v := configMap[g.serversPath].(type) {
        case []interface{}:
                servers = v
        case []map[string]interface{}:
                servers = make([]interface{}, len(v))
                for i, s := range v {
                        servers[i] = s
                }
        default:
                servers = []interface{}{}
        }

        // Find and update existing entry or append new one
        found := false
        for i, s := range servers {
                if serverEntry, ok := s.(map[string]interface{}); ok {
                        if id, exists := serverEntry[g.identifierField]; exists && id == serverName {
                                servers[i] = entryMap
                                found = true
                                break
                        }
                }
        }

        if !found {
                servers = append(servers, entryMap)
        }

        configMap[g.serversPath] = servers
        return nil
}

// RemoveEntry removes an entry based on storage type
func (g *GenericYAMLConverter) RemoveEntry(config interface{}, serverName string) error {
        configMap, ok := config.(map[string]interface{})
        if !ok {
                return fmt.Errorf("invalid config format")
        }

        if configMap[g.serversPath] == nil {
                return nil // Nothing to remove
        }

        if g.storageType == YAMLStorageTypeMap {
                return g.removeMapEntry(configMap, serverName)
        }
        return g.removeArrayEntry(configMap, serverName)
}

// removeMapEntry handles removal from map-based storage
func (g *GenericYAMLConverter) removeMapEntry(configMap map[string]interface{}, serverName string) error {
        servers, ok := configMap[g.serversPath].(map[string]interface{})
        if !ok {
                return fmt.Errorf("invalid servers format")
        }

        delete(servers, serverName)
        return nil
}

// removeArrayEntry handles removal from array-based storage
func (g *GenericYAMLConverter) removeArrayEntry(configMap map[string]interface{}, serverName string) error {
        var servers []interface{}

        // Get the servers array
        switch v := configMap[g.serversPath].(type) {
        case []interface{}:
                servers = v
        case []map[string]interface{}:
                servers = make([]interface{}, len(v))
                for i, s := range v {
                        servers[i] = s
                }
        default:
                return nil // Nothing to remove
        }

        // Filter out the server with matching identifier
        filtered := make([]interface{}, 0, len(servers))
        for _, s := range servers {
                if serverEntry, ok := s.(map[string]interface{}); ok {
                        if name, exists := serverEntry[g.identifierField]; !exists || name != serverName {
                                filtered = append(filtered, s)
                        }
                } else {
                        filtered = append(filtered, s)
                }
        }

        configMap[g.serversPath] = filtered
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "context"
        "fmt"
        "os"
        "os/exec"
        "path/filepath"
        "runtime"

        "github.com/stacklok/toolhive/pkg/config"
        "github.com/stacklok/toolhive/pkg/groups"
)

// ClientManager encapsulates dependencies for client operations
//
//nolint:revive // ClientManager is intentionally named to avoid conflict with existing Manager interface
type ClientManager struct {
        homeDir            string
        groupManager       groups.Manager
        clientIntegrations []clientAppConfig
        configProvider     config.Provider
        lookPath           func(string) (string, error)
}

// HomeDir returns the home directory this ClientManager is rooted at. Exported
// so plugin adapters can resolve user-scoped config-file paths (e.g. Claude
// Code's ~/.claude/settings.json) under the manager's home — which, for tests,
// is a temp dir injected via NewTestClientManagerWithHome.
func (cm *ClientManager) HomeDir() string {
        return cm.homeDir
}

// NewClientManager creates a new ClientManager with default dependencies
func NewClientManager() (*ClientManager, error) {
        home, err := os.UserHomeDir()
        if err != nil {
                return nil, fmt.Errorf("failed to get home directory: %w", err)
        }

        groupManager, err := groups.NewManager()
        if err != nil {
                // If group manager fails to initialize, we'll just skip group checks
                groupManager = nil
        }

        // Copy the static integration list so we can wire per-client detectors
        // without mutating the shared supportedClientIntegrations slice.
        integrations := make([]clientAppConfig, len(supportedClientIntegrations))
        copy(integrations, supportedClientIntegrations)
        codexDetector := newCodexDesktopDetector(home, runtime.GOOS).installed
        for i := range integrations {
                if integrations[i].ClientType == Codex {
                        integrations[i].LLMInstalledDetector = codexDetector
                        break
                }
        }

        return &ClientManager{
                homeDir:            home,
                groupManager:       groupManager,
                clientIntegrations: integrations,
                configProvider:     config.NewDefaultProvider(),
                lookPath:           exec.LookPath,
        }, nil
}

// NewTestClientManager creates a ClientManager with test dependencies. Native
// Codex desktop detection is intentionally omitted to keep tests deterministic.
func NewTestClientManager(
        homeDir string,
        groupManager groups.Manager,
        clientIntegrations []clientAppConfig,
        configProvider config.Provider,
) *ClientManager {
        return &ClientManager{
                homeDir:            homeDir,
                groupManager:       groupManager,
                clientIntegrations: clientIntegrations,
                configProvider:     configProvider,
                lookPath:           exec.LookPath,
        }
}

// ClientAppStatus represents the status of a supported client application
//
//nolint:revive // ClientAppStatus is intentionally named for clarity across packages
type ClientAppStatus struct {
        // ClientType is the type of MCP client
        ClientType ClientApp `json:"client_type"`

        // Installed indicates whether the client is installed on the system
        Installed bool `json:"installed"`

        // Registered indicates whether the client is registered in the ToolHive configuration
        Registered bool `json:"registered"`

        // SupportsSkills indicates whether ToolHive can install skills for this client
        SupportsSkills bool `json:"supports_skills"`
        // SupportsPlugins indicates whether ToolHive can install plugins for this client
        SupportsPlugins bool `json:"supports_plugins"`
}

// IsClientInstalled reports whether the given client appears to be installed on
// the current system. Detection is based on the presence of the client's
// configuration directory (or settings file when neither a relative path nor a
// platform prefix is defined).
func (cm *ClientManager) IsClientInstalled(clientType ClientApp) bool {
        cfg := cm.lookupClientAppConfig(clientType)
        if cfg == nil || cfg.LLMGatewayOnly {
                return false
        }
        pathToCheck := func() string {
                if dir := cm.configDirOverride(cfg); dir != "" {
                        return dir
                }
                if len(cfg.RelPath) == 0 && len(cfg.PlatformPrefix) == 0 {
                        return filepath.Join(cm.homeDir, cfg.SettingsFile)
                }
                return buildConfigDirectoryPath(cfg.RelPath, cfg.PlatformPrefix, []string{cm.homeDir})
        }()
        _, err := os.Stat(pathToCheck)
        return err == nil
}

// GetClientStatus returns the status of all supported MCP clients using this manager's dependencies
func (cm *ClientManager) GetClientStatus(ctx context.Context) ([]ClientAppStatus, error) {
        var statuses []ClientAppStatus

        // Get app configuration to check for registered clients
        appConfig := cm.configProvider.GetConfig()
        registeredClients := make(map[string]bool)

        // Create a map of registered clients for quick lookup from config
        for _, client := range appConfig.Clients.RegisteredClients {
                registeredClients[client] = true
        }

        // Also check for clients registered in groups if group manager is available
        if cm.groupManager != nil {
                allGroups, err := cm.groupManager.List(ctx)
                if err == nil {
                        // Collect clients from all groups
                        for _, group := range allGroups {
                                for _, clientName := range group.RegisteredClients {
                                        registeredClients[clientName] = true
                                }
                        }
                }
        }

        for _, cfg := range cm.clientIntegrations {
                if cfg.LLMGatewayOnly {
                        continue
                }
                status := ClientAppStatus{
                        ClientType:      cfg.ClientType,
                        Installed:       cm.IsClientInstalled(cfg.ClientType),
                        Registered:      registeredClients[string(cfg.ClientType)],
                        SupportsSkills:  cfg.SupportsSkills,
                        SupportsPlugins: cfg.SupportsPlugins,
                }
                statuses = append(statuses, status)
        }

        return statuses, nil
}

// GetClientStatus returns the status of all supported MCP clients using the default config provider
func GetClientStatus(ctx context.Context) ([]ClientAppStatus, error) {
        manager, err := NewClientManager()
        if err != nil {
                return nil, err
        }
        return manager.GetClientStatus(ctx)
}

func buildConfigDirectoryPath(relPath []string, platformPrefix map[Platform][]string, path []string) string {
        if prefix, ok := platformPrefix[Platform(runtime.GOOS)]; ok {
                path = append(path, prefix...)
        }
        path = append(path, relPath...)
        return filepath.Clean(filepath.Join(path...))
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "bytes"
        "fmt"
        "os"
        "path/filepath"
        "strings"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/llmgateway"
)

// ConfigureEnvFile writes the client's LLM env-file entries to the nominated
// dotenv file (e.g. ~/.gemini/.env). Existing entries for other keys are
// preserved; entries managed by thv are added or updated.
//
// Returns the absolute path of the written file, or ("", nil) when the client
// has no env-file entries configured.
func (cm *ClientManager) ConfigureEnvFile(clientType ClientApp, cfg llmgateway.ApplyConfig) (string, error) {
        appCfg := cm.lookupClientAppConfig(clientType)
        if appCfg == nil || len(appCfg.LLMEnvFileKeys) == 0 {
                return "", nil
        }

        path := cm.buildEnvFilePath(appCfg)

        if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
                return "", fmt.Errorf("creating directory for %s: %w", path, err)
        }

        err := fileutils.WithFileLock(path, func() error {
                content, err := readOrInitFile(path, nil)
                if err != nil {
                        return err
                }

                entries := parseEnvFile(content)

                for _, spec := range appCfg.LLMEnvFileKeys {
                        val, err := envFileValueForSpec(spec, cfg)
                        if err != nil {
                                return err
                        }
                        entries = setEnvEntry(entries, spec.Name, val)
                }

                return fileutils.AtomicWriteFile(path, marshalEnvFile(entries), 0o600)
        })
        if err != nil {
                return "", err
        }
        return path, nil
}

// RevertEnvFile removes the entries that ConfigureEnvFile wrote for clientType
// from the nominated dotenv file. Other entries in the file are preserved.
// If the file does not exist the call is a no-op.
func (cm *ClientManager) RevertEnvFile(clientType ClientApp, envFilePath string) error {
        appCfg := cm.lookupClientAppConfig(clientType)
        if appCfg == nil || len(appCfg.LLMEnvFileKeys) == 0 {
                return nil
        }
        if envFilePath == "" {
                return nil
        }

        if _, err := os.Stat(envFilePath); os.IsNotExist(err) {
                return nil
        }

        return fileutils.WithFileLock(envFilePath, func() error {
                content, err := os.ReadFile(envFilePath) // #nosec G304 -- path is a known tool config file
                if err != nil {
                        if os.IsNotExist(err) {
                                return nil
                        }
                        return fmt.Errorf("reading %s: %w", envFilePath, err)
                }

                entries := parseEnvFile(content)
                for _, spec := range appCfg.LLMEnvFileKeys {
                        entries = removeEnvEntry(entries, spec.Name)
                }

                return fileutils.AtomicWriteFile(envFilePath, marshalEnvFile(entries), 0o600)
        })
}

// buildEnvFilePath constructs the absolute path to the .env file for the
// given client, using the same home-dir convention as LLM settings paths.
func (cm *ClientManager) buildEnvFilePath(cfg *clientAppConfig) string {
        return buildConfigFilePath(
                cfg.LLMEnvFileName,
                cfg.LLMEnvFileRelPath,
                nil, // no platform prefix for env files
                []string{cm.homeDir},
        )
}

// HasEnvFileSupport reports whether clientType has .env file entries to manage.
func (cm *ClientManager) HasEnvFileSupport(clientType ClientApp) bool {
        cfg := cm.lookupClientAppConfig(clientType)
        return cfg != nil && len(cfg.LLMEnvFileKeys) > 0
}

// envFileValueForSpec resolves the value to write for a single LLMEnvFileKeySpec.
// Exactly one of spec.Literal or spec.ValueField must be set.
func envFileValueForSpec(spec LLMEnvFileKeySpec, cfg llmgateway.ApplyConfig) (string, error) {
        if spec.Literal != "" && spec.ValueField != "" {
                return "", fmt.Errorf("LLMEnvFileKeySpec for %q has both Literal and ValueField set; exactly one must be set", spec.Name)
        }
        if spec.Literal == "" && spec.ValueField == "" {
                return "", fmt.Errorf("LLMEnvFileKeySpec for %q has neither Literal nor ValueField set; exactly one must be set", spec.Name)
        }
        if spec.Literal != "" {
                return spec.Literal, nil
        }
        if v, ok := resolveApplyConfigField(spec.ValueField, cfg); ok {
                return v, nil
        }
        return "", fmt.Errorf("unknown ValueField %q in LLMEnvFileKeySpec for %q; use Literal for constant values",
                spec.ValueField, spec.Name)
}

// envEntry is an ordered KEY=value entry in a dotenv file. Comments and blank
// lines are stored verbatim so that round-tripping preserves formatting.
type envEntry struct {
        raw   string // raw line for comments / blank lines (when key == "")
        key   string // non-empty for KEY=value lines
        value string
}

// parseEnvFile parses the content of a dotenv file into an ordered list of
// entries. Comments, blank lines, and malformed lines (no "=") are stored as
// raw entries so they survive a round-trip through marshalEnvFile.
func parseEnvFile(content []byte) []envEntry {
        var entries []envEntry
        // bytes.SplitSeq handles arbitrarily long lines since content is already in
        // memory; bufio.Scanner would error on lines longer than its 64 KiB limit.
        // Trim trailing newlines first so an empty or newline-only file produces no
        // entries, and to avoid a spurious empty token after the final newline.
        trimmed := bytes.TrimRight(content, "\r\n")
        if len(trimmed) == 0 {
                return entries
        }
        for rawLine := range bytes.SplitSeq(trimmed, []byte("\n")) {
                // Strip \r to handle CRLF files.
                line := string(bytes.TrimRight(rawLine, "\r"))
                if line == "" || strings.HasPrefix(line, "#") {
                        entries = append(entries, envEntry{raw: line})
                        continue
                }
                key, value, found := strings.Cut(line, "=")
                if !found {
                        // Malformed line — preserve as raw.
                        entries = append(entries, envEntry{raw: line})
                        continue
                }
                entries = append(entries, envEntry{key: key, value: value})
        }
        return entries
}

// marshalEnvFile serialises entries back to a dotenv-formatted byte slice.
// A trailing newline is always appended.
func marshalEnvFile(entries []envEntry) []byte {
        var b bytes.Buffer
        for _, e := range entries {
                if e.key == "" {
                        b.WriteString(e.raw)
                } else {
                        b.WriteString(e.key)
                        b.WriteByte('=')
                        b.WriteString(e.value)
                }
                b.WriteByte('\n')
        }
        return b.Bytes()
}

// setEnvEntry updates the value for key in entries, or appends a new entry if
// key is not present. Returns the (possibly grown) slice.
func setEnvEntry(entries []envEntry, key, value string) []envEntry {
        for i, e := range entries {
                if e.key == key {
                        entries[i].value = value
                        return entries
                }
        }
        return append(entries, envEntry{key: key, value: value})
}

// removeEnvEntry removes the entry with the given key from entries (if any).
func removeEnvEntry(entries []envEntry, key string) []envEntry {
        result := entries[:0:len(entries)]
        for _, e := range entries {
                if e.key != key {
                        result = append(result, e)
                }
        }
        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "errors"
        "slices"

        "github.com/stacklok/toolhive/pkg/groups"
)

// ErrAllClientsRegistered is returned when all available clients are already
// registered for the selected groups.
var ErrAllClientsRegistered = errors.New("all installed clients are already registered for the selected groups")

// FilterClientsAlreadyRegistered returns only clients that are NOT already
// registered in all of the provided groups. A client is excluded only when
// every group in selectedGroups already lists it in RegisteredClients.
func FilterClientsAlreadyRegistered(
        clients []ClientAppStatus,
        selectedGroups []*groups.Group,
) []ClientAppStatus {
        if len(selectedGroups) == 0 {
                return clients
        }

        var filtered []ClientAppStatus
        for _, cli := range clients {
                if !isClientRegisteredInAllGroups(string(cli.ClientType), selectedGroups) {
                        filtered = append(filtered, cli)
                }
        }
        return filtered
}

func isClientRegisteredInAllGroups(clientName string, selectedGroups []*groups.Group) bool {
        for _, group := range selectedGroups {
                if !slices.Contains(group.RegisteredClients, clientName) {
                        return false
                }
        }
        return true
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "runtime"
        "strconv"
        "strings"

        "github.com/tailscale/hujson"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/llmgateway"
)

// llmPlaceholderAPIKey is the static API key written into proxy-mode tool
// configurations. The localhost reverse proxy accepts any non-empty value.
const llmPlaceholderAPIKey = "thv-proxy"

// llmPatchOp is a single RFC 6902 JSON Patch operation, marshaled via
// encoding/json so all string fields are properly escaped.
type llmPatchOp struct {
        Op    string          `json:"op"`
        Path  string          `json:"path"`
        Value json.RawMessage `json:"value,omitempty"`
}

// ConfigureLLMGateway patches the tool's LLM-gateway settings file with cfg
// and returns the absolute path of the patched file.
//
// It uses fileutils.WithFileLock so concurrent calls (e.g. two "thv llm setup"
// invocations) are serialised. Comments and trailing commas in JSONC settings
// files are preserved via hujson. Writes are crash-safe via AtomicWriteFile.
func (cm *ClientManager) ConfigureLLMGateway(clientType ClientApp, cfg llmgateway.ApplyConfig) (string, error) {
        appCfg := cm.lookupClientAppConfig(clientType)
        if appCfg == nil || appCfg.LLMGatewayMode == "" {
                return "", fmt.Errorf("client %q does not support LLM gateway configuration", clientType)
        }

        // Credential-helper clients (Claude Desktop) use a document + selector model
        // that does not fit JSON-key patching; dispatch to the dedicated writer.
        if appCfg.LLMGatewayMode == llmgateway.ModeCredentialHelper {
                return cm.configureCredentialHelper(appCfg, cfg)
        }
        // Codex's config is TOML, not JSON, and authenticates via a command-backed
        // bearer token rather than JSON-Pointer keys; dispatch to its writer.
        if appCfg.LLMGatewayMode == llmgateway.ModeCodexAuth {
                return cm.configureCodexAuth(appCfg, cfg)
        }
        if appCfg.LLMGatewayMode == llmgateway.ModePi {
                return cm.configurePi(appCfg, cfg)
        }
        if appCfg.LLMGatewayMode == llmgateway.ModeVSCode {
                return cm.configureVSCode(appCfg, cfg)
        }

        path := cm.buildLLMSettingsPath(appCfg)

        if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
                return "", fmt.Errorf("creating directory for %s: %w", path, err)
        }

        err := fileutils.WithFileLock(path, func() error {
                content, err := readOrInitFile(path, []byte("{}"))
                if err != nil {
                        return err
                }

                // Parse with hujson first so that JSONC (comments, trailing commas) is
                // handled correctly for all subsequent operations.
                v, err := parseJSONC(content, path)
                if err != nil {
                        return err
                }

                if err := applyLLMGatewayKeys(&v, appCfg.LLMGatewayKeys, cfg, path); err != nil {
                        return err
                }

                formatted, err := hujson.Format(v.Pack())
                if err != nil {
                        return fmt.Errorf("formatting %s: %w", path, err)
                }
                return fileutils.AtomicWriteFile(path, formatted, 0o600)
        })
        if err != nil {
                return "", err
        }
        return path, nil
}

// applyLLMGatewayKeys writes or removes each key spec into v according to cfg.
// Specs with ClearWhenEmpty=true are removed when their resolved value is empty,
// allowing conditional keys (e.g. NODE_TLS_REJECT_UNAUTHORIZED) to be cleaned
// up when the associated flag is cleared.
func applyLLMGatewayKeys(v *hujson.Value, specs []LLMGatewayKeySpec, cfg llmgateway.ApplyConfig, filePath string) error {
        // Resolve all values up front so an unknown ValueField fails fast before
        // any file is modified.
        values := make([]string, len(specs))
        for i, spec := range specs {
                if spec.Literal != "" {
                        values[i] = spec.Literal
                } else {
                        v, err := llmValueForSpec(spec.ValueField, cfg)
                        if err != nil {
                                return err
                        }
                        values[i] = v
                }
        }

        // Ensure ancestors only for specs that will be written (not removed).
        for i, spec := range specs {
                if spec.ClearWhenEmpty && values[i] == "" {
                        continue
                }
                if err := ensureLLMAncestors(v, spec.JSONPointer, filePath); err != nil {
                        return err
                }
        }

        // Standardize once for existence checks in the remove path.
        standardized, err := standardizeJSONC(v, filePath)
        if err != nil {
                return err
        }

        for i, spec := range specs {
                value := values[i]
                if spec.ClearWhenEmpty && value == "" {
                        if err := removeLLMKey(v, spec.JSONPointer, filePath, standardized); err != nil {
                                return err
                        }
                        continue
                }
                if err := addLLMKey(v, spec.JSONPointer, value, filePath); err != nil {
                        return err
                }
        }
        return nil
}

// removeLLMKey removes the key at ptr from v if it exists. standardized is
// pre-computed hujson.Standardize output used for the existence check.
func removeLLMKey(v *hujson.Value, ptr, filePath string, standardized []byte) error {
        if !jsonPointerExists(standardized, ptr) {
                return nil
        }
        patchDoc, err := json.Marshal([]llmPatchOp{{Op: "remove", Path: ptr}})
        if err != nil {
                return fmt.Errorf("marshaling remove patch for %s: %w", ptr, err)
        }
        if err := v.Patch(patchDoc); err != nil {
                return fmt.Errorf("removing %s from %s: %w", ptr, filePath, err)
        }
        return nil
}

// addLLMKey writes value to the key at ptr inside v.
func addLLMKey(v *hujson.Value, ptr, value, filePath string) error {
        valueJSON, err := json.Marshal(value)
        if err != nil {
                return fmt.Errorf("marshaling value for %s: %w", ptr, err)
        }
        patchDoc, err := json.Marshal([]llmPatchOp{{Op: "add", Path: ptr, Value: valueJSON}})
        if err != nil {
                return fmt.Errorf("marshaling patch for %s: %w", ptr, err)
        }
        if err := v.Patch(patchDoc); err != nil {
                return fmt.Errorf("patching %s in %s: %w", ptr, filePath, err)
        }
        return nil
}

// RevertLLMGateway removes the LLM gateway keys from the tool's settings file.
// If the file does not exist the call is a no-op. Comments and trailing commas
// in JSONC settings files are preserved.
func (cm *ClientManager) RevertLLMGateway(clientType ClientApp, configPath string) error {
        appCfg := cm.lookupClientAppConfig(clientType)
        if appCfg == nil || appCfg.LLMGatewayMode == "" {
                return fmt.Errorf("client %q does not support LLM gateway configuration", clientType)
        }

        // Credential-helper clients (Claude Desktop) revert via the dedicated writer.
        if appCfg.LLMGatewayMode == llmgateway.ModeCredentialHelper {
                return cm.revertCredentialHelper(appCfg, configPath)
        }
        // Codex reverts via its dedicated TOML writer.
        if appCfg.LLMGatewayMode == llmgateway.ModeCodexAuth {
                return cm.revertCodexAuth(appCfg, configPath)
        }
        if appCfg.LLMGatewayMode == llmgateway.ModePi {
                return cm.revertPi(appCfg, configPath)
        }
        if appCfg.LLMGatewayMode == llmgateway.ModeVSCode {
                return cm.revertVSCode(appCfg, configPath)
        }

        return revertJSONPointerGateway(appCfg, configPath)
}

// revertJSONPointerGateway removes appCfg.LLMGatewayKeys from configPath, the
// JSON-Pointer-based revert path shared by every LLM-gateway mode except
// ModeCredentialHelper and ModeCodexAuth (which use dedicated writers).
func revertJSONPointerGateway(appCfg *clientAppConfig, configPath string) error {
        // Guard against a missing file (or deleted parent directory) before trying
        // to acquire the lock — WithFileLock creates configPath+".lock", which
        // fails when the directory no longer exists.
        if _, err := os.Stat(configPath); os.IsNotExist(err) {
                return nil
        }

        return fileutils.WithFileLock(configPath, func() error {
                content, err := os.ReadFile(configPath) // #nosec G304 -- path is caller-supplied config file
                if err != nil {
                        if os.IsNotExist(err) {
                                return nil // file removed between stat and lock acquisition
                        }
                        return fmt.Errorf("reading %s: %w", configPath, err)
                }
                if len(content) == 0 {
                        return nil
                }

                v, err := parseJSONC(content, configPath)
                if err != nil {
                        return err
                }

                // Standardize once for all existence checks below.
                standardized, err := standardizeJSONC(&v, configPath)
                if err != nil {
                        return err
                }

                for _, spec := range appCfg.LLMGatewayKeys {
                        // Skip keys that are already absent — avoids brittle error-string matching.
                        if !jsonPointerExists(standardized, spec.JSONPointer) {
                                continue
                        }
                        patchDoc, err := json.Marshal([]llmPatchOp{{Op: "remove", Path: spec.JSONPointer}})
                        if err != nil {
                                return fmt.Errorf("marshaling patch for %s: %w", spec.JSONPointer, err)
                        }
                        if err := v.Patch(patchDoc); err != nil {
                                return fmt.Errorf("reverting %s from %s: %w", spec.JSONPointer, configPath, err)
                        }
                }

                formatted, err := hujson.Format(v.Pack())
                if err != nil {
                        return fmt.Errorf("formatting %s: %w", configPath, err)
                }
                return fileutils.AtomicWriteFile(configPath, formatted, 0o600)
        })
}

// IsLLMGatewaySupported reports whether clientType has LLM gateway support.
func (cm *ClientManager) IsLLMGatewaySupported(clientType ClientApp) bool {
        cfg := cm.lookupClientAppConfig(clientType)
        return cfg != nil && cfg.LLMGatewayMode != ""
}

var promptCacheEnvironmentControls = [...]string{
        "FORCE_PROMPT_CACHING_5M",
        "CLAUDE_CODE_PROMPT_CACHE_TTL",
        "CLAUDE_CODE_SUBAGENT_PROMPT_CACHE_TTL",
}

// PromptCacheConflict reports a locally discoverable setting that may override
// ToolHive's one-hour Claude Code prompt-cache configuration.
func (*ClientManager) PromptCacheConflict(clientType ClientApp) (string, error) {
        if clientType != ClaudeCode {
                return "", nil
        }
        return promptCacheConflict(claudeCodeManagedSettingsDir(runtime.GOOS), os.LookupEnv)
}

type promptCacheControl struct {
        value  string
        source string
}

func promptCacheConflict(
        managedDir string, lookupEnv func(string) (string, bool),
) (string, error) {
        controls := make(map[string]promptCacheControl, 5)
        settingsPaths, err := managedPromptCacheSettingsPaths(managedDir)
        if err != nil {
                return "", err
        }
        for _, path := range settingsPaths {
                if err := mergePromptCacheControls(path, controls); err != nil {
                        return "", err
                }
        }
        for _, name := range promptCacheEnvironmentControls {
                if value, ok := lookupEnv(name); ok {
                        controls[name] = promptCacheControl{strings.TrimSpace(value), "the process environment"}
                }
        }
        return promptCacheConflictDescription(controls), nil
}

func managedPromptCacheSettingsPaths(managedDir string) ([]string, error) {
        if managedDir == "" {
                return nil, nil
        }
        paths := []string{filepath.Join(managedDir, "managed-settings.json")}
        dropInDir := filepath.Join(managedDir, "managed-settings.d")
        dropIns, err := os.ReadDir(dropInDir)
        if err != nil {
                if os.IsNotExist(err) {
                        return paths, nil
                }
                return nil, fmt.Errorf("reading Claude Code managed settings drop-in directory: %w", err)
        }
        for _, entry := range dropIns {
                if entry.IsDir() || strings.HasPrefix(entry.Name(), ".") || filepath.Ext(entry.Name()) != ".json" {
                        continue
                }
                paths = append(paths, filepath.Join(dropInDir, entry.Name()))
        }
        return paths, nil
}

func mergePromptCacheControls(path string, controls map[string]promptCacheControl) error {
        content, err := os.ReadFile(path) // #nosec G304 -- paths are registered client settings locations
        if err != nil {
                if os.IsNotExist(err) {
                        return nil
                }
                return fmt.Errorf("reading %s: %w", path, err)
        }
        if len(content) == 0 {
                return nil
        }

        v, err := parseJSONC(content, path)
        if err != nil {
                return err
        }
        standardized, err := standardizeJSONC(&v, path)
        if err != nil {
                return err
        }
        var settings map[string]any
        if err := json.Unmarshal(standardized, &settings); err != nil {
                return fmt.Errorf("decoding %s: %w", path, err)
        }
        env, _ := settings["env"].(map[string]any)
        for _, name := range promptCacheEnvironmentControls {
                value, ok := env[name].(string)
                if !ok {
                        continue
                }
                controls[name] = promptCacheControl{strings.TrimSpace(value), path}
        }
        for _, name := range []string{"promptCacheTtl", "subagentPromptCacheTtl"} {
                if value, ok := settings[name].(string); ok {
                        controls[name] = promptCacheControl{strings.TrimSpace(value), path}
                }
        }
        return nil
}

func promptCacheConflictDescription(controls map[string]promptCacheControl) string {
        if force, ok := controls["FORCE_PROMPT_CACHING_5M"]; ok && force.value == "1" {
                return formatPromptCacheConflict("FORCE_PROMPT_CACHING_5M", force)
        }

        buckets := []struct {
                environmentVariable string
                setting             string
        }{
                {environmentVariable: "CLAUDE_CODE_PROMPT_CACHE_TTL", setting: "promptCacheTtl"},
                {environmentVariable: "CLAUDE_CODE_SUBAGENT_PROMPT_CACHE_TTL", setting: "subagentPromptCacheTtl"},
        }
        for _, bucket := range buckets {
                if envValue, ok := controls[bucket.environmentVariable]; ok {
                        switch strings.ToLower(envValue.value) {
                        case "5m":
                                return formatPromptCacheConflict(bucket.environmentVariable, envValue)
                        case "1h":
                                continue
                        }
                }
                if setting, ok := controls[bucket.setting]; ok && strings.EqualFold(setting.value, "5m") {
                        return formatPromptCacheConflict(bucket.setting, setting)
                }
        }
        return ""
}

func formatPromptCacheConflict(name string, control promptCacheControl) string {
        return fmt.Sprintf("%s=%s in %s", name, control.value, control.source)
}

func parseJSONC(content []byte, path string) (hujson.Value, error) {
        v, err := hujson.Parse(content)
        if err != nil {
                return hujson.Value{}, fmt.Errorf("parsing %s: %w", path, err)
        }
        return v, nil
}

func standardizeJSONC(v *hujson.Value, path string) ([]byte, error) {
        standardized, err := hujson.Standardize(v.Pack())
        if err != nil {
                return nil, fmt.Errorf("standardizing %s: %w", path, err)
        }
        return standardized, nil
}

func claudeCodeManagedSettingsDir(goos string) string {
        switch goos {
        case "darwin":
                return filepath.Join(string(filepath.Separator), "Library", "Application Support", "ClaudeCode")
        case "linux":
                return filepath.Join(string(filepath.Separator), "etc", "claude-code")
        case "windows":
                programFiles := os.Getenv("ProgramFiles")
                if programFiles == "" {
                        programFiles = `C:\Program Files`
                }
                return filepath.Join(programFiles, "ClaudeCode")
        default:
                return ""
        }
}

// IsManaged reports whether a managed configuration source (an MDM profile on
// macOS, a managed settings file on Linux) is present for the given client.
// When true, the client reads config from the managed source and ignores the
// local config "thv llm setup" writes, so setup warns the user.
func (cm *ClientManager) IsManaged(clientType ClientApp) bool {
        cfg := cm.lookupClientAppConfig(clientType)
        return cfg != nil && managedProfilePresent(cfg)
}

// LLMGatewayModeFor returns the client's configured gateway integration mode,
// or "" when the client is unknown or does not support an LLM gateway.
func (cm *ClientManager) LLMGatewayModeFor(clientType ClientApp) string {
        cfg := cm.lookupClientAppConfig(clientType)
        if cfg == nil {
                return ""
        }
        return cfg.LLMGatewayMode
}

// DetectedLLMGatewayClients returns the subset of LLM-gateway-capable clients
// that are installed on this machine. Most clients require their settings
// directory and, when configured, their binary on $PATH. A client with an
// LLMInstalledDetector hook (e.g. Codex's desktop app) is also detected when
// the hook returns true; CLI and desktop evidence still produce one canonical
// result.
//
// Requiring both CLI signals prevents false positives from leftover config
// directories (e.g. ~/.claude or ~/.gemini) after a tool is uninstalled.
// GUI editors leave LLMBinaryName empty so detection matches the settings-dir
// check used by IsClientInstalled; those apps are often absent from $PATH.
func (cm *ClientManager) DetectedLLMGatewayClients() []ClientApp {
        var result []ClientApp
        for i := range cm.clientIntegrations {
                cfg := &cm.clientIntegrations[i]
                if cfg.LLMGatewayMode == "" {
                        continue
                }
                if cm.isClientInstalled(cfg) {
                        result = append(result, cfg.ClientType)
                }
        }
        return result
}

// isClientInstalled reports whether a single LLM-gateway-capable client appears
// to be installed. The shared check (settings/detection directory exists and,
// when configured, the binary is on $PATH) is augmented by an optional
// per-client LLMInstalledDetector hook (e.g. Codex desktop app detection).
func (cm *ClientManager) isClientInstalled(cfg *clientAppConfig) bool {
        if cm.sharedCLIDetection(cfg) {
                return true
        }
        if cfg.LLMInstalledDetector == nil {
                return false
        }
        installed, err := cfg.LLMInstalledDetector()
        if err != nil {
                // Log the detector failure so it's distinguishable from "not installed"
                // rather than silently treating a detection error as absence.
                slog.Warn("LLM client detection hook failed",
                        "client", cfg.ClientType, "error", err)
                return false
        }
        return installed
}

func (cm *ClientManager) llmDetectDir(cfg *clientAppConfig) string {
        if cfg.LLMDetectRelPath != nil {
                return buildConfigFilePath("", cfg.LLMDetectRelPath, cfg.LLMDetectPlatformPrefix, []string{cm.homeDir})
        }
        return filepath.Dir(cm.buildLLMSettingsPath(cfg))
}

// sharedCLIDetection implements the generic CLI detection shared by all
// LLM-gateway-capable clients: the detection directory (or settings directory)
// must exist, and when LLMBinaryName is set the binary must be found on $PATH.
func (cm *ClientManager) sharedCLIDetection(cfg *clientAppConfig) bool {
        // Detection directory: for GUI apps whose settings directory does not
        // exist until first configured (e.g. Claude Desktop's configLibrary),
        // LLMDetectRelPath points at a directory that exists once the app is
        // installed. Otherwise fall back to the settings directory.
        if _, err := os.Stat(cm.llmDetectDir(cfg)); err != nil {
                return false
        }
        if cfg.LLMBinaryName != "" {
                if cm.lookPath == nil {
                        return false
                }
                if _, err := cm.lookPath(cfg.LLMBinaryName); err != nil {
                        return false
                }
        }
        return true
}

// LLMClientDetectionHint explains why clientType was not detected. Empty when
// there is nothing extra to say (unknown client, already installed, or the
// settings directory itself is missing).
func (cm *ClientManager) LLMClientDetectionHint(clientType ClientApp) string {
        cfg := cm.lookupClientAppConfig(clientType)
        if cfg == nil || cfg.LLMGatewayMode == "" || cfg.LLMBinaryName == "" {
                return ""
        }
        if cm.isClientInstalled(cfg) {
                return ""
        }
        if _, err := os.Stat(cm.llmDetectDir(cfg)); err != nil {
                return ""
        }
        return fmt.Sprintf("settings directory exists but %q was not found on PATH", cfg.LLMBinaryName)
}

// buildLLMSettingsPath resolves the absolute path to the LLM settings file
// for the given client using the same PlatformPrefix logic as MCP config paths,
// honoring cfg.ConfigDirEnvVar when it is set.
func (cm *ClientManager) buildLLMSettingsPath(cfg *clientAppConfig) string {
        if dir := cm.configDirOverride(cfg); dir != "" {
                return filepath.Join(dir, cfg.LLMSettingsFile)
        }
        return buildConfigFilePath(
                cfg.LLMSettingsFile,
                cfg.LLMSettingsRelPath,
                cfg.LLMSettingsPlatformPrefix,
                []string{cm.homeDir},
        )
}

// resolveApplyConfigField maps a ValueField name to its runtime value from cfg.
// Returns (value, true) for any recognized field, ("", false) for unknown ones.
// This is the single source of truth for ValueField semantics shared by both
// LLMGatewayKeySpec (settings-file patching) and LLMEnvFileKeySpec (.env injection).
func resolveApplyConfigField(valueField string, cfg llmgateway.ApplyConfig) (string, bool) {
        switch valueField {
        case "GatewayURL":
                return cfg.GatewayURL, true
        case "AnthropicBaseURL":
                // Use the pre-computed Anthropic base URL when available; fall back to
                // GatewayURL so existing configs continue to work without the prefix.
                if cfg.AnthropicBaseURL != "" {
                        return cfg.AnthropicBaseURL, true
                }
                return cfg.GatewayURL, true
        case "ProxyBaseURL":
                return cfg.ProxyBaseURL, true
        case "ProxyOrigin":
                return llmgateway.ProxyOriginOf(cfg.ProxyBaseURL), true
        case "TokenHelperCommand":
                return cfg.TokenHelperCommand, true
        case "PlaceholderAPIKey":
                return llmPlaceholderAPIKey, true
        case "ClaudeCodeHelperTTLMillis":
                // Constant, independent of cfg: how often Claude Code re-invokes the
                // apiKeyHelper. Kept in sync with the token source's preemptive refresh
                // window (see pkg/llmgateway constants).
                return strconv.FormatInt(llmgateway.ClaudeCodeHelperTTL.Milliseconds(), 10), true
        case "NodeTLSRejectUnauthorized":
                if cfg.TLSSkipVerify {
                        return "0", true
                }
                return "", true
        case "PromptCacheTTL":
                if !cfg.ShortPromptCache {
                        return "1h", true
                }
                return "", true
        case "PromptCache1HLegacy":
                if !cfg.ShortPromptCache {
                        return "1", true
                }
                return "", true
        default:
                return resolveBedrockField(valueField, cfg)
        }
}

// resolveBedrockField resolves the Claude Code bedrock-compat ValueFields. The
// values are written only in bedrock-compat mode; otherwise they resolve to ""
// so their ClearWhenEmpty key specs remove the key. Returns (_, false) for any
// non-bedrock field so the caller can report an unknown ValueField.
func resolveBedrockField(valueField string, cfg llmgateway.ApplyConfig) (string, bool) {
        var v string
        switch valueField {
        case "BedrockDisableExperimentalBetas":
                v = "1"
        case "BedrockHaikuModel":
                v = cfg.BedrockHaikuModel
        case "BedrockOpusModel":
                v = cfg.BedrockOpusModel
        case "BedrockSonnetModel":
                v = cfg.BedrockSonnetModel
        case "BedrockFableModel":
                v = cfg.BedrockFableModel
        default:
                return "", false
        }
        if !cfg.BedrockCompat {
                return "", true
        }
        return v, true
}

// llmValueForSpec returns the config value for a settings-file key spec.
// An unrecognised ValueField is a programming error and returns an error so
// typos are caught at call time rather than silently written into user config.
// Use LLMGatewayKeySpec.Literal for constant values.
func llmValueForSpec(valueField string, cfg llmgateway.ApplyConfig) (string, error) {
        if v, ok := resolveApplyConfigField(valueField, cfg); ok {
                return v, nil
        }
        return "", fmt.Errorf("unknown ValueField %q in LLMGatewayKeySpec; use Literal for constant values", valueField)
}

// ensureLLMAncestors walks every ancestor of ptr from root inward and creates
// any missing intermediate object. For example, for "/a/b/c" it ensures "/a"
// and then "/a/b" exist, so the final "add" patch for "/a/b/c" never fails
// because a parent is missing.
//
// Existence is checked against standardized JSON (hujson.Standardize strips
// JSONC comments and trailing commas) so that JSONC input never produces a
// false "missing" result that would cause an RFC 6902 "add" to replace an
// existing object.
func ensureLLMAncestors(v *hujson.Value, ptr, filePath string) error {
        segments := strings.Split(strings.TrimPrefix(ptr, "/"), "/")
        if len(segments) <= 1 {
                return nil // top-level key — no ancestors to create
        }
        // Standardize once for all existence checks in this call.
        standardized, err := standardizeJSONC(v, filePath)
        if err != nil {
                return err
        }

        ancestor := ""
        needsCreate := false
        for _, seg := range segments[:len(segments)-1] {
                ancestor += "/" + seg
                // Once a missing ancestor is found, all deeper paths are also absent
                // (we just created an empty object), so skip further existence checks.
                if !needsCreate && jsonPointerExists(standardized, ancestor) {
                        continue
                }
                needsCreate = true
                patchDoc, err := json.Marshal([]llmPatchOp{{Op: "add", Path: ancestor, Value: json.RawMessage("{}")}})
                if err != nil {
                        return fmt.Errorf("marshaling ancestor patch for %s in %s: %w", ancestor, filePath, err)
                }
                if err := v.Patch(patchDoc); err != nil {
                        return fmt.Errorf("creating ancestor object %s in %s: %w", ancestor, filePath, err)
                }
        }
        return nil
}

// jsonPointerExists reports whether the JSON Pointer path resolves to a value
// in standard (non-JSONC) JSON data.
// data must already be standardized via hujson.Standardize.
func jsonPointerExists(data []byte, pointer string) bool {
        var root any
        if err := json.Unmarshal(data, &root); err != nil {
                return false
        }
        current := root
        for _, seg := range strings.Split(strings.TrimPrefix(pointer, "/"), "/") {
                m, ok := current.(map[string]any)
                if !ok {
                        return false
                }
                current, ok = m[seg]
                if !ok {
                        return false
                }
        }
        return true
}

// readOrInitFile reads path and returns its content.
// When the file is missing or empty it returns defaultContent instead.
// Returns an error only for real IO failures.
func readOrInitFile(path string, defaultContent []byte) ([]byte, error) {
        data, err := os.ReadFile(path) // #nosec G304 -- path is a known tool config file location
        if err != nil {
                if os.IsNotExist(err) {
                        return defaultContent, nil
                }
                return nil, fmt.Errorf("reading %s: %w", path, err)
        }
        if len(data) == 0 {
                return defaultContent, nil
        }
        return data, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "fmt"
        "os"
        "path/filepath"
        "strings"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/llmgateway"
)

// Codex's LLM gateway config lives in the same config.toml its MCP-server
// registration already uses (~/.codex/config.toml). Unlike the JSON-Pointer
// LLMGatewayKeys clients, Codex authenticates via a command-backed bearer
// token ([model_providers.<id>.auth]), so it gets a dedicated TOML writer
// here rather than going through the hujson JSON-Pointer machinery, which
// cannot write TOML at all.

// codexProviderID is the stable table name ToolHive owns under
// [model_providers.<id>]. Reusing a fixed id (rather than minting one per
// run) keeps repeated "thv llm setup" calls idempotent and avoids orphaned
// tables, mirroring the fixed-name approach configureCredentialHelper uses
// for Claude Desktop's _meta.json entry.
const codexProviderID = "toolhive-gateway"

// codexBaseURL returns gatewayURL with a "/v1" suffix, avoiding a double
// "/v1/v1" if the configured gateway URL already ends in one (e.g. a gateway
// mounted behind its own "/v1" path prefix).
func codexBaseURL(gatewayURL string) string {
        base := strings.TrimSuffix(gatewayURL, "/")
        if strings.HasSuffix(base, "/v1") {
                return base
        }
        return base + "/v1"
}

// configureCodexAuth writes (or updates) ToolHive's model_provider entry in
// Codex's config.toml: a custom provider pointed at the LLM gateway,
// authenticated via a command that invokes "thv llm token". Any other
// model_providers entries and the mcp_servers table (used by the unrelated
// MCP-client feature) are left untouched. Returns the config file path, which
// RevertLLMGateway later passes back to revertCodexAuth.
func (cm *ClientManager) configureCodexAuth(appCfg *clientAppConfig, cfg llmgateway.ApplyConfig) (string, error) {
        if cfg.TokenHelperPath == "" || len(cfg.TokenHelperArgs) == 0 {
                return "", fmt.Errorf("codex-auth requires TokenHelperPath and TokenHelperArgs to be set")
        }

        path := cm.buildLLMSettingsPath(appCfg)
        if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
                return "", fmt.Errorf("creating %s: %w", filepath.Dir(path), err)
        }

        err := fileutils.WithFileLock(path, func() error {
                config, err := readTOMLConfig(path)
                if err != nil {
                        return err
                }

                config["model_provider"] = codexProviderID

                // Refuse to overwrite a pre-existing model_providers that isn't a table
                // (malformed file, or a future Codex schema): clobbering it would silently
                // destroy the user's config. The revert path makes the same ok check.
                providers, ok := config["model_providers"].(map[string]any)
                if existing, present := config["model_providers"]; present && !ok {
                        return fmt.Errorf("existing model_providers in %s is not a table (%T); refusing to overwrite", path, existing)
                }
                if providers == nil {
                        providers = map[string]any{}
                }
                providers[codexProviderID] = map[string]any{
                        "name": "ToolHive Gateway",
                        // Codex's Responses-API client appends "/responses" directly to
                        // base_url (the same convention OpenAI's own base_url="…/v1" uses),
                        // so the gateway's "/v1" prefix must be baked in here rather than
                        // left to the caller like AnthropicBaseURL's optional prefix.
                        "base_url": codexBaseURL(cfg.GatewayURL),
                        "wire_api": "responses",
                        "auth": map[string]any{
                                "command": cfg.TokenHelperPath,
                                "args":    cfg.TokenHelperArgs,
                                // How often Codex re-invokes the token helper. Kept in sync with the
                                // TTL the other clients write (see pkg/llmgateway constants) so the
                                // token source's preemptive-refresh invariant holds for Codex too.
                                "refresh_interval_ms": llmgateway.CodexHelperTTL.Milliseconds(),
                        },
                }
                config["model_providers"] = providers

                return writeTOMLConfig(path, config)
        })
        if err != nil {
                return "", err
        }
        return path, nil
}

// revertCodexAuth undoes configureCodexAuth: it removes ToolHive's
// model_providers.<id> table (leaving any other providers untouched) and
// clears model_provider only if it still points at ToolHive's entry — a value
// changed since setup (by the user or another tool) is left alone, mirroring
// revertCredentialHelper's appliedId check for Claude Desktop. A missing file
// is treated as already-reverted.
//
// Revert removes model_provider rather than restoring whatever it was before
// setup: configureCodexAuth does not stash the prior selection, so a user who
// had e.g. model_provider = "openai" before setup is left with none after
// revert (Codex then falls back to its own default).
func (*ClientManager) revertCodexAuth(_ *clientAppConfig, configPath string) error {
        if configPath == "" || !fileExistsAt(configPath) {
                return nil
        }

        return fileutils.WithFileLock(configPath, func() error {
                config, err := readTOMLConfig(configPath)
                if err != nil {
                        return err
                }

                if providers, ok := config["model_providers"].(map[string]any); ok {
                        delete(providers, codexProviderID)
                        if len(providers) == 0 {
                                delete(config, "model_providers")
                        } else {
                                config["model_providers"] = providers
                        }
                }

                if provider, ok := config["model_provider"].(string); ok && provider == codexProviderID {
                        delete(config, "model_provider")
                }

                return writeTOMLConfig(configPath, config)
        })
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "runtime"
        "strings"

        "github.com/google/uuid"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/llmgateway"
)

// Claude Desktop's "third-party inference" config lives in a configLibrary
// directory: one <id>.json document per saved configuration, plus a _meta.json
// selector that names the active one. This is a different shape from the single
// JSON settings file that direct/proxy clients patch, so it gets a dedicated
// writer here rather than going through the LLMGatewayKeys machinery.

const (
        // claudeDesktopManagedEntryName is the stable display name of the config
        // entry ToolHive owns in _meta.json. Reusing a fixed name (rather than a
        // fresh UUID per run) keeps setup idempotent — repeated "thv llm setup"
        // calls update the same entry instead of accumulating orphans.
        claudeDesktopManagedEntryName = "ToolHive Gateway"

        // claudeDesktopCredentialKind selects the credential-helper auth model
        // (a local executable that prints a token), as opposed to "interactive"
        // (Claude Desktop runs its own OIDC flow).
        claudeDesktopCredentialKind = "helper-script"
)

// claudeDesktopConfig is the <id>.json document written into the configLibrary.
// ToolHive owns this file. ManagedMCPServers holds the connectors written by
// "thv client connector add" (see claude_desktop_connectors.go); setup carries
// them over from the previous document rather than managing them itself.
// inferenceModels is omitted when empty so Claude Desktop falls back to
// gateway-side model auto-discovery once the gateway serves it.
type claudeDesktopConfig struct {
        InferenceProvider                   string   `json:"inferenceProvider"`
        InferenceGatewayBaseURL             string   `json:"inferenceGatewayBaseUrl"`
        InferenceGatewayAuthScheme          string   `json:"inferenceGatewayAuthScheme"`
        InferenceCredentialKind             string   `json:"inferenceCredentialKind"`
        InferenceCredentialHelper           string   `json:"inferenceCredentialHelper"`
        InferenceCredentialHelperTtlSec     int      `json:"inferenceCredentialHelperTtlSec"`
        InferenceCredentialHelperTimeoutSec int      `json:"inferenceCredentialHelperTimeoutSec"`
        InferenceModels                     []string `json:"inferenceModels,omitempty"`
        ManagedMCPServers                   []any    `json:"managedMcpServers,omitempty"`
}

// configureCredentialHelper writes (or updates) ToolHive's Claude Desktop
// configuration: it generates the credential-helper shim, writes the config
// document, and points _meta.json at it. It returns the absolute path of the
// config document, which RevertLLMGateway later uses to undo the change.
//
// The _meta.json update is done under a file lock and preserves every entry and
// key ToolHive does not own, so a user's other saved configurations are left
// intact. Writes are crash-safe via AtomicWriteFile.
func (cm *ClientManager) configureCredentialHelper(appCfg *clientAppConfig, cfg llmgateway.ApplyConfig) (string, error) {
        metaPath := cm.buildLLMSettingsPath(appCfg)
        dir := filepath.Dir(metaPath)
        if err := os.MkdirAll(dir, 0o700); err != nil {
                return "", fmt.Errorf("creating %s: %w", dir, err)
        }

        // Reuse the shared AnthropicBaseURL resolution (falls back to GatewayURL) so
        // this stays in sync with direct-mode clients if the rule ever changes.
        baseURL, _ := resolveApplyConfigField("AnthropicBaseURL", cfg)

        // Write the shim, config document, and _meta.json selector all inside the
        // lock so concurrent setup/teardown runs cannot interleave — e.g. one run's
        // failure-cleanup deleting a shim another run's committed config depends on.
        // _meta.json is written last, so a mid-write failure never leaves Claude
        // Desktop pointing at our config; best-effort cleanup then removes the
        // unreferenced files it created.
        var configPath string
        err := fileutils.WithFileLock(metaPath, func() error {
                // Track whether the shim already existed so failure cleanup does not
                // delete a shim an earlier successful setup still depends on.
                shimExisted := fileExistsAt(cm.credentialHelperShimPath())
                shimPath, err := cm.writeCredentialHelperShim(cfg.TokenHelperPath)
                if err != nil {
                        return err
                }
                cleanup := func(cp string) {
                        if !shimExisted {
                                _ = os.Remove(shimPath)
                        }
                        if cp != "" {
                                _ = os.Remove(cp)
                        }
                }

                meta, err := readClaudeDesktopMeta(metaPath)
                if err != nil {
                        cleanup("")
                        return err
                }
                id := metaEntryID(meta, claudeDesktopManagedEntryName)
                if id == "" {
                        id = uuid.NewString()
                }
                // Upsert by name so an existing (possibly malformed) "ToolHive Gateway"
                // entry is corrected in place rather than leaving a duplicate — preserving
                // the idempotency the stable name is there to provide.
                meta["entries"] = upsertMetaEntry(metaEntries(meta), id, claudeDesktopManagedEntryName)
                meta["appliedId"] = id

                cp := filepath.Join(dir, id+".json")
                doc := claudeDesktopConfig{
                        InferenceProvider:                   "gateway",
                        InferenceGatewayBaseURL:             baseURL,
                        InferenceGatewayAuthScheme:          "bearer",
                        InferenceCredentialKind:             claudeDesktopCredentialKind,
                        InferenceCredentialHelper:           shimPath,
                        InferenceCredentialHelperTtlSec:     int(llmgateway.ClaudeDesktopHelperTTL.Seconds()),
                        InferenceCredentialHelperTimeoutSec: int(llmgateway.ClaudeDesktopHelperTimeout.Seconds()),
                        InferenceModels:                     cfg.Models,
                        // Keep connectors across re-runs of setup, which rewrites the document.
                        ManagedMCPServers: previousManagedMCPServers(cp),
                }
                docBytes, err := json.MarshalIndent(doc, "", "  ")
                if err != nil {
                        cleanup("")
                        return fmt.Errorf("encoding Claude Desktop config: %w", err)
                }
                if err := fileutils.AtomicWriteFile(cp, docBytes, 0o600); err != nil {
                        cleanup(cp)
                        return fmt.Errorf("writing %s: %w", cp, err)
                }
                if err := writeClaudeDesktopMeta(metaPath, meta); err != nil {
                        cleanup(cp)
                        return err
                }
                configPath = cp
                return nil
        })
        if err != nil {
                return "", err
        }
        return configPath, nil
}

// revertCredentialHelper undoes configureCredentialHelper: it removes ToolHive's
// entry from _meta.json (leaving other entries untouched), clears appliedId when
// it pointed at our config, deletes the config document, and removes the shim.
// A missing file at any step is treated as already-reverted.
//
// Ordering matters: the selector (_meta.json) is de-referenced and the config
// document deleted BEFORE the shim, so a mid-revert failure never leaves Claude
// Desktop pointing at a config that references a missing helper executable. Any
// leftover file after a partial failure is unreferenced and harmless.
func (cm *ClientManager) revertCredentialHelper(appCfg *clientAppConfig, configPath string) error {
        // Nothing recorded to revert. Do NOT touch the shim here: without the config
        // path we cannot confirm _meta.json no longer references it, and removing it
        // could break a still-applied config.
        if configPath == "" {
                return nil
        }
        // Derive the configLibrary dir from appCfg (not from configPath) so a
        // tampered stored configPath cannot redirect teardown at an arbitrary
        // directory. The config document must live inside this dir.
        metaPath := cm.buildLLMSettingsPath(appCfg)
        dir := filepath.Dir(metaPath)
        if !fileExistsAt(dir) {
                // Config directory already gone — nothing to revert.
                return nil
        }

        id := metaIDFromConfigPath(configPath)
        if !isSafeConfigID(id) {
                // Symmetric with configureCredentialHelper: a corrupted or hand-edited
                // stored configPath whose id is not a bare filename is treated as
                // already-reverted rather than risk unlinking through it.
                return nil
        }
        // Reconstruct the expected path and require the stored configPath to match
        // it. os.Remove follows symlinks, so a stored path outside configLibrary
        // (or a symlink inside it) must not let teardown unlink an arbitrary file.
        expectedPath := filepath.Join(dir, id+".json")
        if configPath != expectedPath {
                return nil
        }
        shimPath := cm.credentialHelperShimPath()

        // All steps run under the metaPath lock so a concurrent setup/teardown cannot
        // interleave. Ordering: de-reference _meta.json, delete the config document,
        // then remove the shim last — so a mid-revert failure never leaves Claude
        // Desktop pointing at a config whose helper is gone.
        return fileutils.WithFileLock(metaPath, func() error {
                if fileExistsAt(metaPath) {
                        meta, err := readClaudeDesktopMeta(metaPath)
                        if err != nil {
                                return err
                        }
                        meta["entries"] = removeMetaEntry(metaEntries(meta), id)
                        // Only clear appliedId if it still points at our config; leave a
                        // user's own active config selection alone.
                        applied, ok := meta["appliedId"].(string)
                        if !ok && meta["appliedId"] != nil {
                                slog.Warn("Claude Desktop _meta.json has a non-string appliedId; leaving it unchanged",
                                        "path", metaPath)
                        }
                        if applied == id {
                                meta["appliedId"] = ""
                        }
                        if err := writeClaudeDesktopMeta(metaPath, meta); err != nil {
                                return err
                        }
                }
                // The document also holds any connectors "thv client connector add" wrote;
                // they go with it, since Claude Desktop ignores them without the inference
                // settings anyway.
                if names := connectorNamesInDoc(configPath); len(names) > 0 {
                        slog.Warn("removing Claude Desktop connectors along with the LLM gateway configuration",
                                "connectors", names)
                }
                if err := removeIfExists(configPath); err != nil {
                        return err
                }
                if err := os.Remove(shimPath); err != nil && !os.IsNotExist(err) {
                        return fmt.Errorf("removing credential helper shim %s: %w", shimPath, err)
                }
                return nil
        })
}

func credentialHelperShimName() string {
        if runtime.GOOS == "windows" {
                return "claude-desktop-helper.cmd"
        }
        return "claude-desktop-helper.sh"
}

// credentialHelperShimPath is the fixed location of the generated shim.
func (cm *ClientManager) credentialHelperShimPath() string {
        return filepath.Join(cm.homeDir, ".toolhive", "llm", credentialHelperShimName())
}

// writeCredentialHelperShim generates the no-arg executable that Claude Desktop
// invokes as its inferenceCredentialHelper. Claude Desktop requires an absolute
// path to an executable and passes no arguments, so the shim exists to supply
// the "llm token" arguments — and to choose them per invocation context.
//
// tokenHelperPath is the absolute thv path (ApplyConfig.TokenHelperPath). It is
// quoted into the exec lines rather than resolved via PATH at call time:
// Claude Desktop is only ever GUI-launched. On macOS, GUI apps inherit
// launchd's environment, which does not contain thv's install directory, so a
// bare "thv" would never be found. On Windows, explorer does pass the user's
// PATH, but nothing guarantees thv is on it, and the absolute path pins the thv
// that ran setup either way. This is what
// distinguishes Claude Desktop from the direct-mode clients, which deliberately
// keep the bare, PATH-resolved tokenHelperShellCommand (see pkg/llm/setup.go).
// POSIX uses quoteForPOSIXShell; Windows writes a .cmd via quoteForCmd.
//
// Claude Desktop sets CLAUDE_HELPER_CONTEXT on each call. Only "interactive"
// permits an OIDC browser flow; silent contexts (background / setup-test /
// mid-session-refresh) must not hijack the user's browser, so they use
// --skip-browser and rely on the cached/refresh token that "thv llm setup"
// primed up front.
func (cm *ClientManager) writeCredentialHelperShim(tokenHelperPath string) (string, error) {
        // Require an absolute path. Defeating PATH resolution is the entire point of
        // this shim, and a relative path would silently reintroduce it — resolved
        // against whatever working directory Claude Desktop happens to have. A
        // non-absolute path here is a caller bug (os.Executable() is documented to
        // return an absolute path unless it errors), so fail closed rather than
        // emit a shim that cannot reliably locate thv.
        if !filepath.IsAbs(tokenHelperPath) {
                return "", fmt.Errorf(
                        "credential helper shim requires an absolute token-helper path, got %q", tokenHelperPath)
        }
        if runtime.GOOS == "windows" {
                return cm.writeWindowsCredentialHelperShim(tokenHelperPath)
        }
        return cm.writePOSIXCredentialHelperShim(tokenHelperPath)
}

func (cm *ClientManager) writePOSIXCredentialHelperShim(tokenHelperPath string) (string, error) {
        shimPath := cm.credentialHelperShimPath()
        if err := os.MkdirAll(filepath.Dir(shimPath), 0o700); err != nil {
                return "", fmt.Errorf("creating credential helper directory: %w", err)
        }
        quoted := quoteForPOSIXShell(tokenHelperPath)
        script := "#!/bin/sh\n" +
                "# Generated by `thv llm setup` — Claude Desktop credential helper.\n" +
                "# Prints a fresh LLM gateway token. Do not edit; `thv llm teardown` removes it.\n" +
                "if [ \"$CLAUDE_HELPER_CONTEXT\" = \"interactive\" ]; then\n" +
                "  exec " + quoted + " llm token\n" +
                "fi\n" +
                "exec " + quoted + " llm token --skip-browser\n"
        if err := fileutils.AtomicWriteFile(shimPath, []byte(script), 0o700); err != nil {
                return "", fmt.Errorf("writing credential helper shim %s: %w", shimPath, err)
        }
        return shimPath, nil
}

func (cm *ClientManager) writeWindowsCredentialHelperShim(tokenHelperPath string) (string, error) {
        shimPath := cm.credentialHelperShimPath()
        if err := os.MkdirAll(filepath.Dir(shimPath), 0o700); err != nil {
                return "", fmt.Errorf("creating credential helper directory: %w", err)
        }
        quoted := quoteForCmd(tokenHelperPath)
        // goto, not a parenthesized if: a path containing ")" would terminate the
        // block. A direct invoke, not call: TokenHelperPath is thv.exe from
        // os.Executable(), and call would expand % a second time and double ^
        // inside the quotes.
        script := "@echo off\r\n" +
                "REM Generated by `thv llm setup` - Claude Desktop credential helper.\r\n" +
                "REM Prints a fresh LLM gateway token. Do not edit; `thv llm teardown` removes it.\r\n" +
                "if /I \"%CLAUDE_HELPER_CONTEXT%\"==\"interactive\" goto interactive\r\n" +
                quoted + " llm token --skip-browser\r\n" +
                "exit /b %ERRORLEVEL%\r\n" +
                ":interactive\r\n" +
                quoted + " llm token\r\n" +
                "exit /b %ERRORLEVEL%\r\n"
        if err := fileutils.AtomicWriteFile(shimPath, []byte(script), 0o700); err != nil {
                return "", fmt.Errorf("writing credential helper shim %s: %w", shimPath, err)
        }
        return shimPath, nil
}

// quoteForPOSIXShell wraps s in single quotes so a POSIX shell reads every byte
// of it literally, escaping any embedded single quote as the four-character
// sequence quote-backslash-quote-quote (close the quoted run, emit a
// backslash-escaped quote, reopen).
//
// Single-quoting is a total transform: inside a single-quoted run no character
// is special to the shell — not $, backtick, backslash, or even a newline — so
// there is no input this fails on and callers need no metacharacter validation.
func quoteForPOSIXShell(s string) string {
        return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
}

// quoteForCmd wraps s in double quotes for cmd.exe. "%" is doubled first so
// "%VAR%" inside the path is not expanded; quotes inside the path are doubled
// after that. Used only by the Windows .cmd shim.
func quoteForCmd(s string) string {
        return `"` + strings.ReplaceAll(strings.ReplaceAll(s, `%`, `%%`), `"`, `""`) + `"`
}

// managedProfilePresent reports whether a managed configuration source for the
// client is present. A managed source overrides a client's local config, so
// "thv llm setup" warns when one is detected (the local config it writes would
// be ignored). Checks the MDM managed-preferences plist on macOS and the
// managed settings file on Linux; returns false on other platforms or when the
// client declares no managed source for the current platform.
func managedProfilePresent(cfg *clientAppConfig) bool {
        switch runtime.GOOS {
        case "darwin":
                return cfg.LLMManagedProfileDomain != "" &&
                        managedProfileExistsUnder(managedPreferencesRoot, cfg.LLMManagedProfileDomain)
        case "linux":
                return len(cfg.LLMManagedSettingsRelPath) > 0 &&
                        managedSettingsPresentUnder(linuxManagedSettingsRoot, cfg.LLMManagedSettingsRelPath)
        default:
                return false
        }
}

// managedPreferencesRoot is the macOS managed-preferences directory. A package
// variable (not a const) so tests can point it at a temp directory.
var managedPreferencesRoot = "/Library/Managed Preferences"

// linuxManagedSettingsRoot is the directory Linux managed settings files are
// resolved under. A package variable (not a const) so tests can point it at a
// temp directory.
var linuxManagedSettingsRoot = "/etc"

// managedProfileExistsUnder reports whether a managed-preferences plist for
// domain exists directly under root or under a per-user subdirectory
// (root/<user>/domain). Platform-independent so it is unit-testable.
func managedProfileExistsUnder(root, domain string) bool {
        if _, err := os.Stat(filepath.Join(root, domain)); err == nil {
                return true
        }
        matches, _ := filepath.Glob(filepath.Join(root, "*", domain))
        return len(matches) > 0
}

// managedSettingsPresentUnder reports whether anything exists at root/relPath.
// Claude Desktop rejects a managed settings file that is a symlink, not owned
// by root, or group- or world-writable, but it then treats the device as
// managed and disables local settings anyway. Either way the local config is
// ignored, so any entry counts and Desktop's checks are not replicated.
// Platform-independent so it is unit-testable.
func managedSettingsPresentUnder(root string, relPath []string) bool {
        _, err := os.Lstat(filepath.Join(root, filepath.Join(relPath...)))
        return err == nil
}

// ── _meta.json helpers ─────────────────────────────────────────────────────

// readClaudeDesktopMeta reads _meta.json into a generic map so unknown keys are
// preserved on write-back. A missing or empty file yields a fresh selector.
func readClaudeDesktopMeta(path string) (map[string]any, error) {
        data, err := os.ReadFile(path) // #nosec G304 -- known config file location
        if err != nil {
                if os.IsNotExist(err) {
                        return map[string]any{"appliedId": "", "entries": []any{}}, nil
                }
                return nil, fmt.Errorf("reading %s: %w", path, err)
        }
        if len(data) == 0 {
                return map[string]any{"appliedId": "", "entries": []any{}}, nil
        }
        var meta map[string]any
        if err := json.Unmarshal(data, &meta); err != nil {
                return nil, fmt.Errorf("parsing %s: %w", path, err)
        }
        if meta == nil {
                meta = map[string]any{}
        }
        // Guard the "valid JSON, wrong shape" case: if entries is present but not a
        // JSON array, bail rather than silently dropping it — metaEntries would treat
        // it as empty and a subsequent append would overwrite the user's data.
        if e, ok := meta["entries"]; ok && e != nil {
                if _, isArray := e.([]any); !isArray {
                        return nil, fmt.Errorf("parsing %s: %q must be a JSON array", path, "entries")
                }
        }
        return meta, nil
}

// writeClaudeDesktopMeta encodes meta and writes it atomically.
func writeClaudeDesktopMeta(path string, meta map[string]any) error {
        data, err := json.MarshalIndent(meta, "", "  ")
        if err != nil {
                return fmt.Errorf("encoding %s: %w", path, err)
        }
        if err := fileutils.AtomicWriteFile(path, data, 0o600); err != nil {
                return fmt.Errorf("writing %s: %w", path, err)
        }
        return nil
}

// metaEntries returns the entries slice from meta, or an empty slice.
func metaEntries(meta map[string]any) []any {
        entries, _ := meta["entries"].([]any)
        return entries
}

// metaEntryID returns the id of the entry with the given name, or "" if absent.
// _meta.json lives in a directory Claude Desktop and users also write to, so the
// stored id is untrusted: an unsafe value is treated as absent so the caller
// mints a fresh, safe UUID rather than joining a path-traversing id into the
// configLibrary path. See isSafeConfigID.
func metaEntryID(meta map[string]any, name string) string {
        for _, e := range metaEntries(meta) {
                entry, ok := e.(map[string]any)
                if !ok {
                        continue
                }
                if n, _ := entry["name"].(string); n == name {
                        id, _ := entry["id"].(string)
                        if !isSafeConfigID(id) {
                                return ""
                        }
                        return id
                }
        }
        return ""
}

// isSafeConfigID reports whether id is a bare filename safe to join into the
// configLibrary path — non-empty, no path separators, and no "..". Guards
// against a corrupted or hand-edited _meta.json entry escaping configLibrary.
func isSafeConfigID(id string) bool {
        return id != "" && filepath.Base(id) == id && !strings.Contains(id, "..")
}

// upsertMetaEntry sets the id of the entry with the given name if one exists
// (correcting a malformed or stale same-name entry in place), or appends a new
// entry otherwise. Prevents duplicate ToolHive-owned entries.
func upsertMetaEntry(entries []any, id, name string) []any {
        for _, e := range entries {
                entry, ok := e.(map[string]any)
                if !ok {
                        continue
                }
                if n, _ := entry["name"].(string); n == name {
                        entry["id"] = id
                        return entries
                }
        }
        return append(entries, map[string]any{"id": id, "name": name})
}

// removeMetaEntry returns entries with the entry matching id removed.
func removeMetaEntry(entries []any, id string) []any {
        out := make([]any, 0, len(entries))
        for _, e := range entries {
                if entry, ok := e.(map[string]any); ok {
                        if eid, _ := entry["id"].(string); eid == id {
                                continue
                        }
                }
                out = append(out, e)
        }
        return out
}

// metaIDFromConfigPath derives the config id from a "<id>.json" path.
func metaIDFromConfigPath(configPath string) string {
        base := filepath.Base(configPath)
        return base[:len(base)-len(filepath.Ext(base))]
}

// fileExistsAt reports whether a file exists at path.
func fileExistsAt(path string) bool {
        _, err := os.Stat(path)
        return err == nil
}

// removeIfExists deletes path, treating a missing file as success.
func removeIfExists(path string) error {
        if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
                return fmt.Errorf("removing %s: %w", path, err)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "bytes"
        "encoding/json"
        "fmt"
        "os"
        "path/filepath"
        "slices"
        "strings"

        "github.com/tailscale/hujson"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/llmgateway"
)

const piProviderName = "toolhive"
const piAnthropicProviderName = "toolhive-anthropic"

// Pi runs !-prefixed API keys on every request with its own fixed timeout;
// it has no configurable helper TTL or timeout to write into models.json.
const piTokenCommand = "!" + llmgateway.TokenHelperShellCommand //nolint:gosec // command, not a credential

// configurePi adds a named provider without disturbing the user's other providers.
func (cm *ClientManager) configurePi(appCfg *clientAppConfig, cfg llmgateway.ApplyConfig) (string, error) {
        openAIModels, anthropicModels, err := selectPiModels(cfg)
        if err != nil {
                return "", err
        }
        if len(openAIModels) == 0 && len(anthropicModels) == 0 {
                return "", fmt.Errorf("pi requires a model list from authenticated gateway discovery")
        }
        path := cm.buildLLMSettingsPath(appCfg)
        if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
                return "", fmt.Errorf("creating Pi config directory: %w", err)
        }
        err = fileutils.WithFileLock(path, func() error {
                root, providers, err := readPiConfig(path)
                if err != nil {
                        return err
                }
                for _, name := range []string{piProviderName, piAnthropicProviderName} {
                        if existing, ok := providers[name]; ok && !piProviderManaged(existing) {
                                return fmt.Errorf("pi provider %q already exists and is not managed by ToolHive", name)
                        }
                }
                baseURL := strings.TrimRight(cfg.GatewayURL, "/")
                if len(openAIModels) > 0 {
                        provider, err := encodePiGatewayProvider(baseURL+"/v1", "openai-completions", openAIModels,
                                cfg.DiscoveredModelContextWindows, nil, nil)
                        if err != nil {
                                return err
                        }
                        providers[piProviderName] = provider
                } else {
                        delete(providers, piProviderName)
                }
                if len(anthropicModels) > 0 {
                        anthropicBaseURL := baseURL
                        if cfg.AnthropicBaseURL != "" {
                                anthropicBaseURL = strings.TrimRight(cfg.AnthropicBaseURL, "/")
                        }
                        provider, err := encodePiGatewayProvider(anthropicBaseURL, "anthropic-messages", anthropicModels,
                                cfg.AnthropicModelContextWindows, cfg.AnthropicModelReasoning, cfg.AnthropicModelAdaptive)
                        if err != nil {
                                return err
                        }
                        providers[piAnthropicProviderName] = provider
                } else {
                        delete(providers, piAnthropicProviderName)
                }
                return writePiConfig(path, root, providers)
        })
        if err != nil {
                return "", err
        }
        return path, nil
}

func selectPiModels(cfg llmgateway.ApplyConfig) ([]string, []string, error) {
        if len(cfg.Models) == 0 {
                return cfg.DiscoveredModels, cfg.AnthropicModels, nil
        }
        var openAIModels, anthropicModels []string
        for _, id := range cfg.Models {
                openAI := slices.Contains(cfg.DiscoveredModels, id)
                anthropic := slices.Contains(cfg.AnthropicModels, id)
                if !openAI && !anthropic {
                        return nil, nil, fmt.Errorf("pi model %q was not found by authenticated model discovery; cannot determine its API", id)
                }
                if openAI {
                        openAIModels = append(openAIModels, id)
                }
                if anthropic {
                        anthropicModels = append(anthropicModels, id)
                }
        }
        return openAIModels, anthropicModels, nil
}

func piProviderManaged(existing json.RawMessage) bool {
        var provider struct {
                APIKey string `json:"apiKey"`
        }
        return json.Unmarshal(existing, &provider) == nil && provider.APIKey == piTokenCommand
}

func encodePiGatewayProvider(
        baseURL, api string, models []string, contextWindows map[string]int, reasoning, adaptive map[string]bool,
) (json.RawMessage, error) {
        type piCompat struct {
                ForceAdaptiveThinking bool `json:"forceAdaptiveThinking"`
        }
        type piModel struct {
                ID            string    `json:"id"`
                ContextWindow int       `json:"contextWindow,omitempty"`
                Reasoning     bool      `json:"reasoning,omitempty"`
                Compat        *piCompat `json:"compat,omitempty"`
        }
        entries := make([]piModel, 0, len(models))
        for _, id := range models {
                entry := piModel{ID: id, ContextWindow: contextWindows[id], Reasoning: reasoning[id]}
                if adaptive[id] {
                        entry.Compat = &piCompat{ForceAdaptiveThinking: true}
                }
                entries = append(entries, entry)
        }
        provider, err := json.Marshal(struct {
                BaseURL    string    `json:"baseUrl"`
                API        string    `json:"api"`
                APIKey     string    `json:"apiKey"`
                AuthHeader bool      `json:"authHeader"`
                Models     []piModel `json:"models"`
        }{baseURL, api, piTokenCommand, true, entries})
        if err != nil {
                return nil, fmt.Errorf("encoding Pi provider: %w", err)
        }
        return provider, nil
}

func (cm *ClientManager) revertPi(appCfg *clientAppConfig, path string) error {
        if path != cm.buildLLMSettingsPath(appCfg) {
                return fmt.Errorf("refusing to revert Pi config at unexpected path %q", path)
        }
        if _, err := os.Stat(path); os.IsNotExist(err) {
                return nil
        } else if err != nil {
                return fmt.Errorf("checking Pi config: %w", err)
        }
        return fileutils.WithFileLock(path, func() error {
                root, providers, err := readPiConfig(path)
                if err != nil {
                        return err
                }
                owned := false
                for _, name := range []string{piProviderName, piAnthropicProviderName} {
                        existing, ok := providers[name]
                        if !ok {
                                continue
                        }
                        owned = true
                        if !piProviderManaged(existing) {
                                return fmt.Errorf("pi provider %q is no longer managed by ToolHive", name)
                        }
                }
                if !owned {
                        return nil
                }
                delete(providers, piProviderName)
                delete(providers, piAnthropicProviderName)
                return writePiConfig(path, root, providers)
        })
}

type piConfig struct {
        value  hujson.Value
        hasBOM bool
}

var piBOM = []byte{0xef, 0xbb, 0xbf}

func readPiConfig(path string) (piConfig, map[string]json.RawMessage, error) {
        content, err := os.ReadFile(path) // #nosec G304 -- path is the registered Pi config location
        if err != nil && !os.IsNotExist(err) {
                return piConfig{}, nil, fmt.Errorf("reading Pi config: %w", err)
        }
        hasBOM := bytes.HasPrefix(content, piBOM)
        content = bytes.TrimPrefix(content, piBOM)
        if len(content) == 0 {
                content = []byte("{}")
        }
        value, err := parseJSONC(content, path)
        if err != nil {
                return piConfig{}, nil, err
        }
        standardized, err := standardizeJSONC(&value, path)
        if err != nil {
                return piConfig{}, nil, err
        }
        var root map[string]json.RawMessage
        if err := json.Unmarshal(standardized, &root); err != nil || root == nil {
                return piConfig{}, nil, fmt.Errorf("decoding Pi config %s: expected a JSON object: %v", path, err)
        }
        providers := make(map[string]json.RawMessage)
        if raw, ok := root["providers"]; ok {
                if err := json.Unmarshal(raw, &providers); err != nil || providers == nil {
                        return piConfig{}, nil, fmt.Errorf("decoding Pi providers in %s: expected a JSON object: %v", path, err)
                }
        }
        return piConfig{value: value, hasBOM: hasBOM}, providers, nil
}

func writePiConfig(path string, root piConfig, providers map[string]json.RawMessage) error {
        standardized, err := standardizeJSONC(&root.value, path)
        if err != nil {
                return err
        }
        if !jsonPointerExists(standardized, "/providers") && len(providers) != 0 {
                if err := root.value.Patch([]byte(`[{"op":"add","path":"/providers","value":{}}]`)); err != nil {
                        return fmt.Errorf("adding Pi providers in %s: %w", path, err)
                }
        }
        for _, name := range []string{piProviderName, piAnthropicProviderName} {
                ptr := "/providers/" + name
                value, ok := providers[name]
                if !ok {
                        if err := removeLLMKey(&root.value, ptr, path, standardized); err != nil {
                                return err
                        }
                        continue
                }
                patch, err := json.Marshal([]llmPatchOp{{Op: "add", Path: ptr, Value: value}})
                if err != nil {
                        return fmt.Errorf("encoding Pi provider %q: %w", name, err)
                }
                if err := root.value.Patch(patch); err != nil {
                        return fmt.Errorf("patching Pi provider %q in %s: %w", name, path, err)
                }
        }
        content, err := hujson.Format(root.value.Pack())
        if err != nil {
                return fmt.Errorf("formatting Pi config %s: %w", path, err)
        }
        if root.hasBOM {
                content = append(append([]byte{}, piBOM...), content...)
        }
        if err := fileutils.AtomicWriteFile(path, content, 0o600); err != nil {
                return fmt.Errorf("writing Pi config: %w", err)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"

        "github.com/tailscale/hujson"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/llmgateway"
)

const (
        vsCodeProviderName   = "ToolHive"
        vsCodeProviderVendor = "customendpoint"
        // The OpenAI-compatible /v1/models response exposes IDs but not token limits.
        // VS Code requires positive values, so ToolHive supplies conservative defaults.
        vsCodeMaxInputTokens  = 128000
        vsCodeMaxOutputTokens = 16000
)

type vsCodeProviderGroup struct {
        Name   string        `json:"name"`
        Vendor string        `json:"vendor"`
        Models []vsCodeModel `json:"models"`
}

type vsCodeModel struct {
        ID              string            `json:"id"`
        Name            string            `json:"name"`
        URL             string            `json:"url"`
        ToolCalling     bool              `json:"toolCalling"`
        Vision          bool              `json:"vision"`
        MaxInputTokens  int               `json:"maxInputTokens"`
        MaxOutputTokens int               `json:"maxOutputTokens"`
        RequestHeaders  map[string]string `json:"requestHeaders"`
}

func (cm *ClientManager) configureVSCode(appCfg *clientAppConfig, cfg llmgateway.ApplyConfig) (string, error) {
        if len(cfg.DiscoveredModels) == 0 {
                return "", fmt.Errorf("configuring %s: gateway model discovery returned no models", appCfg.ClientType)
        }
        path := cm.buildLLMSettingsPath(appCfg)
        if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
                return "", fmt.Errorf("creating directory for %s: %w", path, err)
        }

        models := make([]vsCodeModel, 0, len(cfg.DiscoveredModels))
        for _, id := range cfg.DiscoveredModels {
                models = append(models, vsCodeModel{
                        ID: id, Name: id, URL: cfg.ProxyBaseURL,
                        ToolCalling: true, Vision: false,
                        MaxInputTokens: vsCodeMaxInputTokens, MaxOutputTokens: vsCodeMaxOutputTokens,
                        RequestHeaders: map[string]string{"Authorization": "Bearer " + llmPlaceholderAPIKey},
                })
        }
        group, err := json.Marshal(vsCodeProviderGroup{
                Name: vsCodeProviderName, Vendor: vsCodeProviderVendor, Models: models,
        })
        if err != nil {
                return "", fmt.Errorf("marshaling VS Code provider group: %w", err)
        }

        err = fileutils.WithFileLock(path, func() error {
                content, err := readOrInitFile(path, []byte("[]"))
                if err != nil {
                        return err
                }
                v, err := hujson.Parse(content)
                if err != nil {
                        return fmt.Errorf("parsing %s: %w", path, err)
                }
                if err := patchVSCodeProviderGroups(&v, path, group); err != nil {
                        return err
                }
                formatted, err := hujson.Format(v.Pack())
                if err != nil {
                        return fmt.Errorf("formatting %s: %w", path, err)
                }
                return fileutils.AtomicWriteFile(path, formatted, 0o600)
        })
        if err != nil {
                return "", err
        }

        // Cleanup is best-effort after the authoritative new configuration is
        // durable. A failure must not make setup lose the path needed by teardown.
        if err := removeLegacyVSCodeSettings(filepath.Join(filepath.Dir(path), "settings.json")); err != nil {
                slog.Warn("Could not remove obsolete VS Code Copilot settings", "error", err)
        }
        return path, nil
}

func (*ClientManager) revertVSCode(_ *clientAppConfig, configPath string) error {
        legacyPath := filepath.Join(filepath.Dir(configPath), "settings.json")
        if filepath.Base(configPath) == "settings.json" {
                legacyPath = configPath
        }
        if err := removeLegacyVSCodeSettings(legacyPath); err != nil {
                return err
        }
        // Older ToolHive versions persisted settings.json as ConfigPath. That file
        // is an object and cannot contain a provider group from the new integration.
        if filepath.Base(configPath) == "settings.json" {
                return nil
        }
        if _, err := os.Stat(configPath); os.IsNotExist(err) {
                return nil
        }
        return fileutils.WithFileLock(configPath, func() error {
                content, err := os.ReadFile(configPath) // #nosec G304 -- persisted configuration path
                if err != nil {
                        if os.IsNotExist(err) {
                                return nil
                        }
                        return fmt.Errorf("reading %s: %w", configPath, err)
                }
                v, err := hujson.Parse(content)
                if err != nil {
                        return fmt.Errorf("parsing %s: %w", configPath, err)
                }
                if err := patchVSCodeProviderGroups(&v, configPath, nil); err != nil {
                        return err
                }
                formatted, err := hujson.Format(v.Pack())
                if err != nil {
                        return fmt.Errorf("formatting %s: %w", configPath, err)
                }
                return fileutils.AtomicWriteFile(configPath, formatted, 0o600)
        })
}

// patchVSCodeProviderGroups removes all ToolHive-owned provider groups and,
// when replacement is non-nil, appends the replacement. Patching the hujson
// syntax tree preserves comments in unrelated groups.
func patchVSCodeProviderGroups(v *hujson.Value, path string, replacement json.RawMessage) error {
        standardized, err := hujson.Standardize(v.Pack())
        if err != nil {
                return fmt.Errorf("standardizing %s: %w", path, err)
        }
        var groups []json.RawMessage
        if err := json.Unmarshal(standardized, &groups); err != nil {
                return fmt.Errorf("parsing %s: expected an array of provider groups: %w", path, err)
        }

        operations := make([]llmPatchOp, 0, len(groups)+1)
        for i := len(groups) - 1; i >= 0; i-- {
                var identity struct {
                        Name   string `json:"name"`
                        Vendor string `json:"vendor"`
                }
                if err := json.Unmarshal(groups[i], &identity); err != nil {
                        return fmt.Errorf("parsing provider group in %s: %w", path, err)
                }
                if identity.Name == vsCodeProviderName && identity.Vendor == vsCodeProviderVendor {
                        operations = append(operations, llmPatchOp{Op: "remove", Path: fmt.Sprintf("/%d", i)})
                }
        }
        if replacement != nil {
                operations = append(operations, llmPatchOp{Op: "add", Path: "/-", Value: replacement})
        }
        if len(operations) == 0 {
                return nil
        }
        patch, err := json.Marshal(operations)
        if err != nil {
                return fmt.Errorf("marshaling provider-group patch for %s: %w", path, err)
        }
        if err := v.Patch(patch); err != nil {
                return fmt.Errorf("patching provider groups in %s: %w", path, err)
        }
        return nil
}

func removeLegacyVSCodeSettings(path string) error {
        legacy := &clientAppConfig{LLMGatewayKeys: []LLMGatewayKeySpec{
                {JSONPointer: "/github.copilot.advanced.serverUrl"},
                {JSONPointer: "/github.copilot.advanced.apiKey"},
        }}
        if err := revertJSONPointerGateway(legacy, path); err != nil {
                return fmt.Errorf("removing obsolete VS Code Copilot settings: %w", err)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "context"
        "errors"
        "fmt"
        "log/slog"

        "github.com/stacklok/toolhive/pkg/config"
        ct "github.com/stacklok/toolhive/pkg/container"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/core"
        "github.com/stacklok/toolhive/pkg/groups"
)

// Client represents a registered ToolHive client.
type Client struct {
        Name ClientApp `json:"name"`
}

// RegisteredClient represents a registered client with its associated groups.
type RegisteredClient struct {
        Name   ClientApp `json:"name"`
        Groups []string  `json:"groups"`
}

// Manager is the interface for managing registered ToolHive clients.
//
//go:generate mockgen -destination=mocks/mock_manager.go -package=mocks -source=manager.go Manager
type Manager interface {
        // ListClients returns a list of all registered clients with their group information.
        ListClients(ctx context.Context) ([]RegisteredClient, error)
        // RegisterClients registers multiple clients with ToolHive for the specified workloads.
        RegisterClients(clients []Client, workloads []core.Workload) error
        // UnregisterClients unregisters multiple clients from ToolHive for the specified workloads.
        UnregisterClients(ctx context.Context, clients []Client, workloads []core.Workload) error
        // AddServerToClients adds an MCP server to the appropriate client configurations.
        AddServerToClients(ctx context.Context, serverName, serverURL, transportType, group string) error
        // RemoveServerFromClients removes an MCP server from the appropriate client configurations.
        RemoveServerFromClients(ctx context.Context, serverName, group string) error
}

type defaultManager struct {
        runtime        rt.Runtime
        groupManager   groups.Manager
        configProvider config.Provider
}

// NewManager creates a new client manager instance.
func NewManager(ctx context.Context) (Manager, error) {
        runtime, err := ct.NewFactory().Create(ctx)
        if err != nil {
                return nil, err
        }

        groupManager, err := groups.NewManager()
        if err != nil {
                return nil, err
        }

        return &defaultManager{
                runtime:        runtime,
                groupManager:   groupManager,
                configProvider: config.NewDefaultProvider(),
        }, nil
}

// NewManagerWithProvider creates a new client manager instance with a custom config provider.
// This is useful for testing to avoid using the singleton config.
func NewManagerWithProvider(ctx context.Context, configProvider config.Provider) (Manager, error) {
        runtime, err := ct.NewFactory().Create(ctx)
        if err != nil {
                return nil, err
        }

        groupManager, err := groups.NewManager()
        if err != nil {
                return nil, err
        }

        return &defaultManager{
                runtime:        runtime,
                groupManager:   groupManager,
                configProvider: configProvider,
        }, nil
}

// SetConfigProvider sets a custom config provider for testing purposes.
// This allows tests to inject a test config provider to avoid modifying the real config file.
func (m *defaultManager) SetConfigProvider(provider config.Provider) {
        m.configProvider = provider
}

func (m *defaultManager) ListClients(ctx context.Context) ([]RegisteredClient, error) {
        cfg := m.configProvider.GetConfig()

        // Get all groups
        allGroups, err := m.groupManager.List(ctx)
        if err != nil {
                return nil, fmt.Errorf("failed to list groups: %w", err)
        }

        clientGroups := make(map[string][]string) // client -> groups
        allRegisteredClients := make(map[string]bool)

        if len(allGroups) > 0 {
                // Collect clients from all groups
                for _, group := range allGroups {
                        for _, clientName := range group.RegisteredClients {
                                allRegisteredClients[clientName] = true
                                clientGroups[clientName] = append(clientGroups[clientName], group.Name)
                        }
                }
        }

        // Add clients from global config that might not be in any group
        for _, clientName := range cfg.Clients.RegisteredClients {
                if !allRegisteredClients[clientName] {
                        allRegisteredClients[clientName] = true
                        if len(allGroups) > 0 {
                                clientGroups[clientName] = []string{} // no groups
                        }
                }
        }

        // Convert to slice for return
        // Initialize as empty slice to ensure JSON encodes as [] instead of null when empty
        registeredClients := make([]RegisteredClient, 0)
        for clientName := range allRegisteredClients {
                registered := RegisteredClient{
                        Name:   ClientApp(clientName),
                        Groups: clientGroups[clientName],
                }
                registeredClients = append(registeredClients, registered)
        }

        return registeredClients, nil
}

// RegisterClients registers multiple clients with ToolHive for the specified workloads.
func (m *defaultManager) RegisterClients(clients []Client, workloads []core.Workload) error {
        for _, client := range clients {
                // Add specified workloads to the client
                if err := m.addWorkloadsToClient(client.Name, workloads); err != nil {
                        return fmt.Errorf("failed to add workloads to client %s: %w", client.Name, err)
                }
        }
        return nil
}

// UnregisterClients unregisters multiple clients from ToolHive for the specified workloads.
func (m *defaultManager) UnregisterClients(_ context.Context, clients []Client, workloads []core.Workload) error {
        for _, client := range clients {
                // Remove specified workloads from the client
                if err := m.removeWorkloadsFromClient(client.Name, workloads); err != nil {
                        return fmt.Errorf("failed to remove workloads from client %s: %w", client.Name, err)
                }
        }
        return nil
}

// AddServerToClients adds an MCP server to the appropriate client configurations.
// If the workload belongs to a group, only clients registered with that group are updated.
// If the workload has no group, all registered clients are updated (backward compatibility).
func (m *defaultManager) AddServerToClients(
        ctx context.Context, serverName, serverURL, transportType, group string,
) error {
        targetClients := m.getTargetClients(ctx, serverName, group)

        if len(targetClients) == 0 {
                slog.Debug("no target clients found for server", "server", serverName)
                return nil
        }

        // Add the server to each target client
        for _, clientName := range targetClients {
                if err := m.updateClientWithServer(clientName, serverName, serverURL, transportType); err != nil {
                        slog.Warn("failed to update client", "client", clientName, "error", err)
                }
        }
        return nil
}

// RemoveServerFromClients removes an MCP server from the appropriate client configurations.
// If the server belongs to a group, only clients registered with that group are updated.
// If the server has no group, all registered clients are updated (backward compatibility).
func (m *defaultManager) RemoveServerFromClients(ctx context.Context, serverName, group string) error {
        targetClients := m.getTargetClients(ctx, serverName, group)

        if len(targetClients) == 0 {
                slog.Debug("no target clients found for server", "server", serverName)
                return nil
        }

        // Remove the server from each target client
        for _, clientName := range targetClients {
                if err := m.removeServerFromClient(ClientApp(clientName), serverName); err != nil && !errors.Is(err, ErrConfigFileNotFound) {
                        slog.Warn("failed to remove server from client", "client", clientName, "error", err)
                }
        }

        return nil
}

// addWorkloadsToClient adds the specified workloads to the client's configuration
func (m *defaultManager) addWorkloadsToClient(clientType ClientApp, workloads []core.Workload) error {
        if len(workloads) == 0 {
                // No workloads to add, nothing more to do
                return nil
        }

        // For each workload, add it to the client configuration
        for _, workload := range workloads {
                // Use the common update function (creates config if needed)
                err := m.updateClientWithServer(
                        string(clientType), workload.Name, workload.URL, string(workload.TransportType),
                )
                if err != nil {
                        return fmt.Errorf("failed to add workload %s to client %s: %w", workload.Name, clientType, err)
                }

                slog.Debug("added mcp server to client", "server", workload.Name, "client", clientType)
        }

        return nil
}

// removeWorkloadsFromClient removes the specified workloads from the client's configuration
func (m *defaultManager) removeWorkloadsFromClient(clientType ClientApp, workloads []core.Workload) error {
        if len(workloads) == 0 {
                // No workloads to remove, nothing to do
                return nil
        }

        // For each workload, remove it from the client configuration
        for _, workload := range workloads {
                err := m.removeServerFromClient(clientType, workload.Name)
                if err != nil {
                        return fmt.Errorf("failed to remove workload %s from client %s: %w", workload.Name, clientType, err)
                }
        }

        return nil
}

// removeServerFromClient removes an MCP server from a single client configuration
func (*defaultManager) removeServerFromClient(clientName ClientApp, serverName string) error {
        clientConfig, err := FindClientConfig(clientName)
        if err != nil {
                return fmt.Errorf("failed to find client configurations: %w", err)
        }

        // Remove the MCP server configuration with locking
        if err := clientConfig.ConfigUpdater.Remove(serverName); err != nil {
                return fmt.Errorf("failed to remove MCP server configuration from %s: %w", clientConfig.Path, err)
        }

        slog.Debug("removed mcp server from client", "server", serverName, "client", clientName)
        return nil
}

// updateClientWithServer updates a single client with an MCP server configuration, creating config if needed
func (*defaultManager) updateClientWithServer(clientName, serverName, serverURL, transportType string) error {
        clientConfig, err := FindClientConfig(ClientApp(clientName))
        if err != nil {
                if errors.Is(err, ErrConfigFileNotFound) {
                        // Create a new client configuration if it doesn't exist
                        clientConfig, err = CreateClientConfig(ClientApp(clientName))
                        if err != nil {
                                return fmt.Errorf("failed to create client configuration for %s: %w", clientName, err)
                        }
                } else {
                        return fmt.Errorf("failed to find client configuration: %w", err)
                }
        }

        slog.Debug("updating client configuration", "path", clientConfig.Path)

        if err := Upsert(*clientConfig, serverName, serverURL, transportType); err != nil {
                return fmt.Errorf("failed to update MCP server configuration in %s: %w", clientConfig.Path, err)
        }

        slog.Debug("successfully updated client configuration", "path", clientConfig.Path)
        return nil
}

// getTargetClients determines which clients should be updated based on workload group
func (m *defaultManager) getTargetClients(ctx context.Context, serverName, groupName string) []string {
        // Server belongs to a group - only update clients registered with that group
        if groupName != "" {
                group, err := m.groupManager.Get(ctx, groupName)
                if err != nil {
                        slog.Warn("failed to get group, skipping client config updates",
                                "group", groupName, "server", serverName, "error", err)
                        return nil
                }

                slog.Debug("server belongs to group, updating registered clients",
                        "server", serverName, "group", group.Name, "count", len(group.RegisteredClients))
                return group.RegisteredClients
        }

        // Server has no group - use backward compatible behavior (update all registered clients)
        appConfig := m.configProvider.GetConfig()
        targetClients := appConfig.Clients.RegisteredClients
        slog.Debug("server has no group, updating globally registered clients for backward compatibility",
                "server", serverName, "count", len(targetClients))
        return targetClients
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

// ReadTOMLConfig reads and parses a TOML config file into a generic map.
// Exported wrapper around the unexported readTOMLConfig so the plugin adapters
// (which live in a separate package) can reuse the same parse path without
// duplicating code.
func ReadTOMLConfig(path string) (map[string]any, error) {
        return readTOMLConfig(path)
}

// WriteTOMLConfig marshals and writes the config map to a TOML file atomically.
// Exported wrapper around the unexported writeTOMLConfig.
func WriteTOMLConfig(path string, config map[string]any) error {
        return writeTOMLConfig(path, config)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "errors"
        "fmt"
        "path/filepath"
        "runtime"
        "sort"

        "github.com/stacklok/toolhive/pkg/plugins"
        "github.com/stacklok/toolhive/pkg/skills"
)

var (
        // ErrPluginsNotSupported is returned when a client does not support plugins.
        ErrPluginsNotSupported = errors.New("client does not support plugins")
        // ErrNoPluginPath is returned when a client has no plugin path for the requested scope.
        ErrNoPluginPath = errors.New("client has no plugin path for the requested scope")
)

// SupportsPlugins returns whether the given client supports plugins.
func (cm *ClientManager) SupportsPlugins(clientType ClientApp) bool {
        cfg := cm.lookupClientAppConfig(clientType)
        return cfg != nil && cfg.SupportsPlugins
}

// ListPluginSupportingClients returns a sorted slice of all clients that support plugins.
func (cm *ClientManager) ListPluginSupportingClients() []ClientApp {
        var clients []ClientApp
        for _, cfg := range cm.clientIntegrations {
                if cfg.SupportsPlugins {
                        clients = append(clients, cfg.ClientType)
                }
        }
        sort.Slice(clients, func(i, j int) bool {
                return clients[i] < clients[j]
        })
        return clients
}

// GetPluginPath resolves the filesystem path for a plugin installation.
//
// For [plugins.ScopeUser], it returns ~/<PluginsGlobalPath>/<pluginName>.
// For [plugins.ScopeProject], it returns <projectRoot>/<PluginsProjectPath>/<pluginName>.
//
// Returns an error if the client doesn't support plugins, the scope has no
// configured path, the project root is empty, or the plugin name would result
// in path traversal outside the plugins directory.
func (cm *ClientManager) GetPluginPath(
        clientType ClientApp, pluginName string, scope plugins.Scope, projectRoot string,
) (string, error) {
        if err := skills.ValidateSkillName(pluginName); err != nil {
                return "", err
        }

        cfg := cm.lookupClientAppConfig(clientType)
        if cfg == nil {
                return "", fmt.Errorf("%w: %s", ErrUnsupportedClientType, clientType)
        }
        if !cfg.SupportsPlugins {
                return "", fmt.Errorf("%w: %s", ErrPluginsNotSupported, clientType)
        }

        switch scope {
        case plugins.ScopeUser:
                return cm.buildPluginsGlobalPath(cfg, pluginName)
        case plugins.ScopeProject:
                return buildPluginsProjectPath(cfg, pluginName, projectRoot)
        default:
                return "", fmt.Errorf("%w: %q", ErrUnknownScope, scope)
        }
}

func (cm *ClientManager) buildPluginsGlobalPath(cfg *clientAppConfig, pluginName string) (string, error) {
        if len(cfg.PluginsGlobalPath) == 0 {
                return "", fmt.Errorf("%w: %s has no global plugin path", ErrNoPluginPath, cfg.ClientType)
        }

        parts := []string{cm.homeDir}
        if prefix, ok := cfg.PluginsPlatformPrefix[Platform(runtime.GOOS)]; ok {
                parts = append(parts, prefix...)
        }
        parts = append(parts, cfg.PluginsGlobalPath...)
        parts = append(parts, pluginName)
        return filepath.Join(parts...), nil
}

func buildPluginsProjectPath(cfg *clientAppConfig, pluginName string, projectRoot string) (string, error) {
        if len(cfg.PluginsProjectPath) == 0 {
                return "", fmt.Errorf("%w: %s has no project plugin path", ErrNoPluginPath, cfg.ClientType)
        }

        if projectRoot == "" {
                return "", ErrProjectRootRequired
        }

        parts := []string{projectRoot}
        parts = append(parts, cfg.PluginsProjectPath...)
        parts = append(parts, pluginName)
        return filepath.Join(parts...), nil
}

// GetConfigPath returns the absolute path to clientType's settings file under
// the manager's home directory. Returns ErrUnsupportedClientType for unknown
// clients. Exported so plugin adapters can locate the Codex config.toml without
// reaching into ClientManager internals.
func (cm *ClientManager) GetConfigPath(clientType ClientApp) (string, error) {
        cfg := cm.lookupClientAppConfig(clientType)
        if cfg == nil {
                return "", fmt.Errorf("%w: %s", ErrUnsupportedClientType, clientType)
        }
        return cm.buildMCPConfigPath(cfg), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

import (
        "errors"
        "fmt"
        "os"
        "path/filepath"
        "runtime"
        "sort"

        "github.com/stacklok/toolhive/pkg/skills"
)

var (
        // ErrSkillsNotSupported is returned when a client does not support skills.
        ErrSkillsNotSupported = errors.New("client does not support skills")
        // ErrNoSkillPath is returned when a client has no skill path for the requested scope.
        ErrNoSkillPath = errors.New("client has no skill path for the requested scope")
        // ErrProjectRootRequired is returned when project root is empty for project-scoped skills.
        ErrProjectRootRequired = errors.New("project root must be provided for project-scoped skills")
        // ErrProjectRootNotFound is returned when no project root can be detected.
        ErrProjectRootNotFound = errors.New("could not detect project root (no .git found)")
        // ErrUnknownScope is returned when an unrecognized skill scope is provided.
        ErrUnknownScope = errors.New("unknown skill scope")
)

// SupportsSkills returns whether the given client supports skills.
func (cm *ClientManager) SupportsSkills(clientType ClientApp) bool {
        cfg := cm.lookupClientAppConfig(clientType)
        return cfg != nil && cfg.SupportsSkills
}

// ListSkillSupportingClients returns a sorted slice of all clients that support skills.
func (cm *ClientManager) ListSkillSupportingClients() []ClientApp {
        var clients []ClientApp
        for _, cfg := range cm.clientIntegrations {
                if cfg.SupportsSkills {
                        clients = append(clients, cfg.ClientType)
                }
        }
        sort.Slice(clients, func(i, j int) bool {
                return clients[i] < clients[j]
        })
        return clients
}

// GetSkillPath resolves the filesystem path for a skill installation.
//
// For [skills.ScopeUser], it returns ~/<SkillsGlobalPath>/<skillName>, or
// <dir>/skills/<skillName> when the client's ConfigDirEnvVar is set.
// For [skills.ScopeProject], it returns <projectRoot>/<SkillsProjectPath>/<skillName>.
//
// Returns an error if the client doesn't support skills, the scope has no
// configured path, the project root is empty, or the skill name would result
// in path traversal outside the skills directory.
func (cm *ClientManager) GetSkillPath(
        clientType ClientApp, skillName string, scope skills.Scope, projectRoot string,
) (string, error) {
        if err := skills.ValidateSkillName(skillName); err != nil {
                return "", err
        }

        cfg := cm.lookupClientAppConfig(clientType)
        if cfg == nil {
                return "", fmt.Errorf("%w: %s", ErrUnsupportedClientType, clientType)
        }
        if !cfg.SupportsSkills {
                return "", fmt.Errorf("%w: %s", ErrSkillsNotSupported, clientType)
        }

        switch scope {
        case skills.ScopeUser:
                return cm.buildSkillsGlobalPath(cfg, skillName)
        case skills.ScopeProject:
                return buildSkillsProjectPath(cfg, skillName, projectRoot)
        default:
                return "", fmt.Errorf("%w: %q", ErrUnknownScope, scope)
        }
}

// DetectProjectRoot walks up from startDir looking for a .git directory or file
// (which indicates a git worktree). If startDir is empty, it uses the current
// working directory.
func DetectProjectRoot(startDir string) (string, error) {
        dir := startDir
        if dir == "" {
                var err error
                dir, err = os.Getwd()
                if err != nil {
                        return "", fmt.Errorf("failed to get working directory: %w", err)
                }
        }

        for {
                gitPath := filepath.Join(dir, ".git")
                if info, err := os.Lstat(gitPath); err == nil {
                        // Accept both directories (regular repos) and files (worktrees/submodules)
                        if info.IsDir() || info.Mode().IsRegular() {
                                return dir, nil
                        }
                }

                parent := filepath.Dir(dir)
                if parent == dir {
                        // Reached filesystem root without finding .git
                        return "", ErrProjectRootNotFound
                }
                dir = parent
        }
}

func (cm *ClientManager) lookupClientAppConfig(clientType ClientApp) *clientAppConfig {
        for i := range cm.clientIntegrations {
                if cm.clientIntegrations[i].ClientType == clientType {
                        return &cm.clientIntegrations[i]
                }
        }
        return nil
}

func (cm *ClientManager) buildSkillsGlobalPath(cfg *clientAppConfig, skillName string) (string, error) {
        if len(cfg.SkillsGlobalPath) == 0 {
                return "", fmt.Errorf("%w: %s has no global skill path", ErrNoSkillPath, cfg.ClientType)
        }

        if dir := cm.configDirOverride(cfg); dir != "" {
                return filepath.Join(dir, skillsDirName, skillName), nil
        }

        parts := []string{cm.homeDir}
        if prefix, ok := cfg.SkillsPlatformPrefix[Platform(runtime.GOOS)]; ok {
                parts = append(parts, prefix...)
        }
        parts = append(parts, cfg.SkillsGlobalPath...)
        parts = append(parts, skillName)
        return filepath.Join(parts...), nil
}

func buildSkillsProjectPath(cfg *clientAppConfig, skillName string, projectRoot string) (string, error) {
        if len(cfg.SkillsProjectPath) == 0 {
                return "", fmt.Errorf("%w: %s has no project skill path", ErrNoSkillPath, cfg.ClientType)
        }

        if projectRoot == "" {
                return "", ErrProjectRootRequired
        }

        parts := []string{projectRoot}
        parts = append(parts, cfg.SkillsProjectPath...)
        parts = append(parts, skillName)
        return filepath.Join(parts...), nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package client

// LLMTestEntry describes a minimal LLM-gateway-capable client integration for
// use in tests. No platform prefix is applied, so settings files resolve as
// homeDir/SettingsDir.../SettingsFile on all platforms.
type LLMTestEntry struct {
        ClientType     ClientApp
        Mode           string   // "direct" or "proxy"
        SettingsDir    []string // path segments from homeDir to the settings directory
        SettingsFile   string   // settings filename
        JSONPointers   []string // RFC 6901 JSON Pointer paths to patch
        ValueFields    []string // value-field names parallel to JSONPointers
        ClearWhenEmpty []bool   // ClearWhenEmpty flags parallel to JSONPointers (optional)
}

// LLMTestIntegrations converts []LLMTestEntry into an internal []clientAppConfig
// slice suitable for passing as the third argument to NewTestClientManager.
// Because clientAppConfig is unexported, callers assign the result via :=
// (type inferred) and pass it directly to NewTestClientManager.
func LLMTestIntegrations(entries []LLMTestEntry) []clientAppConfig {
        cfgs := make([]clientAppConfig, len(entries))
        for i, e := range entries {
                keys := make([]LLMGatewayKeySpec, len(e.JSONPointers))
                for j, ptr := range e.JSONPointers {
                        vf := ""
                        if j < len(e.ValueFields) {
                                vf = e.ValueFields[j]
                        }
                        cwe := false
                        if j < len(e.ClearWhenEmpty) {
                                cwe = e.ClearWhenEmpty[j]
                        }
                        keys[j] = LLMGatewayKeySpec{JSONPointer: ptr, ValueField: vf, ClearWhenEmpty: cwe}
                }
                cfgs[i] = clientAppConfig{
                        ClientType:         e.ClientType,
                        LLMGatewayMode:     e.Mode,
                        LLMSettingsFile:    e.SettingsFile,
                        LLMSettingsRelPath: e.SettingsDir,
                        LLMGatewayKeys:     keys,
                }
        }
        return cfgs
}

// NewTestClientManagerWithHome returns a ClientManager backed by the production
// client integrations (so plugin/skill/config metadata is fully populated) but
// rooted at homeDir instead of the real user home. Exported so adapter tests in
// other packages can resolve plugin paths and config-file paths under a temp
// directory without mutating process-wide HOME (which forbids t.Parallel).
func NewTestClientManagerWithHome(homeDir string) *ClientManager {
        return &ClientManager{
                homeDir:            homeDir,
                clientIntegrations: supportedClientIntegrations,
                lookPath:           nil,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "fmt"
        "strings"
)

// BuildAuthFileSecretPrefix is the prefix used for storing build auth file content in secrets
// #nosec G101 -- This is not a credential, just a prefix for secret names
const BuildAuthFileSecretPrefix = "BUILD_AUTH_FILE_"

// SupportedAuthFiles maps file type names to their target paths in the container
var SupportedAuthFiles = map[string]string{
        "npmrc":  "/root/.npmrc",
        "netrc":  "/root/.netrc",
        "yarnrc": "/root/.yarnrc",
}

// ValidateBuildAuthFileName checks if the file name is supported
func ValidateBuildAuthFileName(name string) error {
        if _, ok := SupportedAuthFiles[name]; !ok {
                supported := make([]string, 0, len(SupportedAuthFiles))
                for k := range SupportedAuthFiles {
                        supported = append(supported, k)
                }
                return fmt.Errorf("unsupported auth file type %q; supported types: %s", name, strings.Join(supported, ", "))
        }
        return nil
}

// BuildAuthFileSecretName returns the secret name for a given auth file type
func BuildAuthFileSecretName(fileType string) string {
        return BuildAuthFileSecretPrefix + fileType
}

// markBuildAuthFileConfigured marks an auth file type as configured in the config.
// The actual content is stored in the secrets provider, not in the config.
func markBuildAuthFileConfigured(p Provider, name string) error {
        if err := ValidateBuildAuthFileName(name); err != nil {
                return err
        }

        return p.UpdateConfig(func(c *Config) error {
                if c.BuildAuthFiles == nil {
                        c.BuildAuthFiles = make(map[string]string)
                }
                // Store only a marker - actual content is in secrets
                c.BuildAuthFiles[name] = "secret:" + BuildAuthFileSecretName(name)
                return nil
        })
}

// isBuildAuthFileConfigured checks if an auth file type is configured
func isBuildAuthFileConfigured(p Provider, name string) bool {
        config := p.GetConfig()
        if config.BuildAuthFiles == nil {
                return false
        }
        _, exists := config.BuildAuthFiles[name]
        return exists
}

// getConfiguredBuildAuthFiles returns a list of configured auth file types.
// Note: This only returns which files are configured, not their content.
// Use the secrets provider to retrieve actual content.
func getConfiguredBuildAuthFiles(p Provider) []string {
        config := p.GetConfig()
        if config.BuildAuthFiles == nil {
                return nil
        }
        result := make([]string, 0, len(config.BuildAuthFiles))
        for k := range config.BuildAuthFiles {
                result = append(result, k)
        }
        return result
}

// unsetBuildAuthFile removes an auth file configuration marker.
// Note: This only removes the config marker. The caller should also delete
// the corresponding secret from the secrets provider.
func unsetBuildAuthFile(p Provider, name string) error {
        return p.UpdateConfig(func(c *Config) error {
                if c.BuildAuthFiles != nil {
                        delete(c.BuildAuthFiles, name)
                }
                return nil
        })
}

// unsetAllBuildAuthFiles removes all auth file configuration markers.
// Note: This only removes the config markers. The caller should also delete
// the corresponding secrets from the secrets provider.
func unsetAllBuildAuthFiles(p Provider) error {
        return p.UpdateConfig(func(c *Config) error {
                c.BuildAuthFiles = nil
                return nil
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "fmt"
        "regexp"
        "strings"
)

// Build environment validation constants
const (
        errInvalidEnvKeyFormat  = "invalid environment variable name: %s (must match pattern %s)"
        errReservedEnvKey       = "environment variable name %s is reserved and cannot be overridden"
        errInvalidEnvValueChars = "environment variable value contains potentially dangerous characters"
)

// envKeyPattern matches valid environment variable names.
// Must start with uppercase letter, followed by uppercase letters, numbers, or underscores.
var envKeyPattern = regexp.MustCompile(`^[A-Z][A-Z0-9_]*$`)

// reservedEnvKeys lists environment variables that cannot be overridden for security reasons.
var reservedEnvKeys = map[string]bool{
        "PATH":            true,
        "HOME":            true,
        "USER":            true,
        "SHELL":           true,
        "PWD":             true,
        "HOSTNAME":        true,
        "TERM":            true,
        "LANG":            true,
        "LC_ALL":          true,
        "LD_PRELOAD":      true,
        "LD_LIBRARY_PATH": true,
}

// ValidateBuildEnvKey validates that an environment variable key follows the required pattern
// and is not a reserved variable.
func ValidateBuildEnvKey(key string) error {
        if !envKeyPattern.MatchString(key) {
                return fmt.Errorf(errInvalidEnvKeyFormat, key, "^[A-Z][A-Z0-9_]*$")
        }

        if reservedEnvKeys[key] {
                return fmt.Errorf(errReservedEnvKey, key)
        }

        return nil
}

// ValidateBuildEnvValue validates that an environment variable value does not contain
// potentially dangerous characters that could enable shell injection in Dockerfiles.
func ValidateBuildEnvValue(value string) error {
        // Check for shell metacharacters that could enable injection
        dangerousPatterns := []string{
                "`",  // Command substitution
                "$(", // Command substitution
                "${", // Variable expansion (could be used for injection)
                "\\", // Escape sequences
                "\n", // Newlines could break Dockerfile syntax
                "\r", // Carriage returns
                "\"", // Double quotes could break ENV syntax
                ";",  // Command separator
                "&&", // Command chaining
                "||", // Command chaining
                "|",  // Pipe
                ">",  // Redirection
                "<",  // Redirection
        }

        for _, pattern := range dangerousPatterns {
                if strings.Contains(value, pattern) {
                        return fmt.Errorf("%s: contains '%s'", errInvalidEnvValueChars, pattern)
                }
        }

        return nil
}

// ValidateBuildEnvEntry validates both the key and value of a build environment variable.
func ValidateBuildEnvEntry(key, value string) error {
        if err := ValidateBuildEnvKey(key); err != nil {
                return err
        }
        return ValidateBuildEnvValue(value)
}

// checkBuildEnvKeyConflict checks if a key is already configured in another source.
func checkBuildEnvKeyConflict(p Provider, key string) error {
        config := p.GetConfig()

        // Check literal values
        if config.BuildEnv != nil {
                if _, exists := config.BuildEnv[key]; exists {
                        return fmt.Errorf("key %s already configured as literal value; unset it first with 'thv config unset-build-env %s'", key, key)
                }
        }

        // Check secret references
        if config.BuildEnvFromSecrets != nil {
                if _, exists := config.BuildEnvFromSecrets[key]; exists {
                        return fmt.Errorf("key %s already configured from secret; unset it first", key)
                }
        }

        // Check shell references
        for _, k := range config.BuildEnvFromShell {
                if k == key {
                        return fmt.Errorf("key %s already configured from shell; unset it first", key)
                }
        }

        return nil
}

// setBuildEnv is a helper function that validates and sets a build environment variable.
func setBuildEnv(p Provider, key, value string) error {
        if err := ValidateBuildEnvEntry(key, value); err != nil {
                return err
        }

        // Check for conflicts with other sources
        if err := checkBuildEnvKeyConflict(p, key); err != nil {
                return err
        }

        return p.UpdateConfig(func(c *Config) error {
                if c.BuildEnv == nil {
                        c.BuildEnv = make(map[string]string)
                }
                c.BuildEnv[key] = value
                return nil
        })
}

// getBuildEnv is a helper function that retrieves a build environment variable.
func getBuildEnv(p Provider, key string) (value string, exists bool) {
        config := p.GetConfig()
        if config.BuildEnv == nil {
                return "", false
        }
        value, exists = config.BuildEnv[key]
        return value, exists
}

// getAllBuildEnv is a helper function that retrieves all build environment variables.
func getAllBuildEnv(p Provider) map[string]string {
        config := p.GetConfig()
        if config.BuildEnv == nil {
                return make(map[string]string)
        }
        // Return a copy to prevent external modifications
        result := make(map[string]string, len(config.BuildEnv))
        for k, v := range config.BuildEnv {
                result[k] = v
        }
        return result
}

// unsetBuildEnv is a helper function that removes a specific build environment variable.
func unsetBuildEnv(p Provider, key string) error {
        return p.UpdateConfig(func(c *Config) error {
                if c.BuildEnv != nil {
                        delete(c.BuildEnv, key)
                }
                return nil
        })
}

// unsetAllBuildEnv is a helper function that removes all build environment variables.
func unsetAllBuildEnv(p Provider) error {
        return p.UpdateConfig(func(c *Config) error {
                c.BuildEnv = nil
                return nil
        })
}

// setBuildEnvFromSecret validates and stores a secret reference for a build environment variable.
func setBuildEnvFromSecret(p Provider, key, secretName string) error {
        // Validate the key follows the pattern
        if err := ValidateBuildEnvKey(key); err != nil {
                return err
        }

        // Check for conflicts with other sources
        if err := checkBuildEnvKeyConflict(p, key); err != nil {
                return err
        }

        return p.UpdateConfig(func(c *Config) error {
                if c.BuildEnvFromSecrets == nil {
                        c.BuildEnvFromSecrets = make(map[string]string)
                }
                c.BuildEnvFromSecrets[key] = secretName
                return nil
        })
}

// getBuildEnvFromSecret retrieves the secret name for a build environment variable.
func getBuildEnvFromSecret(p Provider, key string) (secretName string, exists bool) {
        config := p.GetConfig()
        if config.BuildEnvFromSecrets == nil {
                return "", false
        }
        secretName, exists = config.BuildEnvFromSecrets[key]
        return secretName, exists
}

// getAllBuildEnvFromSecrets returns all build env secret references.
func getAllBuildEnvFromSecrets(p Provider) map[string]string {
        config := p.GetConfig()
        if config.BuildEnvFromSecrets == nil {
                return make(map[string]string)
        }
        result := make(map[string]string, len(config.BuildEnvFromSecrets))
        for k, v := range config.BuildEnvFromSecrets {
                result[k] = v
        }
        return result
}

// unsetBuildEnvFromSecret removes a secret reference.
func unsetBuildEnvFromSecret(p Provider, key string) error {
        return p.UpdateConfig(func(c *Config) error {
                if c.BuildEnvFromSecrets != nil {
                        delete(c.BuildEnvFromSecrets, key)
                }
                return nil
        })
}

// setBuildEnvFromShell adds an environment variable name to read from shell at build time.
func setBuildEnvFromShell(p Provider, key string) error {
        // Validate the key follows the pattern
        if err := ValidateBuildEnvKey(key); err != nil {
                return err
        }

        // Check if already in the list - skip if so
        if getBuildEnvFromShell(p, key) {
                return nil // Already exists, nothing to do
        }

        // Check for conflicts with other sources (not including shell since we checked above)
        if err := checkBuildEnvKeyConflictExcludingShell(p, key); err != nil {
                return err
        }

        return p.UpdateConfig(func(c *Config) error {
                c.BuildEnvFromShell = append(c.BuildEnvFromShell, key)
                return nil
        })
}

// checkBuildEnvKeyConflictExcludingShell checks if a key is already configured in literal or secret sources.
func checkBuildEnvKeyConflictExcludingShell(p Provider, key string) error {
        config := p.GetConfig()

        // Check literal values
        if config.BuildEnv != nil {
                if _, exists := config.BuildEnv[key]; exists {
                        return fmt.Errorf("key %s already configured as literal value; unset it first with 'thv config unset-build-env %s'", key, key)
                }
        }

        // Check secret references
        if config.BuildEnvFromSecrets != nil {
                if _, exists := config.BuildEnvFromSecrets[key]; exists {
                        return fmt.Errorf("key %s already configured from secret; unset it first", key)
                }
        }

        return nil
}

// getBuildEnvFromShell checks if a key is configured to read from shell.
func getBuildEnvFromShell(p Provider, key string) bool {
        config := p.GetConfig()
        for _, k := range config.BuildEnvFromShell {
                if k == key {
                        return true
                }
        }
        return false
}

// getAllBuildEnvFromShell returns all keys configured to read from shell.
func getAllBuildEnvFromShell(p Provider) []string {
        config := p.GetConfig()
        if config.BuildEnvFromShell == nil {
                return []string{}
        }
        result := make([]string, len(config.BuildEnvFromShell))
        copy(result, config.BuildEnvFromShell)
        return result
}

// unsetBuildEnvFromShell removes a key from shell environment list.
func unsetBuildEnvFromShell(p Provider, key string) error {
        return p.UpdateConfig(func(c *Config) error {
                if c.BuildEnvFromShell == nil {
                        return nil
                }
                newList := make([]string, 0, len(c.BuildEnvFromShell))
                for _, k := range c.BuildEnvFromShell {
                        if k != key {
                                newList = append(newList, k)
                        }
                }
                c.BuildEnvFromShell = newList
                return nil
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "fmt"
        "os"

        "github.com/stacklok/toolhive/pkg/certs"
)

// setCACert validates and sets the CA certificate path using the provided provider.
// It performs the following validations:
//   - Verifies the file exists and is readable
//   - Reads the certificate content
//   - Validates the certificate format using pkg/certs.ValidateCACertificate
//   - Cleans the file path
//
// The function returns an error if any validation fails or if updating the configuration fails.
func setCACert(provider Provider, certPath string) error {
        // Validate and clean the file path
        cleanPath, err := validateFilePath(certPath)
        if err != nil {
                return fmt.Errorf("CA certificate %w", err)
        }

        // Read the certificate
        certContent, err := readFile(cleanPath)
        if err != nil {
                return fmt.Errorf("CA certificate %w", err)
        }

        // Validate the certificate format
        if err := certs.ValidateCACertificate(certContent); err != nil {
                return fmt.Errorf("invalid CA certificate: %w", err)
        }

        // Update the configuration
        err = provider.UpdateConfig(func(c *Config) error {
                c.CACertificatePath = cleanPath
                return nil
        })
        if err != nil {
                return fmt.Errorf("failed to update configuration: %w", err)
        }

        return nil
}

// getCACert returns the currently configured CA certificate path and its accessibility status.
// It returns three values:
//   - certPath: the configured certificate path (empty string if not set)
//   - exists: true if a CA certificate is configured in the config
//   - accessible: true if the certificate file exists and is accessible on the filesystem
//
// Note: exists can be true while accessible is false if the file was deleted after configuration.
func getCACert(provider Provider) (certPath string, exists bool, accessible bool) {
        cfg := provider.GetConfig()

        if cfg.CACertificatePath == "" {
                return "", false, false
        }

        certPath = cfg.CACertificatePath
        exists = true

        // Check if the file is still accessible
        if _, err := os.Stat(certPath); err != nil {
                accessible = false
        } else {
                accessible = true
        }

        return certPath, exists, accessible
}

// unsetCACert removes the CA certificate configuration from the config file.
// If no CA certificate is currently configured, this function is a no-op and returns nil.
// Returns an error if updating the configuration fails.
func unsetCACert(provider Provider) error {
        cfg := provider.GetConfig()

        if cfg.CACertificatePath == "" {
                // Already unset, no-op
                return nil
        }

        // Update the configuration
        err := provider.UpdateConfig(func(c *Config) error {
                c.CACertificatePath = ""
                return nil
        })
        if err != nil {
                return fmt.Errorf("failed to update configuration: %w", err)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package config contains the definition of the application config structure
// and logic required to load and update it.
package config

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "os"
        "path"
        "time"

        "github.com/adrg/xdg"
        "go.yaml.in/yaml/v3"

        "github.com/stacklok/toolhive-core/env"
        "github.com/stacklok/toolhive/pkg/container/templates"
        "github.com/stacklok/toolhive/pkg/llm"
        "github.com/stacklok/toolhive/pkg/lockfile"
        "github.com/stacklok/toolhive/pkg/oidc"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// lockTimeout is the maximum time to wait for a file lock
const lockTimeout = 1 * time.Second

// Config represents the configuration of the application.
type Config struct {
        Secrets                      Secrets                             `yaml:"secrets"`
        Clients                      Clients                             `yaml:"clients"`
        RegistryUrl                  string                              `yaml:"registry_url"`
        RegistryApiUrl               string                              `yaml:"registry_api_url"`
        LocalRegistryPath            string                              `yaml:"local_registry_path"`
        AllowPrivateRegistryIp       bool                                `yaml:"allow_private_registry_ip"`
        CACertificatePath            string                              `yaml:"ca_certificate_path,omitempty"`
        OTEL                         OpenTelemetryConfig                 `yaml:"otel,omitempty"`
        DefaultGroupMigration        bool                                `yaml:"default_group_migration,omitempty"`
        TelemetryConfigMigration     bool                                `yaml:"telemetry_config_migration,omitempty"`
        MiddlewareTelemetryMigration bool                                `yaml:"middleware_telemetry_migration,omitempty"`
        SecretScopeMigration         bool                                `yaml:"secret_scope_migration,omitempty"`
        DisableUsageMetrics          bool                                `yaml:"disable_usage_metrics,omitempty"`
        BuildEnv                     map[string]string                   `yaml:"build_env,omitempty"`
        BuildEnvFromSecrets          map[string]string                   `yaml:"build_env_from_secrets,omitempty"`
        BuildEnvFromShell            []string                            `yaml:"build_env_from_shell,omitempty"`
        BuildAuthFiles               map[string]string                   `yaml:"build_auth_files,omitempty"`
        RuntimeConfigs               map[string]*templates.RuntimeConfig `yaml:"runtime_configs,omitempty"`
        RegistryAuth                 RegistryAuth                        `yaml:"registry_auth,omitempty"`
        LLM                          llm.Config                          `yaml:"llm,omitempty"`
}

// RegistryAuthTypeOAuth is the auth type for OAuth/OIDC authentication.
const RegistryAuthTypeOAuth = "oauth"

// RegistryAuth holds authentication configuration for remote registries.
type RegistryAuth struct {
        // Type is the authentication type: RegistryAuthTypeOAuth or "" (none).
        Type string `yaml:"type,omitempty"`

        // OAuth holds OAuth/OIDC authentication configuration.
        OAuth *RegistryOAuthConfig `yaml:"oauth,omitempty"`
}

// RegistryOAuthConfig holds OAuth/OIDC configuration for registry authentication.
// PKCE (S256) is always enforced per OAuth 2.1 requirements for public clients.
//
// This is a type alias for oidc.ClientConfig so that registry and LLM gateway
// authentication always share the same field set and validation logic.
type RegistryOAuthConfig = oidc.ClientConfig

// Secrets contains the settings for secrets management.
type Secrets struct {
        ProviderType   string `yaml:"provider_type"`
        SetupCompleted bool   `yaml:"setup_completed"`
}

// validateProviderType validates and returns the secrets provider type.
func validateProviderType(provider string) (secrets.ProviderType, error) {
        switch provider {
        case string(secrets.EncryptedType):
                return secrets.EncryptedType, nil
        case string(secrets.OnePasswordType):
                return secrets.OnePasswordType, nil
        case string(secrets.EnvironmentType):
                return secrets.EnvironmentType, nil
        default:
                return "", fmt.Errorf("invalid secrets provider type: %s (valid types: %s, %s, %s)",
                        provider,
                        string(secrets.EncryptedType),
                        string(secrets.OnePasswordType),
                        string(secrets.EnvironmentType),
                )
        }
}

// GetProviderType returns the secrets provider type from the environment variable or application config.
// It first checks the TOOLHIVE_SECRETS_PROVIDER environment variable (allowing Kubernetes deployments
// to override without local setup), and falls back to the config file.
// Returns ErrSecretsNotSetup only if the environment variable is not set and secrets have not been configured.
func (s *Secrets) GetProviderType() (secrets.ProviderType, error) {
        return s.GetProviderTypeWithEnv(&env.OSReader{})
}

// GetProviderTypeWithEnv returns the secrets provider type using the provided environment reader.
// This method allows for dependency injection of environment variable access for testing.
//
// Precedence order:
//  1. Environment variable (TOOLHIVE_SECRETS_PROVIDER) - takes highest precedence
//  2. Config file (requires SetupCompleted to be true)
//
// Special handling when SetupCompleted is false:
//   - Only the "environment" provider can be set via env var when SetupCompleted is false
//   - Other providers (encrypted, 1password) require setup and will return ErrSecretsNotSetup
//   - This prevents confusing errors later when trying to create providers that need setup
//
// Why environment provider bypasses SetupCompleted:
//   - In Kubernetes environments, pods don't have config files set up
//   - The operator sets TOOLHIVE_SECRETS_PROVIDER=environment via env vars
//   - The environment provider doesn't require "setup" - it reads directly from env vars
//   - This allows the operator to work without requiring users to run 'thv secret setup'
//
// For CLI users:
//   - If they set TOOLHIVE_SECRETS_PROVIDER=environment, it works without setup
//   - If they set TOOLHIVE_SECRETS_PROVIDER=encrypted/1password without setup, it returns an error
//   - This prevents confusing errors when providers fail to initialize later
func (s *Secrets) GetProviderTypeWithEnv(envReader env.Reader) (secrets.ProviderType, error) {
        // First check the environment variable (takes precedence) - this allows Kubernetes deployments
        // to override the secrets provider without requiring local setup
        envVar := envReader.Getenv(secrets.ProviderEnvVar)
        if envVar != "" {
                providerType, err := validateProviderType(envVar)
                if err != nil {
                        return "", err
                }

                // Special case: Only allow "environment" provider when SetupCompleted is false
                // Other providers (encrypted, 1password) require setup and will fail later when
                // trying to create them (keyring, password, 1Password CLI, etc.)
                if !s.SetupCompleted && providerType != secrets.EnvironmentType {
                        return "", fmt.Errorf(
                                "provider %q requires setup to be completed. "+
                                        "Only 'environment' provider can be used without setup. "+
                                        "Please run 'thv secret setup' or use TOOLHIVE_SECRETS_PROVIDER=environment",
                                providerType,
                        )
                }

                return providerType, nil
        }

        // Check if secrets setup has been completed (required for config file-based provider)
        // Only checked when environment variable is not set
        if !s.SetupCompleted {
                return "", secrets.ErrSecretsNotSetup
        }

        // Fall back to config file
        return validateProviderType(s.ProviderType)
}

// Clients contains settings for client configuration.
type Clients struct {
        RegisteredClients []string `yaml:"registered_clients"`
}

// defaultPathGenerator generates the default config path using xdg
var defaultPathGenerator = func() (string, error) {
        return xdg.ConfigFile("toolhive/config.yaml")
}

// getConfigPath is the current path generator, can be replaced in tests
var getConfigPath = defaultPathGenerator

// createNewConfigWithDefaults creates a new config with default values
func createNewConfigWithDefaults() Config {
        return Config{
                Secrets: Secrets{
                        ProviderType:   "", // No default provider - user must run setup
                        SetupCompleted: false,
                },
                RegistryUrl:                  "",
                RegistryApiUrl:               "",
                AllowPrivateRegistryIp:       false,
                DefaultGroupMigration:        false,
                TelemetryConfigMigration:     false,
                MiddlewareTelemetryMigration: false,
        }
}

// applyBackwardCompatibility applies backward compatibility fixes to existing configs.
// Any migration that needs to be persisted is written back to configPath, the same
// path the config was loaded from, so that path-based loads (e.g. PathProvider) stay
// isolated. An empty configPath falls back to the default path via saveToPath.
func applyBackwardCompatibility(config *Config, configPath string) error {
        // Hack - if the secrets provider type is set to the old `basic` type,
        // just change it to `encrypted`.
        if config.Secrets.ProviderType == "basic" {
                slog.Debug("cleaning up basic secrets provider, migrating to encrypted type")
                // Attempt to cleanup path, treat errors as non fatal.
                oldPath, err := xdg.DataFile("toolhive/secrets")
                if err == nil {
                        _ = os.Remove(oldPath)
                }
                config.Secrets.ProviderType = string(secrets.EncryptedType)
                err = config.saveToPath(configPath)
                if err != nil {
                        return fmt.Errorf("error updating config: %w", err)
                }
        }

        // Handle backward compatibility: if provider is set but setup_completed is false,
        // consider it as setup completed (for existing users)
        if config.Secrets.ProviderType != "" && !config.Secrets.SetupCompleted {
                config.Secrets.SetupCompleted = true
                err := config.saveToPath(configPath)
                if err != nil {
                        return fmt.Errorf("error updating config for backward compatibility: %w", err)
                }
        }

        return nil
}

// LoadOrCreateConfig fetches the application configuration.
// If it does not already exist - it will create a new config file with default values.
func LoadOrCreateConfig() (*Config, error) {
        provider := NewProvider()
        return provider.LoadOrCreateConfig()
}

// LoadOrCreateConfigWithDefaultPath is the internal implementation for loading config with the default path.
// This avoids circular dependency issues.
func LoadOrCreateConfigWithDefaultPath() (*Config, error) {
        configPath, err := getConfigPath()
        if err != nil {
                return nil, fmt.Errorf("unable to fetch config path: %w", err)
        }
        return LoadOrCreateConfigFromPath(configPath)
}

// LoadOrCreateConfigWithPath fetches the application configuration from a specific path.
// If configPath is empty, it uses the default path.
// If it does not already exist - it will create a new config file with default values.
func LoadOrCreateConfigWithPath(configPath string) (*Config, error) {
        if configPath == "" {
                // When no path is specified, use the provider pattern to handle runtime-specific behavior
                return LoadOrCreateConfig()
        }

        return LoadOrCreateConfigFromPath(configPath)
}

// LoadOrCreateConfigFromPath is the core implementation for loading/creating config from a specific path
func LoadOrCreateConfigFromPath(configPath string) (*Config, error) {
        var config Config
        var err error

        // Check to see if the config file already exists.
        configPath = path.Clean(configPath)
        newConfig := false
        // #nosec G304: File path is not configurable at this time.
        _, err = os.Stat(configPath)
        if err != nil {
                if errors.Is(err, os.ErrNotExist) {
                        newConfig = true
                } else {
                        return nil, fmt.Errorf("failed to stat secrets file: %w", err)
                }
        }

        if newConfig {
                // Create a new config with default values.
                config = createNewConfigWithDefaults()

                // Persist the new default to disk using the specific path
                //nolint:gosec // G706: config path is validated and cleaned before use
                slog.Debug("initializing configuration file", "path", configPath)
                err = config.saveToPath(configPath)
                if err != nil {
                        return nil, fmt.Errorf("failed to write default config: %w", err)
                }
        } else {
                // Load the existing config and decode.
                // #nosec G304: File path is not configurable at this time.
                configFile, err := os.ReadFile(configPath)
                if err != nil {
                        return nil, fmt.Errorf("unable to read config file %s: %w", configPath, err)
                }
                err = yaml.Unmarshal(configFile, &config)
                if err != nil {
                        return nil, fmt.Errorf("failed to parse config file yaml: %w", err)
                }

                // Apply backward compatibility fixes, persisting any migration back to the
                // same path the config was loaded from.
                err = applyBackwardCompatibility(&config, configPath)
                if err != nil {
                        return nil, fmt.Errorf("failed to apply backward compatibility fixes: %w", err)
                }
        }

        return &config, nil
}

// saveToPath serializes the config struct and writes it to a specific path.
// If configPath is empty, it uses the default path.
func (c *Config) saveToPath(configPath string) error {
        if configPath == "" {
                var err error
                configPath, err = getConfigPath()
                if err != nil {
                        return fmt.Errorf("unable to fetch config path: %w", err)
                }
        }

        configBytes, err := yaml.Marshal(c)
        if err != nil {
                return fmt.Errorf("error serializing config file: %w", err)
        }

        err = os.WriteFile(configPath, configBytes, 0600)
        if err != nil {
                return fmt.Errorf("error writing config file: %w", err)
        }
        return nil
}

// UpdateConfig locks a separate lock file, reads from disk, applies the changes
// from the anonymous function, writes to disk and unlocks the file.
func UpdateConfig(updateFn func(*Config) error) error {
        provider := NewProvider()
        return provider.UpdateConfig(updateFn)
}

// UpdateConfigAtPath locks a separate lock file, reads from disk, applies the changes
// from the anonymous function, writes to disk and unlocks the file.
// If configPath is empty, it uses the default path.
func UpdateConfigAtPath(configPath string, updateFn func(*Config) error) error {
        if configPath == "" {
                var err error
                configPath, err = getConfigPath()
                if err != nil {
                        return fmt.Errorf("unable to fetch config path: %w", err)
                }
        }

        // Use a separate lock file for cross-platform compatibility
        lockPath := configPath + ".lock"
        fileLock := lockfile.NewTrackedLock(lockPath)
        ctx, cancel := context.WithTimeout(context.Background(), lockTimeout)
        defer cancel()

        // Try and acquire a file lock.
        locked, err := fileLock.TryLockContext(ctx, 100*time.Millisecond)
        if err != nil {
                return fmt.Errorf("failed to acquire lock: %w", err)
        }
        if !locked {
                return fmt.Errorf("failed to acquire lock: timeout after %v", lockTimeout)
        }
        defer lockfile.ReleaseTrackedLock(lockPath, fileLock)

        // Load the config after acquiring the lock to avoid race conditions
        c, err := LoadOrCreateConfigWithPath(configPath)
        if err != nil {
                return fmt.Errorf("failed to load config from disk: %w", err)
        }

        // Apply changes to the config file.
        if err := updateFn(c); err != nil {
                return err
        }

        // Write the updated config to disk.
        err = c.saveToPath(configPath)
        if err != nil {
                return fmt.Errorf("failed to save config: %w", err)
        }

        // Lock is released automatically when the function returns.
        return nil
}

// OpenTelemetryConfig contains the settings for OpenTelemetry configuration.
//
// Fields whose CLI default is true use *bool so that "never set" (nil) is
// distinguishable from "explicitly set to false". Without this, the config
// zero value (false) would silently override the CLI default (true) for
// every user who never touched the field. Plain bool with omitempty is fine
// for fields that default to false on both the CLI and config sides.
type OpenTelemetryConfig struct {
        Endpoint                    string   `yaml:"endpoint,omitempty"`
        SamplingRate                float64  `yaml:"sampling-rate,omitempty"`
        EnvVars                     []string `yaml:"env-vars,omitempty"`
        MetricsEnabled              *bool    `yaml:"metrics-enabled"`
        TracingEnabled              *bool    `yaml:"tracing-enabled"`
        Insecure                    bool     `yaml:"insecure,omitempty"`
        EnablePrometheusMetricsPath bool     `yaml:"enable-prometheus-metrics-path,omitempty"`
        UseLegacyAttributes         *bool    `yaml:"use-legacy-attributes"`
}

// getRuntimeConfig returns the runtime configuration for a given transport type
func getRuntimeConfig(provider Provider, transportType string) (*templates.RuntimeConfig, error) {
        config := provider.GetConfig()
        if config.RuntimeConfigs == nil {
                return nil, nil
        }

        runtimeConfig, exists := config.RuntimeConfigs[transportType]
        if !exists {
                return nil, nil
        }

        return runtimeConfig, nil
}

// setRuntimeConfig sets the runtime configuration for a given transport type.
// It validates the configuration before storing to prevent shell injection
// when values are interpolated into Dockerfile templates.
func setRuntimeConfig(provider Provider, transportType string, runtimeConfig *templates.RuntimeConfig) error {
        if runtimeConfig != nil {
                if err := runtimeConfig.Validate(); err != nil {
                        return fmt.Errorf("invalid runtime config: %w", err)
                }
        }

        return provider.UpdateConfig(func(c *Config) error {
                if c.RuntimeConfigs == nil {
                        c.RuntimeConfigs = make(map[string]*templates.RuntimeConfig)
                }
                c.RuntimeConfigs[transportType] = runtimeConfig
                return nil
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "errors"
        "fmt"
)

var (
        // ErrRegistryTimeout is returned when a registry validation times out
        ErrRegistryTimeout = errors.New("registry validation timed out")

        // ErrRegistryUnreachable is returned when a registry cannot be reached
        ErrRegistryUnreachable = errors.New("registry is unreachable")

        // ErrRegistryValidationFailed is returned when registry validation fails
        ErrRegistryValidationFailed = errors.New("registry validation failed")
)

// RegistryError wraps registry-related errors with additional context
type RegistryError struct {
        // Type is the type of registry (url, api, file)
        Type string
        // URL is the registry URL or path
        URL string
        // Err is the underlying error
        Err error
}

func (e *RegistryError) Error() string {
        return fmt.Sprintf("registry error for %s (%s): %v", e.Type, e.URL, e.Err)
}

func (e *RegistryError) Unwrap() error {
        return e.Err
}

// IsTimeout checks if the error is a timeout error
func (e *RegistryError) IsTimeout() bool {
        return errors.Is(e.Err, ErrRegistryTimeout)
}

// IsUnreachable checks if the error is an unreachable error
func (e *RegistryError) IsUnreachable() bool {
        return errors.Is(e.Err, ErrRegistryUnreachable)
}

// IsValidationFailed checks if the error is a validation error
func (e *RegistryError) IsValidationFailed() bool {
        return errors.Is(e.Err, ErrRegistryValidationFailed)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

// ProviderFactory is a function that optionally creates a Provider.
// Returning nil signals that the caller should fall back to the default provider.
type ProviderFactory func() Provider

// registeredFactory is the package-level factory, nil by default.
var registeredFactory ProviderFactory

// RegisterProviderFactory sets a custom factory to be used by NewProvider.
// It must be called before the first call to NewProvider (typically in main or init).
// Calling it a second time replaces the previously registered factory.
func RegisterProviderFactory(f ProviderFactory) {
        registeredFactory = f
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/container/templates"
)

// Provider defines the interface for configuration operations
//
//go:generate mockgen -destination=mocks/mock_provider.go -package=mocks -source=interface.go Provider
type Provider interface {
        GetConfig() *Config
        UpdateConfig(updateFn func(*Config) error) error
        LoadOrCreateConfig() (*Config, error)

        // Registry operations
        SetRegistryURL(registryURL string, allowPrivateRegistryIp bool) error
        SetRegistryAPI(apiURL string, allowPrivateRegistryIp bool) error
        SetRegistryFile(registryPath string) error
        UnsetRegistry() error
        GetRegistryConfig() (url, localPath string, allowPrivateIP bool, registryType string)

        // CA certificate operations
        SetCACert(certPath string) error
        GetCACert() (certPath string, exists bool, accessible bool)
        UnsetCACert() error

        // Build environment operations
        SetBuildEnv(key, value string) error
        GetBuildEnv(key string) (value string, exists bool)
        GetAllBuildEnv() map[string]string
        UnsetBuildEnv(key string) error
        UnsetAllBuildEnv() error

        // Build environment from secrets operations
        SetBuildEnvFromSecret(key, secretName string) error
        GetBuildEnvFromSecret(key string) (secretName string, exists bool)
        GetAllBuildEnvFromSecrets() map[string]string
        UnsetBuildEnvFromSecret(key string) error

        // Build environment from shell operations
        SetBuildEnvFromShell(key string) error
        GetBuildEnvFromShell(key string) (exists bool)
        GetAllBuildEnvFromShell() []string
        UnsetBuildEnvFromShell(key string) error

        // Build auth file operations (content stored in secrets provider, not config)
        MarkBuildAuthFileConfigured(name string) error
        IsBuildAuthFileConfigured(name string) bool
        GetConfiguredBuildAuthFiles() []string
        UnsetBuildAuthFile(name string) error
        UnsetAllBuildAuthFiles() error

        // Runtime configuration operations
        GetRuntimeConfig(transportType string) (*templates.RuntimeConfig, error)
        SetRuntimeConfig(transportType string, config *templates.RuntimeConfig) error
}

// DefaultProvider implements Provider using the default XDG config path
type DefaultProvider struct{}

// NewDefaultProvider creates a new default config provider
func NewDefaultProvider() *DefaultProvider {
        return &DefaultProvider{}
}

// GetConfig returns the singleton config (for backward compatibility)
func (*DefaultProvider) GetConfig() *Config {
        return getSingletonConfig()
}

// UpdateConfig updates the config using the default path
func (*DefaultProvider) UpdateConfig(updateFn func(*Config) error) error {
        return UpdateConfigAtPath("", updateFn)
}

// LoadOrCreateConfig loads or creates config using the default path
func (*DefaultProvider) LoadOrCreateConfig() (*Config, error) {
        return LoadOrCreateConfigWithDefaultPath()
}

// SetRegistryURL validates and sets a registry URL
func (d *DefaultProvider) SetRegistryURL(registryURL string, allowPrivateRegistryIp bool) error {
        return setRegistryURL(d, registryURL, allowPrivateRegistryIp)
}

// SetRegistryAPI validates and sets an MCP Registry API endpoint
func (d *DefaultProvider) SetRegistryAPI(apiURL string, allowPrivateRegistryIp bool) error {
        return setRegistryAPI(d, apiURL, allowPrivateRegistryIp)
}

// SetRegistryFile validates and sets a local registry file
func (d *DefaultProvider) SetRegistryFile(registryPath string) error {
        return setRegistryFile(d, registryPath)
}

// UnsetRegistry resets registry configuration to defaults
func (d *DefaultProvider) UnsetRegistry() error {
        return unsetRegistry(d)
}

// GetRegistryConfig returns current registry configuration
func (d *DefaultProvider) GetRegistryConfig() (url, localPath string, allowPrivateIP bool, registryType string) {
        return getRegistryConfig(d)
}

// SetCACert validates and sets the CA certificate path
func (d *DefaultProvider) SetCACert(certPath string) error {
        return setCACert(d, certPath)
}

// GetCACert returns the currently configured CA certificate path and its accessibility status
func (d *DefaultProvider) GetCACert() (certPath string, exists bool, accessible bool) {
        return getCACert(d)
}

// UnsetCACert removes the CA certificate configuration
func (d *DefaultProvider) UnsetCACert() error {
        return unsetCACert(d)
}

// SetBuildEnv validates and sets a build environment variable
func (d *DefaultProvider) SetBuildEnv(key, value string) error {
        return setBuildEnv(d, key, value)
}

// GetBuildEnv returns a specific build environment variable
func (d *DefaultProvider) GetBuildEnv(key string) (value string, exists bool) {
        return getBuildEnv(d, key)
}

// GetAllBuildEnv returns all build environment variables
func (d *DefaultProvider) GetAllBuildEnv() map[string]string {
        return getAllBuildEnv(d)
}

// UnsetBuildEnv removes a specific build environment variable
func (d *DefaultProvider) UnsetBuildEnv(key string) error {
        return unsetBuildEnv(d, key)
}

// UnsetAllBuildEnv removes all build environment variables
func (d *DefaultProvider) UnsetAllBuildEnv() error {
        return unsetAllBuildEnv(d)
}

// SetBuildEnvFromSecret validates and sets a secret reference for a build environment variable
func (d *DefaultProvider) SetBuildEnvFromSecret(key, secretName string) error {
        return setBuildEnvFromSecret(d, key, secretName)
}

// GetBuildEnvFromSecret retrieves the secret name for a build environment variable
func (d *DefaultProvider) GetBuildEnvFromSecret(key string) (secretName string, exists bool) {
        return getBuildEnvFromSecret(d, key)
}

// GetAllBuildEnvFromSecrets returns all build env secret references
func (d *DefaultProvider) GetAllBuildEnvFromSecrets() map[string]string {
        return getAllBuildEnvFromSecrets(d)
}

// UnsetBuildEnvFromSecret removes a secret reference
func (d *DefaultProvider) UnsetBuildEnvFromSecret(key string) error {
        return unsetBuildEnvFromSecret(d, key)
}

// SetBuildEnvFromShell adds an environment variable name to read from shell at build time
func (d *DefaultProvider) SetBuildEnvFromShell(key string) error {
        return setBuildEnvFromShell(d, key)
}

// GetBuildEnvFromShell checks if a key is configured to read from shell
func (d *DefaultProvider) GetBuildEnvFromShell(key string) bool {
        return getBuildEnvFromShell(d, key)
}

// GetAllBuildEnvFromShell returns all keys configured to read from shell
func (d *DefaultProvider) GetAllBuildEnvFromShell() []string {
        return getAllBuildEnvFromShell(d)
}

// UnsetBuildEnvFromShell removes a key from shell environment list
func (d *DefaultProvider) UnsetBuildEnvFromShell(key string) error {
        return unsetBuildEnvFromShell(d, key)
}

// MarkBuildAuthFileConfigured marks an auth file type as configured
func (d *DefaultProvider) MarkBuildAuthFileConfigured(name string) error {
        return markBuildAuthFileConfigured(d, name)
}

// IsBuildAuthFileConfigured checks if an auth file type is configured
func (d *DefaultProvider) IsBuildAuthFileConfigured(name string) bool {
        return isBuildAuthFileConfigured(d, name)
}

// GetConfiguredBuildAuthFiles returns list of configured auth file types
func (d *DefaultProvider) GetConfiguredBuildAuthFiles() []string {
        return getConfiguredBuildAuthFiles(d)
}

// UnsetBuildAuthFile removes an auth file configuration
func (d *DefaultProvider) UnsetBuildAuthFile(name string) error {
        return unsetBuildAuthFile(d, name)
}

// UnsetAllBuildAuthFiles removes all auth file configurations
func (d *DefaultProvider) UnsetAllBuildAuthFiles() error {
        return unsetAllBuildAuthFiles(d)
}

// GetRuntimeConfig returns the runtime configuration for a given transport type
func (d *DefaultProvider) GetRuntimeConfig(transportType string) (*templates.RuntimeConfig, error) {
        return getRuntimeConfig(d, transportType)
}

// SetRuntimeConfig sets the runtime configuration for a given transport type
func (d *DefaultProvider) SetRuntimeConfig(transportType string, config *templates.RuntimeConfig) error {
        return setRuntimeConfig(d, transportType, config)
}

// PathProvider implements Provider using a specific config path
type PathProvider struct {
        configPath string
}

// NewPathProvider creates a new config provider with a specific path
func NewPathProvider(configPath string) *PathProvider {
        return &PathProvider{configPath: configPath}
}

// GetConfig loads and returns the config from the specific path
func (p *PathProvider) GetConfig() *Config {
        config, err := LoadOrCreateConfigWithPath(p.configPath)
        if err != nil {
                // Return default config on error, similar to singleton behavior
                defaultConfig := createNewConfigWithDefaults()
                return &defaultConfig
        }
        return config
}

// UpdateConfig updates the config at the specific path
func (p *PathProvider) UpdateConfig(updateFn func(*Config) error) error {
        return UpdateConfigAtPath(p.configPath, updateFn)
}

// LoadOrCreateConfig loads or creates config at the specific path
func (p *PathProvider) LoadOrCreateConfig() (*Config, error) {
        return LoadOrCreateConfigWithPath(p.configPath)
}

// SetRegistryURL validates and sets a registry URL
func (p *PathProvider) SetRegistryURL(registryURL string, allowPrivateRegistryIp bool) error {
        return setRegistryURL(p, registryURL, allowPrivateRegistryIp)
}

// SetRegistryAPI validates and sets an MCP Registry API endpoint
func (p *PathProvider) SetRegistryAPI(apiURL string, allowPrivateRegistryIp bool) error {
        return setRegistryAPI(p, apiURL, allowPrivateRegistryIp)
}

// SetRegistryFile validates and sets a local registry file
func (p *PathProvider) SetRegistryFile(registryPath string) error {
        return setRegistryFile(p, registryPath)
}

// UnsetRegistry resets registry configuration to defaults
func (p *PathProvider) UnsetRegistry() error {
        return unsetRegistry(p)
}

// GetRegistryConfig returns current registry configuration
func (p *PathProvider) GetRegistryConfig() (url, localPath string, allowPrivateIP bool, registryType string) {
        return getRegistryConfig(p)
}

// SetCACert validates and sets the CA certificate path
func (p *PathProvider) SetCACert(certPath string) error {
        return setCACert(p, certPath)
}

// GetCACert returns the currently configured CA certificate path and its accessibility status
func (p *PathProvider) GetCACert() (certPath string, exists bool, accessible bool) {
        return getCACert(p)
}

// UnsetCACert removes the CA certificate configuration
func (p *PathProvider) UnsetCACert() error {
        return unsetCACert(p)
}

// SetBuildEnv validates and sets a build environment variable
func (p *PathProvider) SetBuildEnv(key, value string) error {
        return setBuildEnv(p, key, value)
}

// GetBuildEnv returns a specific build environment variable
func (p *PathProvider) GetBuildEnv(key string) (value string, exists bool) {
        return getBuildEnv(p, key)
}

// GetAllBuildEnv returns all build environment variables
func (p *PathProvider) GetAllBuildEnv() map[string]string {
        return getAllBuildEnv(p)
}

// UnsetBuildEnv removes a specific build environment variable
func (p *PathProvider) UnsetBuildEnv(key string) error {
        return unsetBuildEnv(p, key)
}

// UnsetAllBuildEnv removes all build environment variables
func (p *PathProvider) UnsetAllBuildEnv() error {
        return unsetAllBuildEnv(p)
}

// SetBuildEnvFromSecret validates and sets a secret reference for a build environment variable
func (p *PathProvider) SetBuildEnvFromSecret(key, secretName string) error {
        return setBuildEnvFromSecret(p, key, secretName)
}

// GetBuildEnvFromSecret retrieves the secret name for a build environment variable
func (p *PathProvider) GetBuildEnvFromSecret(key string) (secretName string, exists bool) {
        return getBuildEnvFromSecret(p, key)
}

// GetAllBuildEnvFromSecrets returns all build env secret references
func (p *PathProvider) GetAllBuildEnvFromSecrets() map[string]string {
        return getAllBuildEnvFromSecrets(p)
}

// UnsetBuildEnvFromSecret removes a secret reference
func (p *PathProvider) UnsetBuildEnvFromSecret(key string) error {
        return unsetBuildEnvFromSecret(p, key)
}

// SetBuildEnvFromShell adds an environment variable name to read from shell at build time
func (p *PathProvider) SetBuildEnvFromShell(key string) error {
        return setBuildEnvFromShell(p, key)
}

// GetBuildEnvFromShell checks if a key is configured to read from shell
func (p *PathProvider) GetBuildEnvFromShell(key string) bool {
        return getBuildEnvFromShell(p, key)
}

// GetAllBuildEnvFromShell returns all keys configured to read from shell
func (p *PathProvider) GetAllBuildEnvFromShell() []string {
        return getAllBuildEnvFromShell(p)
}

// UnsetBuildEnvFromShell removes a key from shell environment list
func (p *PathProvider) UnsetBuildEnvFromShell(key string) error {
        return unsetBuildEnvFromShell(p, key)
}

// MarkBuildAuthFileConfigured marks an auth file type as configured
func (p *PathProvider) MarkBuildAuthFileConfigured(name string) error {
        return markBuildAuthFileConfigured(p, name)
}

// IsBuildAuthFileConfigured checks if an auth file type is configured
func (p *PathProvider) IsBuildAuthFileConfigured(name string) bool {
        return isBuildAuthFileConfigured(p, name)
}

// GetConfiguredBuildAuthFiles returns list of configured auth file types
func (p *PathProvider) GetConfiguredBuildAuthFiles() []string {
        return getConfiguredBuildAuthFiles(p)
}

// UnsetBuildAuthFile removes an auth file configuration
func (p *PathProvider) UnsetBuildAuthFile(name string) error {
        return unsetBuildAuthFile(p, name)
}

// UnsetAllBuildAuthFiles removes all auth file configurations
func (p *PathProvider) UnsetAllBuildAuthFiles() error {
        return unsetAllBuildAuthFiles(p)
}

// GetRuntimeConfig returns the runtime configuration for a given transport type
func (p *PathProvider) GetRuntimeConfig(transportType string) (*templates.RuntimeConfig, error) {
        return getRuntimeConfig(p, transportType)
}

// SetRuntimeConfig sets the runtime configuration for a given transport type
func (p *PathProvider) SetRuntimeConfig(transportType string, config *templates.RuntimeConfig) error {
        return setRuntimeConfig(p, transportType, config)
}

// KubernetesProvider is a no-op implementation of Provider for Kubernetes environments.
// In Kubernetes, configuration is managed by the cluster, not by local files.
type KubernetesProvider struct{}

// NewKubernetesProvider creates a new no-op config provider for Kubernetes environments
func NewKubernetesProvider() *KubernetesProvider {
        return &KubernetesProvider{}
}

// GetConfig returns a default config for Kubernetes environments
func (*KubernetesProvider) GetConfig() *Config {
        config := createNewConfigWithDefaults()
        return &config
}

// UpdateConfig is a no-op for Kubernetes environments
func (*KubernetesProvider) UpdateConfig(_ func(*Config) error) error {
        return nil
}

// LoadOrCreateConfig returns a default config for Kubernetes environments
func (*KubernetesProvider) LoadOrCreateConfig() (*Config, error) {
        config := createNewConfigWithDefaults()
        return &config, nil
}

// SetRegistryURL is a no-op for Kubernetes environments
func (*KubernetesProvider) SetRegistryURL(_ string, _ bool) error {
        return nil
}

// SetRegistryAPI is a no-op for Kubernetes environments
func (*KubernetesProvider) SetRegistryAPI(_ string, _ bool) error {
        return nil
}

// SetRegistryFile is a no-op for Kubernetes environments
func (*KubernetesProvider) SetRegistryFile(_ string) error {
        return nil
}

// UnsetRegistry is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetRegistry() error {
        return nil
}

// GetRegistryConfig returns empty registry configuration for Kubernetes environments
func (*KubernetesProvider) GetRegistryConfig() (url, localPath string, allowPrivateIP bool, registryType string) {
        return "", "", false, ""
}

// SetCACert is a no-op for Kubernetes environments
func (*KubernetesProvider) SetCACert(_ string) error {
        return nil
}

// GetCACert returns empty CA cert configuration for Kubernetes environments
func (*KubernetesProvider) GetCACert() (certPath string, exists bool, accessible bool) {
        return "", false, false
}

// UnsetCACert is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetCACert() error {
        return nil
}

// SetBuildEnv is a no-op for Kubernetes environments
func (*KubernetesProvider) SetBuildEnv(_, _ string) error {
        return nil
}

// GetBuildEnv returns empty for Kubernetes environments
func (*KubernetesProvider) GetBuildEnv(_ string) (value string, exists bool) {
        return "", false
}

// GetAllBuildEnv returns empty map for Kubernetes environments
func (*KubernetesProvider) GetAllBuildEnv() map[string]string {
        return make(map[string]string)
}

// UnsetBuildEnv is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetBuildEnv(_ string) error {
        return nil
}

// UnsetAllBuildEnv is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetAllBuildEnv() error {
        return nil
}

// SetBuildEnvFromSecret is a no-op for Kubernetes environments
func (*KubernetesProvider) SetBuildEnvFromSecret(_, _ string) error {
        return nil
}

// GetBuildEnvFromSecret returns empty for Kubernetes environments
func (*KubernetesProvider) GetBuildEnvFromSecret(_ string) (secretName string, exists bool) {
        return "", false
}

// GetAllBuildEnvFromSecrets returns empty map for Kubernetes environments
func (*KubernetesProvider) GetAllBuildEnvFromSecrets() map[string]string {
        return make(map[string]string)
}

// UnsetBuildEnvFromSecret is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetBuildEnvFromSecret(_ string) error {
        return nil
}

// SetBuildEnvFromShell is a no-op for Kubernetes environments
func (*KubernetesProvider) SetBuildEnvFromShell(_ string) error {
        return nil
}

// GetBuildEnvFromShell returns false for Kubernetes environments
func (*KubernetesProvider) GetBuildEnvFromShell(_ string) bool {
        return false
}

// GetAllBuildEnvFromShell returns empty slice for Kubernetes environments
func (*KubernetesProvider) GetAllBuildEnvFromShell() []string {
        return []string{}
}

// UnsetBuildEnvFromShell is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetBuildEnvFromShell(_ string) error {
        return nil
}

// MarkBuildAuthFileConfigured is a no-op for Kubernetes environments
func (*KubernetesProvider) MarkBuildAuthFileConfigured(_ string) error {
        return nil
}

// IsBuildAuthFileConfigured returns false for Kubernetes environments
func (*KubernetesProvider) IsBuildAuthFileConfigured(_ string) bool {
        return false
}

// GetConfiguredBuildAuthFiles returns empty slice for Kubernetes environments
func (*KubernetesProvider) GetConfiguredBuildAuthFiles() []string {
        return []string{}
}

// UnsetBuildAuthFile is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetBuildAuthFile(_ string) error {
        return nil
}

// UnsetAllBuildAuthFiles is a no-op for Kubernetes environments
func (*KubernetesProvider) UnsetAllBuildAuthFiles() error {
        return nil
}

// GetRuntimeConfig returns nil for Kubernetes environments (runtime config not supported)
func (*KubernetesProvider) GetRuntimeConfig(_ string) (*templates.RuntimeConfig, error) {
        return nil, nil
}

// SetRuntimeConfig is a no-op for Kubernetes environments
func (*KubernetesProvider) SetRuntimeConfig(_ string, _ *templates.RuntimeConfig) error {
        return nil
}

// NewProvider creates the appropriate config provider based on the runtime environment.
// If a custom ProviderFactory has been registered via RegisterProviderFactory and it
// returns a non-nil Provider, that provider is used. Otherwise, the built-in selection
// logic applies.
func NewProvider() Provider {
        if registeredFactory != nil {
                if p := registeredFactory(); p != nil {
                        return p
                }
        }
        if runtime.IsKubernetesRuntime() {
                return NewKubernetesProvider()
        }
        return NewDefaultProvider()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net"
        "net/http"
        neturl "net/url"
        "os"
        "path/filepath"
        "strings"
        "time"

        registrytypes "github.com/stacklok/toolhive-core/registry/types"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/registry/legacyhint"
)

const (
        // RegistryTypeFile represents a local file registry
        RegistryTypeFile = "file"
        // RegistryTypeURL represents a remote URL registry
        RegistryTypeURL = "url"
        // RegistryTypeAPI represents an MCP Registry API endpoint
        RegistryTypeAPI = "api"
        // RegistryTypeDefault represents a built-in registry
        RegistryTypeDefault = "default"
)

// DetectRegistryType determines if input is a URL or file path and returns cleaned path
func DetectRegistryType(input string, allowPrivateIPs bool) (registryType string, cleanPath string) {
        // Check for explicit file:// protocol
        if strings.HasPrefix(input, "file://") {
                return RegistryTypeFile, strings.TrimPrefix(input, "file://")
        }

        // Check for HTTP/HTTPS URLs
        if networking.IsURL(input) {
                // If URL ends with .json, treat as static registry file
                if strings.HasSuffix(input, ".json") {
                        return RegistryTypeURL, input
                }

                // For URLs without .json extension, probe to determine the type
                registryType := probeRegistryURL(input, allowPrivateIPs)
                return registryType, input
        }

        // Default: treat as file path
        return RegistryTypeFile, filepath.Clean(input)
}

// probeRegistryURL attempts to determine if a URL is a static JSON file or an API endpoint
// by checking if the MCP Registry API endpoint (/v0.1/servers) exists and returns valid API responses.
// Uses a 5-second timeout for connectivity check.
func probeRegistryURL(url string, allowPrivateIPs bool) string {
        // Create HTTP client for probing with user's private IP preference and 5-second timeout
        // If private IPs are allowed, also allow HTTP (for localhost testing)
        builder := networking.NewHttpClientBuilder().
                WithPrivateIPs(allowPrivateIPs).
                WithTimeout(5 * time.Second)
        if allowPrivateIPs {
                builder = builder.WithInsecureAllowHTTP(true)
        }
        client, err := builder.Build()
        if err != nil {
                // If we can't create a client, default to static JSON
                return RegistryTypeURL
        }

        // Check if the MCP Registry API endpoint exists by trying a lightweight GET request
        // Note: We use GET instead of HEAD because some API implementations don't support HEAD
        apiURL, err := neturl.JoinPath(url, "/v0.1/servers")
        if err == nil {
                // Add query parameters to minimize response size
                params := neturl.Values{}
                params.Add("limit", "1")
                params.Add("version", "latest")
                fullAPIURL := fmt.Sprintf("%s?%s", apiURL, params.Encode())

                resp, err := client.Get(fullAPIURL)
                if err == nil {
                        defer func() {
                                if err := resp.Body.Close(); err != nil {
                                        slog.Debug("failed to close response body", "error", err)
                                }
                        }()
                        // If API endpoint returns 2xx or 401/403 (auth errors), it's an API
                        // 404 means endpoint doesn't exist, 405 means method not supported, 5xx means server error
                        if resp.StatusCode >= 200 && resp.StatusCode < 300 {
                                // Verify the response looks like an API response
                                if isValidAPIResponse(resp) {
                                        return RegistryTypeAPI
                                }
                        } else if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
                                // Auth errors indicate an API endpoint (it exists but requires auth)
                                return RegistryTypeAPI
                        }
                }
        }

        // If no API endpoint found, check if it's valid registry JSON
        if err := isValidRegistryJSON(client, url); err == nil {
                return RegistryTypeURL
        }

        // Default to static JSON file (validation will catch errors later)
        return RegistryTypeURL
}

// isValidAPIResponse checks if an HTTP response contains a valid MCP Registry API response
// by verifying the JSON structure matches the expected API format (ServerListResponse).
func isValidAPIResponse(resp *http.Response) bool {
        // Check Content-Type header
        contentType := resp.Header.Get("Content-Type")
        if !strings.Contains(contentType, "application/json") {
                return false
        }

        // Try to parse as MCP Registry API response structure
        var data map[string]interface{}
        if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
                return false
        }

        // Check for API-specific structure (servers array and metadata object)
        servers, hasServers := data["servers"]
        metadata, hasMetadata := data["metadata"]

        // Valid API response should have both 'servers' (array) and 'metadata' (object)
        if !hasServers || !hasMetadata {
                return false
        }

        // Verify servers is an array
        if _, ok := servers.([]interface{}); !ok {
                return false
        }

        // Verify metadata is an object
        if _, ok := metadata.(map[string]interface{}); !ok {
                return false
        }

        return true
}

// isValidRegistryJSON checks if a URL returns valid ToolHive registry JSON
// by attempting to parse it. Accepts both upstream and legacy formats.
func isValidRegistryJSON(client *http.Client, url string) error {
        resp, err := client.Get(url)
        if err != nil {
                return classifyNetworkError(err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        data, err := io.ReadAll(resp.Body)
        if err != nil {
                return fmt.Errorf("%w: failed to read response body: %v", ErrRegistryValidationFailed, err)
        }

        if err := parseAndValidateUpstreamRegistry(data); err != nil {
                return fmt.Errorf("%w: %v", ErrRegistryValidationFailed, err)
        }
        return nil
}

// parseAndValidateUpstreamRegistry parses data as an UpstreamRegistry and
// verifies that it contains at least one server, skill, or plugin. If the data
// looks like the legacy ToolHive registry format, returns an error containing
// the migration hint message.
func parseAndValidateUpstreamRegistry(data []byte) error {
        if legacyhint.Looks(data) {
                return errors.New(legacyhint.MigrationMessage)
        }
        var upstream registrytypes.UpstreamRegistry
        if err := json.Unmarshal(data, &upstream); err != nil {
                return fmt.Errorf("invalid upstream registry format: %w", err)
        }
        if len(upstream.Data.Servers) == 0 && len(upstream.Data.Skills) == 0 && len(upstream.Data.Plugins) == 0 {
                return errors.New("upstream registry contains no servers, skills, or plugins")
        }
        return nil
}

// classifyNetworkError wraps network errors with appropriate custom error types
func classifyNetworkError(err error) error {
        if err == nil {
                return nil
        }

        // Check for timeout errors
        var netErr net.Error
        if errors.As(err, &netErr) && netErr.Timeout() {
                return fmt.Errorf("%w: %v", ErrRegistryTimeout, err)
        }

        // Check for context deadline exceeded (another form of timeout)
        if errors.Is(err, context.DeadlineExceeded) {
                return fmt.Errorf("%w: %v", ErrRegistryTimeout, err)
        }

        // Check for connection errors
        errStr := err.Error()
        if strings.Contains(errStr, "connection refused") ||
                strings.Contains(errStr, "no route to host") ||
                strings.Contains(errStr, "network is unreachable") ||
                strings.Contains(errStr, networking.ErrPrivateIpAddress) {
                return fmt.Errorf("%w: %v", ErrRegistryUnreachable, err)
        }

        // Check for DNS errors (name resolution failures)
        var dnsErr *net.DNSError
        if errors.As(err, &dnsErr) {
                return fmt.Errorf("%w: %v", ErrRegistryUnreachable, err)
        }

        // Default: return original error
        return err
}

// setRegistryURL validates and sets a registry URL using the provided provider
// Validates connectivity with a 5-second timeout.
func setRegistryURL(provider Provider, registryURL string, allowPrivateRegistryIp bool) error {
        // Validate URL scheme
        _, err := validateURLScheme(registryURL, allowPrivateRegistryIp)
        if err != nil {
                return fmt.Errorf("invalid registry URL: %w", err)
        }

        // Build HTTP client with appropriate security settings and 5-second timeout
        builder := networking.NewHttpClientBuilder().
                WithPrivateIPs(allowPrivateRegistryIp).
                WithTimeout(5 * time.Second)
        if allowPrivateRegistryIp {
                builder = builder.WithInsecureAllowHTTP(true)
        }
        registryClient, err := builder.Build()
        if err != nil {
                return fmt.Errorf("failed to create HTTP client: %w", err)
        }

        // Check for private IP addresses if not allowed
        if !allowPrivateRegistryIp {
                _, err = registryClient.Get(registryURL)
                if err != nil && strings.Contains(fmt.Sprint(err), networking.ErrPrivateIpAddress) {
                        return &RegistryError{
                                Type: RegistryTypeURL,
                                URL:  registryURL,
                                Err:  classifyNetworkError(err),
                        }
                }
        }

        // Validate that the URL returns valid ToolHive registry JSON
        if err := isValidRegistryJSON(registryClient, registryURL); err != nil {
                return &RegistryError{
                        Type: RegistryTypeURL,
                        URL:  registryURL,
                        Err:  err,
                }
        }

        // Update the configuration
        err = provider.UpdateConfig(func(c *Config) error {
                c.RegistryUrl = registryURL
                c.RegistryApiUrl = ""    // Clear API URL when setting static URL
                c.LocalRegistryPath = "" // Clear local path when setting URL
                c.AllowPrivateRegistryIp = allowPrivateRegistryIp
                return nil
        })
        if err != nil {
                return fmt.Errorf("failed to update configuration: %w", err)
        }

        return nil
}

// setRegistryFile validates and sets a local registry file using the provided provider
func setRegistryFile(provider Provider, registryPath string) error {
        // Validate file path exists
        cleanPath, err := validateFilePath(registryPath)
        if err != nil {
                return fmt.Errorf("local registry %w", err)
        }

        // Validate JSON file
        if err := validateJSONFile(cleanPath); err != nil {
                return &RegistryError{
                        Type: RegistryTypeFile,
                        URL:  registryPath,
                        Err:  fmt.Errorf("%w: %v", ErrRegistryValidationFailed, err),
                }
        }

        // Validate registry structure
        if err := validateRegistryFileStructure(cleanPath); err != nil {
                return &RegistryError{
                        Type: RegistryTypeFile,
                        URL:  registryPath,
                        Err:  fmt.Errorf("%w: %v", ErrRegistryValidationFailed, err),
                }
        }

        // Make the path absolute
        absPath, err := makeAbsolutePath(cleanPath)
        if err != nil {
                return fmt.Errorf("registry file: %w", err)
        }

        // Update the configuration
        err = provider.UpdateConfig(func(c *Config) error {
                c.LocalRegistryPath = absPath
                c.RegistryUrl = ""    // Clear URL when setting local path
                c.RegistryApiUrl = "" // Clear API URL when setting local path
                return nil
        })
        if err != nil {
                return fmt.Errorf("failed to update configuration: %w", err)
        }

        return nil
}

// validateRegistryFileStructure checks if a file contains a valid upstream MCP
// registry structure by parsing it into the UpstreamRegistry type.
func validateRegistryFileStructure(path string) error {
        // Read file content
        // #nosec G304: File path is user-provided but validated by caller
        data, err := os.ReadFile(path)
        if err != nil {
                return fmt.Errorf("failed to read file: %w", err)
        }

        return parseAndValidateUpstreamRegistry(data)
}

// setRegistryAPI validates and sets an MCP Registry API URL using the provided provider
// Validates connectivity with a 5-second timeout.
func setRegistryAPI(provider Provider, apiURL string, allowPrivateRegistryIp bool) error {
        parsedURL, err := neturl.Parse(apiURL)
        if err != nil {
                return fmt.Errorf("invalid registry API URL: %w", err)
        }

        if allowPrivateRegistryIp {
                // we validate either https or http URLs
                if parsedURL.Scheme != networking.HttpScheme && parsedURL.Scheme != networking.HttpsScheme {
                        return fmt.Errorf("registry API URL must start with http:// or https:// when allowing private IPs")
                }
        } else {
                // we just allow https
                if parsedURL.Scheme != networking.HttpsScheme {
                        return fmt.Errorf("registry API URL must start with https:// when not allowing private IPs")
                }
        }

        // Validate that the URL is accessible with 5-second timeout
        if !allowPrivateRegistryIp {
                registryClient, err := networking.NewHttpClientBuilder().
                        WithTimeout(5 * time.Second).
                        Build()
                if err != nil {
                        return fmt.Errorf("failed to create HTTP client: %w", err)
                }
                // Just check the base URL is accessible (don't require specific endpoints)
                _, err = registryClient.Head(apiURL)
                if err != nil {
                        return &RegistryError{
                                Type: RegistryTypeAPI,
                                URL:  apiURL,
                                Err:  classifyNetworkError(err),
                        }
                }
        }

        // Update the configuration
        err = provider.UpdateConfig(func(c *Config) error {
                c.RegistryApiUrl = apiURL
                c.RegistryUrl = ""       // Clear static registry URL when setting API URL
                c.LocalRegistryPath = "" // Clear local path when setting API URL
                c.AllowPrivateRegistryIp = allowPrivateRegistryIp
                return nil
        })
        if err != nil {
                return fmt.Errorf("failed to update configuration: %w", err)
        }

        return nil
}

// unsetRegistry resets registry configuration to defaults using the provided provider
func unsetRegistry(provider Provider) error {
        err := provider.UpdateConfig(func(c *Config) error {
                c.RegistryUrl = ""
                c.RegistryApiUrl = ""
                c.LocalRegistryPath = ""
                c.AllowPrivateRegistryIp = false
                return nil
        })
        if err != nil {
                return fmt.Errorf("failed to update configuration: %w", err)
        }
        return nil
}

// getRegistryConfig returns current registry configuration using the provided provider
func getRegistryConfig(provider Provider) (url, localPath string, allowPrivateIP bool, registryType string) {
        cfg := provider.GetConfig()

        // Check API URL first (highest priority for live data)
        if cfg.RegistryApiUrl != "" {
                return cfg.RegistryApiUrl, "", cfg.AllowPrivateRegistryIp, RegistryTypeAPI
        }

        if cfg.RegistryUrl != "" {
                return cfg.RegistryUrl, "", cfg.AllowPrivateRegistryIp, RegistryTypeURL
        }

        if cfg.LocalRegistryPath != "" {
                return "", cfg.LocalRegistryPath, false, RegistryTypeFile
        }

        return "", "", false, "default"
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "log/slog"
        "os"
        "sync"
)

// Singleton value - should only be written to by the getSingletonConfig function.
var appConfig *Config

var lock = &sync.RWMutex{}

// SetSingletonConfig allows tests to pre-initialize the singleton with test data
// This prevents the singleton from loading the real config file during tests
func SetSingletonConfig(cfg *Config) {
        lock.Lock()
        defer lock.Unlock()
        appConfig = cfg
}

// ResetSingleton clears the singleton - useful for test cleanup
func ResetSingleton() {
        lock.Lock()
        defer lock.Unlock()
        appConfig = nil
}

// getSingletonConfig is a Singleton that returns the application configuration.
// This is only used internally by the DefaultProvider
func getSingletonConfig() *Config {
        // First check with read lock for performance
        lock.RLock()
        if appConfig != nil {
                defer lock.RUnlock()
                return appConfig
        }
        lock.RUnlock()

        // If config is nil, acquire write lock and double-check
        lock.Lock()
        defer lock.Unlock()
        if appConfig == nil {
                config, err := LoadOrCreateConfig()
                if err != nil {
                        slog.Error("error loading configuration", "error", err)
                        os.Exit(1)
                }
                appConfig = config
        }
        return appConfig
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "encoding/json"
        "errors"
        "fmt"
        neturl "net/url"
        "os"
        "path/filepath"
        "strings"

        "github.com/stacklok/toolhive/pkg/networking"
)

// Error message templates for consistent error formatting
const (
        errFileNotFound        = "file not found or not accessible: %w"
        errFileRead            = "failed to read file: %w"
        errInvalidJSON         = "invalid JSON format: %w"
        errInvalidURL          = "invalid URL format: %w"
        errInvalidURLScheme    = "URL must start with %s://"
        errJSONExtensionOnly   = "file must be a JSON file (*.json)"
        errAbsolutePathResolve = "failed to resolve absolute path: %w"
)

// validateFilePath validates that a file path exists and is accessible.
// It also cleans the file path using filepath.Clean.
// Returns the cleaned path and an error if the file doesn't exist or isn't accessible.
func validateFilePath(path string) (string, error) {
        cleanPath := filepath.Clean(path)

        if _, err := os.Stat(cleanPath); err != nil {
                return "", fmt.Errorf(errFileNotFound, err)
        }

        return cleanPath, nil
}

// validateFileExists checks if a file exists and is accessible without cleaning the path.
// This is useful when the path has already been cleaned.
func validateFileExists(path string) error {
        if _, err := os.Stat(path); err != nil {
                return fmt.Errorf(errFileNotFound, err)
        }
        return nil
}

// readFile reads the contents of a file and returns the data.
// This is a wrapper around os.ReadFile with consistent error messaging.
func readFile(path string) ([]byte, error) {
        // #nosec G304: File path is user-provided but should be validated by caller
        data, err := os.ReadFile(path)
        if err != nil {
                return nil, fmt.Errorf(errFileRead, err)
        }
        return data, nil
}

// validateJSONFile validates that a file contains valid JSON.
// It checks the file extension and attempts to parse the content.
func validateJSONFile(path string) error {
        // Check file extension
        if !strings.HasSuffix(strings.ToLower(path), ".json") {
                return errors.New(errJSONExtensionOnly)
        }

        // Read and validate JSON content
        data, err := readFile(path)
        if err != nil {
                return err
        }

        // Basic JSON validation - unmarshal into generic map
        var jsonData map[string]interface{}
        if err := json.Unmarshal(data, &jsonData); err != nil {
                return fmt.Errorf(errInvalidJSON, err)
        }

        return nil
}

// validateURLScheme validates that a URL has the correct scheme (http or https).
// If allowInsecure is false, only https is allowed.
// If allowInsecure is true, both http and https are allowed.
func validateURLScheme(rawURL string, allowInsecure bool) (*neturl.URL, error) {
        parsedURL, err := neturl.Parse(rawURL)
        if err != nil {
                return nil, fmt.Errorf(errInvalidURL, err)
        }

        if allowInsecure {
                // Allow both http and https
                if parsedURL.Scheme != networking.HttpScheme && parsedURL.Scheme != networking.HttpsScheme {
                        return nil, fmt.Errorf("URL must start with http:// or https://")
                }
        } else {
                // Only allow https
                if parsedURL.Scheme != networking.HttpsScheme {
                        return nil, fmt.Errorf(errInvalidURLScheme, networking.HttpsScheme)
                }
        }

        return parsedURL, nil
}

// makeAbsolutePath converts a relative path to an absolute path.
func makeAbsolutePath(path string) (string, error) {
        absPath, err := filepath.Abs(path)
        if err != nil {
                return "", fmt.Errorf(errAbsolutePathResolve, err)
        }
        return absPath, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package docker provides Docker-specific implementation of container runtime,
// including creating, starting, stopping, and monitoring containers.
package docker

import (
        "bytes"
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "maps"
        "net/netip"
        "os"
        "path/filepath"
        rt "runtime"
        "slices"
        "strconv"
        "strings"
        "time"

        "github.com/containerd/errdefs"
        "github.com/moby/moby/api/pkg/stdcopy"
        "github.com/moby/moby/api/types/container"
        "github.com/moby/moby/api/types/mount"
        "github.com/moby/moby/api/types/network"
        mobyclient "github.com/moby/moby/client"

        "github.com/stacklok/toolhive-core/permissions"
        "github.com/stacklok/toolhive/pkg/container/docker/sdk"
        "github.com/stacklok/toolhive/pkg/container/images"
        "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/ignore"
        lb "github.com/stacklok/toolhive/pkg/labels"
        "github.com/stacklok/toolhive/pkg/networking"
)

// DnsImage is the default DNS image used for network permissions
const DnsImage = "dockurr/dnsmasq:latest"

// RuntimeName is the name identifier for the Docker runtime
const RuntimeName = "docker"

// IsAvailable checks if Docker is available by attempting to connect to the Docker daemon
func IsAvailable() bool {
        ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
        defer cancel()

        _, err := NewClient(ctx)
        return err == nil
}

// Workloads
const (
        ToolhiveAuxiliaryWorkloadLabel = "toolhive-auxiliary-workload"
        LabelValueTrue                 = "true"
)

// dockerAPI defines the minimal Docker client surface we need for unit-testing
// ListWorkloads/GetWorkloadInfo through an adapter without requiring a live daemon.
// The signatures mirror the methods on *mobyclient.Client so the real client can
// be assigned directly to Client.api.
type dockerAPI interface {
        ContainerList(ctx context.Context, options mobyclient.ContainerListOptions) (mobyclient.ContainerListResult, error)
        ContainerInspect(
                ctx context.Context,
                containerID string,
                options mobyclient.ContainerInspectOptions,
        ) (mobyclient.ContainerInspectResult, error)
        ContainerStop(
                ctx context.Context,
                containerID string,
                options mobyclient.ContainerStopOptions,
        ) (mobyclient.ContainerStopResult, error)
        ContainerCreate(ctx context.Context, options mobyclient.ContainerCreateOptions) (mobyclient.ContainerCreateResult, error)
        ContainerStart(
                ctx context.Context,
                containerID string,
                options mobyclient.ContainerStartOptions,
        ) (mobyclient.ContainerStartResult, error)
        ContainerRemove(
                ctx context.Context,
                containerID string,
                options mobyclient.ContainerRemoveOptions,
        ) (mobyclient.ContainerRemoveResult, error)
}

// deployOps defines the internal operations used by DeployWorkload.
// It allows unit tests to substitute a fake implementation to avoid hitting a real Docker daemon.
type bridgeNetworkAPI interface {
        NetworkList(context.Context, mobyclient.NetworkListOptions) (mobyclient.NetworkListResult, error)
        NetworkInspect(context.Context, string, mobyclient.NetworkInspectOptions) (mobyclient.NetworkInspectResult, error)
}

type deployOps interface {
        createExternalNetworks(ctx context.Context) error
        createNetwork(ctx context.Context, name string, labels map[string]string, internal bool) error
        createDnsContainer(
                ctx context.Context,
                dnsContainerName string,
                attachStdio bool,
                networkName string,
                endpointsConfig map[string]*network.EndpointSettings,
        ) (string, string, error)
        createMcpContainer(
                ctx context.Context,
                name string,
                networkName string,
                image string,
                command []string,
                envVars map[string]string,
                labels map[string]string,
                attachStdio bool,
                permissionConfig *runtime.PermissionConfig,
                additionalDNS string,
                exposedPorts map[string]struct{},
                portBindings map[string][]runtime.PortBinding,
                isolateNetwork bool,
        ) error
        getContainerNetworkIP(ctx context.Context, containerName, networkName string) (string, error)
        discoverBridgeGateways(ctx context.Context) ([]netip.Addr, error)
}

// Client implements the Deployer interface for Docker (and compatible runtimes)
type Client struct {
        runtimeType  runtime.Type
        socketPath   string
        client       *mobyclient.Client
        api          dockerAPI
        imageManager images.ImageManager
        ops          deployOps
        proxy        networkProxy // selected at construction time via newNetworkProxy
        bridgeAPI    bridgeNetworkAPI
        hostAPI      hostDiscoveryAPI
}

// NewClient creates a new container client
func NewClient(ctx context.Context) (*Client, error) {
        dockerClient, socketPath, runtimeType, err := sdk.NewDockerClient(ctx)
        if err != nil {
                return nil, err // there is already enough context in the error.
        }

        imageManager := images.NewRegistryImageManager(dockerClient)

        c := &Client{
                runtimeType:  runtimeType,
                socketPath:   socketPath,
                client:       dockerClient,
                api:          dockerClient,
                imageManager: imageManager,
                bridgeAPI:    dockerClient,
                hostAPI:      dockerClient,
        }
        // Default ops implementation uses the real client methods.
        c.ops = c

        proxy, err := newNetworkProxy(c)
        if err != nil {
                return nil, fmt.Errorf("failed to initialize network proxy: %w", err)
        }
        c.proxy = proxy

        return c, nil
}

// DeployWorkload creates and starts a workload.
// It configures the workload based on the provided permission profile and transport type.
// If options is nil, default options will be used.
//
//nolint:gocyclo // This function has high complexity due to comprehensive workload setup
func (c *Client) DeployWorkload(
        ctx context.Context,
        image,
        name string,
        command []string,
        envVars,
        labels map[string]string,
        permissionProfile *permissions.Profile,
        transportType string,
        options *runtime.DeployWorkloadOptions,
        isolateNetwork bool,
) (int, error) {
        // Get permission config from profile
        var ignoreConfig *ignore.Config
        if options != nil {
                ignoreConfig = options.IgnoreConfig
        }
        permissionConfig, err := c.getPermissionConfigFromProfile(permissionProfile, transportType, ignoreConfig)
        if err != nil {
                return 0, fmt.Errorf("failed to get permission config: %w", err)
        }

        // Determine if we should attach stdio
        attachStdio := options == nil || options.AttachStdio

        // create networks
        var additionalDNS string
        networkName := fmt.Sprintf("toolhive-%s-internal", name)
        externalEndpointsConfig := map[string]*network.EndpointSettings{
                networkName: {},
        }

        // Network isolation is only enforceable in bridge mode. For non-bridge
        // modes (host/none/custom) the sidecars have no route out and the workload
        // bypasses them, so isolation must be dropped. See issue #5775.
        effectiveIsolation := isolateNetwork && networking.IsBridgeMode(permissionConfig.NetworkMode)
        if isolateNetwork && !effectiveIsolation {
                if permissionConfig.NetworkMode == "none" {
                        // "none" is already maximally confined; isolation is merely redundant.
                        //nolint:gosec // G706: network mode from permission config
                        slog.Debug(networking.NetworkIsolationNoneRedundantMsg, "network_mode", permissionConfig.NetworkMode)
                } else {
                        // host (or other non-bridge) is less restrictive than isolation, so
                        // dropping it is a real reduction in confinement.
                        //nolint:gosec // G706: network mode from permission config
                        slog.Warn(networking.NetworkIsolationHostDroppedMsg, "network_mode", permissionConfig.NetworkMode)
                }
        }

        // Only create external networks and add endpoints if we're not using a custom network mode like "host" or "none"
        if networking.IsBridgeMode(permissionConfig.NetworkMode) {
                // Add toolhive-external to endpoints config for default networking modes
                externalEndpointsConfig["toolhive-external"] = &network.EndpointSettings{}

                err = c.ops.createExternalNetworks(ctx)
                if err != nil {
                        return 0, fmt.Errorf("failed to create external networks: %w", err)
                }
        } else {
                //nolint:gosec // G706: network mode from permission config
                slog.Debug("skipping external network creation for custom network mode", "network_mode", permissionConfig.NetworkMode)
        }

        // For non-stdio isolated workloads, extract the upstream port so the ingress
        // proxy can be configured. Done here so a bad exposed-ports config fails fast.
        var upstreamPort int
        if transportType != "stdio" && effectiveIsolation {
                upstreamPort, err = extractFirstPort(options)
                if err != nil {
                        return 0, err // extractFirstPort already wraps the error with context.
                }
        }

        networkIsolation := false
        // pspec and egress are populated in the isolation block and reused by
        // SetupIngress after the MCP container is created.
        var (
                pspec  proxySpec
                egress egressResult
        )
        if effectiveIsolation {
                networkIsolation = true

                internalNetworkLabels := map[string]string{}
                lb.AddNetworkLabels(internalNetworkLabels, networkName)
                err := c.ops.createNetwork(ctx, networkName, internalNetworkLabels, true)
                if err != nil {
                        return 0, fmt.Errorf("failed to create internal network: %w", err)
                }

                // Start or reuse DNS normally before reading Podman's generated hosts.
                dnsContainerName := fmt.Sprintf("%s-dns", name)
                dnsContainerID, dnsContainerIP, err := c.ops.createDnsContainer(
                        ctx, dnsContainerName, attachStdio, networkName, externalEndpointsConfig,
                )
                if err != nil {
                        return 0, fmt.Errorf("failed to create dns container: %w", err)
                }
                additionalDNS = dnsContainerIP

                // Snapshot after DNS startup, which may recreate a missing network.
                // Bridges or host mappings changed later require a workload redeploy.
                gatewayIPs, err := c.ops.discoverBridgeGateways(ctx)
                if err != nil {
                        return 0, fmt.Errorf("failed to discover Docker bridge gateways: %w", err)
                }
                if c.runtimeType == runtime.TypePodman {
                        aliases, err := c.discoverHostAliases(ctx, dnsContainerID)
                        if err != nil {
                                return 0, fmt.Errorf("failed to discover Podman host addresses: %w", err)
                        }
                        gatewayIPs = append(slices.Clone(gatewayIPs), aliases...)
                        slices.SortFunc(gatewayIPs, netip.Addr.Compare)
                        gatewayIPs = slices.Compact(gatewayIPs)
                }

                pspec = proxySpec{
                        WorkloadName:       name,
                        Permissions:        permissionProfile.Network,
                        AllowDockerGateway: options != nil && options.AllowDockerGateway,
                        GatewayIPs:         gatewayIPs,
                        TransportType:      transportType,
                        UpstreamPort:       upstreamPort,
                        AttachStdio:        attachStdio,
                        Endpoints:          externalEndpointsConfig,
                }

                // SetupEgress runs before createMcpContainer so its env vars can be
                // injected into the workload and the egress proxy has a head start.
                egress, err = c.proxy.SetupEgress(ctx, pspec)
                if err != nil {
                        return 0, fmt.Errorf("failed to set up egress proxy: %w", err)
                }
                envVars = mergeEnvVars(envVars, egress.EnvVars)
        } else {
                networkName = ""
        }

        // only remap if is not an auxiliary tool
        newPortBindings, hostPort, err := generatePortBindings(labels, options.PortBindings)
        if err != nil {
                return 0, fmt.Errorf("failed to generate port bindings: %w", err)
        }

        // Add a label to the MCP server indicating network isolation.
        // This allows other methods to determine whether it needs to care
        // about ingress/egress/dns containers.
        lb.AddNetworkIsolationLabel(labels, networkIsolation)

        err = c.ops.createMcpContainer(
                ctx,
                name,
                networkName,
                image,
                command,
                envVars,
                labels,
                attachStdio,
                permissionConfig,
                additionalDNS,
                options.ExposedPorts,
                newPortBindings,
                effectiveIsolation,
        )
        if err != nil {
                return 0, fmt.Errorf("failed to create mcp container: %w", err)
        }

        // Don't try and set up an ingress proxy if the transport type is stdio.
        if transportType == "stdio" {
                return 0, nil
        }

        firstPortInt, err := extractFirstPort(options)
        if err != nil {
                return 0, err // extractFirstPort already wraps the error with context.
        }
        if effectiveIsolation {
                // Point the ingress reverse proxy at the MCP container's IP rather than
                // its name. Ordering ingress after createMcpContainer is not enough: under
                // concurrent startup the container's DNS record can still lag, and the
                // squid cache_peer caches the failed lookup and never recovers within the
                // readiness window (see #6063). An IP has no lookup, so a transient
                // connection failure is retried instead of latching the peer dead.
                upstreamIP, ipErr := c.ops.getContainerNetworkIP(ctx, name, networkName)
                if ipErr != nil {
                        return 0, fmt.Errorf("failed to resolve MCP container IP for ingress: %w", ipErr)
                }
                pspec.UpstreamHost = upstreamIP
                hostPort, err = c.proxy.SetupIngress(ctx, pspec, egress)
                if err != nil {
                        return 0, fmt.Errorf("failed to set up ingress proxy: %w", err)
                }
        }

        // NOTE: this is a hack to get the final port for the workload.
        // The intended behavior is the following
        // * if network is "bridge" (default) and network isolation is not enabled, then
        //   the Proxy should use the Docker assigned port
        // * if network is "bridge" (default) and network isolation is enabled, then
        //   the Proxy should use the egress container port
        // * if network is "host", then the Proxy should use the MCP server port
        //
        // The last case is not supported in VM-based installations like Docker Desktop.
        // Unfortunately, there's no reliable way to know if the user is using a VM-based
        // installation and we assume that Linux installations are Docker Engine installations.
        finalPort := calculateFinalPort(hostPort, firstPortInt, permissionConfig.NetworkMode)

        return finalPort, nil
}

// ListWorkloads lists workloads
func (c *Client) ListWorkloads(ctx context.Context) ([]runtime.ContainerInfo, error) {
        // Create filter for toolhive containers
        filterArgs := mobyclient.Filters{}.Add("label", "toolhive=true")

        // List containers
        listResult, err := c.api.ContainerList(ctx, mobyclient.ContainerListOptions{
                All:     true,
                Filters: filterArgs,
        })
        if err != nil {
                return nil, NewContainerError(err, "", fmt.Sprintf("failed to list containers: %v", err))
        }
        containers := listResult.Items

        // Convert to our ContainerInfo format
        result := make([]runtime.ContainerInfo, 0, len(containers))
        for _, c := range containers {
                // Skip containers that have the auxiliary workload label set to "true"
                if val, ok := c.Labels["toolhive-auxiliary-workload"]; ok && val == "true" {
                        continue
                }

                // Extract container name (remove leading slash)
                name := ""
                if len(c.Names) > 0 {
                        name = c.Names[0]
                        name = strings.TrimPrefix(name, "/")
                }

                // Extract port mappings
                ports := make([]runtime.PortMapping, 0, len(c.Ports))
                for _, p := range c.Ports {
                        ports = append(ports, runtime.PortMapping{
                                ContainerPort: int(p.PrivatePort),
                                HostPort:      int(p.PublicPort),
                                Protocol:      p.Type,
                        })
                }

                // Convert creation time
                created := time.Unix(c.Created, 0)

                result = append(result, runtime.ContainerInfo{
                        Name:    name,
                        Image:   c.Image,
                        Status:  c.Status,
                        State:   dockerToDomainStatus(string(c.State)),
                        Created: created,
                        Labels:  c.Labels,
                        Ports:   ports,
                })
        }

        return result, nil
}

// StopWorkload stops a workload
// If the workload is already stopped, it returns success
func (c *Client) StopWorkload(ctx context.Context, workloadName string) error {
        // Check if the workload is running
        info, err := c.GetWorkloadInfo(ctx, workloadName)
        if err != nil {
                // If the container doesn't exist, that's fine - it's already "stopped"
                if errors.Is(err, ErrContainerNotFound) {
                        return nil
                }
                return err
        }

        // If the container is not running, return success
        if info.State != runtime.WorkloadStatusRunning {
                return nil
        }

        // Use a reasonable timeout
        timeoutSeconds := 30
        _, err = c.api.ContainerStop(ctx, workloadName, mobyclient.ContainerStopOptions{Timeout: &timeoutSeconds})
        if err != nil {
                return NewContainerError(err, workloadName, fmt.Sprintf("failed to stop workload: %v", err))
        }

        // If network isolation is not enabled, then there is nothing else to do.
        // NOTE: This check treats all workloads created before the introduction of
        // this label as having network isolation enabled. This is to ensure that they
        // get cleaned up properly during stop/rm.
        if !lb.HasNetworkIsolation(info.Labels) {
                return nil
        }

        // remove / from container name
        containerName := strings.TrimPrefix(info.Name, "/")
        egressContainerName := fmt.Sprintf("%s-egress", containerName)
        ingressContainerName := fmt.Sprintf("%s-ingress", containerName)
        dnsContainerName := fmt.Sprintf("%s-dns", containerName)

        // Attempt to stop each auxiliary container gracefully.
        // Treat any errors as non-fatal and log them.
        proxyContainers := []string{egressContainerName, ingressContainerName, dnsContainerName}
        for _, name := range proxyContainers {
                c.stopProxyContainer(ctx, name, timeoutSeconds)
        }

        return nil
}

// RemoveWorkload removes a workload
// If the workload doesn't exist, it returns success
func (c *Client) RemoveWorkload(ctx context.Context, workloadName string) error {
        // get container name from ID
        containerResponse, err := c.inspectContainerByName(ctx, workloadName)
        if err != nil {
                slog.Warn("failed to inspect container", "name", workloadName, "error", err)
        }

        // remove the / if it starts with it
        containerName := containerResponse.Name
        containerName = strings.TrimPrefix(containerName, "/")

        // remove the workload containers
        var labels map[string]string
        if containerResponse.Config != nil {
                labels = containerResponse.Config.Labels
        } else {
                labels = make(map[string]string)
        }
        err = c.removeContainer(ctx, containerResponse.ID)
        if err != nil {
                return err // removeContainer already wraps the error with context.
        }

        // Clean up any proxy containers associated with this workload.
        err = c.removeProxyContainers(ctx, containerName, labels)
        if err != nil {
                return err // removeProxyContainers already wraps the error with context.
        }

        // Clear up any networks associated with this workload.
        // This also deletes the external network if no other workloads are using it.
        err = c.deleteNetworks(ctx, containerName)
        if err != nil {
                slog.Warn("failed to delete networks for container", "name", containerName, "error", err)
        }
        return nil
}

// GetWorkloadLogs gets workload logs
func (c *Client) GetWorkloadLogs(ctx context.Context, workloadName string, follow bool, lines int) (string, error) {
        // follow=true means infinite streaming, lines>0 means finite limit - these are contradictory
        if follow && lines > 0 {
                return "", NewContainerError(
                        fmt.Errorf("cannot use both follow and line limit"),
                        workloadName,
                        "follow mode streams logs indefinitely, which conflicts with line limiting",
                )
        }

        // Configure tail option based on lines parameter
        tail := "all"
        if lines > 0 {
                tail = fmt.Sprintf("%d", lines)
        }

        options := mobyclient.ContainerLogsOptions{
                ShowStdout: true,
                ShowStderr: true,
                Follow:     follow,
                Tail:       tail,
        }

        workloadContainer, err := c.inspectContainerByName(ctx, workloadName)
        if err != nil {
                return "", err
        }

        logs, err := c.client.ContainerLogs(ctx, workloadContainer.ID, options)
        if err != nil {
                return "", NewContainerError(err, workloadName, fmt.Sprintf("failed to get workload logs: %v", err))
        }
        defer func() {
                if err := logs.Close(); err != nil {
                        // Non-fatal: log stream cleanup failure
                        slog.Debug("failed to close log stream", "error", err)
                }
        }()

        if follow {
                _, err = stdcopy.StdCopy(os.Stdout, os.Stderr, logs)
                if err != nil && err != io.EOF {
                        slog.Error("error reading workload logs", "error", err)
                        return "", NewContainerError(err, workloadName, fmt.Sprintf("failed to follow workload logs: %v", err))
                }
        }

        // Read logs
        var buf bytes.Buffer
        _, err = stdcopy.StdCopy(&buf, &buf, logs)
        if err != nil {
                return "", NewContainerError(err, workloadName, fmt.Sprintf("failed to read workload logs: %v", err))
        }

        return buf.String(), nil
}

// IsWorkloadRunning checks if a workload is running
func (c *Client) IsWorkloadRunning(ctx context.Context, workloadName string) (bool, error) {
        // Inspect workload
        info, err := c.inspectContainerByName(ctx, workloadName)
        if err != nil {
                // Check if the error is because the workload doesn't exist
                if errdefs.IsNotFound(err) {
                        return false, NewContainerError(ErrContainerNotFound, workloadName, "workload not found")
                }
                return false, NewContainerError(err, workloadName, fmt.Sprintf("failed to inspect workload: %v", err))
        }

        return info.State.Running, nil
}

// GetWorkloadInfo gets workload information
func (c *Client) GetWorkloadInfo(ctx context.Context, workloadName string) (runtime.ContainerInfo, error) {
        // Inspect workload
        info, err := c.inspectContainerByName(ctx, workloadName)
        if err != nil {
                // Check if the error is because the workload doesn't exist
                if errdefs.IsNotFound(err) {
                        return runtime.ContainerInfo{}, NewContainerError(ErrContainerNotFound, workloadName, "workload not found")
                }
                return runtime.ContainerInfo{}, NewContainerError(err, workloadName, fmt.Sprintf("failed to inspect workload: %v", err))
        }

        // Extract port mappings
        ports := make([]runtime.PortMapping, 0)
        for containerPort, bindings := range info.NetworkSettings.Ports {
                for _, binding := range bindings {
                        hostPort := 0
                        if _, err := fmt.Sscanf(binding.HostPort, "%d", &hostPort); err != nil {
                                // If we can't parse the port, just use 0
                                //nolint:gosec // G706: host port from container network settings
                                slog.Warn("failed to parse host port", "host_port", binding.HostPort, "error", err)
                        }

                        ports = append(ports, runtime.PortMapping{
                                ContainerPort: int(containerPort.Num()),
                                HostPort:      hostPort,
                                Protocol:      string(containerPort.Proto()),
                        })
                }
        }

        // Convert creation time
        created, err := time.Parse(time.RFC3339, info.Created)
        if err != nil {
                created = time.Time{} // Use zero time if parsing fails
        }

        // Convert start time
        startedAt, err := time.Parse(time.RFC3339Nano, info.State.StartedAt)
        if err != nil {
                startedAt = time.Time{} // Use zero time if parsing fails
        }

        return runtime.ContainerInfo{
                Name:      strings.TrimPrefix(info.Name, "/"),
                Image:     info.Config.Image,
                Status:    string(info.State.Status),
                State:     dockerToDomainStatus(string(info.State.Status)),
                Created:   created,
                StartedAt: startedAt,
                Labels:    info.Config.Labels,
                Ports:     ports,
        }, nil
}

// AttachToWorkload attaches to a workload
func (c *Client) AttachToWorkload(ctx context.Context, workloadName string) (io.WriteCloser, io.ReadCloser, error) {
        // Check if workload exists and is running
        running, err := c.IsWorkloadRunning(ctx, workloadName)
        if err != nil {
                return nil, nil, err
        }
        if !running {
                return nil, nil, NewContainerError(ErrContainerNotRunning, workloadName, "workload is not running")
        }

        // Attach to workload
        resp, err := c.client.ContainerAttach(ctx, workloadName, mobyclient.ContainerAttachOptions{
                Stream: true,
                Stdin:  true,
                Stdout: true,
                Stderr: true,
        })
        if err != nil {
                return nil, nil, NewContainerError(ErrAttachFailed, workloadName, fmt.Sprintf("failed to attach to workload: %v", err))
        }

        stdoutReader, stdoutWriter := io.Pipe()

        go func() {
                defer func() {
                        if err := stdoutWriter.Close(); err != nil {
                                // Non-fatal: writer cleanup failure
                                slog.Debug("failed to close stdout writer", "error", err)
                        }
                }()
                defer resp.Close()

                // Use stdcopy to demultiplex the container streams
                _, err := stdcopy.StdCopy(stdoutWriter, io.Discard, resp.Reader)
                if err != nil && err != io.EOF {
                        slog.Error("error demultiplexing container streams", "error", err)
                }
        }()

        return resp.Conn, stdoutReader, nil
}

// IsRunning checks the health of the container runtime.
// This is used to verify that the runtime is operational and can manage workloads.
func (c *Client) IsRunning(ctx context.Context) error {
        // Try to ping the Docker server
        _, err := c.client.Ping(ctx, mobyclient.PingOptions{})
        if err != nil {
                return fmt.Errorf("failed to ping Docker server: %w", err)
        }

        return nil
}

// getPermissionConfigFromProfile converts a permission profile to a container permission config
// with transport-specific settings (internal function)
// addReadOnlyMounts adds read-only mounts to the permission config
func (*Client) addReadOnlyMounts(
        config *runtime.PermissionConfig,
        mounts []permissions.MountDeclaration,
        ignoreConfig *ignore.Config,
) {
        for _, mountDecl := range mounts {
                source, target, err := mountDecl.Parse()
                if err != nil {
                        // Skip invalid mounts
                        slog.Warn("skipping invalid mount declaration", "mount", mountDecl, "error", err)
                        continue
                }

                // Skip resource URIs for now (they need special handling)
                if strings.Contains(source, "://") {
                        slog.Warn("resource URI mounts not yet supported", "source", source)
                        continue
                }

                // Convert relative paths to absolute paths
                absPath, ok := convertRelativePathToAbsolute(source, mountDecl)
                if !ok {
                        continue
                }

                config.Mounts = append(config.Mounts, runtime.Mount{
                        Source:   absPath,
                        Target:   target,
                        ReadOnly: true,
                        Type:     runtime.MountTypeBind,
                })

                // Process ignore patterns and add tmpfs overlays
                addIgnoreOverlays(config, absPath, target, ignoreConfig)
        }
}

// addReadWriteMounts adds read-write mounts to the permission config
func (*Client) addReadWriteMounts(
        config *runtime.PermissionConfig,
        mounts []permissions.MountDeclaration,
        ignoreConfig *ignore.Config,
) {
        for _, mountDecl := range mounts {
                source, target, err := mountDecl.Parse()
                if err != nil {
                        // Skip invalid mounts
                        slog.Warn("skipping invalid mount declaration", "mount", mountDecl, "error", err)
                        continue
                }

                // Skip resource URIs for now (they need special handling)
                if strings.Contains(source, "://") {
                        slog.Warn("resource URI mounts not yet supported", "source", source)
                        continue
                }

                // Convert relative paths to absolute paths
                absPath, ok := convertRelativePathToAbsolute(source, mountDecl)
                if !ok {
                        continue
                }

                // Check if the path is already mounted read-only
                alreadyMounted := false
                for i, m := range config.Mounts {
                        if m.Target == target {
                                // Update the mount to be read-write
                                config.Mounts[i].ReadOnly = false
                                alreadyMounted = true
                                break
                        }
                }

                // If not already mounted, add a new mount
                if !alreadyMounted {
                        config.Mounts = append(config.Mounts, runtime.Mount{
                                Source:   absPath,
                                Target:   target,
                                ReadOnly: false,
                                Type:     runtime.MountTypeBind,
                        })
                }

                // Process ignore patterns and add tmpfs overlays
                addIgnoreOverlays(config, absPath, target, ignoreConfig)
        }
}

// addIgnoreOverlays processes ignore patterns for a mount and adds overlay mounts
func addIgnoreOverlays(config *runtime.PermissionConfig, sourceDir, containerPath string, ignoreConfig *ignore.Config) {
        // Skip if no ignore configuration is provided
        if ignoreConfig == nil {
                return
        }

        // Create ignore processor with configuration
        ignoreProcessor := ignore.NewProcessor(ignoreConfig)

        // Load global ignore patterns if enabled
        if ignoreConfig.LoadGlobal {
                if err := ignoreProcessor.LoadGlobal(); err != nil {
                        slog.Debug("failed to load global ignore patterns", "error", err)
                        // Continue without global patterns
                }
        }

        // Load local ignore patterns from the source directory
        if err := ignoreProcessor.LoadLocal(sourceDir); err != nil {
                //nolint:gosec // G706: source directory from mount config
                slog.Debug("failed to load local ignore patterns", "dir", sourceDir, "error", err)
                // Continue without local patterns
        }

        // Get overlay mounts (all using bind mounts now)
        overlayMounts := ignoreProcessor.GetOverlayMounts(sourceDir, containerPath)

        // Add overlay mounts to the configuration
        for _, overlayMount := range overlayMounts {
                // All overlays now use bind mounts (no more tmpfs)
                config.Mounts = append(config.Mounts, runtime.Mount{
                        Source:   overlayMount.HostPath,
                        Target:   overlayMount.ContainerPath,
                        ReadOnly: false,
                        Type:     runtime.MountTypeBind,
                })
                //nolint:gosec // G706: overlay paths from ignore config processing
                slog.Debug("added bind overlay for ignored path",
                        "host_path", overlayMount.HostPath,
                        "container_path", overlayMount.ContainerPath)
        }
}

// convertRelativePathToAbsolute converts a relative path to an absolute path
// Returns the absolute path and a boolean indicating if the conversion was successful
func convertRelativePathToAbsolute(source string, mountDecl permissions.MountDeclaration) (string, bool) {
        // If it's already an absolute path, return it as is
        if filepath.IsAbs(source) {
                return source, true
        }

        // Special case for Windows: expand ~ to user profile directory.
        if rt.GOOS == "windows" && strings.HasPrefix(source, "~") {
                homeDir := os.Getenv("USERPROFILE")
                source = strings.Replace(source, "~", homeDir, 1)
        }

        absPath, err := filepath.Abs(source)
        if err != nil {
                slog.Warn("failed to convert to absolute path", "mount", mountDecl, "error", err)
                return "", false
        }

        //nolint:gosec // G706: file path from mount declaration config
        slog.Debug("converting relative path to absolute", "mount", mountDecl, "abs_path", absPath)
        return absPath, true
}

// getPermissionConfigFromProfile converts a permission profile to a container permission config
func (c *Client) getPermissionConfigFromProfile(
        profile *permissions.Profile,
        transportType string,
        ignoreConfig *ignore.Config,
) (*runtime.PermissionConfig, error) {
        // Get network mode from permission profile
        networkMode := ""
        if profile.Network != nil && profile.Network.Mode != "" {
                networkMode = profile.Network.Mode
        }

        // Start with a default permission config
        config := &runtime.PermissionConfig{
                Mounts:      []runtime.Mount{},
                NetworkMode: networkMode,
                CapDrop:     []string{"ALL"},
                CapAdd:      []string{},
                SecurityOpt: []string{"label:disable"},
                Privileged:  profile.Privileged,
        }

        // Add mounts
        c.addReadOnlyMounts(config, profile.Read, ignoreConfig)
        c.addReadWriteMounts(config, profile.Write, ignoreConfig)

        // Validate transport type
        switch transportType {
        case "sse", "stdio", "inspector", "streamable-http":
                // valid, do nothing
        default:
                return nil, fmt.Errorf("unsupported transport type: %s", transportType)
        }

        return config, nil
}

// findExistingContainer finds a container with the exact name
// Uses container runtime's name filter for efficiency but then verifies exact match to prevent partial matching
func (c *Client) findExistingContainer(ctx context.Context, name string) (string, error) {
        // Use name filter to narrow results, then verify exact match
        listResult, err := c.api.ContainerList(ctx, mobyclient.ContainerListOptions{
                All:     true, // Include stopped containers
                Filters: mobyclient.Filters{}.Add("name", name),
        })
        if err != nil {
                return "", NewContainerError(err, "", fmt.Sprintf("failed to list containers: %v", err))
        }

        // Verify exact name match from the filtered results (name filter does partial matching)
        for _, cont := range listResult.Items {
                for _, containerName := range cont.Names {
                        // Container names in the API have a leading slash
                        if containerName == "/"+name || containerName == name {
                                return cont.ID, nil
                        }
                }
        }

        return "", nil
}

// compareBasicConfig compares basic container configuration (image, command, env vars, labels, stdio settings)
func compareBasicConfig(existing *container.InspectResponse, desired *container.Config) bool {
        // Compare image
        if existing.Config.Image != desired.Image {
                return false
        }

        // Compare command
        if len(existing.Config.Cmd) != len(desired.Cmd) {
                return false
        }
        for i, cmd := range existing.Config.Cmd {
                if i >= len(desired.Cmd) || cmd != desired.Cmd[i] {
                        return false
                }
        }

        // Compare environment variables
        if !compareEnvVars(existing.Config.Env, desired.Env) {
                return false
        }

        // Compare labels
        if !compareLabels(existing.Config.Labels, desired.Labels) {
                return false
        }

        // Compare stdio settings
        if existing.Config.AttachStdin != desired.AttachStdin ||
                existing.Config.AttachStdout != desired.AttachStdout ||
                existing.Config.AttachStderr != desired.AttachStderr ||
                existing.Config.OpenStdin != desired.OpenStdin {
                return false
        }

        return true
}

// compareEnvVars compares environment variables
func compareEnvVars(existingEnv, desiredEnv []string) bool {
        // Convert to maps for easier comparison
        existingMap := envSliceToMap(existingEnv)
        desiredMap := envSliceToMap(desiredEnv)

        // Check if all desired env vars are in existing env with correct values
        for k, v := range desiredMap {
                existingVal, exists := existingMap[k]
                if !exists || existingVal != v {
                        return false
                }
        }

        return true
}

// envSliceToMap converts a slice of environment variables to a map
func envSliceToMap(env []string) map[string]string {
        result := make(map[string]string)
        for _, e := range env {
                parts := strings.SplitN(e, "=", 2)
                if len(parts) == 2 {
                        result[parts[0]] = parts[1]
                }
        }
        return result
}

// compareLabels compares container labels
func compareLabels(existingLabels, desiredLabels map[string]string) bool {
        // Check if all desired labels are in existing labels with correct values
        for k, v := range desiredLabels {
                existingVal, exists := existingLabels[k]
                if !exists || existingVal != v {
                        return false
                }
        }
        return true
}

// compareHostConfig compares host configuration (network mode, capabilities, security options)
func compareHostConfig(existing *container.InspectResponse, desired *container.HostConfig) bool {
        // Compare network mode
        if string(existing.HostConfig.NetworkMode) != string(desired.NetworkMode) {
                return false
        }

        // Compare capabilities
        if !compareStringSlices(existing.HostConfig.CapAdd, desired.CapAdd) {
                return false
        }
        if !compareStringSlices(existing.HostConfig.CapDrop, desired.CapDrop) {
                return false
        }

        // Compare security options
        if !compareStringSlices(existing.HostConfig.SecurityOpt, desired.SecurityOpt) {
                return false
        }

        // Compare privileged mode
        if existing.HostConfig.Privileged != desired.Privileged {
                return false
        }

        // Compare restart policy
        if existing.HostConfig.RestartPolicy.Name != desired.RestartPolicy.Name {
                return false
        }

        return true
}

// compareStringSlices compares two string slices
func compareStringSlices(existing, desired []string) bool {
        if len(existing) != len(desired) {
                return false
        }
        for i, s := range existing {
                if i >= len(desired) || s != desired[i] {
                        return false
                }
        }
        return true
}

// compareMounts compares volume mounts
func compareMounts(existing *container.InspectResponse, desired *container.HostConfig) bool {
        if len(existing.HostConfig.Mounts) != len(desired.Mounts) {
                return false
        }

        // Create maps by target path for easier comparison
        existingMountsMap := make(map[string]mount.Mount)
        for _, m := range existing.HostConfig.Mounts {
                existingMountsMap[m.Target] = m
        }

        // Check if all desired mounts exist in the container with matching source and read-only flag
        for _, desiredMount := range desired.Mounts {
                existingMount, exists := existingMountsMap[desiredMount.Target]
                if !exists || existingMount.Source != desiredMount.Source || existingMount.ReadOnly != desiredMount.ReadOnly {
                        return false
                }
        }

        return true
}

// comparePortConfig compares port configuration (exposed ports and port bindings)
func comparePortConfig(existing *container.InspectResponse, desired *container.Config, desiredHost *container.HostConfig) bool {
        // Compare exposed ports
        if len(existing.Config.ExposedPorts) != len(desired.ExposedPorts) {
                return false
        }
        for port := range desired.ExposedPorts {
                if _, exists := existing.Config.ExposedPorts[port]; !exists {
                        return false
                }
        }

        // Compare port bindings
        if len(existing.HostConfig.PortBindings) != len(desiredHost.PortBindings) {
                return false
        }
        for port, bindings := range desiredHost.PortBindings {
                existingBindings, exists := existing.HostConfig.PortBindings[port]
                if !exists || len(existingBindings) != len(bindings) {
                        return false
                }
                for i, binding := range bindings {
                        if i >= len(existingBindings) ||
                                existingBindings[i].HostIP != binding.HostIP ||
                                existingBindings[i].HostPort != binding.HostPort {
                                return false
                        }
                }
        }

        return true
}

// compareContainerConfig compares an existing container's configuration with the desired configuration
func compareContainerConfig(
        existing *container.InspectResponse,
        desired *container.Config,
        desiredHost *container.HostConfig,
) bool {
        // Compare basic configuration
        if !compareBasicConfig(existing, desired) {
                return false
        }

        // Compare host configuration
        if !compareHostConfig(existing, desiredHost) {
                return false
        }

        // Compare mounts
        if !compareMounts(existing, desiredHost) {
                return false
        }

        // Compare port configuration
        if !comparePortConfig(existing, desired, desiredHost) {
                return false
        }

        // All checks passed, configurations match
        return true
}

// handleExistingContainer checks if an existing container's configuration matches the desired configuration
// Returns true if the container can be reused, false if it was removed and needs to be recreated
func (c *Client) handleExistingContainer(
        ctx context.Context,
        containerID string,
        desiredConfig *container.Config,
        desiredHostConfig *container.HostConfig,
) (bool, error) {
        // Get container info
        inspectResult, err := c.api.ContainerInspect(ctx, containerID, mobyclient.ContainerInspectOptions{})
        if err != nil {
                return false, NewContainerError(err, containerID, fmt.Sprintf("failed to inspect container: %v", err))
        }
        info := inspectResult.Container

        // Compare configurations
        if compareContainerConfig(&info, desiredConfig, desiredHostConfig) {
                // Configurations match, container can be reused

                // Check if the container is running
                if !info.State.Running {
                        // Container exists but is not running, start it
                        _, err = c.api.ContainerStart(ctx, containerID, mobyclient.ContainerStartOptions{})
                        if err != nil {
                                return false, NewContainerError(err, containerID, fmt.Sprintf("failed to start existing container: %v", err))
                        }
                }

                return true, nil
        }

        // Configurations don't match, we need to recreate the container.
        // Remove only this container, leave any associated networks and containers intact
        // Any proxy containers (like ingress/egress) will have already recreated themselves at this point
        if err := c.removeContainer(ctx, containerID); err != nil {
                return false, err
        }

        // Container was removed and needs to be recreated
        return false, nil
}

// CreateNetwork creates a network following configuration.
func (c *Client) createNetwork(
        ctx context.Context,
        name string,
        labels map[string]string,
        internal bool,
) error {
        // Check if the network already exists
        // Use name filter for efficiency but verify exact match to avoid partial matching
        networks, err := c.client.NetworkList(ctx, mobyclient.NetworkListOptions{
                Filters: mobyclient.Filters{}.Add("name", name),
        })
        if err != nil {
                return fmt.Errorf("failed to list networks: %w", err)
        }
        // Verify exact name match from filtered results
        for _, n := range networks.Items {
                if n.Name == name {
                        // Network already exists
                        return nil
                }
        }

        networkCreate := mobyclient.NetworkCreateOptions{
                Driver:   "bridge",
                Internal: internal,
                Labels:   labels,
        }

        _, err = c.client.NetworkCreate(ctx, name, networkCreate)
        if err != nil {
                // The existence check above is not atomic with the create: another
                // process (e.g. a second `thv run` sharing the "toolhive-external"
                // network) can create the same network in between. Docker rejects a
                // duplicate name with a conflict error, so treat that as success — the
                // network is already in the state we wanted. Without this, workloads
                // that start concurrently race here and all but one fail to launch.
                if errdefs.IsConflict(err) {
                        return nil
                }
                return err
        }
        return nil
}

// getContainerNetworkIP returns a container's IP address on the named network.
// Used to give the ingress proxy a DNS-free upstream address (see proxySpec.UpstreamHost).
func (c *Client) getContainerNetworkIP(ctx context.Context, containerName, networkName string) (string, error) {
        inspectResult, err := c.client.ContainerInspect(ctx, containerName, mobyclient.ContainerInspectOptions{})
        if err != nil {
                return "", fmt.Errorf("failed to inspect container %s: %w", containerName, err)
        }
        netSettings, ok := inspectResult.Container.NetworkSettings.Networks[networkName]
        if !ok {
                return "", fmt.Errorf("network %s not found in container %s network settings", networkName, containerName)
        }
        // EndpointSettings.IPAddress is a netip.Addr; render it to string form.
        if !netSettings.IPAddress.IsValid() {
                return "", fmt.Errorf("container %s has no IP address on network %s", containerName, networkName)
        }
        return netSettings.IPAddress.String(), nil
}

// discoverBridgeGateways returns every gateway configured on every bridge
// network. Networks are inspected by immutable ID; a concurrent deletion causes
// one bounded relist before failing closed.
func (c *Client) discoverBridgeGateways(ctx context.Context) ([]netip.Addr, error) {
        api := c.bridgeAPI
        if api == nil {
                if c.client == nil {
                        return nil, errors.New("docker client is unavailable")
                }
                api = c.client
        }

        for attempt := 0; attempt < 2; attempt++ {
                networks, err := api.NetworkList(ctx, mobyclient.NetworkListOptions{})
                if err != nil {
                        return nil, fmt.Errorf("list Docker networks: %w", err)
                }

                gateways := make(map[netip.Addr]struct{})
                relist := false
                bridgeCount := 0
                for _, summary := range networks.Items {
                        if summary.Driver != "bridge" {
                                continue
                        }
                        if summary.ID == "" {
                                return nil, errors.New("docker daemon returned a bridge network without an ID")
                        }
                        bridgeCount++
                        networkGateways, err := inspectBridgeNetworkGateways(
                                ctx,
                                api,
                                summary,
                                c.runtimeType == runtime.TypePodman,
                        )
                        if err != nil {
                                if errdefs.IsNotFound(err) && attempt == 0 {
                                        relist = true
                                        break
                                }
                                return nil, err
                        }
                        for _, gateway := range networkGateways {
                                gateways[gateway] = struct{}{}
                        }
                }
                if relist {
                        continue
                }
                if bridgeCount == 0 || len(gateways) == 0 {
                        return nil, errors.New("docker daemon returned no bridge gateways")
                }

                result := slices.Collect(maps.Keys(gateways))
                slices.SortFunc(result, netip.Addr.Compare)
                return result, nil
        }
        return nil, errors.New("docker bridge networks changed during discovery")
}

func inspectBridgeNetworkGateways(
        ctx context.Context,
        api bridgeNetworkAPI,
        summary network.Summary,
        allowAbsentGateways bool,
) ([]netip.Addr, error) {
        inspected, err := api.NetworkInspect(ctx, summary.ID, mobyclient.NetworkInspectOptions{})
        if err != nil {
                return nil, fmt.Errorf("inspect Docker bridge network %s: %w", summary.ID, err)
        }
        if inspected.Network.ID != summary.ID || inspected.Network.Driver != "bridge" {
                return nil, fmt.Errorf("docker bridge network %s inspection returned inconsistent identity", summary.ID)
        }
        if len(inspected.Network.IPAM.Config) == 0 {
                return nil, fmt.Errorf("docker bridge network %s has no IPAM configuration", summary.ID)
        }

        gateways := make([]netip.Addr, 0, len(inspected.Network.IPAM.Config))
        for _, config := range inspected.Network.IPAM.Config {
                if !config.Gateway.IsValid() {
                        if allowAbsentGateways {
                                // The Podman compatibility transport maps its explicit "<nil>"
                                // sentinel to an invalid address. Malformed non-empty values fail
                                // during JSON decoding before discovery reaches this case.
                                continue
                        }
                        return nil, fmt.Errorf("docker bridge network %s has an invalid IPAM gateway", summary.ID)
                }
                gateways = append(gateways, config.Gateway.Unmap())
        }
        return gateways, nil
}

// DeleteNetwork deletes a network by name.
func (c *Client) deleteNetwork(ctx context.Context, name string) error {
        // find the network by name using filter for efficiency but verify exact match
        networks, err := c.client.NetworkList(ctx, mobyclient.NetworkListOptions{
                Filters: mobyclient.Filters{}.Add("name", name),
        })
        if err != nil {
                return err
        }

        // Verify exact name match from filtered results
        var networkToRemove *network.Summary
        for _, n := range networks.Items {
                if n.Name == name {
                        networkToRemove = &n
                        break
                }
        }

        // If the network does not exist, there is nothing to do here.
        if networkToRemove == nil {
                slog.Debug("network not found, nothing to delete", "name", name)
                return nil
        }

        if _, err := c.client.NetworkRemove(ctx, networkToRemove.ID, mobyclient.NetworkRemoveOptions{}); err != nil {
                return fmt.Errorf("failed to remove network %s: %w", name, err)
        }
        return nil
}

// removeContainer removes a container by ID, without removing any associated networks or proxy containers.
func (c *Client) removeContainer(ctx context.Context, containerID string) error {
        _, err := c.api.ContainerRemove(ctx, containerID, mobyclient.ContainerRemoveOptions{
                Force: true,
        })
        if err != nil {
                // If the workload doesn't exist, that's fine - it's already removed
                if errdefs.IsNotFound(err) {
                        return nil
                }
                return NewContainerError(err, containerID, fmt.Sprintf("failed to remove container: %v", err))
        }

        return nil
}

// removeProxyContainers removes the MCP server container and any proxy containers.
func (c *Client) removeProxyContainers(
        ctx context.Context,
        containerName string,
        workloadLabels map[string]string,
) error {
        // remove the / if it starts with it
        containerName = strings.TrimPrefix(containerName, "/")

        // If network isolation is not enabled, then there is nothing else to do.
        // NOTE: This check treats all workloads created before the introduction of
        // this label as having network isolation enabled. This is to ensure that they
        // get cleaned up properly during stop/rm. There may be some spurious warnings
        // from the following code, but they can be ignored.
        if !lb.HasNetworkIsolation(workloadLabels) {
                return nil
        }

        // remove egress, ingress, and dns containers
        suffixes := []string{"egress", "ingress", "dns"}

        for _, suffix := range suffixes {
                containerName := fmt.Sprintf("%s-%s", containerName, suffix)
                containerId, err := c.findExistingContainer(ctx, containerName)
                if err != nil {
                        slog.Debug("failed to find container", "type", suffix, "name", containerName, "error", err)
                        continue
                }
                if containerId == "" {
                        continue
                }

                _, err = c.client.ContainerRemove(ctx, containerId, mobyclient.ContainerRemoveOptions{
                        Force: true,
                })
                if err != nil {
                        if errdefs.IsNotFound(err) {
                                continue
                        }
                        return NewContainerError(err, containerId, fmt.Sprintf("failed to remove %s container: %v", suffix, err))
                }
        }

        return nil
}

// CreateContainer creates a container without starting it
// If options is nil, default options will be used
// convertEnvVars converts a map of environment variables to a slice
func convertEnvVars(envVars map[string]string) []string {
        env := make([]string, 0, len(envVars))
        for k, v := range envVars {
                env = append(env, fmt.Sprintf("%s=%s", k, v))
        }
        return env
}

// convertMounts converts internal mount format to Docker mount format
func convertMounts(mounts []runtime.Mount) []mount.Mount {
        result := make([]mount.Mount, 0, len(mounts))
        for _, m := range mounts {
                // All mounts are now bind mounts (removed tmpfs support for overlays)
                result = append(result, mount.Mount{
                        Type:     mount.TypeBind,
                        Source:   m.Source,
                        Target:   m.Target,
                        ReadOnly: m.ReadOnly,
                })
        }
        return result
}

// setupExposedPorts configures exposed ports for a container
func setupExposedPorts(config *container.Config, exposedPorts map[string]struct{}) error {
        if len(exposedPorts) == 0 {
                return nil
        }

        config.ExposedPorts = network.PortSet{}
        for port := range exposedPorts {
                networkPort, err := network.ParsePort(strings.Split(port, "/")[0])
                if err != nil {
                        return fmt.Errorf("failed to parse port: %w", err)
                }
                config.ExposedPorts[networkPort] = struct{}{}
        }

        return nil
}

// setupPortBindings configures port bindings for a container
func setupPortBindings(hostConfig *container.HostConfig, portBindings map[string][]runtime.PortBinding) error {
        if len(portBindings) == 0 {
                return nil
        }

        hostConfig.PortBindings = network.PortMap{}
        for port, bindings := range portBindings {
                networkPort, err := network.ParsePort(strings.Split(port, "/")[0])
                if err != nil {
                        return fmt.Errorf("failed to parse port: %w", err)
                }

                networkBindings := make([]network.PortBinding, len(bindings))
                for i, binding := range bindings {
                        hostIP, err := parseHostIP(binding.HostIP)
                        if err != nil {
                                return fmt.Errorf("failed to parse host IP %q: %w", binding.HostIP, err)
                        }
                        networkBindings[i] = network.PortBinding{
                                HostIP:   hostIP,
                                HostPort: binding.HostPort,
                        }
                }
                hostConfig.PortBindings[networkPort] = networkBindings
        }

        return nil
}

// parseHostIP converts a host IP string into a netip.Addr. An empty string
// (meaning "all interfaces") maps to the zero netip.Addr, preserving the
// previous behavior where an empty HostIP string was passed through verbatim.
func parseHostIP(hostIP string) (netip.Addr, error) {
        if hostIP == "" {
                return netip.Addr{}, nil
        }
        return netip.ParseAddr(hostIP)
}

func (c *Client) createContainer(
        ctx context.Context,
        containerName string,
        config *container.Config,
        hostConfig *container.HostConfig,
        endpointsConfig map[string]*network.EndpointSettings,
) (string, error) {
        existingID, err := c.findExistingContainer(ctx, containerName)
        if err != nil {
                return "", err
        }

        // If container exists, check if we need to recreate it
        if existingID != "" {
                canReuse, err := c.handleExistingContainer(ctx, existingID, config, hostConfig)
                if err != nil {
                        return "", err
                }

                if canReuse {
                        // Container exists with the right configuration, return its ID
                        return existingID, nil
                }
                // Container was removed and needs to be recreated
        }

        // network config
        networkConfig := &network.NetworkingConfig{
                EndpointsConfig: endpointsConfig,
        }

        id, err := c.createAndStartContainer(ctx, containerName, config, hostConfig, networkConfig)
        if err != nil && errdefs.IsNotFound(err) {
                if _, sharesExternalNetwork := endpointsConfig["toolhive-external"]; sharesExternalNetwork {
                        // "toolhive-external" is shared, concurrency-created infrastructure
                        // reused by every workload; deleteNetworks removes it opportunistically
                        // once it looks like no workload needs it anymore. Under concurrent
                        // `thv` processes that teardown can race this container's own creation,
                        // so the network can vanish between our create-time existence check and
                        // this container actually attaching to it, surfacing as a not-found
                        // error from either Create or Start. Recreate it and retry once instead
                        // of failing the whole workload for what is a transient condition.
                        if recreateErr := c.createExternalNetworks(ctx); recreateErr == nil {
                                id, err = c.createAndStartContainer(ctx, containerName, config, hostConfig, networkConfig)
                        }
                }
        }
        return id, err
}

// createAndStartContainer creates and starts a container, wrapping any
// Docker API error with contextual information.
func (c *Client) createAndStartContainer(
        ctx context.Context,
        containerName string,
        config *container.Config,
        hostConfig *container.HostConfig,
        networkConfig *network.NetworkingConfig,
) (string, error) {
        resp, err := c.api.ContainerCreate(ctx, mobyclient.ContainerCreateOptions{
                Config:           config,
                HostConfig:       hostConfig,
                NetworkingConfig: networkConfig,
                Platform:         nil,
                Name:             containerName,
        })
        if err != nil {
                return "", NewContainerError(err, "", fmt.Sprintf("failed to create container: %v", err))
        }
        if resp.Warnings != nil {
                //nolint:gosec // G706: warnings from container API response
                slog.Debug("container creation warnings", "warnings", resp.Warnings)
        }

        // Start the container
        _, err = c.api.ContainerStart(ctx, resp.ID, mobyclient.ContainerStartOptions{})
        if err != nil {
                // Docker leaves the container behind in a "created" state on a failed
                // start (e.g. its network vanished underneath it). Clean it up so a
                // caller-driven retry under the same name doesn't collide with it.
                if _, rmErr := c.api.ContainerRemove(ctx, resp.ID, mobyclient.ContainerRemoveOptions{Force: true}); rmErr != nil {
                        slog.Debug("failed to remove container after failed start", "id", resp.ID, "error", rmErr)
                }
                return "", NewContainerError(err, resp.ID, fmt.Sprintf("failed to start container: %v", err))
        }

        return resp.ID, nil
}

func (c *Client) createDnsContainer(ctx context.Context, dnsContainerName string,
        attachStdio bool, networkName string, endpointsConfig map[string]*network.EndpointSettings) (string, string, error) {
        slog.Debug("setting up DNS container", "name", dnsContainerName, "image", DnsImage)
        dnsLabels := map[string]string{}
        lb.AddStandardLabels(dnsLabels, dnsContainerName, dnsContainerName, "stdio", 80)
        dnsLabels[ToolhiveAuxiliaryWorkloadLabel] = LabelValueTrue

        // pull the dns image if it is not already pulled
        err := c.imageManager.PullImage(ctx, DnsImage)
        if err != nil {
                // Check if the DNS image exists locally before failing
                _, inspectErr := c.client.ImageInspect(ctx, DnsImage)
                if inspectErr == nil {
                        slog.Debug("dns image exists locally, continuing despite pull failure", "image", DnsImage)
                } else {
                        return "", "", fmt.Errorf("failed to pull DNS image: %w", err)
                }
        }

        configDns := &container.Config{
                Image:        DnsImage,
                Cmd:          nil,
                Env:          nil,
                Labels:       dnsLabels,
                AttachStdin:  attachStdio,
                AttachStdout: attachStdio,
                AttachStderr: attachStdio,
                OpenStdin:    attachStdio,
                Tty:          false,
        }

        dnsHostConfig := &container.HostConfig{
                Mounts:      nil,
                NetworkMode: container.NetworkMode("bridge"),
                CapAdd:      nil,
                CapDrop:     nil,
                SecurityOpt: []string{"label:disable"},
                RestartPolicy: container.RestartPolicy{
                        Name: "unless-stopped",
                },
        }

        // now create the dns container
        dnsContainerId, err := c.createContainer(ctx, dnsContainerName, configDns, dnsHostConfig, endpointsConfig)
        if err != nil {
                return "", "", fmt.Errorf("failed to create dns container: %w", err)
        }

        dnsInspectResult, err := c.client.ContainerInspect(ctx, dnsContainerId, mobyclient.ContainerInspectOptions{})
        if err != nil {
                return "", "", fmt.Errorf("failed to inspect DNS container: %w", err)
        }

        dnsNetworkSettings, ok := dnsInspectResult.Container.NetworkSettings.Networks[networkName]
        if !ok {
                return "", "", fmt.Errorf("network %s not found in container's network settings", networkName)
        }
        // EndpointSettings.IPAddress is a netip.Addr; render it back to the string
        // form the callers expect, treating the zero value as an empty address.
        dnsContainerIP := ""
        if dnsNetworkSettings.IPAddress.IsValid() {
                dnsContainerIP = dnsNetworkSettings.IPAddress.String()
        }

        return dnsContainerId, dnsContainerIP, nil
}

func (c *Client) createMcpContainer(
        ctx context.Context,
        name string,
        networkName string,
        image string,
        command []string,
        envVars map[string]string,
        labels map[string]string,
        attachStdio bool,
        permissionConfig *runtime.PermissionConfig,
        additionalDNS string,
        exposedPorts map[string]struct{},
        portBindings map[string][]runtime.PortBinding,
        isolateNetwork bool,
) error {
        // Create container configuration
        config := &container.Config{
                Image:        image,
                Cmd:          command,
                Env:          convertEnvVars(envVars),
                Labels:       labels,
                AttachStdin:  attachStdio,
                AttachStdout: attachStdio,
                AttachStderr: attachStdio,
                OpenStdin:    attachStdio,
                Tty:          false,
        }

        // Create host configuration
        hostConfig := &container.HostConfig{
                Mounts:      convertMounts(permissionConfig.Mounts),
                NetworkMode: container.NetworkMode(permissionConfig.NetworkMode),
                CapAdd:      permissionConfig.CapAdd,
                CapDrop:     permissionConfig.CapDrop,
                SecurityOpt: permissionConfig.SecurityOpt,
                Privileged:  permissionConfig.Privileged,
                RestartPolicy: container.RestartPolicy{
                        Name: "unless-stopped",
                },
        }
        if additionalDNS != "" {
                dnsAddr, err := netip.ParseAddr(additionalDNS)
                if err != nil {
                        return NewContainerError(err, "", fmt.Sprintf("invalid additional DNS address %q: %v", additionalDNS, err))
                }
                hostConfig.DNS = []netip.Addr{dnsAddr}
        }

        // Configure ports if options are provided
        // Setup exposed ports
        if err := setupExposedPorts(config, exposedPorts); err != nil {
                return NewContainerError(err, "", err.Error())
        }

        // Setup port bindings
        if err := setupPortBindings(hostConfig, portBindings); err != nil {
                return NewContainerError(err, "", err.Error())
        }

        // create mcp container
        internalEndpointsConfig := map[string]*network.EndpointSettings{}

        // Check if we have a custom network mode (e.g., "host", "none", etc.)
        if !networking.IsBridgeMode(permissionConfig.NetworkMode) {
                // For custom network modes like "host", "none", etc., don't add any endpoint configurations
                // The NetworkMode in hostConfig will handle the networking
                //nolint:gosec // G706: network mode from permission config
                slog.Debug("using custom network mode", "network_mode", permissionConfig.NetworkMode)
                // Leave internalEndpointsConfig as empty map
        } else if isolateNetwork {
                internalEndpointsConfig[networkName] = &network.EndpointSettings{
                        NetworkID: networkName,
                }
        } else {
                // for other workloads such as inspector, add to external network
                internalEndpointsConfig["toolhive-external"] = &network.EndpointSettings{
                        NetworkID: "toolhive-external",
                }
        }
        _, err := c.createContainer(ctx, name, config, hostConfig, internalEndpointsConfig)
        if err != nil {
                return fmt.Errorf("failed to create container: %w", err)
        }

        return nil

}

// addEgressEnvVars returns a new map containing all entries from envVars plus
// the HTTP_PROXY/HTTPS_PROXY/NO_PROXY variables for the given egress container.
// The caller's map is never mutated.
func addEgressEnvVars(envVars map[string]string, egressContainerName string) map[string]string {
        egressHost := fmt.Sprintf("http://%s:3128", egressContainerName)
        result := maps.Clone(envVars)
        if result == nil {
                result = make(map[string]string)
        }
        result["HTTP_PROXY"] = egressHost
        result["HTTPS_PROXY"] = egressHost
        result["http_proxy"] = egressHost
        result["https_proxy"] = egressHost
        result["NO_PROXY"] = "localhost,127.0.0.1,::1"
        result["no_proxy"] = "localhost,127.0.0.1,::1"
        return result
}

// mergeEnvVars returns a new map containing all entries from base with all
// entries from extra added (extra values overwrite base on conflict). Neither
// input map is mutated.
func mergeEnvVars(base, extra map[string]string) map[string]string {
        result := make(map[string]string, len(base)+len(extra))
        for k, v := range base {
                result[k] = v
        }
        for k, v := range extra {
                result[k] = v
        }
        return result
}

// setupIngressContainer creates the ingress Squid reverse-proxy container for
// the workload and returns the host-side port it is bound on.
func (c *Client) setupIngressContainer(ctx context.Context, containerName, upstreamHost string, upstreamPort int,
        attachStdio bool, externalEndpointsConfig map[string]*network.EndpointSettings,
        networkPermissions *permissions.NetworkPermissions) (int, error) {
        // A random port avoids every same-image workload racing to bind the same
        // preferred port when starting concurrently (see #6063).
        squidPort, err := networking.FindOrUsePort(0)
        if err != nil {
                return 0, fmt.Errorf("failed to find an available port: %w", err)
        }
        squidExposedPorts := map[string]struct{}{
                fmt.Sprintf("%d/tcp", squidPort): {},
        }
        squidPortBindings := map[string][]runtime.PortBinding{
                fmt.Sprintf("%d/tcp", squidPort): {
                        {
                                HostIP:   "127.0.0.1",
                                HostPort: fmt.Sprintf("%d", squidPort),
                        },
                },
        }
        ingressContainerName := fmt.Sprintf("%s-ingress", containerName)
        _, err = createIngressSquidContainer(
                ctx,
                c,
                containerName,
                upstreamHost,
                ingressContainerName,
                attachStdio,
                upstreamPort,
                squidPort,
                squidExposedPorts,
                externalEndpointsConfig,
                squidPortBindings,
                networkPermissions,
        )
        if err != nil {
                return 0, fmt.Errorf("failed to create ingress container: %w", err)
        }
        return squidPort, nil

}

func extractFirstPort(options *runtime.DeployWorkloadOptions) (int, error) {
        var firstPort string
        if len(options.ExposedPorts) == 0 {
                return 0, fmt.Errorf("no exposed ports specified in options.ExposedPorts")
        }
        for port := range options.ExposedPorts {
                firstPort = port

                // need to strip the protocol
                firstPort = strings.Split(firstPort, "/")[0]
                break // take only the first one
        }
        firstPortInt, err := strconv.Atoi(firstPort)
        if err != nil {
                return 0, fmt.Errorf("failed to convert port %s to int: %w", firstPort, err)
        }
        return firstPortInt, nil
}

func (c *Client) createExternalNetworks(ctx context.Context) error {
        externalNetworkLabels := map[string]string{}
        lb.AddNetworkLabels(externalNetworkLabels, "toolhive-external")
        err := c.createNetwork(ctx, "toolhive-external", externalNetworkLabels, false)
        if err != nil {
                return err
        }
        return nil
}

func generatePortBindings(labels map[string]string,
        portBindings map[string][]runtime.PortBinding) (map[string][]runtime.PortBinding, int, error) {
        // Clone portBindings so we never mutate the caller's map.
        portBindings = maps.Clone(portBindings)
        var hostPort int
        // check if we need to map to a random port of not
        if _, ok := labels[ToolhiveAuxiliaryWorkloadLabel]; ok && labels[ToolhiveAuxiliaryWorkloadLabel] == LabelValueTrue {
                // find first port
                var err error
                for _, bindings := range portBindings {
                        if len(bindings) > 0 {
                                hostPortStr := bindings[0].HostPort
                                hostPort, err = strconv.Atoi(hostPortStr)
                                if err != nil {
                                        return nil, 0, fmt.Errorf("failed to convert host port %s to int: %w", hostPortStr, err)
                                }
                                break
                        }
                }
        } else {
                // first port binding needs to map to the host port
                // For consistency, we only use FindAvailable for the primary port if it's not already set
                for key, bindings := range portBindings {
                        if len(bindings) > 0 {
                                hostPortStr := bindings[0].HostPort
                                if hostPortStr == "" || hostPortStr == "0" {
                                        hostPort = networking.FindAvailable()
                                        if hostPort == 0 {
                                                return nil, 0, fmt.Errorf("could not find an available port")
                                        }
                                        bindings[0].HostPort = fmt.Sprintf("%d", hostPort)
                                        portBindings[key] = bindings
                                } else {
                                        var err error
                                        hostPort, err = strconv.Atoi(hostPortStr)
                                        if err != nil {
                                                return nil, 0, fmt.Errorf("failed to convert host port %s to int: %w", hostPortStr, err)
                                        }
                                }
                                break
                        }
                }
        }

        return portBindings, hostPort, nil
}

func (c *Client) stopProxyContainer(ctx context.Context, containerName string, timeoutSeconds int) {
        containerId, err := c.findExistingContainer(ctx, containerName)
        if err != nil {
                slog.Debug("failed to find internal container", "name", containerName, "error", err)
                return
        }
        if containerId == "" {
                return
        }
        if _, err := c.api.ContainerStop(ctx, containerId, mobyclient.ContainerStopOptions{Timeout: &timeoutSeconds}); err != nil {
                slog.Debug("failed to stop internal container", "name", containerName, "error", err)
        }
}

func (c *Client) deleteNetworks(ctx context.Context, containerName string) error {
        // Delete networks if there are no containers using them.
        toolHiveContainers, err := c.client.ContainerList(ctx, mobyclient.ContainerListOptions{
                All:     true,
                Filters: mobyclient.Filters{}.Add("label", "toolhive=true"),
        })
        if err != nil {
                return fmt.Errorf("failed to list containers: %w", err)
        }

        // Delete associated internal network
        networkName := fmt.Sprintf("toolhive-%s-internal", containerName)
        if err := c.deleteNetwork(ctx, networkName); err != nil {
                // just log the error and continue
                slog.Warn("failed to delete network", "name", networkName, "error", err)
        }

        if len(toolHiveContainers.Items) == 0 {
                // remove external network
                if err := c.deleteNetwork(ctx, "toolhive-external"); err != nil {
                        // just log the error and continue
                        slog.Warn("failed to delete network", "name", "toolhive-external", "error", err)
                }
        }
        return nil
}

func dockerToDomainStatus(status string) runtime.WorkloadStatus {
        // Reference: https://docs.docker.com/reference/cli/docker/container/ls/#status
        switch status {
        case "running":
                return runtime.WorkloadStatusRunning
        case "created", "restarting":
                return runtime.WorkloadStatusStarting
        case "paused", "exited", "dead":
                return runtime.WorkloadStatusStopped
        case "removing": // TODO: add handling new workload creation
                return runtime.WorkloadStatusRemoving
        }
        // We should not reach here.
        return runtime.WorkloadStatusUnknown
}

// findContainerByLabel finds a container by the base name label.
// Returns the container ID if found, empty string otherwise.
func (c *Client) findContainerByLabel(ctx context.Context, workloadName string) (string, error) {
        filterArgs := mobyclient.Filters{}.
                Add("label", "toolhive=true").
                Add("label", fmt.Sprintf("toolhive-basename=%s", workloadName))

        listResult, err := c.api.ContainerList(ctx, mobyclient.ContainerListOptions{
                All:     true,
                Filters: filterArgs,
        })
        if err != nil {
                return "", NewContainerError(err, "", fmt.Sprintf("failed to list containers: %v", err))
        }
        containers := listResult.Items

        if len(containers) == 0 {
                return "", nil
        }

        // If multiple containers have the same base name, prefer the running one
        var containerID string
        for _, cont := range containers {
                if cont.State == "running" {
                        containerID = cont.ID
                        break
                }
        }
        // If no running container found, use the first one
        if containerID == "" {
                containerID = containers[0].ID
        }

        return containerID, nil
}

// findContainerByExactName finds a container by exact name matching.
// Returns the container ID if found, empty string otherwise.
// Uses container runtime's name filter for efficiency but then verifies exact match to prevent partial matching
func (c *Client) findContainerByExactName(ctx context.Context, workloadName string) (string, error) {
        filterArgs := mobyclient.Filters{}.
                Add("label", "toolhive=true").
                Add("name", workloadName) // Use name filter for efficiency

        listResult, err := c.api.ContainerList(ctx, mobyclient.ContainerListOptions{
                All:     true,
                Filters: filterArgs,
        })
        if err != nil {
                return "", NewContainerError(err, "", fmt.Sprintf("failed to list containers: %v", err))
        }
        containers := listResult.Items

        if len(containers) == 0 {
                return "", nil
        }

        // Verify exact name match from the filtered results (name filter does partial matching)
        // The name in the API has a leading slash, so we need to search for that.
        prefixedName := "/" + workloadName
        for _, cont := range containers {
                // Check if any of the container names match exactly
                if slices.Contains(cont.Names, prefixedName) || slices.Contains(cont.Names, workloadName) {
                        return cont.ID, nil
                }
        }

        return "", nil
}

// inspectContainerByName finds a container by the workload name and inspects it.
// It first tries to find by base name label, then falls back to exact name matching.
func (c *Client) inspectContainerByName(ctx context.Context, workloadName string) (container.InspectResponse, error) {
        empty := container.InspectResponse{}

        // First try to find container by base name label
        containerID, err := c.findContainerByLabel(ctx, workloadName)
        if err != nil {
                return empty, err
        }
        if containerID != "" {
                result, err := c.api.ContainerInspect(ctx, containerID, mobyclient.ContainerInspectOptions{})
                return result.Container, err
        }

        // Fall back to exact name matching for backward compatibility
        containerID, err = c.findContainerByExactName(ctx, workloadName)
        if err != nil {
                return empty, err
        }
        if containerID == "" {
                return empty, NewContainerError(runtime.ErrWorkloadNotFound, workloadName, "no containers found")
        }

        result, err := c.api.ContainerInspect(ctx, containerID, mobyclient.ContainerInspectOptions{})
        return result.Container, err
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

//go:build linux

package docker

func calculateFinalPort(hostPort int, firstPortInt int, networkName string) int {
        if networkName == "host" {
                return firstPortInt
        }
        return hostPort
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package docker

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "net/netip"
        "os"
        "regexp"
        "slices"
        "strconv"
        "strings"

        "github.com/moby/moby/api/types/container"
        "k8s.io/utils/ptr"

        "github.com/stacklok/toolhive/pkg/container/runtime"
        lb "github.com/stacklok/toolhive/pkg/labels"
        "github.com/stacklok/toolhive/pkg/networking"
)

const (
        // defaultEnvoyImage is pinned by tag+digest for supply-chain integrity.
        // TOOLHIVE_ENVOY_IMAGE overrides must use Envoy v1.39+ with DFP
        // resolved_address_filter support.
        defaultEnvoyImage = "envoyproxy/envoy:distroless-v1.39.1@sha256:eb2c01c13125d1629637cb4e4cce7207009fb7cc2c8027f9742758549d15b6f4"

        // Protobuf type URLs required by Envoy's protobuf-JSON bootstrap format.
        // Every typed_config field must carry an @type URL or Envoy will reject the
        // config on startup.
        typeHCM        = "type.googleapis.com/envoy.extensions.filters.network.http_connection_manager.v3.HttpConnectionManager"
        typeHTTPRBAC   = "type.googleapis.com/envoy.extensions.filters.http.rbac.v3.RBAC"
        typeDFPFilter  = "type.googleapis.com/envoy.extensions.filters.http.dynamic_forward_proxy.v3.FilterConfig"
        typeRouter     = "type.googleapis.com/envoy.extensions.filters.http.router.v3.Router"
        typeDFPCluster = "type.googleapis.com/envoy.extensions.clusters.dynamic_forward_proxy.v3.ClusterConfig"

        // dfpCacheName is the shared DNS cache config name used by both the DFP
        // HTTP filter and the DFP cluster extension.
        dfpCacheName = "dynamic_forward_proxy_cache_config"

        // dfpClusterName is the cluster referenced by the egress route config.
        dfpClusterName = "dynamic_forward_proxy_cluster"

        // ingressClusterName is the cluster referenced by the ingress route config.
        ingressClusterName = "ingress_upstream"
)

func getEnvoyImage() string {
        if img := os.Getenv("TOOLHIVE_ENVOY_IMAGE"); img != "" {
                return img
        }
        return defaultEnvoyImage
}

// ── Bootstrap ────────────────────────────────────────────────────────────────

// envoyBootstrap is the top-level Envoy bootstrap configuration.
// The admin interface is intentionally omitted: Envoy does not start an admin
// server when the field is absent, eliminating the attack surface at runtime.
type envoyBootstrap struct {
        StaticResources envoyStaticResources `json:"static_resources"`
}

// envoyStaticResources holds the static listeners and clusters.
type envoyStaticResources struct {
        Listeners []envoyListener `json:"listeners,omitempty"`
        Clusters  []envoyCluster  `json:"clusters,omitempty"`
}

// ── Address types ────────────────────────────────────────────────────────────

// envoyAddress wraps a socket address for Envoy config.
type envoyAddress struct {
        SocketAddress envoySocketAddress `json:"socket_address"`
}

// envoySocketAddress is an IP + port pair used throughout Envoy config.
type envoySocketAddress struct {
        Address   string `json:"address"`
        PortValue int    `json:"port_value"`
}

// ── Listener / filter chain ──────────────────────────────────────────────────

// envoyListener is an Envoy listener binding on a socket address.
type envoyListener struct {
        Name         string             `json:"name"`
        Address      envoyAddress       `json:"address"`
        FilterChains []envoyFilterChain `json:"filter_chains"`
}

// envoyFilterChain is a sequence of network-level filters applied to matching
// connections.
type envoyFilterChain struct {
        Filters []envoyNetworkFilter `json:"filters"`
}

// envoyNetworkFilter is a named network filter whose typed config is an
// arbitrary protobuf-JSON object (e.g. envoyHCM). The any type lets us embed
// concrete structs directly so that JSON serialisation includes @type.
type envoyNetworkFilter struct {
        Name        string `json:"name"`
        TypedConfig any    `json:"typed_config"`
}

// ── HttpConnectionManager ───────────────────────────────────────────────────

// envoyHCM is the typed config for
// envoy.filters.network.http_connection_manager.
type envoyHCM struct {
        Type       string           `json:"@type"`
        StatPrefix string           `json:"stat_prefix"`
        AccessLog  []envoyAccessLog `json:"access_log,omitempty"`
        // StreamIdleTimeout caps how long a stream may be idle. Envoy defaults it to
        // 5m, which would reap a sparse-but-open SSE stream; "0s" disables it.
        StreamIdleTimeout string               `json:"stream_idle_timeout,omitempty"`
        UpgradeConfigs    []envoyUpgradeConfig `json:"upgrade_configs,omitempty"`
        HTTPFilters       []envoyHTTPFilter    `json:"http_filters"`
        RouteConfig       *envoyRouteConfig    `json:"route_config,omitempty"`
}

// envoyAccessLog configures request access logging for an HCM.
type envoyAccessLog struct {
        Name        string `json:"name"`
        TypedConfig any    `json:"typed_config"`
}

// envoyStdoutAccessLog is the typed config for envoy.access_loggers.stdout.
type envoyStdoutAccessLog struct {
        Type string `json:"@type"`
}

const typeStdoutAccessLog = "type.googleapis.com/envoy.extensions.access_loggers.stream.v3.StdoutAccessLog"

func stdoutAccessLog() []envoyAccessLog {
        return []envoyAccessLog{{
                Name:        "envoy.access_loggers.stdout",
                TypedConfig: &envoyStdoutAccessLog{Type: typeStdoutAccessLog},
        }}
}

// envoyUpgradeConfig enables HTTP upgrade protocols (e.g. CONNECT) in HCM.
type envoyUpgradeConfig struct {
        UpgradeType   string              `json:"upgrade_type"`
        ConnectConfig *envoyConnectConfig `json:"connect_config,omitempty"`
}

// envoyHTTPFilter is a named HTTP-layer filter embedded inside HCM.
type envoyHTTPFilter struct {
        Name        string `json:"name"`
        TypedConfig any    `json:"typed_config"`
}

// ── RBAC ─────────────────────────────────────────────────────────────────────

// envoyHTTPRBAC is the typed config for envoy.filters.http.rbac.
type envoyHTTPRBAC struct {
        Type  string         `json:"@type"`
        Rules envoyRBACRules `json:"rules"`
}

// envoyRBACRules holds the RBAC action and policy map.
// CRITICAL: Policies must NOT have omitempty. An empty map serializes as {} and
// is intentional — an absent field would silently turn deny-all into allow-all.
type envoyRBACRules struct {
        Action   string                     `json:"action"`
        Policies map[string]envoyRBACPolicy `json:"policies"`
}

// envoyRBACPolicy pairs permissions with principals for a single RBAC policy.
type envoyRBACPolicy struct {
        Permissions []envoyPermission `json:"permissions"`
        Principals  []envoyPrincipal  `json:"principals"`
}

// envoyPermission matches a request by various criteria. Exactly one field
// should be set per permission entry.
type envoyPermission struct {
        Any     *bool               `json:"any,omitempty"`
        Header  *envoyHeaderMatcher `json:"header,omitempty"`
        OrRules *envoyOrRules       `json:"or_rules,omitempty"`
}

// envoyOrRules composes multiple permissions with logical OR.
type envoyOrRules struct {
        Rules []envoyPermission `json:"rules"`
}

// envoyHeaderMatcher matches an HTTP header value.
type envoyHeaderMatcher struct {
        Name        string            `json:"name"`
        StringMatch *envoyStringMatch `json:"string_match,omitempty"`
}

// envoyStringMatch matches a string by exact value, prefix, suffix, or regex.
type envoyStringMatch struct {
        Exact     string          `json:"exact,omitempty"`
        Prefix    string          `json:"prefix,omitempty"`
        Suffix    string          `json:"suffix,omitempty"`
        SafeRegex *envoySafeRegex `json:"safe_regex,omitempty"`
}

// envoySafeRegex is an RE2 regex matcher. Envoy anchors the pattern fully
// (implicit ^ and $), so the pattern must match the entire input.
type envoySafeRegex struct {
        Regex string `json:"regex"`
}

// envoyPrincipal matches a downstream principal. Any:true is a wildcard.
type envoyPrincipal struct {
        Any bool `json:"any"`
}

// ── DFP filter ───────────────────────────────────────────────────────────────

// envoyDFPFilter is the typed config for
// envoy.filters.http.dynamic_forward_proxy.
type envoyDFPFilter struct {
        Type           string        `json:"@type"`
        DNSCacheConfig envoyDNSCache `json:"dns_cache_config"`
}

// envoyDNSCache is the shared DNS cache config referenced by both the DFP
// HTTP filter and the DFP cluster extension.
type envoyDNSCache struct {
        Name                  string               `json:"name"`
        DNSLookupFamily       string               `json:"dns_lookup_family"`
        ResolvedAddressFilter *envoyAddressMatcher `json:"resolved_address_filter,omitempty"`
}

// envoyAddressMatcher filters resolved DNS addresses against exact gateway CIDRs.
type envoyAddressMatcher struct {
        Ranges []envoyCIDRRange `json:"ranges"`
}

type envoyCIDRRange struct {
        AddressPrefix string `json:"address_prefix"`
        PrefixLen     int    `json:"prefix_len"`
}

func buildDFPDNSCache(spec proxySpec) envoyDNSCache {
        cache := envoyDNSCache{Name: dfpCacheName, DNSLookupFamily: "V4_ONLY"}
        if spec.AllowDockerGateway {
                return cache
        }
        ranges := make([]envoyCIDRRange, len(spec.GatewayIPs))
        for i, gateway := range spec.GatewayIPs {
                gateway = gateway.Unmap()
                ranges[i] = envoyCIDRRange{AddressPrefix: gateway.String(), PrefixLen: gateway.BitLen()}
        }
        cache.ResolvedAddressFilter = &envoyAddressMatcher{Ranges: ranges}
        return cache
}

// ── Router ────────────────────────────────────────────────────────────────────

// envoyRouter is the typed config for envoy.filters.http.router.
type envoyRouter struct {
        Type string `json:"@type"`
}

// ── Route config ─────────────────────────────────────────────────────────────

// envoyRouteConfig holds the list of virtual hosts for a listener.
type envoyRouteConfig struct {
        VirtualHosts []envoyVirtualHost `json:"virtual_hosts"`
}

// envoyVirtualHost matches requests by domain and dispatches to a cluster.
type envoyVirtualHost struct {
        Name    string       `json:"name"`
        Domains []string     `json:"domains"`
        Routes  []envoyRoute `json:"routes"`
}

// envoyRoute matches a request prefix and forwards it to a cluster.
type envoyRoute struct {
        Match envoyRouteMatch   `json:"match"`
        Route *envoyRouteAction `json:"route,omitempty"`
}

// envoyRouteMatch matches incoming requests by URI prefix or CONNECT method.
// Exactly one of Prefix or ConnectMatcher must be set.
type envoyRouteMatch struct {
        Prefix         string               `json:"prefix,omitempty"`
        ConnectMatcher *envoyConnectMatcher `json:"connect_matcher,omitempty"`
}

// envoyConnectMatcher matches HTTP CONNECT requests (used for HTTPS tunneling).
type envoyConnectMatcher struct{}

// envoyRouteAction forwards matched requests to an upstream cluster.
type envoyRouteAction struct {
        Cluster string `json:"cluster"`
        // Timeout is the route's total request→response cap. Envoy defaults it to
        // 15s, which truncates long-lived MCP streams (SSE, streamable-http); set
        // "0s" to disable it for proxied traffic.
        Timeout        string               `json:"timeout,omitempty"`
        UpgradeConfigs []envoyUpgradeConfig `json:"upgrade_configs,omitempty"`
}

// envoyConnectConfig is the per-route CONNECT tunnel configuration.
type envoyConnectConfig struct{}

// ── Cluster ───────────────────────────────────────────────────────────────────

// envoyCluster is an Envoy upstream cluster definition.
type envoyCluster struct {
        Name           string `json:"name"`
        ConnectTimeout string `json:"connect_timeout"`
        LbPolicy       string `json:"lb_policy,omitempty"`
        Type           string `json:"type,omitempty"`
        // DnsLookupFamily restricts DNS resolution to IPv4 only. Without this,
        // STRICT_DNS clusters attempt AAAA lookups first, which adds latency in
        // environments where IPv6 is unavailable or times out.
        DnsLookupFamily string               `json:"dns_lookup_family,omitempty"`
        ClusterType     *envoyClusterType    `json:"cluster_type,omitempty"`
        LoadAssignment  *envoyLoadAssignment `json:"load_assignment,omitempty"`
}

// envoyClusterType is the custom cluster discovery extension (e.g. DFP).
type envoyClusterType struct {
        Name        string `json:"name"`
        TypedConfig any    `json:"typed_config"`
}

// envoyDFPClusterConfig is the typed config for
// envoy.clusters.dynamic_forward_proxy.
type envoyDFPClusterConfig struct {
        Type           string        `json:"@type"`
        DNSCacheConfig envoyDNSCache `json:"dns_cache_config"`
}

// envoyLoadAssignment is an EDS-style static load assignment.
type envoyLoadAssignment struct {
        ClusterName string          `json:"cluster_name"`
        Endpoints   []envoyEndpoint `json:"endpoints"`
}

// envoyEndpoint is a group of LB endpoints.
type envoyEndpoint struct {
        LBEndpoints []envoyLBEndpoint `json:"lb_endpoints"`
}

// envoyLBEndpoint is a single upstream endpoint.
type envoyLBEndpoint struct {
        Endpoint envoyEndpointAddress `json:"endpoint"`
}

// envoyEndpointAddress wraps an address for an LB endpoint.
type envoyEndpointAddress struct {
        Address envoyAddress `json:"address"`
}

// ── Builder functions ─────────────────────────────────────────────────────────

// buildEgressListener builds the Envoy listener config for outbound traffic.
//
// The resulting HCM HTTP filter chain is:
//  1. (when !spec.AllowDockerGateway) RBAC DENY on gateway IP and hostnames
//  2. RBAC ALLOW allowlist — empty policies map is Envoy's deny-all
//  3. dynamic_forward_proxy HTTP filter — removes resolved gateway addresses
//  4. router
//
// CONNECT upgrades are enabled so that HTTPS CONNECT tunnels pass through.
func buildEgressListener(spec proxySpec) envoyListener {
        httpFilters := buildEgressHTTPFilters(spec)

        hcm := &envoyHCM{
                Type:              typeHCM,
                StatPrefix:        "egress_http",
                AccessLog:         stdoutAccessLog(),
                StreamIdleTimeout: "0s", // don't reap sparse long-lived outbound streams
                UpgradeConfigs:    []envoyUpgradeConfig{{UpgradeType: "CONNECT"}},
                HTTPFilters:       httpFilters,
                RouteConfig: &envoyRouteConfig{
                        VirtualHosts: []envoyVirtualHost{
                                {
                                        Name:    "local_service",
                                        Domains: []string{"*"},
                                        Routes: []envoyRoute{
                                                // CONNECT match must come first: handles HTTPS tunneling.
                                                // Without this, CONNECT requests don't match prefix "/" and get 404.
                                                {
                                                        Match: envoyRouteMatch{ConnectMatcher: &envoyConnectMatcher{}},
                                                        Route: &envoyRouteAction{
                                                                Cluster: dfpClusterName,
                                                                UpgradeConfigs: []envoyUpgradeConfig{
                                                                        {UpgradeType: "CONNECT", ConnectConfig: &envoyConnectConfig{}},
                                                                },
                                                        },
                                                },
                                                // Prefix match handles plain HTTP requests. Timeout "0s"
                                                // disables Envoy's 15s default so long-lived plain-HTTP
                                                // outbound streams aren't truncated (CONNECT tunnels are
                                                // unaffected by the default, but disabling is harmless there).
                                                {
                                                        Match: envoyRouteMatch{Prefix: "/"},
                                                        Route: &envoyRouteAction{Cluster: dfpClusterName, Timeout: "0s"},
                                                },
                                        },
                                },
                        },
                },
        }

        return envoyListener{
                Name: fmt.Sprintf("%s-egress", spec.WorkloadName),
                Address: envoyAddress{
                        SocketAddress: envoySocketAddress{
                                Address:   "0.0.0.0",
                                PortValue: 3128,
                        },
                },
                FilterChains: []envoyFilterChain{
                        {
                                Filters: []envoyNetworkFilter{
                                        {
                                                Name:        "envoy.filters.network.http_connection_manager",
                                                TypedConfig: hcm,
                                        },
                                },
                        },
                },
        }
}

// buildEgressHTTPFilters constructs the ordered list of HTTP filters for the
// egress HCM. Gateway deny rules (when !spec.AllowDockerGateway) are placed
// first so they are evaluated before the allowlist.
func buildEgressHTTPFilters(spec proxySpec) []envoyHTTPFilter {
        var filters []envoyHTTPFilter

        if !spec.AllowDockerGateway {
                filters = append(filters, envoyHTTPFilter{
                        Name:        "envoy.filters.http.rbac",
                        TypedConfig: buildGatewayDenyRBAC(spec.GatewayIPs),
                })
        }

        filters = append(filters,
                envoyHTTPFilter{
                        Name:        "envoy.filters.http.rbac",
                        TypedConfig: buildAllowlistRBAC(spec),
                },
                envoyHTTPFilter{
                        Name: "envoy.filters.http.dynamic_forward_proxy",
                        TypedConfig: &envoyDFPFilter{
                                Type:           typeDFPFilter,
                                DNSCacheConfig: buildDFPDNSCache(spec),
                        },
                },
                envoyHTTPFilter{
                        Name:        "envoy.filters.http.router",
                        TypedConfig: &envoyRouter{Type: typeRouter},
                },
        )

        return filters
}

// buildGatewayDenyRBAC builds an RBAC DENY filter that blocks the Docker bridge
// gateway — both the resolved gateway IP and the Docker-internal hostnames —
// matched on the :authority header. This filter must precede the allowlist.
//
// Matching is on :authority (not destination_ip): for a forward proxy the RBAC
// destination_ip resolves to Envoy's own listener socket, not the proxied
// target, so an L3 CIDR rule never matches the upstream. The forward-proxy
// target is carried in :authority for both plain HTTP and HTTPS CONNECT, so
// authority matching — including the gateway IP literal — is what actually
// blocks it (mirroring Squid's `dst`/`dstdomain` denies).
func buildGatewayDenyRBAC(gatewayIPs []netip.Addr) *envoyHTTPRBAC {
        return &envoyHTTPRBAC{
                Type: typeHTTPRBAC,
                Rules: envoyRBACRules{
                        Action: "DENY",
                        Policies: map[string]envoyRBACPolicy{
                                "gateway": {
                                        Permissions: gatewayAuthorityPermissions(gatewayIPs),
                                        Principals:  []envoyPrincipal{{Any: true}},
                                },
                        },
                },
        }
}

// gatewayAuthorityPermissions matches the protected Docker/Podman host IPs
// and shared host aliases on the :authority header.
// Matches are anchored, case-insensitive, and tolerate an optional port; IPv6
// includes every legal compressed, expanded, and dotted-suffix representation.
// Shared by the deny path (gateway blocked)
// and the allow path (--allow-docker-gateway).
func gatewayAuthorityPermissions(gatewayIPs []netip.Addr) []envoyPermission {
        patterns := make([]string, 0, len(gatewayIPs)+len(hostGatewayNames))
        for _, gateway := range gatewayIPs {
                patterns = append(patterns, gatewayAuthorityRegexes(gateway)...)
        }
        for _, hostname := range hostGatewayNames {
                patterns = append(patterns, hostMatchRegex(hostname))
        }
        rules := make([]envoyPermission, 0, len(patterns))
        for _, re := range patterns {
                rules = append(rules, envoyPermission{
                        Header: &envoyHeaderMatcher{
                                Name:        ":authority",
                                StringMatch: &envoyStringMatch{SafeRegex: &envoySafeRegex{Regex: re}},
                        },
                })
        }
        return []envoyPermission{{OrRules: &envoyOrRules{Rules: rules}}}
}

func gatewayAuthorityRegexes(addr netip.Addr) []string {
        addr = addr.Unmap()
        if addr.Is4() {
                return append([]string{hostMatchRegex(addr.String())}, ipv4MappedAuthorityRegexes(addr.As4())...)
        }
        return ipv6AuthorityRegexes(addr)
}

func ipv4MappedAuthorityRegexes(v4 [4]byte) []string {
        hexTail := `0*` + strconv.FormatUint(uint64(uint16(v4[0])<<8|uint16(v4[1])), 16) +
                `:0*` + strconv.FormatUint(uint64(uint16(v4[2])<<8|uint16(v4[3])), 16)
        dottedTail := strings.Join([]string{
                strconv.Itoa(int(v4[0])), strconv.Itoa(int(v4[1])),
                strconv.Itoa(int(v4[2])), strconv.Itoa(int(v4[3])),
        }, `\.`)
        prefixes := []string{
                `(?:0+:){5}0*ffff:`,
                `(?:0+(?::0+){0,3})?::(?:0+:){0,4}0*ffff:`,
        }
        patterns := make([]string, 0, 4)
        for _, prefix := range prefixes {
                patterns = append(patterns,
                        `(?i)\[`+prefix+hexTail+`\](:[0-9]+)?`,
                        `(?i)\[`+prefix+dottedTail+`\](:[0-9]+)?`,
                )
        }
        return patterns
}

func ipv6AuthorityRegexes(addr netip.Addr) []string {
        bytes := addr.As16()
        words := make([]uint16, 8)
        for i := range words {
                words[i] = uint16(bytes[2*i])<<8 | uint16(bytes[2*i+1])
        }
        patterns := ipv6TextPatterns(words, "")
        dotted := strings.Join([]string{
                strconv.Itoa(int(bytes[12])), strconv.Itoa(int(bytes[13])),
                strconv.Itoa(int(bytes[14])), strconv.Itoa(int(bytes[15])),
        }, `\.`)
        patterns = append(patterns, ipv6TextPatterns(words[:6], dotted)...)
        for i, pattern := range patterns {
                patterns[i] = `(?i)\[` + pattern + `\](:[0-9]+)?`
        }
        return patterns
}

func ipv6TextPatterns(words []uint16, suffix string) []string {
        tokens := make([]string, len(words))
        for i, word := range words {
                if word == 0 {
                        tokens[i] = `0{1,4}`
                } else {
                        tokens[i] = `0{0,3}` + strconv.FormatUint(uint64(word), 16)
                }
        }
        patterns := []string{strings.Join(append(slices.Clone(tokens), suffixIfSet(suffix)...), ":")}
        for start := range words {
                for end := start + 1; end <= len(words); end++ {
                        allZero := true
                        for _, word := range words[start:end] {
                                allZero = allZero && word == 0
                        }
                        if !allZero {
                                break
                        }
                        left := strings.Join(tokens[:start], ":")
                        rightTokens := slices.Clone(tokens[end:])
                        rightTokens = append(rightTokens, suffixIfSet(suffix)...)
                        right := strings.Join(rightTokens, ":")
                        patterns = append(patterns, left+"::"+right)
                }
        }
        return patterns
}

func suffixIfSet(s string) []string {
        if s == "" {
                return nil
        }
        return []string{s}
}

// buildAllowlistRBAC builds an RBAC ALLOW filter encoding the outbound
// allowlist. An empty Policies map is Envoy's deny-all under ALLOW action.
func buildAllowlistRBAC(spec proxySpec) *envoyHTTPRBAC {
        return &envoyHTTPRBAC{
                Type: typeHTTPRBAC,
                Rules: envoyRBACRules{
                        Action:   "ALLOW",
                        Policies: buildAllowlistPolicies(spec),
                },
        }
}

// buildAllowlistPolicies returns the policy map for the egress RBAC ALLOW
// filter. An empty map (deny-all) is returned when:
//   - spec.Permissions is nil
//   - spec.Permissions.Outbound is nil
//
// InsecureAllowAll produces a single wildcard policy. Each AllowHost entry
// becomes a policy matching the :authority header via hostMatchRegex (Squid
// dstdomain semantics). When AllowPort is also set, each host policy AND-s a
// port-suffix permission so that host AND port must both match — mirroring
// Squid's "allowed_ports AND allowed_dsts" combination.
//
// Known divergence from Squid: plain-HTTP requests omit the port from
// :authority (e.g. "example.com" not "example.com:80"). A port-suffix matcher
// for ":80" will not match the portless authority, so AllowPort:[80] with plain
// HTTP is more restrictive than Squid's allowed_ports ACL. This is fail-closed
// (over-blocks) rather than a bypass.
//
// When spec.AllowDockerGateway is set, an explicit ALLOW policy for the gateway
// is added: omitting the deny filter alone is not enough, because the allowlist
// is default-deny, so the gateway would still be blocked under a non-permissive
// profile. This mirrors Squid, which emits explicit allow rules for the flag.
func buildAllowlistPolicies(spec proxySpec) map[string]envoyRBACPolicy {
        policies := make(map[string]envoyRBACPolicy)

        // --allow-docker-gateway: grant the gateway explicitly (see doc comment).
        // Harmless under InsecureAllowAll (allow-all already covers it).
        if spec.AllowDockerGateway {
                policies["gateway"] = envoyRBACPolicy{
                        Permissions: gatewayAuthorityPermissions(spec.GatewayIPs),
                        Principals:  []envoyPrincipal{{Any: true}},
                }
        }

        if spec.Permissions == nil || spec.Permissions.Outbound == nil {
                return policies
        }
        out := spec.Permissions.Outbound
        if out.InsecureAllowAll {
                policies["allow-all"] = envoyRBACPolicy{
                        Permissions: []envoyPermission{{Any: ptr.To(true)}},
                        Principals:  []envoyPrincipal{{Any: true}},
                }
                return policies
        }
        if len(out.AllowHost) == 0 && len(out.AllowPort) > 0 {
                // AllowPort only — any host, but only on the listed ports.
                // Use a single regex so both the explicit-port ("host:80") and the
                // bare-hostname ("host", implying port 80) forms are handled correctly.
                policies["any-host-allowed-ports"] = envoyRBACPolicy{
                        Permissions: []envoyPermission{{
                                Header: &envoyHeaderMatcher{
                                        Name:        ":authority",
                                        StringMatch: &envoyStringMatch{SafeRegex: &envoySafeRegex{Regex: anyHostPortRegex(out.AllowPort)}},
                                },
                        }},
                        Principals: []envoyPrincipal{{Any: true}},
                }
                return policies
        }

        for _, host := range out.AllowHost {
                var regex string
                if len(out.AllowPort) > 0 {
                        // AllowHost + AllowPort: build a single regex that encodes both the
                        // host pattern and the allowed ports. Using a combined regex rather
                        // than separate AND-d permissions means plain-HTTP requests (where
                        // :authority omits the default port 80) are handled correctly — a
                        // bare "example.com" is permitted when port 80 is in the list.
                        regex = hostPortMatchRegex(host, out.AllowPort)
                } else {
                        // AllowHost only — any port permitted.
                        regex = hostMatchRegex(host)
                }
                policies[host] = envoyRBACPolicy{
                        Permissions: []envoyPermission{{
                                Header: &envoyHeaderMatcher{
                                        Name:        ":authority",
                                        StringMatch: &envoyStringMatch{SafeRegex: &envoySafeRegex{Regex: regex}},
                                },
                        }},
                        Principals: []envoyPrincipal{{Any: true}},
                }
        }
        return policies
}

// portGroupRegex builds the ":(?:port1|port2|…)" alternation for a list of
// allowed ports. Returns (group, includes80) where includes80 indicates whether
// port 80 is present (caller uses this to decide whether the port suffix is
// optional). Shared by hostPortMatchRegex and anyHostPortRegex.
func portGroupRegex(ports []int) (group string, includes80 bool) {
        alts := make([]string, 0, len(ports))
        for _, p := range ports {
                alts = append(alts, strconv.Itoa(p))
                if p == 80 {
                        includes80 = true
                }
        }
        return ":(?:" + strings.Join(alts, "|") + ")", includes80
}

// hostBasePattern returns the escaped, case-insensitive host portion of an
// :authority regex — without any trailing port group. It mirrors hostMatchRegex
// semantics: leading dot / *. = subdomain wildcard, optional trailing dot.
// Shared by hostPortMatchRegex and hostMatchRegex.
func hostBasePattern(host string) string {
        subdomains := false
        switch {
        case strings.HasPrefix(host, "*."):
                host = host[2:]
                subdomains = true
        case strings.HasPrefix(host, "."):
                host = host[1:]
                subdomains = true
        }
        p := regexp.QuoteMeta(host)
        if subdomains {
                p = `(.*\.)?` + p
        }
        return p
}

// hostPortMatchRegex builds a single anchored RE2 pattern that encodes both a
// host allowlist entry and a port allowlist, so that AllowHost+AllowPort can be
// expressed as one safe_regex permission rather than AND-ing two separate matchers.
//
// Port 80 is treated as the HTTP default: when 80 is in the ports list the port
// suffix is optional, so a bare "example.com" authority (plain HTTP) is permitted
// alongside "example.com:80". For all other ports the port suffix is required.
//
// The host portion mirrors hostMatchRegex semantics (leading dot / *. = subdomain
// wildcard, case-insensitive, optional trailing dot, no port in the base pattern).
func hostPortMatchRegex(host string, ports []int) string {
        portGroup, includes80 := portGroupRegex(ports)
        hostPattern := hostBasePattern(host)
        if includes80 {
                // Port is optional: bare "example.com" counts as port 80.
                return `(?i)` + hostPattern + `\.?(` + portGroup + `)?`
        }
        // Port is required: bare hostname is not accepted.
        return `(?i)` + hostPattern + `\.?` + portGroup
}

// anyHostPortRegex builds a regex matching any hostname on the given ports.
// When port 80 is in the list the port suffix is optional (bare hostname implies
// port 80 for plain HTTP). For other ports the port suffix is required.
//
// [^:]+ matches any hostname without an explicit port; IPv6 bracket addresses
// (which contain colons) are outside the supported scope for this proxy backend.
func anyHostPortRegex(ports []int) string {
        portGroup, includes80 := portGroupRegex(ports)
        if includes80 {
                return `(?i)[^:]+(` + portGroup + `)?`
        }
        return `(?i)[^:]+` + portGroup
}

// hostMatchRegex builds an anchored, case-insensitive RE2 pattern that matches
// an AllowHost entry against the HTTP :authority header. It tolerates an
// optional ":port" suffix because HTTPS CONNECT authorities include the port
// (e.g. "example.com:443"), and it mirrors Squid's dstdomain semantics:
//
//   - ".example.com" or "*.example.com" — matches the apex "example.com" AND
//     any subdomain ("www.example.com", "a.b.example.com").
//   - "example.com" — matches that host exactly (no subdomains).
//
// The domain is regex-escaped so dots are literal, and the pattern is anchored
// by Envoy so it cannot match "example.com.attacker.com".
func hostMatchRegex(host string) string {
        pattern := hostBasePattern(host)
        // (?i) case-insensitive (hostnames are); \.? tolerates a trailing-dot FQDN
        // ("host.docker.internal." resolves identically and would otherwise bypass a
        // deny); (:[0-9]+)? optional port.
        return `(?i)` + pattern + `\.?(:[0-9]+)?`
}

// buildEgressCluster returns the dynamic_forward_proxy cluster required by the
// egress listener's route config.
func buildEgressCluster(spec proxySpec) envoyCluster {
        return envoyCluster{
                Name:           dfpClusterName,
                ConnectTimeout: "10s",
                LbPolicy:       "CLUSTER_PROVIDED",
                ClusterType: &envoyClusterType{
                        Name: "envoy.clusters.dynamic_forward_proxy",
                        TypedConfig: &envoyDFPClusterConfig{
                                Type:           typeDFPCluster,
                                DNSCacheConfig: buildDFPDNSCache(spec),
                        },
                },
        }
}

// buildIngressListener builds the Envoy listener config for inbound (ingress)
// traffic. It binds on 0.0.0.0:hostPort inside the container so that Docker's
// port forwarding (host:127.0.0.1:<hostPort> → container:<hostPort>) can reach
// it. The host-side HostIP:"127.0.0.1" in the port binding provides the
// localhost-only restriction; the container-side address must be 0.0.0.0 or
// Docker's bridge forwarding cannot deliver traffic to the listener.
//
// The virtual host domain is always "*". Inbound host filtering is enforced
// by the egress RBAC `:authority` matcher (hostMatchRegex), not the ingress
// virtual host list — the transparent proxy sends "127.0.0.1:<port>" as the
// Host header (port included), which would not match bare hostnames in a
// domain list.
func buildIngressListener(spec proxySpec, hostPort int) envoyListener {
        domains := ingressDomains(spec)

        hcm := &envoyHCM{
                Type:       typeHCM,
                StatPrefix: "ingress_http",
                AccessLog:  stdoutAccessLog(),
                // Disable the idle timer so a sparse-but-open SSE stream isn't reaped.
                StreamIdleTimeout: "0s",
                HTTPFilters: []envoyHTTPFilter{
                        {
                                Name:        "envoy.filters.http.router",
                                TypedConfig: &envoyRouter{Type: typeRouter},
                        },
                },
                RouteConfig: &envoyRouteConfig{
                        VirtualHosts: []envoyVirtualHost{
                                {
                                        Name:    "ingress_service",
                                        Domains: domains,
                                        Routes: []envoyRoute{
                                                {
                                                        Match: envoyRouteMatch{Prefix: "/"},
                                                        // Timeout "0s" disables Envoy's 15s default total-response
                                                        // cap, which would otherwise truncate SSE / streamable-http
                                                        // MCP streams. This is the primary transport path.
                                                        Route: &envoyRouteAction{Cluster: ingressClusterName, Timeout: "0s"},
                                                },
                                        },
                                },
                        },
                },
        }

        return envoyListener{
                Name: fmt.Sprintf("%s-ingress", spec.WorkloadName),
                Address: envoyAddress{
                        SocketAddress: envoySocketAddress{
                                Address:   "0.0.0.0", // must be 0.0.0.0; Docker port forwarding targets the bridge IP, not container loopback
                                PortValue: hostPort,
                        },
                },
                FilterChains: []envoyFilterChain{
                        {
                                Filters: []envoyNetworkFilter{
                                        {
                                                Name:        "envoy.filters.network.http_connection_manager",
                                                TypedConfig: hcm,
                                        },
                                },
                        },
                },
        }
}

// ingressDomains returns the virtual host domain list for the ingress listener.
// Always returns a wildcard: the Inbound.AllowHost list contains bare hostnames
// (e.g. "localhost", "127.0.0.1") but the transparent proxy sends a Host header
// with a port suffix ("127.0.0.1:22354"), which Envoy would not match against
// a bare hostname. The inbound access restriction is instead enforced by the
// host-side port binding to 127.0.0.1, which already limits the ingress to
// local connections only.
func ingressDomains(_ proxySpec) []string {
        return []string{"*"}
}

// buildIngressCluster returns the STRICT_DNS upstream cluster for the ingress
// listener, pointing at spec.WorkloadName:spec.UpstreamPort.
func buildIngressCluster(spec proxySpec) envoyCluster {
        // Prefer the resolved upstream IP (no DNS dependency); fall back to the
        // workload name. See proxySpec.UpstreamHost.
        upstreamHost := spec.UpstreamHost
        if upstreamHost == "" {
                upstreamHost = spec.WorkloadName
        }
        return envoyCluster{
                Name:            ingressClusterName,
                ConnectTimeout:  "10s",
                Type:            "STRICT_DNS",
                DnsLookupFamily: "V4_ONLY",
                LoadAssignment: &envoyLoadAssignment{
                        ClusterName: ingressClusterName,
                        Endpoints: []envoyEndpoint{
                                {
                                        LBEndpoints: []envoyLBEndpoint{
                                                {
                                                        Endpoint: envoyEndpointAddress{
                                                                Address: envoyAddress{
                                                                        SocketAddress: envoySocketAddress{
                                                                                Address:   upstreamHost,
                                                                                PortValue: spec.UpstreamPort,
                                                                        },
                                                                },
                                                        },
                                                },
                                        },
                                },
                        },
                },
        }
}

// writeEnvoyBootstrap marshals b to JSON and writes it to a temporary file at
// mode 0644, returning the path. The bootstrap contains no secrets and must be
// readable by the nonroot Envoy container. On success the file must outlive this call —
// the Envoy container bind-mounts it read-only — so it is not removed here; the
// caller removes it if setup fails afterward. On any error during writing, the
// partially-created file is removed before returning.
func writeEnvoyBootstrap(b envoyBootstrap) (string, error) {
        data, err := json.Marshal(b)
        if err != nil {
                return "", fmt.Errorf("failed to marshal envoy bootstrap: %w", err)
        }
        tmpFile, err := os.CreateTemp("", "envoy-bootstrap-*.json")
        if err != nil {
                return "", fmt.Errorf("failed to create envoy bootstrap temp file: %w", err)
        }
        created := tmpFile.Name()
        defer func() {
                if cerr := tmpFile.Close(); cerr != nil {
                        slog.Warn("failed to close envoy bootstrap temp file", "error", cerr)
                }
        }()
        if _, err := tmpFile.Write(data); err != nil {
                _ = os.Remove(created)
                return "", fmt.Errorf("failed to write envoy bootstrap: %w", err)
        }
        // 0o644: world-readable so the Envoy distroless container (UID 101) can
        // read the bind-mounted file. On Linux Docker Engine, strict POSIX
        // permissions apply — 0o600 prevents the container user from reading the
        // file, causing Envoy to crash-loop. The bootstrap contains no secrets
        // (only network topology: hostnames, ports, RBAC rules).
        if err := tmpFile.Chmod(0o644); err != nil {
                _ = os.Remove(created)
                return "", fmt.Errorf("failed to set envoy bootstrap file permissions: %w", err)
        }
        return created, nil
}

// ── envoyProxy ────────────────────────────────────────────────────────────────

// envoyProxy implements networkProxy using Envoy as the proxy backend.
// It creates a single Envoy container that handles both egress (forward proxy
// on :3128) and ingress (reverse proxy) as separate listeners, reducing aux
// container count from 3 (Squid: egress + ingress + dns) to 2 (Envoy: combined
// + dns).
type envoyProxy struct {
        client *Client
}

// SetupEgress implements networkProxy for the Envoy backend.
//
// Only the egress proxy env vars are returned here — no container is created
// yet. Container creation is deferred to SetupIngress (which runs after
// createMcpContainer) so that the MCP container's hostname is resolvable the
// moment the Envoy STRICT_DNS ingress cluster first probes it. Creating Envoy
// before the MCP container caused the ingress cluster to cache a negative DNS
// response on Linux Docker Engine, preventing the server from ever becoming
// ready within the readiness window (see #5922).
func (*envoyProxy) SetupEgress(_ context.Context, spec proxySpec) (egressResult, error) {
        egressContainerName := fmt.Sprintf("%s-egress", spec.WorkloadName)
        return egressResult{EnvVars: addEgressEnvVars(nil, egressContainerName)}, nil
}

// SetupIngress implements networkProxy for the Envoy backend.
//
// Creates the Envoy container after the MCP container exists, with both the
// egress forward-proxy listener and (for non-stdio transports) the ingress
// reverse-proxy listener. Running after createMcpContainer ensures the
// STRICT_DNS upstream cluster resolves the MCP hostname on the first probe,
// avoiding the Linux Docker Engine readiness failure described in #5922.
func (e *envoyProxy) SetupIngress(ctx context.Context, spec proxySpec, _ egressResult) (int, error) {
        // The container is named <name>-egress (not <name>-proxy or similar) to
        // preserve parity with the Squid backend: addEgressEnvVars and the suffix-
        // based cleanup loop both key off this name.
        egressContainerName := fmt.Sprintf("%s-egress", spec.WorkloadName)

        bootstrap := envoyBootstrap{
                StaticResources: envoyStaticResources{
                        Listeners: []envoyListener{buildEgressListener(spec)},
                        Clusters:  []envoyCluster{buildEgressCluster(spec)},
                },
        }

        var ingressPort int
        if spec.TransportType != "stdio" && spec.UpstreamPort > 0 {
                // A random port avoids every same-image workload racing to bind the
                // same preferred port when starting concurrently (see #6063).
                port, err := networking.FindOrUsePort(0)
                if err != nil {
                        return 0, fmt.Errorf("failed to find ingress port: %w", err)
                }
                ingressPort = port
                bootstrap.StaticResources.Listeners = append(
                        bootstrap.StaticResources.Listeners,
                        buildIngressListener(spec, ingressPort),
                )
                bootstrap.StaticResources.Clusters = append(
                        bootstrap.StaticResources.Clusters,
                        buildIngressCluster(spec),
                )
        }

        configPath, err := writeEnvoyBootstrap(bootstrap)
        if err != nil {
                return 0, err
        }
        success := false
        defer func() {
                if !success {
                        _ = os.Remove(configPath)
                }
        }()

        envoyImage := getEnvoyImage()
        //nolint:gosec // G706: envoy image name from config
        slog.Debug("setting up envoy container", "name", egressContainerName, "image", envoyImage)

        if err := e.client.imageManager.PullImage(ctx, envoyImage); err != nil {
                exists, inspectErr := e.client.imageManager.ImageExists(ctx, envoyImage)
                if inspectErr != nil || !exists {
                        return 0, fmt.Errorf("failed to pull envoy image: %w", err)
                }
                //nolint:gosec // G706: envoy image name from config
                slog.Debug("envoy image exists locally, continuing despite pull failure", "image", envoyImage)
        }

        envoyLabels := map[string]string{}
        lb.AddStandardLabels(envoyLabels, egressContainerName, egressContainerName, "stdio", 80)
        envoyLabels[ToolhiveAuxiliaryWorkloadLabel] = LabelValueTrue

        containerConfig := &container.Config{
                Image:  envoyImage,
                Cmd:    []string{"-c", "/etc/envoy/envoy.json"},
                Labels: envoyLabels,
        }

        mounts := []runtime.Mount{
                {Source: configPath, Target: "/etc/envoy/envoy.json", ReadOnly: true},
        }

        var exposedPorts map[string]struct{}
        var portBindings map[string][]runtime.PortBinding
        if ingressPort > 0 {
                portStr := strconv.Itoa(ingressPort)
                portKey := portStr + "/tcp"
                exposedPorts = map[string]struct{}{portKey: {}}
                portBindings = map[string][]runtime.PortBinding{
                        portKey: {{HostIP: "127.0.0.1", HostPort: portStr}},
                }
        }

        hostConfig := &container.HostConfig{
                Mounts:      convertMounts(mounts),
                NetworkMode: container.NetworkMode("bridge"),
                SecurityOpt: []string{"label:disable"},
                CapDrop:     []string{"ALL"},
                RestartPolicy: container.RestartPolicy{
                        Name: "unless-stopped",
                },
        }
        if portBindings != nil {
                if err := setupPortBindings(hostConfig, portBindings); err != nil {
                        return 0, fmt.Errorf("failed to setup port bindings: %w", err)
                }
        }
        if err := setupExposedPorts(containerConfig, exposedPorts); err != nil {
                return 0, fmt.Errorf("failed to setup exposed ports: %w", err)
        }

        if _, err := e.client.createContainer(ctx, egressContainerName, containerConfig, hostConfig, spec.Endpoints); err != nil {
                return 0, fmt.Errorf("failed to create envoy container: %w", err)
        }

        success = true
        return ingressPort, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package docker

import (
        "archive/tar"
        "bufio"
        "bytes"
        "context"
        "errors"
        "fmt"
        "io"
        "net/netip"
        "slices"
        "strings"
        "time"

        mobyclient "github.com/moby/moby/client"
)

var hostGatewayNames = [...]string{dockerGatewayHostname, dockerAltGatewayHostname, podmanGatewayHostname}

// CopyFromContainer uses the runtime's archive endpoint, not the client's host
// filesystem. Podman only exposes its generated hosts file after container start.
type hostDiscoveryAPI interface {
        CopyFromContainer(context.Context, string, mobyclient.CopyFromContainerOptions) (mobyclient.CopyFromContainerResult, error)
}

// discoverHostAliases reads the running DNS sidecar's runtime-generated hosts.
// ToolHive's DNS and proxy sidecars use bridge networking without hosts overrides;
// all bridge gateways separately cover Podman's rootful gateway selection.
func (c *Client) discoverHostAliases(ctx context.Context, dnsContainerID string) (_ []netip.Addr, retErr error) {
        ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
        defer cancel()
        archive, err := c.hostAPI.CopyFromContainer(ctx, dnsContainerID, mobyclient.CopyFromContainerOptions{SourcePath: "/etc/hosts"})
        if err != nil {
                return nil, fmt.Errorf("read running DNS hosts: %w", err)
        }
        defer func() { retErr = errors.Join(retErr, archive.Content.Close()) }()
        return parseHostAliasArchive(archive.Content)
}

// Bound the entire archive as well as the member, so extended headers, padding,
// and unexpected files cannot cause unbounded reads. Never extract to disk.
func parseHostAliasArchive(r io.Reader) ([]netip.Addr, error) {
        const maxHostsSize = 1 << 20
        const maxArchiveSize = 2 << 20
        data, err := io.ReadAll(io.LimitReader(r, maxArchiveSize+1))
        if err != nil {
                return nil, fmt.Errorf("read hosts archive: %w", err)
        }
        if len(data) > maxArchiveSize {
                return nil, errors.New("hosts archive exceeds size limit")
        }
        tr := tar.NewReader(bytes.NewReader(data))
        header, err := tr.Next()
        if err != nil {
                return nil, fmt.Errorf("read hosts archive header: %w", err)
        }
        if header.Name != "hosts" || header.Typeflag != tar.TypeReg || header.Size > maxHostsSize {
                return nil, errors.New("hosts archive must contain one bounded regular hosts file")
        }
        addresses, err := parseHostAliases(tr)
        if err != nil {
                return nil, err
        }
        if _, err := tr.Next(); !errors.Is(err, io.EOF) {
                return nil, errors.New("hosts archive contains extra or malformed entries")
        }
        return addresses, nil
}

func parseHostAliases(r io.Reader) ([]netip.Addr, error) {
        var foundPodmanAlias bool
        var addresses []netip.Addr
        scanner := bufio.NewScanner(r)
        scanner.Buffer(make([]byte, 4096), (1<<20)+1)
        for scanner.Scan() {
                line, _, _ := strings.Cut(scanner.Text(), "#")
                fields := strings.Fields(line)
                if len(fields) == 0 {
                        continue
                }
                addr, err := netip.ParseAddr(fields[0])
                if err != nil || addr.Zone() != "" || len(fields) < 2 {
                        return nil, errors.New("malformed hosts entry")
                }
                for _, name := range fields[1:] {
                        name = strings.TrimSuffix(strings.ToLower(name), ".")
                        if slices.Contains(hostGatewayNames[:], name) {
                                foundPodmanAlias = foundPodmanAlias || name == podmanGatewayHostname
                                addresses = append(addresses, addr.Unmap())
                        }
                }
        }
        if err := scanner.Err(); err != nil {
                return nil, fmt.Errorf("read hosts: %w", err)
        }
        // Only Podman's own alias is required: older Podman releases do not write
        // host.docker.internal, and an absent alias cannot resolve to the host.
        if !foundPodmanAlias {
                return nil, errors.New("hosts missing " + podmanGatewayHostname + "; check Podman hosts configuration")
        }
        slices.SortFunc(addresses, netip.Addr.Compare)
        return slices.Compact(addresses), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package docker

import (
        "context"
        "fmt"
        "net/netip"
        "os"

        "github.com/moby/moby/api/types/network"

        "github.com/stacklok/toolhive-core/permissions"
)

// networkProxy is the single enforcement point for outbound allowlisting,
// docker-gateway blocking, and inbound reverse-proxying for isolated workloads.
//
// What it enables: all egress and ingress proxy containers are created through
// this interface, ensuring a consistent policy-enforcement seam that can be
// swapped at startup via TOOLHIVE_NETWORK_PROXY.
//
// What it does NOT solve: non-cooperative egress is contained by the
// Internal:true network blackhole created by createNetwork, not by this proxy.
// True L3/L4 traffic interception is a separate Phase 2 concern and is not
// addressed here.
type networkProxy interface {
        // SetupEgress provisions egress enforcement BEFORE the MCP container is
        // created and returns the environment variables to inject into the workload
        // (HTTP_PROXY etc.).
        //
        // Squid creates its egress container here. The Envoy backend creates no
        // container in SetupEgress — it returns only env vars — and defers all
        // container creation (both egress and ingress listeners, one container) to
        // SetupIngress so the STRICT_DNS ingress cluster can resolve the MCP hostname
        // on its first probe.
        //
        // It must run before createMcpContainer so the returned env vars land in the
        // workload's environment.
        SetupEgress(ctx context.Context, spec proxySpec) (egressResult, error)

        // SetupIngress creates the ingress proxy AFTER the MCP container exists and
        // returns the host-side ingress port (0 for stdio / UpstreamPort==0).
        //
        // Squid creates its ingress container here. The Envoy backend creates its
        // single dual-listener container here (both egress forward-proxy and ingress
        // reverse-proxy listeners). Running after createMcpContainer ensures the MCP
        // container's hostname resolves on first probe, avoiding a cached negative DNS
        // lookup that would leave the ingress permanently unable to reach the upstream.
        //
        // It must run after createMcpContainer.
        SetupIngress(ctx context.Context, spec proxySpec, egress egressResult) (int, error)
}

// proxySpec contains all the parameters needed to set up proxy containers for
// an isolated workload.
type proxySpec struct {
        // WorkloadName is the base name of the MCP container (e.g. "myserver").
        WorkloadName string
        // UpstreamHost is the address the ingress reverse proxy connects to for the
        // MCP container. It is the container's resolved IP on the internal network
        // (not its name) so the ingress proxy has no DNS dependency: under
        // concurrent startup the container's DNS record can lag, and a hostname-based
        // Squid cache_peer caches the failed lookup and never recovers within the
        // readiness window (see #6063). Empty falls back to WorkloadName.
        UpstreamHost string
        // Permissions holds the network permission profile from the workload's
        // permission profile, governing what outbound traffic is allowed.
        Permissions *permissions.NetworkPermissions
        // AllowDockerGateway explicitly allows Docker/Podman host IPs and aliases
        // on any port, replacing the default deny rules.
        AllowDockerGateway bool
        // GatewayIPs is the canonical set of bridge gateways and runtime host-alias addresses.
        GatewayIPs []netip.Addr
        // TransportType is the MCP transport in use (e.g. "stdio", "sse",
        // "streamable-http"). A value of "stdio" suppresses ingress proxy creation.
        TransportType string
        // UpstreamPort is the container port the MCP server listens on. Ignored
        // when TransportType is "stdio" or the value is 0.
        UpstreamPort int
        // AttachStdio controls whether the proxy containers attach stdio streams.
        AttachStdio bool
        // Endpoints is the set of network endpoints the proxy containers should
        // join, keyed by network name.
        Endpoints map[string]*network.EndpointSettings
}

// egressResult is the output of a successful SetupEgress call. It is passed to
// SetupIngress so backends can carry any state set up during egress forward
// without holding per-workload state on the shared proxy.
type egressResult struct {
        // EnvVars contains environment variables that must be merged into the MCP
        // container's environment (e.g. HTTP_PROXY, HTTPS_PROXY).
        EnvVars map[string]string
}

// newNetworkProxy reads the TOOLHIVE_NETWORK_PROXY environment variable and
// returns the corresponding networkProxy implementation. An empty value or
// "squid" selects the default squid-based proxy. Any other value returns an
// error so that misconfiguration is caught at startup.
func newNetworkProxy(c *Client) (networkProxy, error) {
        val := os.Getenv("TOOLHIVE_NETWORK_PROXY")
        switch val {
        case "", "squid":
                return &squidProxy{client: c}, nil
        case "envoy":
                return &envoyProxy{client: c}, nil
        default:
                return nil, fmt.Errorf("unknown TOOLHIVE_NETWORK_PROXY value %q: supported values are \"squid\" (default), \"envoy\"", val)
        }
}

// Compile-time assertion that squidProxy satisfies networkProxy.
var _ networkProxy = (*squidProxy)(nil)

// Compile-time assertion that envoyProxy satisfies networkProxy.
var _ networkProxy = (*envoyProxy)(nil)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package docker

import (
        "context"

        "github.com/stacklok/toolhive/pkg/container/runtime"
)

func init() {
        runtime.RegisterRuntime(&runtime.Info{
                Name:     RuntimeName,
                Priority: 100,
                Initializer: func(ctx context.Context) (runtime.Runtime, error) {
                        return NewClient(ctx)
                },
                AutoDetector: func() bool {
                        return IsAvailable()
                },
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

//go:build !windows

package sdk

import (
        "context"
        "fmt"
        "log/slog"
        "net"
        "net/http"
        "os"
        "path/filepath"

        mobyclient "github.com/moby/moby/client"

        "github.com/stacklok/toolhive/pkg/container/runtime"
)

// ErrRuntimeNotFound is returned when a container runtime is not found
var ErrRuntimeNotFound = fmt.Errorf("container runtime not found")

// newPlatformClient creates a Docker client using Unix sockets
func newPlatformClient(socketPath string) (*http.Client, []mobyclient.Opt) {
        // Create a custom HTTP client that uses the Unix socket
        httpClient := &http.Client{
                Transport: &http.Transport{
                        DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
                                return net.Dial("unix", socketPath)
                        },
                },
        }

        // Create Docker client options. API-version negotiation is enabled by
        // default in the new client, so it no longer needs to be requested
        // explicitly.
        opts := []mobyclient.Opt{
                // Configure the host before installing a potentially wrapped Podman transport.
                mobyclient.WithHost("unix://" + socketPath),
                mobyclient.WithHTTPClient(httpClient),
        }

        return httpClient, opts
}

// findPlatformContainerSocket finds container socket paths on Unix systems.
// Returns every candidate path that exists on disk, in priority order, so the
// caller can fall back to the next candidate when an earlier one is found but
// cannot be connected to (e.g. /var/run/docker.sock present but EACCES, while
// Docker Desktop's per-user socket would have worked).
func findPlatformContainerSocket(rt runtime.Type) ([]string, runtime.Type, error) {
        // First check for custom socket paths via environment variables. The
        // override is explicit, so we don't add other fallbacks alongside it.
        if customSocketPath := os.Getenv(PodmanSocketEnv); customSocketPath != "" {
                //nolint:gosec // G706: socket path from trusted environment variable
                slog.Debug("using Podman socket from env", "path", customSocketPath)
                if _, err := os.Stat(customSocketPath); err != nil { //nolint:gosec // G703: socket path from trusted environment variable
                        return nil, runtime.TypePodman, fmt.Errorf("invalid Podman socket path: %w", err)
                }
                return []string{customSocketPath}, runtime.TypePodman, nil
        }

        if customSocketPath := os.Getenv(DockerSocketEnv); customSocketPath != "" {
                //nolint:gosec // G706: socket path from trusted environment variable
                slog.Debug("using Docker socket from env", "path", customSocketPath)
                if _, err := os.Stat(customSocketPath); err != nil { //nolint:gosec // G703: socket path from trusted environment variable
                        return nil, runtime.TypeDocker, fmt.Errorf("invalid Docker socket path: %w", err)
                }
                return []string{customSocketPath}, runtime.TypeDocker, nil
        }

        if customSocketPath := os.Getenv(ColimaSocketEnv); customSocketPath != "" {
                //nolint:gosec // G706: socket path from trusted environment variable
                slog.Debug("using Colima socket from env", "path", customSocketPath)
                if _, err := os.Stat(customSocketPath); err != nil { //nolint:gosec // G703: socket path from trusted environment variable
                        return nil, runtime.TypeDocker, fmt.Errorf("invalid Colima socket path: %w", err)
                }
                return []string{customSocketPath}, runtime.TypeDocker, nil
        }

        if rt == runtime.TypePodman {
                if paths := findPodmanSocket(); len(paths) > 0 {
                        return paths, runtime.TypePodman, nil
                }
        }

        if rt == runtime.TypeDocker {
                if paths := findDockerSocket(); len(paths) > 0 {
                        return paths, runtime.TypeDocker, nil
                }
        }

        if rt == runtime.TypeColima {
                if paths := findColimaSocket(); len(paths) > 0 {
                        return paths, runtime.TypeColima, nil
                }
        }

        return nil, "", ErrRuntimeNotFound
}

// findPodmanSocket attempts to locate Podman sockets. Returns every candidate
// path that exists on disk in priority order; an empty slice means no Podman
// socket was found.
func findPodmanSocket() []string {
        var paths []string

        if _, err := os.Stat(PodmanSocketPath); err == nil {
                slog.Debug("found Podman socket", "path", PodmanSocketPath)
                paths = append(paths, PodmanSocketPath)
        } else {
                slog.Debug("failed to check Podman socket", "path", PodmanSocketPath, "error", err)
        }

        if xdgRuntimeDir := os.Getenv("XDG_RUNTIME_DIR"); xdgRuntimeDir != "" {
                xdgSocketPath := filepath.Join(xdgRuntimeDir, PodmanXDGRuntimeSocketPath)
                if _, err := os.Stat(xdgSocketPath); err == nil { //nolint:gosec // G703: path from trusted env + constant
                        //nolint:gosec // G706: socket path derived from XDG_RUNTIME_DIR env var
                        slog.Debug("found Podman socket", "path", xdgSocketPath)
                        paths = append(paths, xdgSocketPath)
                } else {
                        //nolint:gosec // G706: socket path derived from XDG_RUNTIME_DIR env var
                        slog.Debug("failed to check Podman socket", "path", xdgSocketPath, "error", err)
                }
        }

        if home := os.Getenv("HOME"); home != "" {
                userSocketPath := filepath.Join(home, ".local/share/containers/podman/machine/podman.sock")
                if _, err := os.Stat(userSocketPath); err == nil { //nolint:gosec // G703: path from trusted env + constant
                        //nolint:gosec // G706: socket path derived from HOME env var
                        slog.Debug("found Podman socket", "path", userSocketPath)
                        paths = append(paths, userSocketPath)
                } else {
                        //nolint:gosec // G706: socket path derived from HOME env var
                        slog.Debug("failed to check Podman socket", "path", userSocketPath, "error", err)
                }
        }

        // Check TMPDIR for Podman Machine API sockets (macOS).
        // The socket path follows the pattern: $TMPDIR/podman/<machine-name>-api.sock
        if tmpDir := os.Getenv("TMPDIR"); tmpDir != "" {
                podmanTmpDir := filepath.Join(tmpDir, "podman")
                if _, err := os.Stat(podmanTmpDir); err == nil { //nolint:gosec // G703: path from trusted env
                        matches, err := filepath.Glob(filepath.Join(podmanTmpDir, "*-api.sock"))
                        if err == nil && len(matches) > 0 {
                                for _, socketPath := range matches {
                                        //nolint:gosec // G706: socket path discovered from TMPDIR
                                        slog.Debug("found Podman machine API socket", "path", socketPath)
                                        paths = append(paths, socketPath)
                                }
                        } else {
                                //nolint:gosec // G706: directory path from TMPDIR env var
                                slog.Debug("no Podman machine API sockets found", "dir", podmanTmpDir)
                        }
                } else {
                        //nolint:gosec // G706: directory path from TMPDIR env var
                        slog.Debug("podman temp directory not found", "dir", podmanTmpDir, "error", err)
                }
        }

        return paths
}

// systemDockerSocketPath is the system-wide Docker socket path probed by
// findDockerSocket. It defaults to DockerSocketPath and is package-private so
// tests can redirect the system check to a sandbox path.
var systemDockerSocketPath = DockerSocketPath

// findDockerSocket attempts to locate Docker sockets. Returns every candidate
// path that exists on disk in priority order; an empty slice means no Docker
// socket was found.
//
// Returning every candidate lets callers fall back to a per-user socket like
// ~/.docker/desktop/docker.sock when the system socket /var/run/docker.sock
// exists but is not connectable (e.g. the running user is not in the docker
// group). Stat-OK on a path does not imply the daemon will accept a connect.
func findDockerSocket() []string {
        var paths []string

        if _, err := os.Stat(systemDockerSocketPath); err == nil {
                slog.Debug("found Docker socket", "path", systemDockerSocketPath)
                paths = append(paths, systemDockerSocketPath)
        } else {
                slog.Debug("failed to check Docker socket", "path", systemDockerSocketPath, "error", err)
        }

        home := os.Getenv("HOME")
        if home == "" {
                return paths
        }

        userCandidates := []struct {
                label    string
                relative string
        }{
                {"Docker Desktop socket", DockerDesktopMacSocketPath},
                {"Docker Desktop socket", DockerDesktopLinuxSocketPath},
                {"Rancher Desktop socket", RancherDesktopMacSocketPath},
                {"OrbStack socket", OrbStackMacSocketPath},
        }
        for _, c := range userCandidates {
                path := filepath.Join(home, c.relative)
                if _, err := os.Stat(path); err == nil { // #nosec G703 -- path is built from HOME + constant socket path
                        //nolint:gosec // G706: socket path derived from HOME env var
                        slog.Debug("found "+c.label, "path", path)
                        paths = append(paths, path)
                } else {
                        //nolint:gosec // G706: socket path derived from HOME env var
                        slog.Debug("failed to check "+c.label, "path", path, "error", err)
                }
        }

        return paths
}

// findColimaSocket attempts to locate Colima sockets. Returns every candidate
// path that exists on disk in priority order; an empty slice means no Colima
// socket was found.
func findColimaSocket() []string {
        var paths []string

        if _, err := os.Stat(ColimaDesktopMacSocketPath); err == nil {
                slog.Debug("found Colima socket", "path", ColimaDesktopMacSocketPath)
                paths = append(paths, ColimaDesktopMacSocketPath)
        } else {
                slog.Debug("failed to check Colima socket", "path", ColimaDesktopMacSocketPath, "error", err)
        }

        if home := os.Getenv("HOME"); home != "" {
                userSocketPath := filepath.Join(home, ColimaDesktopMacSocketPath)
                if _, err := os.Stat(userSocketPath); err == nil { // #nosec G703 -- path is built from HOME + constant socket path
                        //nolint:gosec // G706: socket path derived from HOME env var
                        slog.Debug("found Colima socket", "path", userSocketPath)
                        paths = append(paths, userSocketPath)
                } else {
                        //nolint:gosec // G706: socket path derived from HOME env var
                        slog.Debug("failed to check Colima socket", "path", userSocketPath, "error", err)
                }
        }

        return paths
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package sdk provides a factory method for creating a Docker client.
package sdk

import (
        "context"
        "errors"
        "fmt"
        "io/fs"
        "log/slog"
        "strings"

        mobyclient "github.com/moby/moby/client"

        "github.com/stacklok/toolhive/pkg/container/runtime"
)

/*
 * This file contains logic for instantiating the Docker SDK client.
 * In future, this may be centralized behind a single client wrapper.
 */

// Environment variable names
const (
        // DockerSocketEnv is the environment variable for custom Docker socket path
        DockerSocketEnv = "TOOLHIVE_DOCKER_SOCKET"
        // PodmanSocketEnv is the environment variable for custom Podman socket path
        PodmanSocketEnv = "TOOLHIVE_PODMAN_SOCKET"
        // ColimaSocketEnv is the environment variable for custom Colima socket path
        ColimaSocketEnv = "TOOLHIVE_COLIMA_SOCKET"
)

// Common socket paths
const (
        // PodmanSocketPath is the default Podman socket path
        PodmanSocketPath = "/var/run/podman/podman.sock"
        // PodmanXDGRuntimeSocketPath is the XDG runtime Podman socket path
        PodmanXDGRuntimeSocketPath = "podman/podman.sock"
        // DockerSocketPath is the default Docker socket path
        DockerSocketPath = "/var/run/docker.sock"
        // DockerDesktopMacSocketPath is the Docker Desktop socket path on macOS
        DockerDesktopMacSocketPath = ".docker/run/docker.sock"
        // DockerDesktopLinuxSocketPath is the Docker Desktop socket path on Linux
        // (relative to $HOME). Docker Desktop on Linux registers a "desktop-linux"
        // Docker context that points to this socket.
        DockerDesktopLinuxSocketPath = ".docker/desktop/docker.sock"
        // RancherDesktopMacSocketPath is the Docker socket path for Rancher Desktop on macOS
        RancherDesktopMacSocketPath = ".rd/docker.sock"
        // OrbStackMacSocketPath is the Docker socket path for OrbStack on macOS
        OrbStackMacSocketPath = ".orbstack/run/docker.sock"
        // ColimaDesktopMacSocketPath is the Docker socket path for Colima on macOS
        ColimaDesktopMacSocketPath = ".colima/default/docker.sock"
)

var supportedSocketPaths = []runtime.Type{runtime.TypePodman, runtime.TypeDocker, runtime.TypeColima}

// NewDockerClient creates a new container client.
//
// For each runtime (Podman, Docker, Colima) it asks the platform helper for
// every candidate socket that exists on disk and then tries to connect to each
// in turn. This matters in mixed setups — e.g. /var/run/docker.sock is present
// but the running user is not in the docker group, while Docker Desktop's
// per-user socket at ~/.docker/desktop/docker.sock would work. Without
// per-runtime fallback the first stat-OK socket short-circuits discovery and
// we'd surface "no container runtime available" even though a usable Docker
// daemon is reachable through a different socket.
func NewDockerClient(ctx context.Context) (*mobyclient.Client, string, runtime.Type, error) {
        var lastErr error

        for _, sp := range supportedSocketPaths {
                socketPaths, runtimeType, err := findContainerSocket(sp)
                if err != nil {
                        //nolint:gosec // G706: runtime type from internal config
                        slog.Debug("failed to find socket", "runtime", sp, "error", err)
                        lastErr = err
                        continue
                }

                for _, socketPath := range socketPaths {
                        c, err := newClientWithSocketPath(ctx, socketPath, runtimeType)
                        if err != nil {
                                lastErr = err
                                //nolint:gosec // G706: runtime type from internal config
                                slog.Debug("failed to create client", "runtime", sp, "socket", socketPath, "error", err)
                                continue
                        }

                        //nolint:gosec // G706: runtime type from internal detection
                        slog.Debug("successfully connected to runtime", "runtime", runtimeType, "socket", socketPath)
                        return c, socketPath, runtimeType, nil
                }
        }

        if lastErr != nil {
                return nil, "", "", fmt.Errorf("no supported container runtime available%s: %w", dockerPermissionHint(lastErr), lastErr)
        }
        return nil, "", "", fmt.Errorf("no supported container runtime found/running")
}

// dockerPermissionHint returns an actionable hint when the underlying socket
// connect was rejected with a permission error. The most common cause on
// Linux is the calling user not being a member of the docker group, which the
// raw "permission denied" error does not call out.
func dockerPermissionHint(err error) string {
        if err == nil {
                return ""
        }
        // errors.Is handles syscall.EACCES/EPERM (which map to fs.ErrPermission via
        // the syscall.Errno.Is shim); the Docker client wraps the underlying
        // *os.SyscallError in its own message, so also fall back to a substring
        // match on the rendered text for robustness across error wrappings.
        if !errors.Is(err, fs.ErrPermission) && !strings.Contains(err.Error(), "permission denied") {
                return ""
        }
        return " (hint: the Docker socket rejected the connection with 'permission denied' — " +
                "is your user in the 'docker' group? On Linux: 'sudo usermod -aG docker $USER', " +
                "then log out and back in. See https://docs.docker.com/engine/install/linux-postinstall/.)"
}

// NewClientWithSocketPath creates a new container client with a specific socket path
func newClientWithSocketPath(ctx context.Context, socketPath string, runtimeType runtime.Type) (*mobyclient.Client, error) {
        // Create platform-specific client
        httpClient, opts := newPlatformClient(socketPath)
        if runtimeType == runtime.TypePodman {
                httpClient.Transport = newPodmanNetworkCompatRoundTripper(httpClient.Transport)
        }

        // Create Docker client with the custom HTTP client
        dockerClient, err := mobyclient.New(opts...)
        if err != nil {
                return nil, fmt.Errorf("failed to create client: %w", err)
        }

        // Make sure we can ping the server.
        _, err = dockerClient.Ping(ctx, mobyclient.PingOptions{})
        if err != nil {
                return nil, fmt.Errorf("failed to ping Docker server at %s: %w", socketPath, err)
        }

        return dockerClient, nil
}

// findContainerSocket finds candidate container socket paths for the given
// runtime, preferring Podman over Docker. The returned slice is ordered from
// highest to lowest priority and may contain multiple entries on Unix when
// several known socket locations exist on disk simultaneously.
func findContainerSocket(rt runtime.Type) ([]string, runtime.Type, error) {
        // Use platform-specific implementation
        return findPlatformContainerSocket(rt)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package sdk

import (
        "bytes"
        "encoding/json"
        "fmt"
        "io"
        "net/http"
        "strconv"
        "strings"

        "github.com/stacklok/toolhive/pkg/networking"
)

// podmanNetworkCompatRoundTripper normalizes Podman's Docker-compatible
// representation of a gatewayless network before the Moby client decodes it.
// Podman serializes a missing net.IP as the string "<nil>", which netip.Addr
// correctly rejects as malformed. Only that exact sentinel is treated as an
// absent gateway; all other values are left for Moby's strict decoder.
type podmanNetworkCompatRoundTripper struct {
        base http.RoundTripper
}

func newPodmanNetworkCompatRoundTripper(base http.RoundTripper) http.RoundTripper {
        if base == nil {
                base = http.DefaultTransport
        }
        return &podmanNetworkCompatRoundTripper{base: base}
}

func (t *podmanNetworkCompatRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
        resp, err := t.base.RoundTrip(req)
        if err != nil || resp.Body == nil || req.Method != http.MethodGet ||
                resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices ||
                !isNetworkAPIPath(req.URL.Path) {
                return resp, err
        }

        body, readErr := io.ReadAll(resp.Body)
        closeErr := resp.Body.Close()
        if readErr != nil {
                return nil, fmt.Errorf("read Podman network response: %w", readErr)
        }
        if closeErr != nil {
                return nil, fmt.Errorf("close Podman network response: %w", closeErr)
        }

        body = normalizePodmanNetworkResponse(body)
        resp.Body = io.NopCloser(bytes.NewReader(body))
        resp.ContentLength = int64(len(body))
        if resp.Header != nil {
                resp.Header.Set("Content-Length", strconv.Itoa(len(body)))
        }
        return resp, nil
}

func (t *podmanNetworkCompatRoundTripper) CloseIdleConnections() {
        networking.ForwardCloseIdle(t.base)
}

var _ networking.IdleConnectionCloser = (*podmanNetworkCompatRoundTripper)(nil)

func normalizePodmanNetworkResponse(body []byte) []byte {
        trimmed := bytes.TrimSpace(body)
        if len(trimmed) == 0 {
                return body
        }

        switch trimmed[0] {
        case '[':
                var networks []json.RawMessage
                if err := json.Unmarshal(body, &networks); err != nil {
                        return body
                }

                changed := false
                for i, rawNetwork := range networks {
                        normalized, networkChanged := normalizePodmanNetwork(rawNetwork)
                        if networkChanged {
                                networks[i] = normalized
                                changed = true
                        }
                }
                if !changed {
                        return body
                }
                normalized, err := json.Marshal(networks)
                if err != nil {
                        return body
                }
                return normalized
        case '{':
                normalized, changed := normalizePodmanNetwork(body)
                if changed {
                        return normalized
                }
        }
        return body
}

func normalizePodmanNetwork(rawNetwork json.RawMessage) (json.RawMessage, bool) {
        var networkFields map[string]json.RawMessage
        if err := json.Unmarshal(rawNetwork, &networkFields); err != nil {
                return rawNetwork, false
        }

        rawIPAM, ok := networkFields["IPAM"]
        if !ok {
                return rawNetwork, false
        }
        normalizedIPAM, changed := normalizePodmanIPAM(rawIPAM)
        if !changed {
                return rawNetwork, false
        }

        networkFields["IPAM"] = normalizedIPAM
        normalizedNetwork, err := json.Marshal(networkFields)
        if err != nil {
                return rawNetwork, false
        }
        return normalizedNetwork, true
}

func normalizePodmanIPAM(rawIPAM json.RawMessage) (json.RawMessage, bool) {
        var ipamFields map[string]json.RawMessage
        if err := json.Unmarshal(rawIPAM, &ipamFields); err != nil {
                return rawIPAM, false
        }

        rawConfigs, ok := ipamFields["Config"]
        if !ok {
                return rawIPAM, false
        }
        var configs []json.RawMessage
        if err := json.Unmarshal(rawConfigs, &configs); err != nil {
                return rawIPAM, false
        }

        changed := false
        for i, rawConfig := range configs {
                var configFields map[string]json.RawMessage
                if err := json.Unmarshal(rawConfig, &configFields); err != nil {
                        return rawIPAM, false
                }
                rawGateway, ok := configFields["Gateway"]
                if !ok {
                        continue
                }
                var gateway string
                if err := json.Unmarshal(rawGateway, &gateway); err != nil || gateway != "<nil>" {
                        continue
                }

                configFields["Gateway"] = json.RawMessage("null")
                normalizedConfig, err := json.Marshal(configFields)
                if err != nil {
                        return rawIPAM, false
                }
                configs[i] = normalizedConfig
                changed = true
        }
        if !changed {
                return rawIPAM, false
        }

        normalizedConfigs, err := json.Marshal(configs)
        if err != nil {
                return rawIPAM, false
        }
        ipamFields["Config"] = normalizedConfigs
        normalizedIPAM, err := json.Marshal(ipamFields)
        if err != nil {
                return rawIPAM, false
        }
        return normalizedIPAM, true
}

func isNetworkAPIPath(path string) bool {
        parts := strings.Split(strings.Trim(path, "/"), "/")
        if len(parts) > 0 && strings.HasPrefix(parts[0], "v") {
                parts = parts[1:]
        }
        return len(parts) > 0 && parts[0] == "networks"
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package docker

import (
        "context"
        "fmt"
        "log/slog"
        "net/netip"
        "os"
        "strconv"
        "strings"

        "github.com/moby/moby/api/types/container"
        "github.com/moby/moby/api/types/network"

        "github.com/stacklok/toolhive-core/permissions"
        "github.com/stacklok/toolhive/pkg/container/runtime"
        lb "github.com/stacklok/toolhive/pkg/labels"
)

const defaultSquidImage = "ghcr.io/stacklok/toolhive/egress-proxy:latest"

// These Docker/Podman host aliases are blocked by default to prevent
// containers from reaching host services unintentionally.
const (
        podmanGatewayHostname    = "host.containers.internal"
        dockerGatewayHostname    = "host.docker.internal"
        dockerAltGatewayHostname = "gateway.docker.internal"
)

type proxyDirection int

const (
        proxyIngress proxyDirection = iota
        proxyEgress
)

// createIngressSquidContainer creates an instance of the squid proxy for ingress traffic.
func createIngressSquidContainer(
        ctx context.Context,
        c *Client,
        containerName string,
        upstreamHost string,
        squidContainerName string,
        attachStdio bool,
        upstreamPort int,
        squidPort int,
        exposedPorts map[string]struct{},
        endpointsConfig map[string]*network.EndpointSettings,
        portBindings map[string][]runtime.PortBinding,
        networkPermissions *permissions.NetworkPermissions,
) (string, error) {
        squidConfPath, err := createTempIngressSquidConf(containerName, upstreamHost, upstreamPort, squidPort, networkPermissions)
        if err != nil {
                return "", fmt.Errorf("failed to create temporary squid.conf: %w", err)
        }

        return createSquidContainer(
                ctx,
                c,
                squidContainerName,
                attachStdio,
                exposedPorts,
                endpointsConfig,
                portBindings,
                squidConfPath,
        )
}

// createEgressSquidContainer creates an instance of the squid proxy for egress traffic.
func createEgressSquidContainer(
        ctx context.Context,
        c *Client,
        containerName string,
        squidContainerName string,
        attachStdio bool,
        exposedPorts map[string]struct{},
        endpointsConfig map[string]*network.EndpointSettings,
        perm *permissions.NetworkPermissions,
        allowDockerGateway bool,
        gatewayIPs []netip.Addr,
) (string, error) {
        squidConfPath, err := createTempEgressSquidConf(perm, containerName, allowDockerGateway, gatewayIPs)
        if err != nil {
                return "", fmt.Errorf("failed to create temporary squid.conf: %w", err)
        }

        return createSquidContainer(
                ctx,
                c,
                squidContainerName,
                attachStdio,
                exposedPorts,
                endpointsConfig,
                nil,
                squidConfPath,
        )
}

// createSquidContainer contains the shared logic for creating a squid container.
func createSquidContainer(
        ctx context.Context,
        c *Client, // TODO: refactor the methods we need from docker.Client into a lower level interface.
        squidContainerName string,
        attachStdio bool,
        exposedPorts map[string]struct{},
        endpointsConfig map[string]*network.EndpointSettings,
        portBindings map[string][]runtime.PortBinding, // used for ingress only
        squidConfPath string,
) (string, error) {

        //nolint:gosec // G706: squid container name and image from config
        slog.Debug("setting up squid container", "name", squidContainerName, "image", getSquidImage())
        squidLabels := map[string]string{}
        lb.AddStandardLabels(squidLabels, squidContainerName, squidContainerName, "stdio", 80)
        squidLabels[ToolhiveAuxiliaryWorkloadLabel] = LabelValueTrue

        // pull the squid image if it is not already pulled
        squidImage := getSquidImage()
        // TODO: Move these down into an image operations layer.
        err := c.imageManager.PullImage(ctx, squidImage)
        if err != nil {
                // Check if the squid image exists locally before failing
                _, inspectErr := c.imageManager.ImageExists(ctx, squidImage)
                if inspectErr == nil {
                        //nolint:gosec // G706: squid image name from config
                        slog.Debug("squid image exists locally, continuing despite pull failure", "image", squidImage)
                } else {
                        return "", fmt.Errorf("failed to pull squid image: %w", err)
                }
        }

        // Create container options
        config := &container.Config{
                Image:        getSquidImage(),
                Cmd:          nil,
                Env:          nil,
                Labels:       squidLabels,
                AttachStdin:  attachStdio,
                AttachStdout: attachStdio,
                AttachStderr: attachStdio,
                OpenStdin:    attachStdio,
                Tty:          false,
        }

        mounts := []runtime.Mount{}
        mounts = append(mounts, runtime.Mount{
                Source:   squidConfPath,
                Target:   "/etc/squid/squid.conf",
                ReadOnly: true,
        })

        // Create squid host configuration
        squidHostConfig := &container.HostConfig{
                Mounts:      convertMounts(mounts),
                NetworkMode: container.NetworkMode("bridge"),
                CapAdd:      []string{"CAP_SETUID", "CAP_SETGID"},
                CapDrop:     nil,
                SecurityOpt: []string{"label:disable"},
                RestartPolicy: container.RestartPolicy{
                        Name: "unless-stopped",
                },
        }

        // Setup port bindings
        if portBindings != nil {
                if err := setupPortBindings(squidHostConfig, portBindings); err != nil {
                        return "", NewContainerError(err, "", err.Error())
                }
        }

        // Setup port bindings
        if err := setupExposedPorts(config, exposedPorts); err != nil {
                return "", NewContainerError(err, "", err.Error())
        }

        // Create squid container itself
        squidContainerId, err := c.createContainer(ctx, squidContainerName, config, squidHostConfig, endpointsConfig)
        if err != nil {
                return "", fmt.Errorf("failed to create egress container: %w", err)
        }

        return squidContainerId, nil
}

// writeDockerGatewayDenyRules emits Squid ACL definitions and http_access deny
// rules that block the Docker gateway addresses. These rules MUST be written
// before any http_access allow rules: Squid evaluates access control in
// first-match-wins order, so a deny placed after an allow is never reached.
//
// gatewayIPs contains every bridge gateway and discovered Podman host-alias IP.
// hostGatewayNames covers hostname-based access; the dst rule also covers direct
// IPs and other names resolving to a protected address.
func writeDockerGatewayDenyRules(sb *strings.Builder, gatewayIPs []netip.Addr) {
        sb.WriteString(
                "# Block Docker gateway addresses — opt in with --allow-docker-gateway\n" +
                        "acl docker_gateway_hosts dstdomain " +
                        strings.Join(hostGatewayNames[:], " ") + "\n" +
                        "acl docker_gateway_ip dst " + gatewayAddrList(gatewayIPs) + "\n" +
                        "http_access deny docker_gateway_hosts\n" +
                        "http_access deny docker_gateway_ip\n\n",
        )
}

// writeDockerGatewayAllowRules emits Squid ACL definitions and http_access allow
// rules that grant access to the Docker gateway addresses. They are written when
// --allow-docker-gateway is set so the flag alone makes the host reachable — in
// allowlist mode, suppressing the deny is not enough, since Squid is
// first-match-wins and only a later allow can let the request through.
//
// The allow rules are emitted as standalone http_access lines that reference only
// the host/IP ACLs, NOT the port ACL. On a single http_access line Squid AND-s
// the ACLs together, so folding the gateway host into the port-restricted
// allowed_ports/allowed_dsts rule would block it on the arbitrary ports that
// host-local services use (5432, 3306, 6379, 8080, …). Gateway access is
// therefore port-independent — a deliberate relaxation, since reaching the host
// is already a separate, more-privileged opt-in.
//
// Like the deny rules, these MUST precede any later http_access allow/deny so
// they are reached under Squid's first-match-wins evaluation.
//
// See writeDockerGatewayDenyRules for the meaning of gatewayIPs and the gateway
// hostnames.
func writeDockerGatewayAllowRules(sb *strings.Builder, gatewayIPs []netip.Addr) {
        sb.WriteString(
                "# Grant Docker gateway access — enabled via --allow-docker-gateway\n" +
                        "acl docker_gateway_hosts dstdomain " +
                        strings.Join(hostGatewayNames[:], " ") + "\n" +
                        "acl docker_gateway_ip dst " + gatewayAddrList(gatewayIPs) + "\n" +
                        "http_access allow docker_gateway_hosts\n" +
                        "http_access allow docker_gateway_ip\n\n",
        )
}

func gatewayAddrList(gateways []netip.Addr) string {
        parts := make([]string, len(gateways))
        for i, gateway := range gateways {
                parts[i] = gateway.String()
        }
        return strings.Join(parts, " ")
}

func createTempEgressSquidConf(
        networkPermissions *permissions.NetworkPermissions,
        serverHostname string,
        allowDockerGateway bool,
        gatewayIPs []netip.Addr,
) (string, error) {
        var sb strings.Builder

        writeCommonConfig(&sb, serverHostname, proxyEgress)

        // Handle Docker gateway access. MUST precede any http_access allow — Squid
        // is first-match-wins. When --allow-docker-gateway is set, emit explicit
        // allow rules so the flag alone grants host access; otherwise block the
        // gateway addresses.
        if allowDockerGateway {
                writeDockerGatewayAllowRules(&sb, gatewayIPs)
        } else {
                writeDockerGatewayDenyRules(&sb, gatewayIPs)
        }

        if networkPermissions == nil || (networkPermissions.Outbound != nil && networkPermissions.Outbound.InsecureAllowAll) {
                sb.WriteString("# Allow all traffic\nhttp_access allow all\n")
        } else {
                writeOutboundACLs(&sb, networkPermissions.Outbound)
                writeHttpAccessRules(&sb, networkPermissions.Outbound)
        }

        sb.WriteString("http_access deny all\n")

        tmpFile, err := os.CreateTemp("", "squid-*.conf")
        if err != nil {
                return "", err
        }
        defer func() {
                if err := tmpFile.Close(); err != nil {
                        // Non-fatal: temp file cleanup failure
                        slog.Warn("failed to close temp file", "error", err)
                }
        }()

        if _, err := tmpFile.WriteString(sb.String()); err != nil {
                return "", fmt.Errorf("failed to write to temporary file: %w", err)
        }

        // Set file permissions to be readable by all users (including squid user in container)
        if err := tmpFile.Chmod(0644); err != nil {
                return "", fmt.Errorf("failed to set file permissions: %w", err)
        }

        return tmpFile.Name(), nil
}

func writeCommonConfig(sb *strings.Builder, hostnameBase string, direction proxyDirection) {
        var serverHostname string

        if direction == proxyEgress {
                serverHostname = hostnameBase + "-egress"
                sb.WriteString("http_port 3128\n")
        } else {
                serverHostname = hostnameBase + "-ingress"
        }

        sb.WriteString(
                "visible_hostname " + serverHostname + "\n" +
                        "access_log stdio:/dev/stdout squid\n" +
                        "pid_filename none\n" +
                        "# Avoid allocation errors caused by max_filedescriptors inference\n" +
                        "max_filedescriptors 1024\n" +
                        "# Disable memory and disk caching\n" +
                        "cache deny all\n" +
                        "cache_mem 0 MB\n" +
                        "maximum_object_size 0 KB\n" +
                        "maximum_object_size_in_memory 0 KB\n" +
                        "# Don't use cache directories\n" +
                        "cache_store_log none\n\n")
}

func writeOutboundACLs(sb *strings.Builder, outbound *permissions.OutboundNetworkPermissions) {
        if len(outbound.AllowPort) > 0 {
                sb.WriteString("# Define allowed ports\nacl allowed_ports port")
                for _, port := range outbound.AllowPort {
                        sb.WriteString(" " + strconv.Itoa(port))
                }
                sb.WriteString("\n")
        }

        if len(outbound.AllowHost) > 0 {
                sb.WriteString("# Define allowed destinations\nacl allowed_dsts dstdomain")
                for _, host := range outbound.AllowHost {
                        sb.WriteString(" " + host)
                }
                sb.WriteString("\n")
        }
}

func writeHttpAccessRules(sb *strings.Builder, outbound *permissions.OutboundNetworkPermissions) {
        var conditions []string
        if len(outbound.AllowPort) > 0 {
                conditions = append(conditions, "allowed_ports")
        }
        if len(outbound.AllowHost) > 0 {
                conditions = append(conditions, "allowed_dsts")
        }
        if len(conditions) > 0 {
                sb.WriteString("\n# Define http_access rules\n")
                sb.WriteString("http_access allow " + strings.Join(conditions, " ") + "\n")
        }
}

func getSquidImage() string {
        if egressImage := os.Getenv("TOOLHIVE_EGRESS_IMAGE"); egressImage != "" {
                return egressImage
        }
        return defaultSquidImage
}

// squidProxy is the default networkProxy implementation. It creates egress and
// ingress Squid-based proxy containers for isolated workloads.
type squidProxy struct {
        client *Client
}

// SetupEgress creates the egress Squid container before the MCP container is
// created and returns the proxy env vars to inject into the workload.
func (s *squidProxy) SetupEgress(ctx context.Context, spec proxySpec) (egressResult, error) {
        egressContainerName := fmt.Sprintf("%s-egress", spec.WorkloadName)
        _, err := createEgressSquidContainer(
                ctx, s.client, spec.WorkloadName, egressContainerName,
                spec.AttachStdio, nil, spec.Endpoints, spec.Permissions,
                spec.AllowDockerGateway, spec.GatewayIPs,
        )
        if err != nil {
                return egressResult{}, fmt.Errorf("failed to create egress container: %w", err)
        }
        // ingressPort stays 0: squid creates and binds the ingress container later,
        // in SetupIngress, once the MCP container's hostname resolves.
        return egressResult{EnvVars: addEgressEnvVars(nil, egressContainerName)}, nil
}

// SetupIngress creates the ingress Squid container after the MCP container
// exists. Creating it here (rather than before the MCP container) ensures the
// cache_peer hostname resolves on first probe; a Squid ingress created against a
// not-yet-existent upstream caches the negative DNS lookup and never recovers
// within the workload's readiness window.
func (s *squidProxy) SetupIngress(ctx context.Context, spec proxySpec, _ egressResult) (int, error) {
        if spec.TransportType == "stdio" || spec.UpstreamPort == 0 {
                return 0, nil
        }
        // Prefer the resolved upstream IP so the cache_peer has no DNS dependency;
        // fall back to the workload name when unset (see proxySpec.UpstreamHost).
        upstreamHost := spec.UpstreamHost
        if upstreamHost == "" {
                upstreamHost = spec.WorkloadName
        }
        ingressPort, err := s.client.setupIngressContainer(
                ctx, spec.WorkloadName, upstreamHost, spec.UpstreamPort, spec.AttachStdio, spec.Endpoints, spec.Permissions,
        )
        if err != nil {
                return 0, err
        }
        return ingressPort, nil
}

func createTempIngressSquidConf(
        serverHostname string,
        upstreamHost string,
        upstreamPort int,
        squidPort int,
        networkPermissions *permissions.NetworkPermissions,
) (string, error) {
        var sb strings.Builder

        writeCommonConfig(&sb, serverHostname, proxyIngress)

        writeIngressProxyConfig(&sb, serverHostname, upstreamHost, upstreamPort, squidPort, networkPermissions)
        sb.WriteString("http_access deny all\n")

        tmpFile, err := os.CreateTemp("", "squid-*.conf")
        if err != nil {
                return "", err
        }
        defer func() {
                if err := tmpFile.Close(); err != nil {
                        // Non-fatal: temp file cleanup failure
                        slog.Warn("failed to close temp file", "error", err)
                }
        }()

        if _, err := tmpFile.WriteString(sb.String()); err != nil {
                return "", fmt.Errorf("failed to write to temporary file: %w", err)
        }

        // Set file permissions to be readable by all users (including squid user in container)
        if err := tmpFile.Chmod(0644); err != nil {
                return "", fmt.Errorf("failed to set file permissions: %w", err)
        }

        return tmpFile.Name(), nil
}

func writeIngressProxyConfig(
        sb *strings.Builder,
        serverHostname string,
        upstreamHost string,
        upstreamPort int,
        squidPort int,
        networkPermissions *permissions.NetworkPermissions,
) {
        portNum := strconv.Itoa(upstreamPort)
        squidPortNum := strconv.Itoa(squidPort)
        // cache_peer targets upstreamHost (the MCP container's IP), not its name, so
        // the peer has no DNS lookup to cache-and-latch on a cold-start miss (#6063).
        // defaultsite keeps the name for the origin Host header. standby=2 keeps warm
        // idle connections open to the upstream so the first request after a cold
        // start (notably a long-lived GET SSE stream that a server-initiated request
        // rides on) is forwarded without paying inline TCP connect latency. Without
        // it, the cold first GET races behind a later POST that reuses a warmed path,
        // reordering server->client streams.
        sb.WriteString(
                "\n# Reverse proxy setup for port " + portNum + "\n" +
                        "http_port 0.0.0.0:" + squidPortNum + " accel defaultsite=" + serverHostname + "\n" +
                        "cache_peer " + upstreamHost + " parent " + portNum + " 0 no-query originserver name=origin_" +
                        portNum + " connect-timeout=5 connect-fail-limit=5 standby=2\n")

        // Check if inbound network permissions are configured
        if networkPermissions != nil && networkPermissions.Inbound != nil && len(networkPermissions.Inbound.AllowHost) > 0 {
                // Use only the configured allowed hosts
                sb.WriteString("acl allowed_hosts dstdomain")
                for _, host := range networkPermissions.Inbound.AllowHost {
                        sb.WriteString(" " + host)
                }
                sb.WriteString("\n")
                sb.WriteString("http_access allow allowed_hosts\n")
        } else {
                // Default: Allow container hostname, localhost, and 127.0.0.1
                sb.WriteString("acl site_" + portNum + " dstdomain " + serverHostname + "\n" +
                        "acl local_dst dst 127.0.0.1\n" +
                        "acl local_domain dstdomain localhost\n" +
                        "http_access allow site_" + portNum + "\n" +
                        "http_access allow local_dst\n" +
                        "http_access allow local_domain\n")
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package container provides utilities for managing containers,
// including creating, starting, stopping, and monitoring containers.
package container

import (
        "context"
        "fmt"
        "os"
        "sort"
        "strings"
        "sync"

        "github.com/stacklok/toolhive/pkg/container/runtime"
)

// Factory creates container runtimes with pluggable runtime support
type Factory struct {
        mu       sync.RWMutex
        runtimes map[string]*runtime.Info
}

// NewFactory creates a new container factory seeded from the DefaultRegistry.
// Runtimes register themselves via init() in their respective packages;
// the container/runtimes.go collector file ensures all built-in runtimes are imported.
func NewFactory() *Factory {
        return NewFactoryFromRegistry(runtime.DefaultRegistry)
}

// NewFactoryFromRegistry creates a new container factory seeded from the given registry.
// This is useful for testing with isolated registry instances.
func NewFactoryFromRegistry(reg *runtime.Registry) *Factory {
        f := &Factory{
                runtimes: make(map[string]*runtime.Info),
        }

        for _, info := range reg.All() {
                f.runtimes[info.Name] = info
        }

        return f
}

// Register registers a new runtime with the factory
func (f *Factory) Register(info *runtime.Info) error {
        if info == nil {
                return fmt.Errorf("runtime info cannot be nil")
        }
        if info.Name == "" {
                return fmt.Errorf("runtime name cannot be empty")
        }
        if info.Initializer == nil {
                return fmt.Errorf("runtime initializer cannot be nil")
        }

        f.mu.Lock()
        defer f.mu.Unlock()

        f.runtimes[info.Name] = info
        return nil
}

// Unregister removes a runtime from the factory
func (f *Factory) Unregister(name string) {
        f.mu.Lock()
        defer f.mu.Unlock()
        delete(f.runtimes, name)
}

// GetRuntime retrieves a runtime info by name
func (f *Factory) GetRuntime(name string) (*runtime.Info, bool) {
        f.mu.RLock()
        defer f.mu.RUnlock()
        info, exists := f.runtimes[name]
        return info, exists
}

// ListRuntimes returns all registered runtimes
func (f *Factory) ListRuntimes() map[string]*runtime.Info {
        f.mu.RLock()
        defer f.mu.RUnlock()

        result := make(map[string]*runtime.Info, len(f.runtimes))
        for name, info := range f.runtimes {
                result[name] = info
        }
        return result
}

// ListAvailableRuntimes returns all runtimes that are currently available
func (f *Factory) ListAvailableRuntimes() map[string]*runtime.Info {
        f.mu.RLock()
        defer f.mu.RUnlock()

        result := make(map[string]*runtime.Info)
        for name, info := range f.runtimes {
                if info.AutoDetector == nil || info.AutoDetector() {
                        result[name] = info
                }
        }
        return result
}

// Create creates a container runtime
// It first checks the TOOLHIVE_RUNTIME environment variable for a specific runtime,
// otherwise falls back to auto-detection
func (f *Factory) Create(ctx context.Context) (runtime.Runtime, error) {
        return f.CreateWithRuntimeName(ctx, f.getRuntimeFromEnv())
}

// CreateWithRuntimeName creates a container runtime with a specific runtime name
// If runtimeName is empty, it falls back to auto-detection
func (f *Factory) CreateWithRuntimeName(ctx context.Context, runtimeName string) (runtime.Runtime, error) {
        var runtimeInfo *runtime.Info
        var selectedRuntimeName string

        if runtimeName != "" {
                // Use specified runtime
                info, exists := f.GetRuntime(runtimeName)
                if !exists {
                        available := f.ListRuntimes()
                        var availableNames []string
                        for n := range available {
                                availableNames = append(availableNames, n)
                        }
                        return nil, fmt.Errorf("runtime %q not found. Available runtimes: %s",
                                runtimeName, strings.Join(availableNames, ", "))
                }

                // Check if the runtime is available
                if info.AutoDetector != nil && !info.AutoDetector() {
                        return nil, fmt.Errorf("runtime %q is not available on this system", runtimeName)
                }

                selectedRuntimeName = runtimeName
                runtimeInfo = info
        } else {
                // Auto-detect runtime
                detectedName, detectedInfo := f.autoDetectRuntime()
                if detectedInfo == nil {
                        return nil, fmt.Errorf("no available runtime found")
                }
                selectedRuntimeName = detectedName
                runtimeInfo = detectedInfo
        }

        rt, err := runtimeInfo.Initializer(ctx)
        if err != nil {
                return nil, fmt.Errorf("failed to initialize %s runtime: %w", selectedRuntimeName, err)
        }

        return rt, nil
}

// NewMonitor creates a new container monitor
func NewMonitor(rt runtime.Runtime, containerName string) runtime.Monitor {
        return runtime.NewMonitor(rt, containerName)
}

// CheckRuntimeAvailable checks if any container runtime is available
// and returns a user-friendly error message if none are found
func CheckRuntimeAvailable() error {
        factory := NewFactory()
        available := factory.ListAvailableRuntimes()

        if len(available) == 0 {
                registered := runtime.RegisteredRuntimesByPriority()
                names := make([]string, 0, len(registered))
                for _, r := range registered {
                        names = append(names, r.Name)
                }
                return fmt.Errorf("no container runtime available. ToolHive requires a Docker-compatible container runtime "+
                        "(Docker, Podman, or Colima) or Kubernetes to run MCP servers. Registered runtimes: [%s]",
                        strings.Join(names, ", "))
        }

        return nil
}

// autoDetectRuntime returns the first available runtime based on auto-detection.
// Runtimes are tried in priority order (lowest first); the first one whose
// AutoDetector is nil or returns true is selected.
func (f *Factory) autoDetectRuntime() (string, *runtime.Info) {
        f.mu.RLock()
        ordered := make([]*runtime.Info, 0, len(f.runtimes))
        for _, info := range f.runtimes {
                ordered = append(ordered, info)
        }
        f.mu.RUnlock()

        sort.SliceStable(ordered, func(i, j int) bool {
                if ordered[i].Priority != ordered[j].Priority {
                        return ordered[i].Priority < ordered[j].Priority
                }
                return ordered[i].Name < ordered[j].Name
        })

        for _, info := range ordered {
                if info.AutoDetector == nil || info.AutoDetector() {
                        return info.Name, info
                }
        }

        return "", nil
}

// Clear removes all registered runtimes
// This is useful for testing or when you want to start with a clean slate
func (f *Factory) Clear() {
        f.mu.Lock()
        defer f.mu.Unlock()

        f.runtimes = make(map[string]*runtime.Info)
}

// getRuntimeFromEnv gets the runtime name from the TOOLHIVE_RUNTIME environment variable
// This is separated for easier testing
func (*Factory) getRuntimeFromEnv() string {
        return strings.TrimSpace(os.Getenv("TOOLHIVE_RUNTIME"))
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package images handles container image management operations.
package images

import (
        "context"
        "log/slog"

        "github.com/stacklok/toolhive/pkg/container/docker/sdk"
        "github.com/stacklok/toolhive/pkg/container/runtime"
)

// ImageManager defines the interface for managing container images.
// It has been extracted from the runtime interface as part of
// ongoing refactoring. It may be merged into a more general container
// management interface in future.
type ImageManager interface {
        // ImageExists checks if an image exists locally
        ImageExists(ctx context.Context, image string) (bool, error)

        // PullImage pulls an image from a registry
        PullImage(ctx context.Context, image string) error

        // BuildImage builds a Docker image from a Dockerfile in the specified context directory
        BuildImage(ctx context.Context, contextDir, imageName string) error
}

// NewImageManager creates an instance of ImageManager appropriate
// for the current environment, or returns an error if it is not supported.
func NewImageManager(ctx context.Context) ImageManager {
        // Check if we are running in a Kubernetes environment
        if runtime.IsKubernetesRuntime() {
                slog.Debug("running in Kubernetes environment, using no-op image manager")
                return &NoopImageManager{}
        }

        // Check if we are running in a Docker or compatible environment
        dockerClient, _, _, err := sdk.NewDockerClient(ctx)
        if err != nil {
                slog.Debug("no docker runtime found, using no-op image manager")
                return &NoopImageManager{}
        }

        return NewRegistryImageManager(dockerClient)
}

// NoopImageManager is a no-op implementation of ImageManager.
type NoopImageManager struct{}

// ImageExists always returns false for the no-op implementation.
func (*NoopImageManager) ImageExists(_ context.Context, _ string) (bool, error) {
        return false, nil
}

// PullImage does nothing for the no-op implementation.
func (*NoopImageManager) PullImage(_ context.Context, _ string) error {
        return nil
}

// BuildImage does nothing for the no-op implementation.
func (*NoopImageManager) BuildImage(_ context.Context, _, _ string) error {
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package images

import (
        "github.com/google/go-containerregistry/pkg/authn"

        coreimages "github.com/stacklok/toolhive-core/container/images"
)

// NewCompositeKeychain creates a keychain that tries environment variables first,
// then falls back to the default keychain.
// Deprecated: use github.com/stacklok/toolhive-core/container/images.NewCompositeKeychain.
func NewCompositeKeychain() authn.Keychain {
        return coreimages.NewCompositeKeychain()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package images

import (
        "archive/tar"
        "context"
        "encoding/base64"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "os"
        "path/filepath"
        "runtime"

        "github.com/google/go-containerregistry/pkg/authn"
        "github.com/google/go-containerregistry/pkg/name"
        v1 "github.com/google/go-containerregistry/pkg/v1"
        "github.com/google/go-containerregistry/pkg/v1/daemon"
        "github.com/google/go-containerregistry/pkg/v1/remote"
        mobyclient "github.com/moby/moby/client"
)

// RegistryImageManager implements the ImageManager interface using go-containerregistry
// for direct registry operations without requiring a Docker daemon.
// However, for building images from Dockerfiles, it still uses the Docker client.
type RegistryImageManager struct {
        keychain authn.Keychain
        platform *v1.Platform
        // dockerClient is used for building images from Dockerfiles and, because the
        // moby/moby client already satisfies the go-containerregistry daemon.Client
        // interface, for daemon.Image/daemon.Write operations as well.
        dockerClient *mobyclient.Client
}

// NewRegistryImageManager creates a new RegistryImageManager instance
func NewRegistryImageManager(dockerClient *mobyclient.Client) *RegistryImageManager {
        return &RegistryImageManager{
                keychain:     NewCompositeKeychain(), // Use composite keychain (env vars + default)
                platform:     getDefaultPlatform(),   // Use a default platform based on host architecture
                dockerClient: dockerClient,
        }
}

// getDefaultPlatform returns the default platform for containers
// Uses host architecture
func getDefaultPlatform() *v1.Platform {
        return &v1.Platform{
                Architecture: runtime.GOARCH,
                OS:           "linux", // TODO: Should we support Windows too?
        }
}

// ImageExists checks if an image exists locally in the daemon or remotely in the registry
func (r *RegistryImageManager) ImageExists(_ context.Context, imageName string) (bool, error) {
        // Parse the image reference
        ref, err := name.ParseReference(imageName)
        if err != nil {
                return false, fmt.Errorf("failed to parse image reference %q: %w", imageName, err)
        }

        // First check if image exists locally in daemon
        if _, err := daemon.Image(ref, daemon.WithClient(r.dockerClient)); err != nil {
                // Image does not exist locally
                return false, nil
        }
        // Image exists locally
        return true, nil
}

// PullImage pulls an image from a registry and saves it to the local daemon.
//
// For digest-pinned references (tag@digest or digest-only), the Docker daemon's
// native ImagePull is used. The daemon registers the OCI manifest digest so that
// subsequent container creates using a tag@digest reference can resolve the image.
// go-containerregistry's daemon.Write stores images under a tag only, without
// registering the manifest digest; Docker's tag@digest cross-check at
// container-create time then fails with "No such image".
//
// For plain tag references, go-containerregistry is used for cross-platform
// support and custom keychain auth.
func (r *RegistryImageManager) PullImage(ctx context.Context, imageName string) error {
        //nolint:gosec // G706: image name from user/config input
        slog.Info("pulling image", "image", imageName)

        ref, err := name.ParseReference(imageName)
        if err != nil {
                return fmt.Errorf("failed to parse image reference %q: %w", imageName, err)
        }

        if _, ok := ref.(name.Digest); ok {
                return r.pullDigestViaDockerDaemon(ctx, imageName, ref)
        }

        tag, ok := ref.(name.Tag)
        if !ok {
                return fmt.Errorf("unsupported image reference type %T", ref)
        }

        remoteOpts := []remote.Option{
                remote.WithAuthFromKeychain(r.keychain),
                remote.WithContext(ctx),
        }
        if r.platform != nil {
                remoteOpts = append(remoteOpts, remote.WithPlatform(*r.platform))
        }

        img, err := remote.Image(ref, remoteOpts...)
        if err != nil {
                return fmt.Errorf("failed to pull image from registry: %w", err)
        }

        response, err := daemon.Write(tag, img, daemon.WithClient(r.dockerClient))
        if err != nil {
                return fmt.Errorf("failed to write image to daemon: %w", err)
        }

        if _, err := fmt.Fprintf(os.Stdout, "Successfully pulled %s\n", imageName); err != nil {
                slog.Debug("failed to write success message", "error", err)
        }
        //nolint:gosec // G706: image name and response from registry pull
        slog.Debug("pull complete", "image", imageName, "response", response)

        return nil
}

// pullDigestViaDockerDaemon pulls a digest-pinned image (tag@digest or
// digest-only form) using the Docker daemon client so the OCI manifest digest
// is stored and Docker can resolve the reference at container-create time.
func (r *RegistryImageManager) pullDigestViaDockerDaemon(ctx context.Context, imageName string, ref name.Reference) error {
        // Resolve auth from the keychain and encode for the daemon API.
        var registryAuth string
        if auth, err := r.keychain.Resolve(ref.Context().Registry); err == nil && auth != authn.Anonymous {
                if authCfg, err := auth.Authorization(); err == nil {
                        if authJSON, err := json.Marshal(authCfg); err == nil {
                                registryAuth = base64.URLEncoding.EncodeToString(authJSON)
                        }
                }
        }

        resp, err := r.dockerClient.ImagePull(ctx, imageName, mobyclient.ImagePullOptions{
                RegistryAuth: registryAuth,
        })
        if err != nil {
                return fmt.Errorf("failed to pull image: %w", err)
        }
        if err := resp.Wait(ctx); err != nil {
                return fmt.Errorf("failed waiting for image pull: %w", err)
        }

        if _, err := fmt.Fprintf(os.Stdout, "Successfully pulled %s\n", imageName); err != nil {
                slog.Debug("failed to write success message", "error", err)
        }

        return nil
}

// BuildImage builds a Docker image from a Dockerfile in the specified context directory
func (r *RegistryImageManager) BuildImage(ctx context.Context, contextDir, imageName string) error {
        return buildDockerImage(ctx, r.dockerClient, contextDir, imageName)
}

// WithKeychain sets the keychain for authentication
func (r *RegistryImageManager) WithKeychain(keychain authn.Keychain) *RegistryImageManager {
        r.keychain = keychain
        return r
}

// WithPlatform sets the platform for the RegistryImageManager
func (r *RegistryImageManager) WithPlatform(platform *v1.Platform) *RegistryImageManager {
        r.platform = platform
        return r
}

// buildDockerImage builds a Docker image using the Docker client API
func buildDockerImage(ctx context.Context, dockerClient *mobyclient.Client, contextDir, imageName string) error {
        //nolint:gosec // G706: image name and context dir from config
        slog.Debug("building image", "image", imageName, "context_dir", contextDir)

        // Create a tar archive of the context directory
        tarFile, err := os.CreateTemp("", "docker-build-context-*.tar")
        if err != nil {
                return fmt.Errorf("failed to create temporary tar file: %w", err)
        }
        defer func() {
                // #nosec G703 -- tarFile.Name() is from os.CreateTemp, not user input
                if err := os.Remove(tarFile.Name()); err != nil {
                        // Non-fatal: temp file cleanup failure
                        //nolint:gosec // G706: temp file path from os.CreateTemp
                        slog.Debug("failed to remove temporary file", "path", tarFile.Name(), "error", err)
                }
        }()
        defer func() {
                if err := tarFile.Close(); err != nil {
                        // Docker client closes the reader on success, so ignore "already closed" errors
                        if !errors.Is(err, os.ErrClosed) {
                                // Non-fatal: file cleanup failure
                                slog.Debug("failed to close tar file", "error", err)
                        }
                }
        }()

        // Create a tar archive of the context directory
        if err := createTarFromDir(contextDir, tarFile); err != nil {
                return fmt.Errorf("failed to create tar archive: %w", err)
        }

        // Reset the file pointer to the beginning of the file
        if _, err := tarFile.Seek(0, 0); err != nil {
                return fmt.Errorf("failed to reset tar file pointer: %w", err)
        }

        // Build the image
        buildOptions := mobyclient.ImageBuildOptions{
                Tags:       []string{imageName},
                Dockerfile: "Dockerfile",
                Remove:     true,
        }

        response, err := dockerClient.ImageBuild(ctx, tarFile, buildOptions)
        if err != nil {
                return fmt.Errorf("failed to build image: %w", err)
        }
        defer func() {
                if err := response.Body.Close(); err != nil {
                        // Non-fatal: response body cleanup failure
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        // Parse and log the build output
        if err := parseBuildOutput(response.Body, os.Stdout); err != nil {
                return fmt.Errorf("failed to process build output: %w", err)
        }

        return nil
}

// createTarFromDir creates a tar archive from a directory
func createTarFromDir(srcDir string, writer io.Writer) error {
        // Create a new tar writer
        tw := tar.NewWriter(writer)
        defer func() {
                if err := tw.Close(); err != nil {
                        // Non-fatal: tar writer cleanup failure
                        slog.Debug("failed to close tar writer", "error", err)
                }
        }()

        // Walk through the directory and add files to the tar archive
        return filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
                if err != nil {
                        return err
                }

                // Get the relative path
                relPath, err := filepath.Rel(srcDir, path)
                if err != nil {
                        return fmt.Errorf("failed to get relative path: %w", err)
                }

                // Skip the root directory
                if relPath == "." {
                        return nil
                }

                // Create a tar header
                header, err := tar.FileInfoHeader(info, "")
                if err != nil {
                        return fmt.Errorf("failed to create tar header: %w", err)
                }

                // Set the name to the relative path
                header.Name = relPath

                // Write the header
                if err := tw.WriteHeader(header); err != nil {
                        return fmt.Errorf("failed to write tar header: %w", err)
                }

                // If it's a regular file, write the contents
                if !info.IsDir() {
                        // #nosec G304 - This is safe because we're only opening files within the specified context directory
                        file, err := os.Open(path) //nolint:gosec // G122 - path from filepath.Walk within validated source directory
                        if err != nil {
                                return fmt.Errorf("failed to open file: %w", err)
                        }
                        defer func() {
                                if err := file.Close(); err != nil {
                                        // Non-fatal: file cleanup failure
                                        slog.Debug("failed to close file", "error", err)
                                }
                        }()

                        if _, err := io.Copy(tw, file); err != nil {
                                return fmt.Errorf("failed to copy file contents: %w", err)
                        }
                }

                return nil
        })
}

// parseBuildOutput parses the Docker image build output and formats it in a more readable way
func parseBuildOutput(reader io.Reader, writer io.Writer) error {
        decoder := json.NewDecoder(reader)
        for {
                var buildOutput struct {
                        Stream string `json:"stream,omitempty"`
                        Error  string `json:"error,omitempty"`
                }

                if err := decoder.Decode(&buildOutput); err != nil {
                        if err == io.EOF {
                                break
                        }
                        return fmt.Errorf("failed to decode build output: %w", err)
                }

                // Check for errors
                if buildOutput.Error != "" {
                        return fmt.Errorf("build error: %s", buildOutput.Error)
                }

                // Print the stream output
                if buildOutput.Stream != "" {
                        if _, err := fmt.Fprint(writer, buildOutput.Stream); err != nil {
                                slog.Debug("failed to write build output", "error", err)
                        }
                }
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package kubernetes provides a client for the Kubernetes runtime
// including creating, starting, stopping, and retrieving container information.
package kubernetes

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "os"
        "sort"
        "strconv"
        "strings"
        "time"

        "github.com/cenkalti/backoff/v7"
        appsv1 "k8s.io/api/apps/v1"
        corev1 "k8s.io/api/core/v1"
        apierrors "k8s.io/apimachinery/pkg/api/errors"
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/util/intstr"
        apimwatch "k8s.io/apimachinery/pkg/watch"
        appsv1apply "k8s.io/client-go/applyconfigurations/apps/v1"
        corev1apply "k8s.io/client-go/applyconfigurations/core/v1"
        metav1apply "k8s.io/client-go/applyconfigurations/meta/v1"
        "k8s.io/client-go/kubernetes"
        "k8s.io/client-go/kubernetes/scheme"
        "k8s.io/client-go/rest"
        "k8s.io/client-go/tools/remotecommand"
        "k8s.io/client-go/tools/watch"
        "k8s.io/utils/ptr"

        "github.com/stacklok/toolhive-core/permissions"
        "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/k8s"
        transtypes "github.com/stacklok/toolhive/pkg/transport/types"
)

// Constants for container status
const (
        // UnknownStatus represents an unknown container status
        UnknownStatus = "unknown"
        // mcpContainerName is the name of the MCP container. This is a known constant.
        mcpContainerName = "mcp"
        // defaultNamespace is the default Kubernetes namespace
        defaultNamespace = "default"
        // serviceFieldManager is the field manager name for server-side apply operations
        serviceFieldManager = "toolhive-container-manager"

        // KubectlRestartedAtAnnotation is the pod-template annotation written by
        // `kubectl rollout restart`. Server-side apply must re-assert it or the
        // operator/runner clobber the bounce (#6344). Exported so operator
        // controllers share the same key.
        KubectlRestartedAtAnnotation = "kubectl.kubernetes.io/restartedAt"

        // RunConfigMCPServerGenerationAnnotation carries the MCPServer .metadata.generation that
        // produced the RunConfig applied to this StatefulSet. Used as a monotonic version stamp
        // to prevent stale proxyrunner pods (from an old Deployment ReplicaSet) from clobbering
        // a newer RunConfig's apply. The gate only becomes effective once proxyrunner is upgraded
        // to a version that reads this annotation; operator-only upgrades leave the race window
        // in place until proxyrunner is also rolled. Exported because it forms a wire contract
        // that external readers (operator, diagnostic tooling) may consume.
        //
        // The operator also stamps this same annotation on the proxyrunner Deployment's
        // pod template and projects it into the proxyrunner container as the env var
        // EnvVarMCPServerGeneration via the downward API. That projection freezes the
        // generation per pod at creation time, so two coexisting proxyrunner pods cannot
        // converge on the same generation by re-reading the live-mounted RunConfig
        // ConfigMap (issue #5360).
        RunConfigMCPServerGenerationAnnotation = "toolhive.stacklok.dev/mcpserver-generation"

        // EnvVarMCPServerGeneration is the env var name through which the proxyrunner
        // container receives its frozen-per-pod MCPServer generation. Sourced via the
        // downward API from the pod-template annotation RunConfigMCPServerGenerationAnnotation,
        // it overrides the value read from /etc/runconfig/runconfig.json (which would
        // otherwise live-update across all proxyrunner pods during a helm upgrade and
        // defeat the apply-gate). See issue #5360.
        EnvVarMCPServerGeneration = "THV_MCPSERVER_GENERATION"
)

// RuntimeName is the name identifier for the Kubernetes runtime
const RuntimeName = "kubernetes"

// Retry configuration for kubectl attach operations
const (
        // attachRetryTimeout is the maximum time to retry kubectl attach before giving up
        // This accommodates typical pod restart times in both local and CI environments,
        // including container image pulls and startup delays
        attachRetryTimeout = 90 * time.Second

        // attachMaxRetryInterval is the maximum delay between individual retry attempts
        attachMaxRetryInterval = 15 * time.Second

        // attachInitialRetryInterval is the initial delay before the first retry
        attachInitialRetryInterval = 1 * time.Second
)

// Client implements the Deployer interface for container operations
type Client struct {
        runtimeType      runtime.Type
        client           kubernetes.Interface
        config           *rest.Config
        platformDetector PlatformDetector
        // waitForStatefulSetReadyFunc is used for testing to mock the waitForStatefulSetReady function
        waitForStatefulSetReadyFunc func(
                ctx context.Context,
                clientset kubernetes.Interface,
                namespace, name string,
                desiredGeneration int64,
        ) error
        // namespaceFunc is used for testing to override namespace detection
        namespaceFunc func() string
        // exitFunc is used for testing to override os.Exit behavior
        exitFunc func(code int)
}

// NewClient creates a new container client
func NewClient(_ context.Context) (*Client, error) {
        // Get kubernetes client and config using the common package
        clientset, config, err := k8s.NewClient()
        if err != nil {
                return nil, err
        }

        return NewClientWithConfig(clientset, config), nil
}

// NewClientWithConfig creates a new container client with a provided config
// This is primarily used for testing with fake clients
func NewClientWithConfig(clientset kubernetes.Interface, config *rest.Config) *Client {
        return &Client{
                runtimeType:      runtime.TypeKubernetes,
                client:           clientset,
                config:           config,
                platformDetector: NewDefaultPlatformDetector(),
        }
}

// NewClientWithConfigAndPlatformDetector creates a new container client with a provided config and platform detector
// This is primarily used for testing with fake clients and mock platform detectors
func NewClientWithConfigAndPlatformDetector(
        clientset kubernetes.Interface,
        config *rest.Config,
        platformDetector PlatformDetector,
) *Client {
        return &Client{
                runtimeType:      runtime.TypeKubernetes,
                client:           clientset,
                config:           config,
                platformDetector: platformDetector,
        }
}

// AttachToWorkload implements runtime.Runtime.
// It establishes a kubectl attach connection to the MCP server pod.
//
// Connection Failure Handling:
// If the connection fails permanently (after retries with exponential backoff),
// this function causes the process to exit with code 1. This triggers a Kubernetes
// restart, allowing the proxy to establish a fresh connection to the current pod.
// This is critical for handling StatefulSet pod restarts - when the MCP pod restarts,
// the old kubectl attach connection becomes stale and cannot be reused. Exiting allows
// Kubernetes to restart the proxy, which then attaches to the new pod.
//
// The retry configuration (see attachRetryTimeout constant) accommodates typical pod
// restart times in both local and CI environments, while still failing fast enough
// for truly unavailable pods.
func (c *Client) AttachToWorkload(ctx context.Context, workloadName string) (io.WriteCloser, io.ReadCloser, error) {
        // AttachToWorkload attaches to a workload in Kubernetes
        // This is a more complex operation in Kubernetes compared to Docker/Podman
        // as it requires setting up an exec session to the pod

        // First, we need to find the pod associated with the workloadID (which is actually the statefulset name)
        namespace := c.getCurrentNamespace()
        pods, err := c.client.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
                LabelSelector: fmt.Sprintf("app=%s", workloadName),
        })
        if err != nil {
                return nil, nil, fmt.Errorf("failed to find pod for workload %s: %w", workloadName, err)
        }

        if len(pods.Items) == 0 {
                return nil, nil, fmt.Errorf("%w: no pods found for workload %s", runtime.ErrWorkloadNotFound, workloadName)
        }

        // Use the first pod found
        podName := pods.Items[0].Name

        attachOpts := &corev1.PodAttachOptions{
                Container: mcpContainerName,
                Stdin:     true,
                Stdout:    true,
                Stderr:    true,
                TTY:       false,
        }

        // Set up the attach request URL (used to create fresh SPDY executors for each retry)
        req := c.client.CoreV1().RESTClient().Post().
                Resource("pods").
                Name(podName).
                Namespace(c.getCurrentNamespace()).
                SubResource("attach").
                VersionedParams(attachOpts, scheme.ParameterCodec)
        attachURL := req.URL()

        slog.Info("attaching to pod", "pod", podName, "workload", workloadName)

        stdinReader, stdinWriter := io.Pipe()
        stdoutReader, stdoutWriter := io.Pipe()
        go func() {
                // Close pipes when this goroutine exits to signal the transport layer.
                // This ensures processStdout() sees EOF and can attempt re-attachment or exit.
                defer func() {
                        if err := stdoutWriter.Close(); err != nil {
                                slog.Debug("error closing stdout writer", "error", err)
                        }
                        if err := stdinReader.Close(); err != nil {
                                slog.Debug("error closing stdin reader", "error", err)
                        }
                }()

                // Create exponential backoff with extended retry window to handle pod restarts
                // in both local and CI environments.
                expBackoff := backoff.NewExponentialBackOff()
                expBackoff.MaxInterval = attachMaxRetryInterval
                expBackoff.InitialInterval = attachInitialRetryInterval

                _, err := backoff.Retry(ctx, func() (any, error) {
                        // Create a fresh SPDY executor for each retry attempt.
                        // This is critical because the SPDY connection state becomes corrupted
                        // after certain failures (e.g., EOF from idle timeout), and reusing
                        // a stale executor prevents recovery.
                        exec, execErr := remotecommand.NewSPDYExecutor(c.config, "POST", attachURL)
                        if execErr != nil {
                                return nil, fmt.Errorf("failed to create SPDY executor: %w", execErr)
                        }

                        return nil, exec.StreamWithContext(ctx, remotecommand.StreamOptions{
                                Stdin:  stdinReader,
                                Stdout: stdoutWriter,
                                Stderr: stdoutWriter,
                                Tty:    false,
                        })
                },
                        backoff.WithBackOff(expBackoff),
                        backoff.WithMaxElapsedTime(attachRetryTimeout),
                        backoff.WithNotify(func(err error, duration time.Duration) {
                                slog.Error("error attaching to workload, retrying", "workload", workloadName, "error", err, "retry_in", duration)
                        }),
                )
                if err != nil {
                        var statusErr *apierrors.StatusError
                        if errors.As(err, &statusErr) {
                                slog.Error("kubernetes API error",
                                        "status", statusErr.ErrStatus.Status,
                                        "message", statusErr.ErrStatus.Message,
                                        "reason", statusErr.ErrStatus.Reason,
                                        "code", statusErr.ErrStatus.Code)

                                // Note: statuscode 0 with empty message indicates the connection was closed
                                // unexpectedly (e.g., container terminated or doesn't read from stdin)
                                if statusErr.ErrStatus.Code == 0 && statusErr.ErrStatus.Message == "" {
                                        slog.Error("connection closed unexpectedly, pod likely terminated or restarted", "workload", workloadName)
                                }
                        } else {
                                slog.Error("non-status error", "error", err)
                        }

                        // Exit the process to trigger a restart by Kubernetes.
                        // This allows the proxy to establish a fresh connection to the current pod
                        // after a pod restart, rather than maintaining a stale connection.
                        //
                        // Note: We call os.Exit(1) directly (bypassing deferred cleanup) because:
                        // 1. The proxy is in a permanently broken state with stale stdin/stdout pipes
                        // 2. Any cleanup of these broken resources would likely fail or hang
                        // 3. We want Kubernetes to perform a complete container restart with fresh state
                        // 4. Deferred cleanup is designed for graceful shutdown, not recovery from broken state
                        slog.Error("kubectl attach failed after all retries, exiting to allow restart", "workload", workloadName)
                        exitFunc := c.exitFunc
                        if exitFunc == nil {
                                exitFunc = os.Exit
                        }
                        exitFunc(1)
                }
        }()

        return stdinWriter, stdoutReader, nil
}

// GetWorkloadLogs implements runtime.Runtime.
func (c *Client) GetWorkloadLogs(ctx context.Context, workloadName string, follow bool, lines int) (string, error) {
        // follow=true means infinite streaming, lines>0 means finite limit - these are contradictory
        if follow && lines > 0 {
                return "", fmt.Errorf(
                        "cannot use both follow and line limit: follow mode streams logs indefinitely, " +
                                "which conflicts with line limiting",
                )
        }

        // In Kubernetes, workloadID is the statefulset name
        namespace := c.getCurrentNamespace()

        // Get the pods associated with this statefulset
        pods, err := c.client.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
                LabelSelector: "toolhive=true",
                FieldSelector: fmt.Sprintf("metadata.name=%s", workloadName),
        })
        if err != nil {
                return "", fmt.Errorf("failed to list pods for statefulset %s: %w", workloadName, err)
        }

        if len(pods.Items) == 0 {
                return "", fmt.Errorf("%w: no pods found for statefulset %s", runtime.ErrWorkloadNotFound, workloadName)
        }

        // Use the first pod
        podName := pods.Items[0].Name

        // Configure tail lines based on lines parameter
        var tailLines *int64
        if lines > 0 {
                tailLinesVal := int64(lines)
                tailLines = &tailLinesVal
        }

        // Get logs from the pod
        logOptions := &corev1.PodLogOptions{
                Container:  mcpContainerName,
                Follow:     follow,
                Previous:   false,
                Timestamps: true,
                TailLines:  tailLines,
        }

        req := c.client.CoreV1().Pods(namespace).GetLogs(podName, logOptions)
        podLogs, err := req.Stream(ctx)
        if err != nil {
                return "", fmt.Errorf("failed to get logs for pod %s: %w", podName, err)
        }
        defer func() {
                if err := podLogs.Close(); err != nil {
                        // Non-fatal: pod logs cleanup failure
                        slog.Debug("failed to close pod logs", "error", err)
                }
        }()

        // Read logs
        logBytes, err := io.ReadAll(podLogs)
        if err != nil {
                return "", fmt.Errorf("failed to read logs for pod %s: %w", podName, err)
        }

        return string(logBytes), nil
}

// DeployWorkload implements runtime.Runtime.
func (c *Client) DeployWorkload(ctx context.Context,
        image string,
        containerName string,
        command []string,
        envVars map[string]string,
        containerLabels map[string]string,
        _ *permissions.Profile, // TODO: Implement permission profile support for Kubernetes
        transportType string,
        options *runtime.DeployWorkloadOptions,
        _ bool,
) (int, error) {
        namespace := c.getCurrentNamespace()
        containerLabels["app"] = containerName
        containerLabels["toolhive"] = "true"

        attachStdio := options == nil || options.AttachStdio

        envVarList := buildSortedEnvVarList(envVars)

        // Create a pod template spec
        podTemplateSpec := ensureObjectMetaApplyConfigurationExists(corev1apply.PodTemplateSpec())

        // Apply the patch if provided
        if options != nil && options.K8sPodTemplatePatch != "" {
                var err error
                podTemplateSpec, err = applyPodTemplatePatch(podTemplateSpec, options.K8sPodTemplatePatch)
                if err != nil {
                        return 0, fmt.Errorf("failed to apply pod template patch: %w", err)
                }
        }

        // Ensure the pod template has required configuration (labels, etc.)
        // Get a config to talk to the apiserver
        cfg := c.config

        // Detect platform type
        platformDetector := c.platformDetector
        if platformDetector == nil {
                platformDetector = NewDefaultPlatformDetector()
        }
        platform, err := platformDetector.DetectPlatform(cfg)
        if err != nil {
                return 0, fmt.Errorf("can't determine api server type: %w", err)
        }

        podTemplateSpec = ensurePodTemplateConfig(podTemplateSpec, containerLabels, platform)

        // Configure the MCP container
        err = configureMCPContainer(
                podTemplateSpec,
                image,
                command,
                attachStdio,
                envVarList,
                transportType,
                options,
                platform,
        )
        if err != nil {
                return 0, err
        }

        ourGen := runConfigGeneration(options)
        existing, err := c.getExistingStatefulSet(ctx, namespace, containerName)
        if err != nil {
                return 0, err
        }
        if skipStatefulSetApply(existing, ourGen) {
                // Intentionally skip ensureBackendServices in the gated path: this pod's RunConfig
                // is stale, so reconciling services here would clobber port/config fields set by
                // the newer-generation pod under the same field manager + Force: true — the same
                // race this gate prevents for the StatefulSet. The newer pod already reconciled
                // services; if that failed, it returns an error and retries on its own.
                return 0, nil
        }

        createdStatefulSet, err := c.applyStatefulSet(
                ctx, namespace, containerName, containerLabels, podTemplateSpec, options, ourGen, existing,
        )
        if err != nil {
                return 0, err
        }

        err = c.ensureBackendServices(
                ctx, containerName, namespace, containerLabels, transportType, options, createdStatefulSet)
        if err != nil {
                return 0, err
        }

        // Wait for the statefulset to be ready
        // Pass the generation from the Apply call to ensure we wait for the controller
        // to process this specific spec version
        waitFunc := waitForStatefulSetReady
        if c.waitForStatefulSetReadyFunc != nil {
                waitFunc = c.waitForStatefulSetReadyFunc
        }
        err = waitFunc(ctx, c.client, namespace, createdStatefulSet.Name, createdStatefulSet.Generation)
        if err != nil {
                return 0, fmt.Errorf("statefulset applied but failed to become ready: %w", err)
        }

        return 0, nil
}

// buildSortedEnvVarList converts an env var map to Kubernetes apply configurations
// with deterministically ordered keys. Go map iteration is randomized, and any
// ordering shift changes the pod template hash and triggers an unnecessary
// StatefulSet rollout (#5063).
func buildSortedEnvVarList(envVars map[string]string) []*corev1apply.EnvVarApplyConfiguration {
        envKeys := make([]string, 0, len(envVars))
        for k := range envVars {
                envKeys = append(envKeys, k)
        }
        sort.Strings(envKeys)
        envVarList := make([]*corev1apply.EnvVarApplyConfiguration, 0, len(envKeys))
        for _, k := range envKeys {
                envVarList = append(envVarList, corev1apply.EnvVar().WithName(k).WithValue(envVars[k]))
        }
        return envVarList
}

// runConfigGeneration extracts the RunConfig MCPServer generation from options,
// returning 0 when options is nil (backward-compat / non-operator callers).
func runConfigGeneration(options *runtime.DeployWorkloadOptions) int64 {
        if options == nil {
                return 0
        }
        return options.RunConfigMCPServerGeneration
}

// applyStatefulSet stamps the MCPServer generation annotation when non-zero,
// builds the StatefulSet apply configuration, and performs the server-side apply.
func (c *Client) applyStatefulSet(
        ctx context.Context,
        namespace, containerName string,
        containerLabels map[string]string,
        podTemplateSpec *corev1apply.PodTemplateSpecApplyConfiguration,
        options *runtime.DeployWorkloadOptions,
        ourGen int64,
        existing *appsv1.StatefulSet,
) (*appsv1.StatefulSet, error) {
        if ourGen > 0 {
                podTemplateSpec = podTemplateSpec.WithAnnotations(map[string]string{
                        RunConfigMCPServerGenerationAnnotation: strconv.FormatInt(ourGen, 10),
                })
        }
        if existing != nil && existing.Spec.Template.Annotations != nil {
                if restartedAt := existing.Spec.Template.Annotations[KubectlRestartedAtAnnotation]; restartedAt != "" {
                        podTemplateSpec = podTemplateSpec.WithAnnotations(map[string]string{
                                KubectlRestartedAtAnnotation: restartedAt,
                        })
                }
        }
        statefulSetApply := appsv1apply.StatefulSet(containerName, namespace).
                WithLabels(containerLabels).
                WithSpec(buildStatefulSetSpec(containerName, podTemplateSpec, options))
        createdStatefulSet, err := c.client.AppsV1().StatefulSets(namespace).
                Apply(ctx, statefulSetApply, metav1.ApplyOptions{
                        FieldManager: serviceFieldManager,
                        Force:        true,
                })
        if err != nil {
                return nil, fmt.Errorf("failed to apply statefulset: %w", err)
        }
        slog.Debug("applied statefulset", "name", createdStatefulSet.Name)
        return createdStatefulSet, nil
}

func (c *Client) getExistingStatefulSet(ctx context.Context, namespace, name string) (*appsv1.StatefulSet, error) {
        existing, err := c.client.AppsV1().StatefulSets(namespace).Get(ctx, name, metav1.GetOptions{})
        if err != nil {
                if apierrors.IsNotFound(err) {
                        return nil, nil
                }
                return nil, fmt.Errorf("failed to get existing statefulset: %w", err)
        }
        return existing, nil
}

// skipStatefulSetApply returns true when the existing StatefulSet is already
// stamped with a strictly greater MCPServer generation than ours, meaning a newer
// proxyrunner pod has already reconciled the workload and ours would be a regression.
// Returns false (apply as normal) when ourGen is zero or negative, when the StatefulSet
// does not yet exist, when the annotation is absent, or when the annotation is unparsable.
func skipStatefulSetApply(existing *appsv1.StatefulSet, ourGen int64) bool {
        if ourGen <= 0 || existing == nil {
                return false
        }
        if existing.Spec.Template.Annotations == nil {
                return false
        }
        theirs := existing.Spec.Template.Annotations[RunConfigMCPServerGenerationAnnotation]
        if theirs == "" {
                return false
        }
        theirsGen, parseErr := strconv.ParseInt(theirs, 10, 64)
        if parseErr != nil {
                slog.Warn("unparsable mcpserver-generation annotation; proceeding with apply",
                        "sts", existing.Name, "value", theirs, "err", parseErr)
                return false
        }
        if theirsGen > ourGen {
                slog.Debug("skipping StatefulSet apply; newer MCPServer generation already applied",
                        "sts", existing.Name, "ours", ourGen, "theirs", theirsGen)
                return true
        }
        return false
}

// buildStatefulSetSpec constructs the StatefulSet spec apply configuration.
// WithReplicas is only included when BackendReplicas is explicitly set; omitting
// the field lets the existing field manager (e.g. HPA or kubectl) retain control
// of scaling, satisfying the nil-omission invariant from RC-11.
func buildStatefulSetSpec(
        containerName string,
        podTemplateSpec *corev1apply.PodTemplateSpecApplyConfiguration,
        options *runtime.DeployWorkloadOptions,
) *appsv1apply.StatefulSetSpecApplyConfiguration {
        spec := appsv1apply.StatefulSetSpec().
                WithSelector(metav1apply.LabelSelector().
                        WithMatchLabels(map[string]string{"app": containerName})).
                WithServiceName(containerName).
                WithTemplate(podTemplateSpec)
        if options != nil && options.ScalingConfig != nil && options.ScalingConfig.BackendReplicas != nil {
                spec = spec.WithReplicas(*options.ScalingConfig.BackendReplicas)
        }
        return spec
}

// ensureBackendServices creates the headless and ClusterIP services needed by
// HTTP-based transports (SSE, streamable-http). Both services are owned by the
// StatefulSet so Kubernetes GC can clean them up automatically.
func (c *Client) ensureBackendServices(
        ctx context.Context,
        containerName, namespace string,
        containerLabels map[string]string,
        transportType string,
        options *runtime.DeployWorkloadOptions,
        sts *appsv1.StatefulSet,
) error {
        if !transportTypeRequiresBackendServices(transportType) || options == nil {
                return nil
        }

        stsOwner := &metav1.OwnerReference{
                APIVersion:         appsv1.SchemeGroupVersion.String(),
                Kind:               "StatefulSet",
                Name:               sts.Name,
                UID:                sts.UID,
                BlockOwnerDeletion: ptr.To(true),
                Controller:         ptr.To(true),
        }

        // Create a headless service for DNS discovery
        if err := c.createHeadlessService(ctx, containerName, namespace, containerLabels, options, stsOwner); err != nil {
                return fmt.Errorf("failed to create headless service: %w", err)
        }

        // Create a regular ClusterIP service with session affinity for the proxy-runner target
        if err := c.createMCPService(ctx, containerName, namespace, containerLabels, options, stsOwner); err != nil {
                return fmt.Errorf("failed to create MCP service: %w", err)
        }

        return nil
}

// GetWorkloadInfo implements runtime.Runtime.
func (c *Client) GetWorkloadInfo(ctx context.Context, workloadName string) (runtime.ContainerInfo, error) {
        // In Kubernetes, workloadID is the statefulset name
        namespace := c.getCurrentNamespace()

        // Get the statefulset
        statefulset, err := c.client.AppsV1().StatefulSets(namespace).Get(ctx, workloadName, metav1.GetOptions{})
        if err != nil {
                if apierrors.IsNotFound(err) {
                        return runtime.ContainerInfo{}, fmt.Errorf("%w: statefulset %s not found", runtime.ErrWorkloadNotFound, workloadName)
                }
                return runtime.ContainerInfo{}, fmt.Errorf("failed to get statefulset %s: %w", workloadName, err)
        }

        // Get the pods associated with this workload.
        pods, err := c.client.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
                LabelSelector: "toolhive=true",
                FieldSelector: fmt.Sprintf("metadata.name=%s", workloadName),
        })
        if err != nil {
                return runtime.ContainerInfo{}, fmt.Errorf("failed to list pods for statefulset %s: %w", workloadName, err)
        }

        // Extract port mappings from pods
        ports := make([]runtime.PortMapping, 0)
        if len(pods.Items) > 0 {
                ports = extractPortMappingsFromPod(&pods.Items[0])
        }

        // Get ports from associated service (for SSE transport)
        service, err := c.client.CoreV1().Services(namespace).Get(ctx, workloadName, metav1.GetOptions{})
        if err == nil {
                // Service exists, add its ports
                ports = extractPortMappingsFromService(service, ports)
        }

        // Determine status and state
        var status string
        var state runtime.WorkloadStatus
        if statefulset.Status.ReadyReplicas > 0 {
                status = "Running"
                state = runtime.WorkloadStatusRunning
        } else if statefulset.Status.Replicas > 0 {
                status = "Pending"
                state = runtime.WorkloadStatusStarting
        } else {
                // NOTE: Not clear if this is correct since the stop operation is a no-op.
                status = "Stopped"
                state = runtime.WorkloadStatusStopped
        }

        // Get the image from the pod template
        image := ""
        if len(statefulset.Spec.Template.Spec.Containers) > 0 {
                image = statefulset.Spec.Template.Spec.Containers[0].Image
        }

        return runtime.ContainerInfo{
                Name:    statefulset.Name,
                Image:   image,
                Status:  status,
                State:   state,
                Created: statefulset.CreationTimestamp.Time,
                Labels:  statefulset.Labels,
                Ports:   ports,
        }, nil
}

// IsWorkloadRunning implements runtime.Runtime.
func (c *Client) IsWorkloadRunning(ctx context.Context, workloadName string) (bool, error) {
        // In Kubernetes, workloadID is the statefulset name
        namespace := c.getCurrentNamespace()

        // Get the statefulset
        statefulset, err := c.client.AppsV1().StatefulSets(namespace).Get(ctx, workloadName, metav1.GetOptions{})
        if err != nil {
                if apierrors.IsNotFound(err) {
                        return false, fmt.Errorf("%w: statefulset %s not found", runtime.ErrWorkloadNotFound, workloadName)
                }
                return false, fmt.Errorf("failed to get statefulset %s: %w", workloadName, err)
        }

        // Check if the statefulset has at least one ready replica
        return statefulset.Status.ReadyReplicas > 0, nil
}

// ListWorkloads implements runtime.Runtime.
func (c *Client) ListWorkloads(ctx context.Context) ([]runtime.ContainerInfo, error) {
        // Create label selector for toolhive containers
        // Only show main MCP server pods (not proxy pods) by requiring toolhive-tool-type label
        labelSelector := "toolhive=true,toolhive-tool-type"

        // Determine namespace to search in
        var namespace string
        if strings.TrimSpace(os.Getenv("TOOLHIVE_KUBERNETES_ALL_NAMESPACES")) != "" {
                // Search in all namespaces
                namespace = ""
        } else {
                // Search in current namespace only
                namespace = c.getCurrentNamespace()
        }

        // List pods with the toolhive label
        pods, err := c.client.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
                LabelSelector: labelSelector,
        })
        if err != nil {
                return nil, fmt.Errorf("failed to list pods: %w", err)
        }

        // Convert to our ContainerInfo format
        result := make([]runtime.ContainerInfo, 0, len(pods.Items))
        for _, pod := range pods.Items {
                // Extract port mappings from pod
                ports := extractPortMappingsFromPod(&pod)

                // Get ports from associated service (for SSE transport)
                service, err := c.client.CoreV1().Services(namespace).Get(ctx, pod.Name, metav1.GetOptions{})
                if err == nil {
                        // Service exists, add its ports
                        ports = extractPortMappingsFromService(service, ports)
                }

                // Get container status
                status := UnknownStatus
                state := runtime.WorkloadStatusUnknown
                if len(pod.Status.ContainerStatuses) > 0 {
                        containerStatus := pod.Status.ContainerStatuses[0]
                        if containerStatus.State.Running != nil {
                                state = runtime.WorkloadStatusRunning
                                status = "Running"
                        } else if containerStatus.State.Waiting != nil {
                                state = runtime.WorkloadStatusStarting
                                status = containerStatus.State.Waiting.Reason
                        } else if containerStatus.State.Terminated != nil {
                                state = runtime.WorkloadStatusRemoving
                                status = containerStatus.State.Terminated.Reason
                        }
                }

                result = append(result, runtime.ContainerInfo{
                        Name:    pod.Name,
                        Image:   pod.Spec.Containers[0].Image,
                        Status:  status,
                        State:   state,
                        Created: pod.CreationTimestamp.Time,
                        Labels:  pod.Labels,
                        Ports:   ports,
                })
        }

        return result, nil
}

// RemoveWorkload implements runtime.Runtime.
func (c *Client) RemoveWorkload(ctx context.Context, workloadName string) error {
        // In Kubernetes, we remove a workload by deleting the statefulset
        namespace := c.getCurrentNamespace()

        // Delete the statefulset
        deleteOptions := metav1.DeleteOptions{}
        err := c.client.AppsV1().StatefulSets(namespace).Delete(ctx, workloadName, deleteOptions)
        if err != nil {
                if apierrors.IsNotFound(err) {
                        // If the statefulset doesn't exist, that's fine
                        slog.Info("statefulset not found, nothing to remove", "name", workloadName)
                        return nil
                }
                return fmt.Errorf("failed to delete statefulset %s: %w", workloadName, err)
        }

        slog.Info("deleted statefulset", "name", workloadName)
        return nil
}

// StopWorkload implements runtime.Runtime.
func (*Client) StopWorkload(_ context.Context, _ string) error {
        return nil
}

// IsRunning checks the health of the container runtime.
// This is used to verify that the runtime is operational and can manage workloads.
func (c *Client) IsRunning(ctx context.Context) error {
        // Use /readyz endpoint to check if the Kubernetes API server is ready.
        var status int
        result := c.client.Discovery().RESTClient().Get().AbsPath("/readyz").Do(ctx)
        if result.StatusCode(&status); status != 200 {
                return fmt.Errorf("kubernetes API server is not ready, status code: %d", status)
        }

        return nil
}

// isStatefulSetReady checks if a StatefulSet is ready after an update.
// It requires the desiredGeneration from the Apply call to ensure
// the controller has processed our spec before considering it ready.
//
// The check requires all three conditions to be true:
// 1. ObservedGeneration >= desiredGeneration (controller has processed our spec)
// 2. UpdatedReplicas == Replicas (all pods are on the new spec)
// 3. ReadyReplicas == Replicas (all pods are ready)
func isStatefulSetReady(desiredGeneration int64, currentSts *appsv1.StatefulSet) bool {
        if currentSts == nil || currentSts.Spec.Replicas == nil {
                return false
        }

        return currentSts.Status.ObservedGeneration >= desiredGeneration &&
                currentSts.Status.UpdatedReplicas == *currentSts.Spec.Replicas &&
                currentSts.Status.ReadyReplicas == *currentSts.Spec.Replicas
}

// waitForStatefulSetReady waits for a statefulset to be ready using the watch API.
// The desiredGeneration parameter is the generation from the Apply call (createdStatefulSet.Generation)
// which is used to ensure the controller has processed our specific spec version.
func waitForStatefulSetReady(
        ctx context.Context,
        clientset kubernetes.Interface,
        namespace, name string,
        desiredGeneration int64,
) error {
        // Create a field selector to watch only this specific statefulset
        fieldSelector := fmt.Sprintf("metadata.name=%s", name)

        // Set up the watch
        watcher, err := clientset.AppsV1().StatefulSets(namespace).Watch(ctx, metav1.ListOptions{
                FieldSelector: fieldSelector,
                Watch:         true,
        })
        if err != nil {
                return fmt.Errorf("error watching statefulset: %w", err)
        }

        // Define the condition function that checks if the statefulset is ready
        isStatefulSetReady := func(event apimwatch.Event) (bool, error) {
                // Check if the event is a statefulset
                statefulSet, ok := event.Object.(*appsv1.StatefulSet)
                if !ok {
                        return false, fmt.Errorf("unexpected object type: %T", event.Object)
                }

                if isStatefulSetReady(desiredGeneration, statefulSet) {
                        return true, nil
                }

                slog.Info("waiting for statefulset to be ready",
                        "name", name,
                        "ready_replicas", statefulSet.Status.ReadyReplicas,
                        "desired_replicas", *statefulSet.Spec.Replicas,
                        "observed_generation", statefulSet.Status.ObservedGeneration,
                        "desired_generation", desiredGeneration)
                return false, nil
        }

        // Create a context with timeout
        timeoutCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
        defer cancel()

        // Wait for the statefulset to be ready
        _, err = watch.UntilWithoutRetry(timeoutCtx, watcher, isStatefulSetReady)
        if err != nil {
                return fmt.Errorf("error waiting for statefulset to be ready: %w", err)
        }

        return nil
}

// parsePortString parses a port string in the format "port/protocol" and returns the port number
func parsePortString(portStr string) (int, error) {
        // Split the port string to get just the port number
        port := strings.Split(portStr, "/")[0]
        portNum, err := strconv.Atoi(port)
        if err != nil {
                return 0, fmt.Errorf("failed to parse port %s: %w", port, err)
        }
        return portNum, nil
}

// configureContainerPorts adds port configurations to a container for SSE transport
func configureContainerPorts(
        containerConfig *corev1apply.ContainerApplyConfiguration,
        options *runtime.DeployWorkloadOptions,
) (*corev1apply.ContainerApplyConfiguration, error) {
        if options == nil {
                return containerConfig, nil
        }

        // Use a map to track which ports have been added
        portMap := make(map[int32]bool)
        var containerPorts []*corev1apply.ContainerPortApplyConfiguration

        // Process exposed ports
        for portStr := range options.ExposedPorts {
                portNum, err := parsePortString(portStr)
                if err != nil {
                        return nil, err
                }

                // Check for integer overflow
                if portNum < 0 || portNum > 65535 {
                        return nil, fmt.Errorf("port number %d is out of valid range (0-65535)", portNum)
                }

                // Add port if not already in the map
                portInt32 := int32(portNum)
                if !portMap[portInt32] {
                        containerPorts = append(containerPorts, corev1apply.ContainerPort().
                                WithContainerPort(portInt32).
                                WithProtocol(corev1.ProtocolTCP))
                        portMap[portInt32] = true
                }
        }

        // Process port bindings
        for portStr := range options.PortBindings {
                portNum, err := parsePortString(portStr)
                if err != nil {
                        return nil, err
                }

                // Check for integer overflow
                if portNum < 0 || portNum > 65535 {
                        return nil, fmt.Errorf("port number %d is out of valid range (0-65535)", portNum)
                }

                // Add port if not already in the map
                portInt32 := int32(portNum)
                if !portMap[portInt32] {
                        containerPorts = append(containerPorts, corev1apply.ContainerPort().
                                WithContainerPort(portInt32).
                                WithProtocol(corev1.ProtocolTCP))
                        portMap[portInt32] = true
                }
        }

        // Add ports to container config
        if len(containerPorts) > 0 {
                containerConfig = containerConfig.WithPorts(containerPorts...)
        }

        return containerConfig, nil
}

// validatePortNumber checks if a port number is within the valid range
func validatePortNumber(portNum int) error {
        if portNum < 0 || portNum > 65535 {
                return fmt.Errorf("port number %d is out of valid range (0-65535)", portNum)
        }
        return nil
}

// createServicePortConfig creates a service port configuration for a given port number
func createServicePortConfig(portNum int) *corev1apply.ServicePortApplyConfiguration {
        //nolint:gosec // G115: Safe int->int32 conversion, range is checked in validatePortNumber
        portInt32 := int32(portNum)
        return corev1apply.ServicePort().
                WithName(fmt.Sprintf("port-%d", portInt32)).
                WithPort(portInt32).
                WithTargetPort(intstr.FromInt32(portInt32)).
                WithProtocol(corev1.ProtocolTCP)
}

// processExposedPorts processes exposed ports and adds them to the port map
func processExposedPorts(
        options *runtime.DeployWorkloadOptions,
        portMap map[int32]*corev1apply.ServicePortApplyConfiguration,
) error {
        for portStr := range options.ExposedPorts {
                portNum, err := parsePortString(portStr)
                if err != nil {
                        return err
                }

                if err := validatePortNumber(portNum); err != nil {
                        return err
                }

                //nolint:gosec // G115: Safe int->int32 conversion, range is checked in validatePortNumber
                portInt32 := int32(portNum)
                // Add port if not already in the map
                if _, exists := portMap[portInt32]; !exists {
                        portMap[portInt32] = createServicePortConfig(portNum)
                }
        }
        return nil
}

// createServicePorts creates service port configurations from container options
func createServicePorts(options *runtime.DeployWorkloadOptions) ([]*corev1apply.ServicePortApplyConfiguration, error) {
        if options == nil {
                return nil, nil
        }

        // Use a map to track which ports have been added
        portMap := make(map[int32]*corev1apply.ServicePortApplyConfiguration)

        // Process exposed ports
        if err := processExposedPorts(options, portMap); err != nil {
                return nil, err
        }

        // Process port bindings
        for portStr, bindings := range options.PortBindings {
                portNum, err := parsePortString(portStr)
                if err != nil {
                        return nil, err
                }

                if err := validatePortNumber(portNum); err != nil {
                        return nil, err
                }

                //nolint:gosec // G115: Safe int->int32 conversion, range is checked in validatePortNumber
                portInt32 := int32(portNum)
                servicePort := portMap[portInt32]
                if servicePort == nil {
                        // Create new service port if not in map
                        servicePort = createServicePortConfig(portNum)
                }

                // If there are bindings with a host port, use the first one as node port
                if len(bindings) > 0 && bindings[0].HostPort != "" {
                        hostPort, err := strconv.Atoi(bindings[0].HostPort)
                        if err == nil && hostPort >= 30000 && hostPort <= 32767 {
                                // NodePort must be in range 30000-32767
                                // Safe to convert to int32 since we've verified the range (30000-32767)
                                // which is well within int32 range (-2,147,483,648 to 2,147,483,647)
                                //nolint:gosec // G109: Safe int->int32 conversion, range is checked above
                                nodePort := int32(hostPort)
                                servicePort = servicePort.WithNodePort(nodePort)
                        }
                }

                //nolint:gosec // G115: Safe int->int32 conversion, range is checked above
                portMap[int32(portNum)] = servicePort
        }

        // Convert map to slice
        var servicePorts []*corev1apply.ServicePortApplyConfiguration
        for _, port := range portMap {
                servicePorts = append(servicePorts, port)
        }

        return servicePorts, nil
}

// serviceConfig holds the configuration for creating a Kubernetes service via applyService.
type serviceConfig struct {
        // nameSuffix is appended to "mcp-<containerName>" to form the service name.
        // Use "-headless" for the headless service or "" for the MCP service.
        nameSuffix string
        // headless makes the service a headless service (ClusterIP: None).
        headless bool
        // sessionAffinity enables ClientIP session affinity with the given timeout.
        sessionAffinity bool
        // sessionAffinityTimeoutSeconds sets the timeout for ClientIP session affinity.
        // Only used when sessionAffinity is true. Kubernetes defaults to 10800s (3h) if unset.
        sessionAffinityTimeoutSeconds int32
}

// applyService creates or updates a Kubernetes service using server-side apply.
// If owner is non-nil, it is set as an owner reference so Kubernetes garbage-collects
// the service when the owner is deleted.
func (c *Client) applyService(
        ctx context.Context,
        containerName string,
        namespace string,
        labels map[string]string,
        options *runtime.DeployWorkloadOptions,
        cfg serviceConfig,
        owner *metav1.OwnerReference,
) (string, error) {
        servicePorts, err := createServicePorts(options)
        if err != nil {
                return "", err
        }

        if len(servicePorts) == 0 {
                slog.Debug("no ports configured, skipping service creation")
                return "", nil
        }

        svcName := fmt.Sprintf("mcp-%s%s", containerName, cfg.nameSuffix)

        // Determine service type based on whether any ports have NodePort set.
        // Headless services (ClusterIP: None) cannot be NodePort, so skip the
        // promotion for those — Kubernetes rejects clusterIP=None + type=NodePort.
        serviceType := corev1.ServiceTypeClusterIP
        if !cfg.headless {
                for _, sp := range servicePorts {
                        if sp.NodePort != nil {
                                serviceType = corev1.ServiceTypeNodePort
                                break
                        }
                }
        }

        spec := corev1apply.ServiceSpec().
                WithSelector(map[string]string{
                        "app": containerName,
                }).
                WithPorts(servicePorts...).
                WithType(serviceType)

        if cfg.headless {
                spec = spec.WithClusterIP("None")
        }

        if cfg.sessionAffinity {
                spec = spec.
                        WithSessionAffinity(corev1.ServiceAffinityClientIP).
                        WithSessionAffinityConfig(corev1apply.SessionAffinityConfig().
                                WithClientIP(corev1apply.ClientIPConfig().
                                        WithTimeoutSeconds(cfg.sessionAffinityTimeoutSeconds)))
        }

        serviceApply := corev1apply.Service(svcName, namespace).
                WithLabels(labels).
                WithSpec(spec)

        if owner != nil {
                serviceApply = serviceApply.WithOwnerReferences(metav1apply.OwnerReference().
                        WithAPIVersion(owner.APIVersion).
                        WithKind(owner.Kind).
                        WithName(owner.Name).
                        WithUID(owner.UID).
                        WithBlockOwnerDeletion(true).
                        WithController(true))
        }

        _, err = c.client.CoreV1().Services(namespace).
                Apply(ctx, serviceApply, metav1.ApplyOptions{
                        FieldManager: serviceFieldManager,
                        Force:        true,
                })
        if err != nil {
                return "", fmt.Errorf("failed to apply service %s: %w", svcName, err)
        }

        slog.Debug("applied service", "name", svcName)
        return svcName, nil
}

// createHeadlessService creates a headless Kubernetes service for the StatefulSet
func (c *Client) createHeadlessService(
        ctx context.Context,
        containerName string,
        namespace string,
        labels map[string]string,
        options *runtime.DeployWorkloadOptions,
        owner *metav1.OwnerReference,
) error {
        _, err := c.applyService(ctx, containerName, namespace, labels, options, serviceConfig{
                nameSuffix: "-headless",
                headless:   true,
        }, owner)
        return err
}

// mcpServiceSessionAffinityTimeout is the timeout in seconds for ClientIP session affinity
// on the MCP service. This controls how long kube-proxy pins a client IP to the same backend pod.
// Note: this provides proxy-runner-level stickiness (L4), not per-MCP-session stickiness (L7).
// True per-session routing would require Mcp-Session-Id-based routing at the proxy layer.
const mcpServiceSessionAffinityTimeout int32 = 1800

// createMCPService creates a regular ClusterIP service with SessionAffinity for the MCP server StatefulSet.
// This service provides load balancing with client-IP-based session stickiness, which the proxy-runner
// uses as its target host. The headless service is retained for DNS discovery purposes.
func (c *Client) createMCPService(
        ctx context.Context,
        containerName string,
        namespace string,
        labels map[string]string,
        options *runtime.DeployWorkloadOptions,
        owner *metav1.OwnerReference,
) error {
        svcName, err := c.applyService(ctx, containerName, namespace, labels, options, serviceConfig{
                sessionAffinity:               true,
                sessionAffinityTimeoutSeconds: mcpServiceSessionAffinityTimeout,
        }, owner)
        if err != nil {
                return err
        }
        options.MCPServiceName = svcName
        return nil
}

// extractPortMappingsFromPod extracts port mappings from a pod's containers
func extractPortMappingsFromPod(pod *corev1.Pod) []runtime.PortMapping {
        ports := make([]runtime.PortMapping, 0)

        for _, container := range pod.Spec.Containers {
                for _, port := range container.Ports {
                        ports = append(ports, runtime.PortMapping{
                                ContainerPort: int(port.ContainerPort),
                                HostPort:      int(port.HostPort),
                                Protocol:      string(port.Protocol),
                        })
                }
        }

        return ports
}

// transportTypeRequiresBackendServices returns true if the transport type requires backend services
func transportTypeRequiresBackendServices(transportType string) bool {
        return transportType == string(transtypes.TransportTypeSSE) || transportType == string(transtypes.TransportTypeStreamableHTTP)
}

// extractPortMappingsFromService extracts port mappings from a Kubernetes service
func extractPortMappingsFromService(service *corev1.Service, existingPorts []runtime.PortMapping) []runtime.PortMapping {
        // Create a map of existing ports for easy lookup and updating
        portMap := make(map[int]runtime.PortMapping)
        for _, p := range existingPorts {
                portMap[p.ContainerPort] = p
        }

        // Update or add ports from the service
        for _, port := range service.Spec.Ports {
                containerPort := int(port.Port)
                hostPort := 0
                if port.NodePort > 0 {
                        hostPort = int(port.NodePort)
                }

                // Update existing port or add new one
                portMap[containerPort] = runtime.PortMapping{
                        ContainerPort: containerPort,
                        HostPort:      hostPort,
                        Protocol:      string(port.Protocol),
                }
        }

        // Convert map back to slice
        result := make([]runtime.PortMapping, 0, len(portMap))
        for _, p := range portMap {
                result = append(result, p)
        }

        return result
}

// applyPodTemplatePatch applies a JSON patch to a pod template spec
func applyPodTemplatePatch(
        baseTemplate *corev1apply.PodTemplateSpecApplyConfiguration,
        patchJSON string,
) (*corev1apply.PodTemplateSpecApplyConfiguration, error) {
        // Check if the base template is nil
        if baseTemplate == nil {
                return nil, fmt.Errorf("base template is nil")
        }

        // Parse the patch JSON
        patchedSpec, err := createPodTemplateFromPatch(patchJSON)
        if err != nil {
                return nil, err
        }

        // Check if the patched spec is nil
        if patchedSpec == nil {
                return baseTemplate, nil
        }

        // Copy fields from the patched spec to our template
        if patchedSpec.ObjectMetaApplyConfiguration != nil && len(patchedSpec.Labels) > 0 {
                baseTemplate = baseTemplate.WithLabels(patchedSpec.Labels)
        }

        // Copy annotations from the patched spec to our template
        if patchedSpec.ObjectMetaApplyConfiguration != nil && len(patchedSpec.Annotations) > 0 {
                baseTemplate = baseTemplate.WithAnnotations(patchedSpec.Annotations)
        }

        if patchedSpec.Spec != nil {
                // Ensure baseTemplate.Spec is not nil
                if baseTemplate.Spec == nil {
                        baseTemplate = baseTemplate.WithSpec(corev1apply.PodSpec())
                }
                // Copy the spec
                baseTemplate = baseTemplate.WithSpec(patchedSpec.Spec)
        }

        return baseTemplate, nil
}

// createPodTemplateFromPatch creates a pod template spec from a JSON string
func createPodTemplateFromPatch(patchJSON string) (*corev1apply.PodTemplateSpecApplyConfiguration, error) {
        // Ensure the patch is valid JSON
        var patchMap map[string]interface{}
        if err := json.Unmarshal([]byte(patchJSON), &patchMap); err != nil {
                return nil, fmt.Errorf("invalid JSON patch: %w", err)
        }

        var podTemplateSpec corev1apply.PodTemplateSpecApplyConfiguration
        if err := json.Unmarshal([]byte(patchJSON), &podTemplateSpec); err != nil {
                return nil, fmt.Errorf("failed to unmarshal patch into pod template spec: %w", err)
        }

        // Ensure the pod template spec is not nil
        return ensureObjectMetaApplyConfigurationExists(&podTemplateSpec), nil
}

// ensurePodTemplateConfig ensures the pod template has required configuration
//
//nolint:gocyclo // Complex but necessary for platform-aware security context configuration
func ensurePodTemplateConfig(
        podTemplateSpec *corev1apply.PodTemplateSpecApplyConfiguration,
        containerLabels map[string]string,
        platform Platform,
) *corev1apply.PodTemplateSpecApplyConfiguration {
        podTemplateSpec = ensureObjectMetaApplyConfigurationExists(podTemplateSpec)
        // Ensure the pod template has labels
        if podTemplateSpec.Labels == nil {
                podTemplateSpec = podTemplateSpec.WithLabels(containerLabels)
        } else {
                // Merge with required labels
                for k, v := range containerLabels {
                        podTemplateSpec.Labels[k] = v
                }
        }

        // Ensure the pod template has a spec
        if podTemplateSpec.Spec == nil {
                podTemplateSpec = podTemplateSpec.WithSpec(corev1apply.PodSpec())
        }

        // Ensure the pod template has a restart policy
        if podTemplateSpec.Spec.RestartPolicy == nil {
                podTemplateSpec.Spec = podTemplateSpec.Spec.WithRestartPolicy(corev1.RestartPolicyAlways)
        }

        // Add pod-level security context using SecurityContextBuilder
        if podTemplateSpec.Spec.SecurityContext == nil {
                securityBuilder := NewSecurityContextBuilder(platform)
                podTemplateSpec.Spec = podTemplateSpec.Spec.WithSecurityContext(
                        securityBuilder.BuildPodSecurityContextApplyConfiguration(),
                )
        } else {
                // If the pod-level security context already exists, merge with platform-aware defaults
                securityBuilder := NewSecurityContextBuilder(platform)
                platformContext := securityBuilder.BuildPodSecurityContextApplyConfiguration()

                // Merge existing context with platform-aware settings
                if podTemplateSpec.Spec.SecurityContext.RunAsNonRoot == nil && platformContext.RunAsNonRoot != nil {
                        podTemplateSpec.Spec.SecurityContext = podTemplateSpec.Spec.SecurityContext.WithRunAsNonRoot(*platformContext.RunAsNonRoot)
                }

                if podTemplateSpec.Spec.SecurityContext.RunAsUser == nil && platformContext.RunAsUser != nil {
                        podTemplateSpec.Spec.SecurityContext = podTemplateSpec.Spec.SecurityContext.WithRunAsUser(*platformContext.RunAsUser)
                }

                if podTemplateSpec.Spec.SecurityContext.RunAsGroup == nil && platformContext.RunAsGroup != nil {
                        podTemplateSpec.Spec.SecurityContext = podTemplateSpec.Spec.SecurityContext.WithRunAsGroup(*platformContext.RunAsGroup)
                }

                if podTemplateSpec.Spec.SecurityContext.FSGroup == nil && platformContext.FSGroup != nil {
                        podTemplateSpec.Spec.SecurityContext = podTemplateSpec.Spec.SecurityContext.WithFSGroup(*platformContext.FSGroup)
                }

                if podTemplateSpec.Spec.SecurityContext.SeccompProfile == nil && platformContext.SeccompProfile != nil {
                        podTemplateSpec.Spec.SecurityContext = podTemplateSpec.Spec.SecurityContext.WithSeccompProfile(platformContext.SeccompProfile)
                }

                // For OpenShift, override certain fields even if they exist
                if platform == PlatformOpenShift {
                        if podTemplateSpec.Spec.SecurityContext.RunAsUser != nil {
                                podTemplateSpec.Spec.SecurityContext.RunAsUser = nil
                        }
                        if podTemplateSpec.Spec.SecurityContext.RunAsGroup != nil {
                                podTemplateSpec.Spec.SecurityContext.RunAsGroup = nil
                        }
                        if podTemplateSpec.Spec.SecurityContext.FSGroup != nil {
                                podTemplateSpec.Spec.SecurityContext.FSGroup = nil
                        }
                }
        }

        return podTemplateSpec
}

// getMCPContainer finds the "mcp" container in the pod template if it exists.
// Returns nil if the container doesn't exist.
func getMCPContainer(
        podTemplateSpec *corev1apply.PodTemplateSpecApplyConfiguration,
) *corev1apply.ContainerApplyConfiguration {
        // Ensure the pod template has a spec
        if podTemplateSpec.Spec == nil {
                podTemplateSpec = podTemplateSpec.WithSpec(corev1apply.PodSpec())
        }

        // Check if the container already exists
        if podTemplateSpec.Spec.Containers != nil {
                for i := range podTemplateSpec.Spec.Containers {
                        // Get a pointer to the container in the slice
                        container := &podTemplateSpec.Spec.Containers[i]
                        if container.Name != nil && *container.Name == "mcp" {
                                return container
                        }
                }
        }

        // Container doesn't exist
        return nil
}

func ensureObjectMetaApplyConfigurationExists(
        podTemplateSpec *corev1apply.PodTemplateSpecApplyConfiguration,
) *corev1apply.PodTemplateSpecApplyConfiguration {
        if podTemplateSpec.ObjectMetaApplyConfiguration == nil {
                podTemplateSpec.ObjectMetaApplyConfiguration = &metav1apply.ObjectMetaApplyConfiguration{}
        }

        return podTemplateSpec
}

// configureContainer configures a container with the given settings
//
//nolint:gocyclo // Complex but necessary for platform-aware security context configuration
func configureContainer(
        container *corev1apply.ContainerApplyConfiguration,
        image string,
        command []string,
        attachStdio bool,
        envVars []*corev1apply.EnvVarApplyConfiguration,
        platform Platform,
) {
        //nolint:gosec // G706: container name and image from config
        slog.Debug("configuring container", "name", *container.Name, "image", image)
        //nolint:gosec // G706: command args from config
        slog.Debug("container command", "args", command)
        slog.Debug("container stdio", "attach_stdio", attachStdio)
        for _, envVar := range envVars {
                //nolint:gosec // G706: env var names from config
                slog.Debug("container env var", "name", *envVar.Name, "value", *envVar.Value)
        }

        container.WithImage(image).
                WithArgs(command...).
                WithStdin(attachStdio).
                WithTTY(false).
                WithEnv(envVars...)

        // Add container security context using SecurityContextBuilder
        securityBuilder := NewSecurityContextBuilder(platform)
        if container.SecurityContext == nil {
                container.WithSecurityContext(securityBuilder.BuildContainerSecurityContextApplyConfiguration())
        } else {
                // If the container security context already exists, merge with platform-aware defaults
                platformContext := securityBuilder.BuildContainerSecurityContextApplyConfiguration()

                // Merge existing context with platform-aware settings
                if container.SecurityContext.Privileged == nil && platformContext.Privileged != nil {
                        container.SecurityContext = container.SecurityContext.WithPrivileged(*platformContext.Privileged)
                }

                if container.SecurityContext.RunAsNonRoot == nil && platformContext.RunAsNonRoot != nil {
                        container.SecurityContext = container.SecurityContext.WithRunAsNonRoot(*platformContext.RunAsNonRoot)
                }

                if container.SecurityContext.RunAsUser == nil && platformContext.RunAsUser != nil {
                        container.SecurityContext = container.SecurityContext.WithRunAsUser(*platformContext.RunAsUser)
                }

                if container.SecurityContext.RunAsGroup == nil && platformContext.RunAsGroup != nil {
                        container.SecurityContext = container.SecurityContext.WithRunAsGroup(*platformContext.RunAsGroup)
                }

                if container.SecurityContext.AllowPrivilegeEscalation == nil && platformContext.AllowPrivilegeEscalation != nil {
                        container.SecurityContext = container.SecurityContext.WithAllowPrivilegeEscalation(*platformContext.AllowPrivilegeEscalation)
                }

                if container.SecurityContext.ReadOnlyRootFilesystem == nil && platformContext.ReadOnlyRootFilesystem != nil {
                        container.SecurityContext = container.SecurityContext.WithReadOnlyRootFilesystem(*platformContext.ReadOnlyRootFilesystem)
                }

                if container.SecurityContext.SeccompProfile == nil && platformContext.SeccompProfile != nil {
                        container.SecurityContext = container.SecurityContext.WithSeccompProfile(platformContext.SeccompProfile)
                }

                if container.SecurityContext.Capabilities == nil && platformContext.Capabilities != nil {
                        container.SecurityContext = container.SecurityContext.WithCapabilities(platformContext.Capabilities)
                }

                // For OpenShift, override certain fields even if they exist
                if platform == PlatformOpenShift {
                        slog.Info("setting OpenShift security context requirements", "container", *container.Name)

                        if container.SecurityContext.RunAsUser != nil {
                                container.SecurityContext.RunAsUser = nil
                        }
                        if container.SecurityContext.RunAsGroup != nil {
                                container.SecurityContext.RunAsGroup = nil
                        }
                }
        }
}

// configureMCPContainer configures the MCP container in the pod template
func configureMCPContainer(
        podTemplateSpec *corev1apply.PodTemplateSpecApplyConfiguration,
        image string,
        command []string,
        attachStdio bool,
        envVarList []*corev1apply.EnvVarApplyConfiguration,
        transportType string,
        options *runtime.DeployWorkloadOptions,
        platform Platform,
) error {
        // Get the "mcp" container if it exists
        mcpContainer := getMCPContainer(podTemplateSpec)

        // If the container doesn't exist, create a new one
        if mcpContainer == nil {
                mcpContainer = corev1apply.Container().WithName("mcp")

                // Configure the container
                configureContainer(mcpContainer, image, command, attachStdio, envVarList, platform)

                // Configure ports if needed
                if options != nil && transportType == string(transtypes.TransportTypeSSE) {
                        var err error
                        mcpContainer, err = configureContainerPorts(mcpContainer, options)
                        if err != nil {
                                return err
                        }
                }

                // Add the fully configured container to the pod template
                podTemplateSpec.Spec.WithContainers(mcpContainer)
        } else {
                // Configure the existing container
                configureContainer(mcpContainer, image, command, attachStdio, envVarList, platform)

                // Configure ports if needed
                if options != nil && transportType == string(transtypes.TransportTypeSSE) {
                        var err error
                        _, err = configureContainerPorts(mcpContainer, options)
                        if err != nil {
                                return err
                        }
                }
        }

        return nil
}

// getCurrentNamespace returns the namespace the pod is running in.
// It tries multiple methods in order:
// 1. Reading from the service account token file (when running inside a pod)
// 2. Getting the namespace from environment variables
// 3. Getting the namespace from the current kubectl context
// 4. Falling back to "default" if all methods fail
func (c *Client) getCurrentNamespace() string {
        // If a custom namespace function is set (for testing), use it
        if c.namespaceFunc != nil {
                return c.namespaceFunc()
        }

        return k8s.GetCurrentNamespace()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package kubernetes

import (
        "log/slog"
        "os"
        "strings"
        "sync"
        "time"

        "k8s.io/apimachinery/pkg/runtime/schema"
        "k8s.io/apimachinery/pkg/util/wait"
        "k8s.io/client-go/discovery"
        "k8s.io/client-go/rest"
)

// Platform represents the Kubernetes platform type
type Platform int

const (
        // PlatformKubernetes represents standard Kubernetes
        PlatformKubernetes Platform = iota
        // PlatformOpenShift represents OpenShift
        PlatformOpenShift
)

// String returns the string representation of the Platform
func (p Platform) String() string {
        switch p {
        case PlatformKubernetes:
                return "Kubernetes"
        case PlatformOpenShift:
                return "OpenShift"
        default:
                return "Unknown"
        }
}

// PlatformDetector defines the interface for detecting the Kubernetes platform type
type PlatformDetector interface {
        DetectPlatform(config *rest.Config) (Platform, error)
}

// DefaultPlatformDetector implements PlatformDetector using the existing OpenShift detection logic
type DefaultPlatformDetector struct {
        once     sync.Once
        platform Platform
        err      error
}

// extra kinds
const (
        // defaultRetries is the number of times a resource discovery is retried
        defaultRetries = 10

        // defaultRetryInterval is the maximum interval between retries for resource discovery (used as cap in exponential backoff)
        defaultRetryInterval = 3 * time.Second
)

// DetectPlatform implements the PlatformDetector interface
func (d *DefaultPlatformDetector) DetectPlatform(config *rest.Config) (Platform, error) {
        d.once.Do(func() {
                // Check if we are running on OpenShift via environment variable override
                value, ok := os.LookupEnv("OPERATOR_OPENSHIFT")
                if ok {
                        //nolint:gosec // G706: env var value from trusted OPERATOR_OPENSHIFT
                        slog.Info("openshift set by env var", "env", "OPERATOR_OPENSHIFT", "value", value)
                        if strings.ToLower(value) == "true" {
                                d.platform = PlatformOpenShift
                        } else {
                                d.platform = PlatformKubernetes
                        }
                        return
                }

                // Check for OpenShift by attempting to discover the Route resource
                discoveryClient, err := discovery.NewDiscoveryClientForConfig(config)
                if err != nil {
                        d.err = err
                        return
                }

                var isOpenShiftResourcePresent bool
                err = wait.ExponentialBackoff(wait.Backoff{
                        Duration: time.Second,          // Initial delay
                        Factor:   2.0,                  // Backoff factor
                        Jitter:   0.1,                  // Add some randomness
                        Steps:    defaultRetries,       // Maximum number of retries
                        Cap:      defaultRetryInterval, // Maximum delay between retries
                }, func() (bool, error) {
                        isOpenShiftResourcePresent, err = discovery.IsResourceEnabled(discoveryClient,
                                schema.GroupVersionResource{
                                        Group:    "route.openshift.io",
                                        Version:  "v1",
                                        Resource: "routes",
                                })

                        if err != nil {
                                // Return false to continue retrying, don't return the error yet
                                return false, nil
                        }

                        // Success - stop retrying
                        return true, nil
                })

                if err != nil {
                        d.err = err
                        return
                }

                if isOpenShiftResourcePresent {
                        slog.Info("OpenShift detected by route resource check")
                        d.platform = PlatformOpenShift
                } else {
                        d.platform = PlatformKubernetes
                }
        })

        return d.platform, d.err
}

// NewDefaultPlatformDetector creates a new DefaultPlatformDetector
func NewDefaultPlatformDetector() PlatformDetector {
        return &DefaultPlatformDetector{}
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package kubernetes

import (
        "context"
        "fmt"
        "log/slog"
        "strings"

        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/client-go/kubernetes"

        "github.com/stacklok/toolhive/pkg/k8s"
)

// RunConfigMapReader defines the interface for reading RunConfig from ConfigMaps
// This interface allows for easy mocking in tests
//
//go:generate mockgen -destination=mocks/mock_configmap.go -package=mocks -source=configmap.go RunConfigMapReader
type RunConfigMapReader interface {
        // GetRunConfigMap retrieves the runconfig.json from a ConfigMap
        // configMapRef should be in the format "namespace/configmap-name"
        // Returns the runconfig.json content as a string
        GetRunConfigMap(ctx context.Context, configMapRef string) (string, error)
}

// ConfigMapReader implements RunConfigMapReader using real Kubernetes API
type ConfigMapReader struct {
        clientset kubernetes.Interface
}

// NewConfigMapReaderWithClient creates a new ConfigMapReader with the provided clientset
// This is useful for testing with a mock clientset
func NewConfigMapReaderWithClient(clientset kubernetes.Interface) *ConfigMapReader {
        return &ConfigMapReader{
                clientset: clientset,
        }
}

// NewConfigMapReader creates a new ConfigMapReader using in-cluster configuration
// This is the standard way to create a reader in production
// Note: This function is not unit tested as it requires a real Kubernetes cluster.
// The business logic is tested via NewConfigMapReaderWithClient with mock clients.
func NewConfigMapReader() (*ConfigMapReader, error) {
        config, err := k8s.GetConfig()
        if err != nil {
                return nil, fmt.Errorf("failed to get kubernetes config: %w", err)
        }

        clientset, err := k8s.NewClientWithConfig(config)
        if err != nil {
                return nil, fmt.Errorf("failed to create Kubernetes client: %w", err)
        }

        return NewConfigMapReaderWithClient(clientset), nil
}

// GetRunConfigMap retrieves the runconfig.json from a ConfigMap
func (c *ConfigMapReader) GetRunConfigMap(ctx context.Context, configMapRef string) (string, error) {
        // Parse the ConfigMap reference
        namespace, name, err := parseConfigMapRef(configMapRef)
        if err != nil {
                return "", fmt.Errorf("invalid configmap reference: %w", err)
        }

        slog.Info("loading runconfig.json from ConfigMap", "namespace", namespace, "name", name)

        // Get the ConfigMap
        configMap, err := c.clientset.CoreV1().ConfigMaps(namespace).Get(ctx, name, metav1.GetOptions{})
        if err != nil {
                return "", fmt.Errorf("failed to get ConfigMap '%s/%s': %w", namespace, name, err)
        }

        // Get the runconfig.json data
        data, ok := configMap.Data["runconfig.json"]
        if !ok {
                return "", fmt.Errorf("ConfigMap '%s/%s' does not contain 'runconfig.json' key", namespace, name)
        }

        slog.Info("successfully loaded runconfig.json from ConfigMap", "bytes", len(data), "namespace", namespace, "name", name)

        return data, nil
}

// parseConfigMapRef parses a ConfigMap reference in the format "namespace/configmap-name"
func parseConfigMapRef(ref string) (namespace, name string, err error) {
        parts := strings.SplitN(ref, "/", 2)
        if len(parts) != 2 {
                return "", "", fmt.Errorf("expected format 'namespace/configmap-name', got '%s'", ref)
        }

        namespace = strings.TrimSpace(parts[0])
        name = strings.TrimSpace(parts[1])

        if namespace == "" {
                return "", "", fmt.Errorf("namespace cannot be empty")
        }
        if name == "" {
                return "", "", fmt.Errorf("configmap name cannot be empty")
        }

        return namespace, name, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package kubernetes

import (
        "context"

        "github.com/stacklok/toolhive/pkg/container/runtime"
)

func init() {
        runtime.RegisterRuntime(&runtime.Info{
                Name:     RuntimeName,
                Priority: 200,
                Initializer: func(ctx context.Context) (runtime.Runtime, error) {
                        return NewClient(ctx)
                },
                AutoDetector: func() bool {
                        return runtime.IsKubernetesRuntime()
                },
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package kubernetes

import (
        "log/slog"

        corev1 "k8s.io/api/core/v1"
        corev1apply "k8s.io/client-go/applyconfigurations/core/v1"
        "k8s.io/utils/ptr"
)

// SecurityContextBuilder provides platform-aware security context configuration
type SecurityContextBuilder struct {
        platform Platform
}

// NewSecurityContextBuilder creates a new SecurityContextBuilder for the given platform
func NewSecurityContextBuilder(platform Platform) *SecurityContextBuilder {
        return &SecurityContextBuilder{
                platform: platform,
        }
}

// BuildPodSecurityContext creates a platform-appropriate pod security context
func (b *SecurityContextBuilder) BuildPodSecurityContext() *corev1.PodSecurityContext {
        // Start with base security context. The seccomp profile is set unconditionally
        // so the pod satisfies the restricted Pod Security Standard on all platforms.
        podSecurityContext := &corev1.PodSecurityContext{
                RunAsNonRoot: ptr.To(true),
                RunAsUser:    ptr.To(int64(1000)),
                RunAsGroup:   ptr.To(int64(1000)),
                FSGroup:      ptr.To(int64(1000)),
                SeccompProfile: &corev1.SeccompProfile{
                        Type: corev1.SeccompProfileTypeRuntimeDefault,
                },
        }

        // Apply platform-specific modifications
        if b.platform == PlatformOpenShift {
                slog.Info("configuring pod security context for OpenShift")
                // OpenShift uses Security Context Constraints (SCCs) to manage user/group assignments
                // Setting these to nil allows OpenShift to assign them dynamically
                podSecurityContext.RunAsUser = nil
                podSecurityContext.RunAsGroup = nil
                podSecurityContext.FSGroup = nil
        } else {
                slog.Info("configuring pod security context for Kubernetes")
        }

        return podSecurityContext
}

// BuildContainerSecurityContext creates a platform-appropriate container security context
func (b *SecurityContextBuilder) BuildContainerSecurityContext() *corev1.SecurityContext {
        // Start with base security context. The seccomp profile and dropped capabilities
        // are set unconditionally so the container satisfies the restricted Pod Security
        // Standard on all platforms.
        containerSecurityContext := &corev1.SecurityContext{
                Privileged:               ptr.To(false),
                RunAsNonRoot:             ptr.To(true),
                RunAsUser:                ptr.To(int64(1000)),
                RunAsGroup:               ptr.To(int64(1000)),
                AllowPrivilegeEscalation: ptr.To(false),
                ReadOnlyRootFilesystem:   ptr.To(true),
                SeccompProfile: &corev1.SeccompProfile{
                        Type: corev1.SeccompProfileTypeRuntimeDefault,
                },
                Capabilities: &corev1.Capabilities{
                        Drop: []corev1.Capability{"ALL"},
                },
        }

        // Apply platform-specific modifications
        if b.platform == PlatformOpenShift {
                slog.Info("configuring container security context for OpenShift")
                // OpenShift uses Security Context Constraints (SCCs) to manage user/group assignments
                // Setting these to nil allows OpenShift to assign them dynamically
                containerSecurityContext.RunAsUser = nil
                containerSecurityContext.RunAsGroup = nil
        } else {
                slog.Info("configuring container security context for Kubernetes")
        }

        return containerSecurityContext
}

// BuildPodSecurityContextApplyConfiguration creates a platform-appropriate pod security context
// using the ApplyConfiguration types used by the client
func (b *SecurityContextBuilder) BuildPodSecurityContextApplyConfiguration() *corev1apply.PodSecurityContextApplyConfiguration {
        baseContext := b.BuildPodSecurityContext()

        applyConfig := corev1apply.PodSecurityContext()

        if baseContext.RunAsNonRoot != nil {
                applyConfig = applyConfig.WithRunAsNonRoot(*baseContext.RunAsNonRoot)
        }

        if baseContext.RunAsUser != nil {
                applyConfig = applyConfig.WithRunAsUser(*baseContext.RunAsUser)
        }

        if baseContext.RunAsGroup != nil {
                applyConfig = applyConfig.WithRunAsGroup(*baseContext.RunAsGroup)
        }

        if baseContext.FSGroup != nil {
                applyConfig = applyConfig.WithFSGroup(*baseContext.FSGroup)
        }

        if baseContext.SeccompProfile != nil {
                applyConfig = applyConfig.WithSeccompProfile(
                        corev1apply.SeccompProfile().WithType(baseContext.SeccompProfile.Type))
        }

        return applyConfig
}

// BuildContainerSecurityContextApplyConfiguration creates a platform-appropriate container security context
// using the ApplyConfiguration types used by the client
func (b *SecurityContextBuilder) BuildContainerSecurityContextApplyConfiguration() *corev1apply.SecurityContextApplyConfiguration { //nolint:lll
        baseContext := b.BuildContainerSecurityContext()

        applyConfig := corev1apply.SecurityContext()

        if baseContext.Privileged != nil {
                applyConfig = applyConfig.WithPrivileged(*baseContext.Privileged)
        }

        if baseContext.RunAsNonRoot != nil {
                applyConfig = applyConfig.WithRunAsNonRoot(*baseContext.RunAsNonRoot)
        }

        if baseContext.RunAsUser != nil {
                applyConfig = applyConfig.WithRunAsUser(*baseContext.RunAsUser)
        }

        if baseContext.RunAsGroup != nil {
                applyConfig = applyConfig.WithRunAsGroup(*baseContext.RunAsGroup)
        }

        if baseContext.AllowPrivilegeEscalation != nil {
                applyConfig = applyConfig.WithAllowPrivilegeEscalation(*baseContext.AllowPrivilegeEscalation)
        }

        if baseContext.ReadOnlyRootFilesystem != nil {
                applyConfig = applyConfig.WithReadOnlyRootFilesystem(*baseContext.ReadOnlyRootFilesystem)
        }

        if baseContext.SeccompProfile != nil {
                applyConfig = applyConfig.WithSeccompProfile(
                        corev1apply.SeccompProfile().WithType(baseContext.SeccompProfile.Type))
        }

        if baseContext.Capabilities != nil {
                capabilities := corev1apply.Capabilities()
                if len(baseContext.Capabilities.Drop) > 0 {
                        capabilities = capabilities.WithDrop(baseContext.Capabilities.Drop...)
                }
                if len(baseContext.Capabilities.Add) > 0 {
                        capabilities = capabilities.WithAdd(baseContext.Capabilities.Add...)
                }
                applyConfig = applyConfig.WithCapabilities(capabilities)
        }

        return applyConfig
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package container

import (
        "fmt"
        "strings"
        "time"
)

// GetOrGenerateContainerName generates a container name if not provided.
// It returns both the container name and the base name.
// If containerName is not empty, it will be used as both the container name and base name.
// If containerName is empty, a name will be generated based on the image.
func GetOrGenerateContainerName(containerName, image string) (string, string) {
        var baseName string

        if containerName == "" {
                // Generate a container name from the image
                baseName = generateContainerBaseName(image)
                containerName = appendTimestamp(baseName)
        } else {
                // If container name is provided, use it as the base name
                baseName = containerName
        }

        return containerName, baseName
}

// generateContainerBaseName generates a base name for a container from the image name
func generateContainerBaseName(image string) string {
        // Find last '/' and last ':' to distinguish port from tag
        lastSlash := strings.LastIndex(image, "/")
        lastColon := strings.LastIndex(image, ":")

        imageWithoutTag := image
        if lastColon > lastSlash {
                imageWithoutTag = image[:lastColon]
        }

        // Split by '/'
        parts := strings.Split(imageWithoutTag, "/")
        var registryOrNamespace, name string
        switch len(parts) {
        case 1:
                name = parts[0]
        case 2:
                registryOrNamespace = parts[0]
                name = parts[1]
        default:
                registryOrNamespace = parts[len(parts)-2]
                name = parts[len(parts)-1]
        }
        // Strip the port from registryOrNamespace if it looks like host:port
        if registryOrNamespace != "" && strings.Contains(registryOrNamespace, ":") {
                registryOrNamespace = strings.SplitN(registryOrNamespace, ":", 2)[0]
        }

        // Construct the base name using the sanitized registryOrNamespace and name
        var base string
        if registryOrNamespace != "" {
                base = registryOrNamespace + "-" + name
        } else {
                base = name
        }

        // Sanitize: allow alphanumeric and dashes
        var sanitizedName strings.Builder
        for _, c := range base {
                if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '-' {
                        sanitizedName.WriteRune(c)
                } else {
                        sanitizedName.WriteRune('-')
                }
        }

        return sanitizedName.String()
}

// appendTimestamp appends a timestamp to a base name to ensure uniqueness
func appendTimestamp(baseName string) string {
        timestamp := time.Now().Unix()
        return fmt.Sprintf("%s-%d", baseName, timestamp)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runtime

import (
        "errors"
        "fmt"
        "net/http"

        "github.com/stacklok/toolhive-core/httperr"
)

// Error types for container operations
var (
        // ErrContainerNotFound is returned when a container is not found
        ErrContainerNotFound = httperr.WithCode(fmt.Errorf("container not found"), http.StatusNotFound)

        // ErrContainerNotRunning is returned when a container is not running
        ErrContainerNotRunning = httperr.WithCode(fmt.Errorf("container not running"), http.StatusBadRequest)

        // ErrContainerExited is returned when a container has exited unexpectedly
        ErrContainerExited = httperr.WithCode(fmt.Errorf("container exited unexpectedly"), http.StatusBadRequest)

        // ErrContainerRestarted is returned when a container has been restarted
        // (e.g., by Docker restart policy). The container is still running.
        // This is distinct from ErrContainerExited.
        ErrContainerRestarted = httperr.WithCode(fmt.Errorf("container restarted"), http.StatusBadRequest)

        // ErrContainerRemoved is returned when a container has been removed
        ErrContainerRemoved = httperr.WithCode(fmt.Errorf("container removed"), http.StatusBadRequest)
)

// ContainerError represents an error related to container operations
type ContainerError struct {
        // Err is the underlying error
        Err error
        // ContainerID is the ID of the container
        ContainerID string
        // Message is an optional error message
        Message string
}

// Error returns the error message
func (e *ContainerError) Error() string {
        if e.Message != "" {
                if e.ContainerID != "" {
                        return fmt.Sprintf("%s: %s (container: %s)", e.Err, e.Message, e.ContainerID)
                }
                return fmt.Sprintf("%s: %s", e.Err, e.Message)
        }

        if e.ContainerID != "" {
                return fmt.Sprintf("%s (container: %s)", e.Err, e.ContainerID)
        }

        return e.Err.Error()
}

// Unwrap returns the underlying error
func (e *ContainerError) Unwrap() error {
        return e.Err
}

// NewContainerError creates a new container error
func NewContainerError(err error, containerID, message string) *ContainerError {
        return &ContainerError{
                Err:         err,
                ContainerID: containerID,
                Message:     message,
        }
}

// IsContainerNotFound checks if the error is a container not found error
func IsContainerNotFound(err error) bool {
        return errors.Is(err, ErrContainerNotFound)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runtime

import (
        "context"
        "fmt"
        "sync"
        "time"
)

// WorkloadMonitor watches a workload's state and reports when it exits
type WorkloadMonitor struct {
        runtime          Runtime
        containerName    string
        stopCh           chan struct{}
        errorCh          chan error
        wg               sync.WaitGroup
        running          bool
        mutex            sync.Mutex
        initialStartTime time.Time // Track container start time to detect restarts
}

// NewMonitor creates a new workload monitor
func NewMonitor(rt Runtime, containerName string) Monitor {
        return &WorkloadMonitor{
                runtime:       rt,
                containerName: containerName,
                stopCh:        make(chan struct{}),
                errorCh:       make(chan error, 1), // Buffered to prevent blocking
        }
}

// StartMonitoring starts monitoring the workload
func (m *WorkloadMonitor) StartMonitoring(ctx context.Context) (<-chan error, error) {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        if m.running {
                return m.errorCh, nil // Already monitoring
        }

        // Check if the workload exists and is running
        running, err := m.runtime.IsWorkloadRunning(ctx, m.containerName)
        if err != nil {
                return nil, err
        }
        if !running {
                return nil, NewContainerError(ErrContainerNotRunning, m.containerName, "container is not running")
        }

        // Get initial container info to track start time
        info, err := m.runtime.GetWorkloadInfo(ctx, m.containerName)
        if err != nil {
                return nil, NewContainerError(err, m.containerName, fmt.Sprintf("failed to get container info: %v", err))
        }
        m.initialStartTime = info.StartedAt

        m.running = true
        m.wg.Add(1)

        // Start monitoring in a goroutine
        go m.monitor(ctx)

        return m.errorCh, nil
}

// StopMonitoring stops monitoring the workload
func (m *WorkloadMonitor) StopMonitoring() {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        if !m.running {
                return // Not monitoring
        }

        close(m.stopCh)
        m.wg.Wait()
        m.running = false
}

// monitor checks the workload status periodically
func (m *WorkloadMonitor) monitor(ctx context.Context) {
        defer m.wg.Done()

        // Check interval
        checkInterval := 5 * time.Second

        ticker := time.NewTicker(checkInterval)
        defer ticker.Stop()

        for {
                select {
                case <-ctx.Done():
                        // Context canceled
                        return
                case <-m.stopCh:
                        // Monitoring stopped
                        return
                case <-ticker.C:
                        // Check if the container is still running
                        // Create a short timeout context for this check, derived from parent
                        // to respect parent cancellation while having its own timeout
                        checkCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
                        running, err := m.runtime.IsWorkloadRunning(checkCtx, m.containerName)
                        cancel() // Always cancel the context to avoid leaks
                        if err != nil {
                                // If the container is not found, it has been removed
                                if IsContainerNotFound(err) {
                                        removeErr := NewContainerError(
                                                ErrContainerRemoved,
                                                m.containerName,
                                                fmt.Sprintf("Container %s not found, it has been removed", m.containerName),
                                        )
                                        m.errorCh <- removeErr
                                        return
                                }

                                // For other errors, log and continue
                                continue
                        }

                        if !running {
                                // Container has exited, get logs and info
                                // Create a short timeout context for these operations, derived from parent
                                infoCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
                                // Get last 50 lines of logs for error reporting
                                logs, _ := m.runtime.GetWorkloadLogs(infoCtx, m.containerName, false, 50)
                                info, _ := m.runtime.GetWorkloadInfo(infoCtx, m.containerName)
                                cancel() // Always cancel the context to avoid leaks

                                exitErr := NewContainerError(
                                        ErrContainerExited,
                                        m.containerName,
                                        fmt.Sprintf("Container %s (%s) exited unexpectedly. Status: %s. Last logs:\n%s",
                                                m.containerName, m.containerName, info.Status, logs),
                                )
                                m.errorCh <- exitErr
                                return
                        }

                        // Container is running - check if it was restarted (different start time)
                        // Derive timeout from parent context to respect cancellation
                        infoCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
                        info, err := m.runtime.GetWorkloadInfo(infoCtx, m.containerName)
                        cancel()
                        if err == nil && !info.StartedAt.IsZero() && !info.StartedAt.Equal(m.initialStartTime) {
                                // Container was restarted (has a different start time)
                                restartErr := NewContainerError(
                                        ErrContainerRestarted,
                                        m.containerName,
                                        fmt.Sprintf("Container %s was restarted (start time changed from %s to %s)",
                                                m.containerName, m.initialStartTime.Format(time.RFC3339), info.StartedAt.Format(time.RFC3339)),
                                )
                                m.errorCh <- restartErr
                                return
                        }
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runtime

import (
        "fmt"
        "sort"
        "sync"
)

// Registry holds a set of registered runtimes. It is safe for concurrent use.
// Tests should create isolated instances via NewRegistry() rather than using
// the DefaultRegistry, so they can run fully in parallel.
type Registry struct {
        mu       sync.RWMutex
        runtimes map[string]*Info
}

// NewRegistry creates a new, empty Registry.
func NewRegistry() *Registry {
        return &Registry{
                runtimes: make(map[string]*Info),
        }
}

// Register adds a runtime to the registry.
// It panics if info is nil, has an empty name, has a nil initializer,
// or if a runtime with the same name is already registered.
func (r *Registry) Register(info *Info) {
        if info == nil {
                panic("runtime info cannot be nil")
        }
        if info.Name == "" {
                panic("runtime name cannot be empty")
        }
        if info.Initializer == nil {
                panic("runtime initializer cannot be nil")
        }
        if info.Priority < 0 {
                panic("runtime priority must be non-negative")
        }

        r.mu.Lock()
        defer r.mu.Unlock()

        if _, exists := r.runtimes[info.Name]; exists {
                panic(fmt.Sprintf("runtime already registered: %s", info.Name))
        }
        r.runtimes[info.Name] = info
}

// Get returns a copy of the Info for the given name, or nil if not found.
func (r *Registry) Get(name string) *Info {
        r.mu.RLock()
        defer r.mu.RUnlock()

        info, ok := r.runtimes[name]
        if !ok {
                return nil
        }
        cp := *info
        return &cp
}

// IsRegistered returns true if a runtime with the given name is registered.
func (r *Registry) IsRegistered(name string) bool {
        r.mu.RLock()
        defer r.mu.RUnlock()

        _, exists := r.runtimes[name]
        return exists
}

// All returns copies of all registered runtimes as a slice.
func (r *Registry) All() []*Info {
        r.mu.RLock()
        defer r.mu.RUnlock()

        result := make([]*Info, 0, len(r.runtimes))
        for _, info := range r.runtimes {
                cp := *info
                result = append(result, &cp)
        }
        return result
}

// ByPriority returns all registered runtimes sorted by priority (ascending).
// When two runtimes share the same priority, they are sorted by name for
// deterministic ordering.
func (r *Registry) ByPriority() []*Info {
        runtimes := r.All()
        sort.SliceStable(runtimes, func(i, j int) bool {
                if runtimes[i].Priority != runtimes[j].Priority {
                        return runtimes[i].Priority < runtimes[j].Priority
                }
                return runtimes[i].Name < runtimes[j].Name
        })
        return runtimes
}

// DefaultRegistry is the global registry used by init()-time self-registration.
// Production code should use this (typically via the package-level convenience
// functions). Tests should create isolated registries with NewRegistry().
var DefaultRegistry = NewRegistry()

// RegisterRuntime registers an Info in the DefaultRegistry.
// This is typically called from an init() function in each runtime package.
func RegisterRuntime(info *Info) {
        DefaultRegistry.Register(info)
}

// RegisteredRuntimesByPriority returns all runtimes from the DefaultRegistry
// sorted by priority.
func RegisteredRuntimesByPriority() []*Info {
        return DefaultRegistry.ByPriority()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package runtime provides interfaces and types for container runtimes,
// including creating, starting, stopping, and monitoring containers.
package runtime

import (
        "context"
        "errors"
        "io"
        "net/http"
        "strings"
        "time"

        "github.com/stacklok/toolhive-core/env"
        "github.com/stacklok/toolhive-core/httperr"
        "github.com/stacklok/toolhive-core/permissions"
        "github.com/stacklok/toolhive/pkg/ignore"
)

// WorkloadStatus is an enum representing the possible statuses of a workload.
type WorkloadStatus string

const (
        // WorkloadStatusRunning indicates that the workload is currently running.
        WorkloadStatusRunning WorkloadStatus = "running"
        // WorkloadStatusStopped indicates that the workload is stopped.
        WorkloadStatusStopped WorkloadStatus = "stopped"
        // WorkloadStatusError indicates that the workload has encountered an error
        // during creation/stop/restart/delete.
        WorkloadStatusError WorkloadStatus = "error"
        // WorkloadStatusStarting indicates that the workload is being started.
        WorkloadStatusStarting WorkloadStatus = "starting"
        // WorkloadStatusStopping indicates that the workload is being stopped.
        WorkloadStatusStopping WorkloadStatus = "stopping"
        // WorkloadStatusUnhealthy indicates that the workload is running, but is
        // in an inconsistent state which prevents normal operation.
        WorkloadStatusUnhealthy WorkloadStatus = "unhealthy"
        // WorkloadStatusRemoving indicates that the workload is being removed.
        WorkloadStatusRemoving WorkloadStatus = "removing"
        // WorkloadStatusUnknown indicates that the workload status is unknown.
        WorkloadStatusUnknown WorkloadStatus = "unknown"
        // WorkloadStatusUnauthenticated indicates that the workload is running but
        // cannot authenticate with the remote MCP server (e.g., expired refresh token).
        WorkloadStatusUnauthenticated WorkloadStatus = "unauthenticated"
        // WorkloadStatusAuthRetrying indicates that background authentication
        // refresh is failing transiently and the monitor is still retrying.
        // The workload may recover on its own (→ Running) or be marked
        // unauthenticated when the configured ceiling is exceeded
        // (→ Unauthenticated). Hot requests fail fast with 503+Retry-After
        // while in this state.
        WorkloadStatusAuthRetrying WorkloadStatus = "auth_retrying"
        // WorkloadStatusPolicyStopped indicates that the workload was stopped by
        // policy enforcement. The StatusContext field carries the human-readable reason.
        WorkloadStatusPolicyStopped WorkloadStatus = "policy_stopped"
)

// ContainerInfo represents information about a container
// TODO: Consider merging this with workloads.Workload
type ContainerInfo struct {
        // Name is the container name
        Name string
        // Image is the container image
        Image string
        // Status is the container status
        // This is usually some human-readable context.
        Status string
        // State is the container state
        State WorkloadStatus
        // Created is the container creation timestamp
        Created time.Time
        // StartedAt is when the container was last started (changes on restart)
        StartedAt time.Time
        // Labels is the container labels
        Labels map[string]string
        // Ports is the container port mappings
        Ports []PortMapping
}

// IsRunning returns true if the container is currently running.
func (c *ContainerInfo) IsRunning() bool {
        return c.State == WorkloadStatusRunning
}

// PortMapping represents a port mapping for a container
type PortMapping struct {
        // ContainerPort is the port inside the container
        ContainerPort int
        // HostPort is the port on the host
        HostPort int
        // Protocol is the protocol (tcp, udp)
        Protocol string
}

// Deployer contains the methods to start and stop a workload.
// This is intended as a subset of the Runtime interface for
// the runner code.
type Deployer interface {
        // DeployWorkload creates and starts a complete workload deployment.
        // This includes the primary container, any required sidecars, networking setup,
        // volume mounts, and service configuration. The workload is started as part
        // of this operation, making it immediately available for use.
        //
        // Parameters:
        // - image: The primary container image to deploy
        // - name: The workload name (used for identification and networking)
        // - command: Command to run in the primary container
        // - envVars: Environment variables for the primary container
        // - labels: Labels to apply to all workload components
        // - permissionProfile: Security and permission configuration
        // - transportType: Communication transport (sse, stdio, etc.)
        // - options: Additional deployment options (ports, sidecars, etc.)
        //
        // Returns the workload ID for subsequent operations.
        // If options is nil, default options will be used.
        //todo: make args a struct to reduce number of args (linter going crazy)
        DeployWorkload(
                ctx context.Context,
                image, name string,
                command []string,
                envVars, labels map[string]string,
                permissionProfile *permissions.Profile,
                transportType string,
                options *DeployWorkloadOptions,
                isolateNetwork bool,
        ) (int, error)

        // StopWorkload gracefully stops a running workload and all its components.
        // This includes stopping the primary container, sidecars, and cleaning up
        // any associated network resources. The workload remains available for restart.
        StopWorkload(ctx context.Context, workloadName string) error

        // AttachToWorkload establishes a direct connection to the primary container
        // of the workload for interactive communication. This is typically used
        // for stdio transport where direct input/output streaming is required.
        AttachToWorkload(ctx context.Context, workloadName string) (io.WriteCloser, io.ReadCloser, error)

        // IsWorkloadRunning checks if a workload is currently running and healthy.
        // This verifies that the primary container is running and that any
        // required sidecars are also operational.
        IsWorkloadRunning(ctx context.Context, workloadName string) (bool, error)
}

// Runtime defines the interface for container runtimes that manage workloads.
//
// A workload in ToolHive represents a complete deployment unit that may consist of:
// - Primary MCP server container
// - Sidecar containers (for logging, monitoring, proxying, etc.)
// - Network configurations and port mappings
// - Volume mounts and storage
// - Service discovery and load balancing components
// - Security policies and permission profiles
//
// This is a departure from simple container management, as modern deployments
// often require orchestrating multiple interconnected components that work
// together to provide a complete service.
//
//go:generate mockgen -destination=mocks/mock_runtime.go -package=mocks -source=types.go Runtime
type Runtime interface {
        Deployer

        // ListWorkloads lists all deployed workloads managed by this runtime.
        // Returns information about each workload including its components,
        // status, and resource usage.
        ListWorkloads(ctx context.Context) ([]ContainerInfo, error)

        // RemoveWorkload completely removes a workload and all its components.
        // This includes removing containers, cleaning up networks, volumes,
        // and any other resources associated with the workload. This operation
        // is irreversible.
        RemoveWorkload(ctx context.Context, workloadName string) error

        // GetWorkloadLogs retrieves logs from the primary container of the workload.
        // If follow is true, the logs will be streamed continuously.
        // The lines parameter specifies the maximum number of lines to return from the end of the logs.
        // If lines is 0, all logs are returned.
        // For workloads with multiple containers, this returns logs from the
        // main MCP server container.
        GetWorkloadLogs(ctx context.Context, workloadName string, follow bool, lines int) (string, error)

        // GetWorkloadInfo retrieves detailed information about a workload.
        // This includes status, resource usage, network configuration,
        // and metadata about all components in the workload.
        GetWorkloadInfo(ctx context.Context, workloadName string) (ContainerInfo, error)

        // IsRunning checks the health of the container runtime.
        // This is used to verify that the runtime is operational and can manage workloads.
        IsRunning(ctx context.Context) error
}

// Monitor defines the interface for container monitoring
type Monitor interface {
        // StartMonitoring starts monitoring the container
        // Returns a channel that will receive an error if the container exits unexpectedly
        StartMonitoring(ctx context.Context) (<-chan error, error)

        // StopMonitoring stops monitoring the container
        StopMonitoring()
}

// Type represents the type of container runtime
type Type string

const (
        // TypePodman represents the Podman runtime
        TypePodman Type = "podman"
        // TypeDocker represents the Docker runtime
        TypeDocker Type = "docker"
        // TypeKubernetes represents the Kubernetes runtime
        TypeKubernetes Type = "kubernetes"
        // TypeColima represents the Colima runtime
        TypeColima Type = "colima"
)

// MountType represents the type of mount
type MountType string

const (
        // MountTypeBind represents a bind mount
        MountTypeBind MountType = "bind"
        // MountTypeTmpfs represents a tmpfs mount
        MountTypeTmpfs MountType = "tmpfs"
)

// String returns the string representation of the mount type
func (mt MountType) String() string {
        return string(mt)
}

// PermissionConfig represents container permission configuration
type PermissionConfig struct {
        // Mounts is the list of volume mounts
        Mounts []Mount
        // NetworkMode is the network mode
        NetworkMode string
        // CapDrop is the list of capabilities to drop
        CapDrop []string
        // CapAdd is the list of capabilities to add
        CapAdd []string
        // SecurityOpt is the list of security options
        SecurityOpt []string
        // Privileged indicates whether the container should run in privileged mode
        Privileged bool
}

// DeployWorkloadOptions represents configuration options for deploying a workload.
// These options control how the workload is deployed, including networking,
// platform-specific configurations, and communication settings.
type DeployWorkloadOptions struct {
        // ExposedPorts is a map of container ports to expose
        // The key is in the format "port/protocol" (e.g., "8080/tcp")
        // The value is an empty struct (not used)
        ExposedPorts map[string]struct{}

        // PortBindings is a map of container ports to host ports
        // The key is in the format "port/protocol" (e.g., "8080/tcp")
        // The value is a slice of host port bindings
        PortBindings map[string][]PortBinding

        // AttachStdio indicates whether to attach stdin/stdout/stderr
        // This is typically set to true for stdio transport
        AttachStdio bool

        // K8sPodTemplatePatch is a JSON string to patch the Kubernetes pod template
        // Only applicable when using Kubernetes runtime
        K8sPodTemplatePatch string

        // MCPServiceName is the name of the Kubernetes ClusterIP service used as the
        // proxy-runner target for MCP server workloads.
        // Only applicable when using Kubernetes runtime with SSE or streamable-http transport.
        MCPServiceName string

        // IgnoreConfig contains configuration for ignore patterns and tmpfs overlays
        // Used to filter bind mount contents by hiding sensitive files
        IgnoreConfig *ignore.Config

        // ScalingConfig contains scaling-related configuration for the workload.
        // Only applicable to Kubernetes deployments.
        ScalingConfig *ScalingConfig

        // AllowDockerGateway permits outbound connections to Docker/Podman host addresses
        // (all daemon bridge gateways, runtime host-alias IPs, host.docker.internal,
        // gateway.docker.internal, and host.containers.internal). These are blocked by
        // default in the egress proxy even when InsecureAllowAll is set.
        // Only applicable to Docker/Podman deployments with network isolation enabled.
        // Addresses are snapshotted at deployment; redeploy after runtime network changes.
        // Gateway access is port-independent: it ignores the permission profile's
        // allowed ports, so once enabled the gateway is reachable on any port.
        AllowDockerGateway bool

        // RunConfigMCPServerGeneration is the monotonic version stamp from the source RunConfig
        // (the MCPServer .metadata.generation). K8s runtime uses it to refuse apply when the
        // StatefulSet is already stamped with a strictly greater value.
        RunConfigMCPServerGeneration int64
}

// ScalingConfig holds horizontal-scaling knobs threaded from RunConfig down to
// the Kubernetes client. Fields mirror runner.ScalingConfig but are defined here
// to avoid an import cycle between pkg/runner and pkg/container/runtime.
type ScalingConfig struct {
        // BackendReplicas is the desired StatefulSet replica count.
        // When nil, the replicas field is omitted from the server-side apply spec so that
        // HPA or kubectl retains control of scaling.
        BackendReplicas *int32
}

// PortBinding represents a host port binding
type PortBinding struct {
        // HostIP is the host IP to bind to (empty for all interfaces)
        HostIP string
        // HostPort is the host port to bind to (empty for random port)
        HostPort string
}

// NewDeployWorkloadOptions creates a new DeployWorkloadOptions with default values.
// This provides a baseline configuration suitable for most workload deployments,
// with empty port mappings and standard communication settings.
func NewDeployWorkloadOptions() *DeployWorkloadOptions {
        return &DeployWorkloadOptions{
                ExposedPorts:        make(map[string]struct{}),
                PortBindings:        make(map[string][]PortBinding),
                AttachStdio:         false,
                K8sPodTemplatePatch: "",
        }
}

// Mount represents a volume mount
type Mount struct {
        // Source is the source path on the host
        Source string
        // Target is the target path in the container
        Target string
        // ReadOnly indicates if the mount is read-only
        ReadOnly bool
        // Type is the mount type (bind or tmpfs)
        Type MountType
}

// IsKubernetesRuntime checks if the current runtime is Kubernetes
// This checks the TOOLHIVE_RUNTIME environment variable first, then falls back to
// checking if we're in a Kubernetes environment
func IsKubernetesRuntime() bool {
        return IsKubernetesRuntimeWithEnv(&env.OSReader{})
}

// IsKubernetesRuntimeWithEnv checks if the current runtime is Kubernetes using the provided environment reader.
// This allows for dependency injection of environment variable access for testing.
func IsKubernetesRuntimeWithEnv(envReader env.Reader) bool {
        // Check if TOOLHIVE_RUNTIME is explicitly set to kubernetes
        if runtimeEnv := strings.TrimSpace(envReader.Getenv("TOOLHIVE_RUNTIME")); runtimeEnv == "kubernetes" {
                return true
        }

        // Fall back to checking if we're in a Kubernetes environment
        return envReader.Getenv("KUBERNETES_SERVICE_HOST") != ""
}

// Initializer is a function that creates a new runtime instance.
type Initializer func(ctx context.Context) (Runtime, error)

// Info contains metadata about a runtime, including its initializer
// and auto-detection logic. This is used by the global runtime registry
// to manage available runtimes.
type Info struct {
        // Name is the runtime name (e.g., "docker", "kubernetes")
        Name string
        // Priority determines the order in which runtimes are tried during auto-detection.
        // Lower values are tried first. Convention:
        //   100 = Docker (default local runtime)
        //   200 = Kubernetes (detected via env vars)
        //   50-99 = External runtimes preferred before Docker
        //   101-199 = External runtimes between Docker and Kubernetes
        //   201+ = External fallback runtimes
        Priority int
        // Initializer is the function to create the runtime instance
        Initializer Initializer
        // AutoDetector is an optional function to detect if this runtime is available.
        // If nil, the runtime is always considered available.
        AutoDetector func() bool
}

// Common errors
var (
        // ErrWorkloadNotFound indicates that the specified workload was not found.
        ErrWorkloadNotFound = httperr.WithCode(
                errors.New("workload not found"),
                http.StatusNotFound,
        )
)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package templates

import (
        "errors"
        "fmt"
        "maps"
        "regexp"
        "slices"
        "strings"

        nameref "github.com/google/go-containerregistry/pkg/name"
)

// maxPackageNameLength is the maximum allowed length for a package name.
const maxPackageNameLength = 128

// packageNamePattern matches valid Alpine/Debian package names.
// Must start with an alphanumeric character, followed by alphanumeric characters,
// dots, underscores, plus signs, or hyphens.
var packageNamePattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._+\-]*$`)

// buildWithPattern matches a safe subset of PEP 508 requirement specifiers for
// BuildWith entries: a package name (optionally with extras) followed by version
// specifiers, e.g. "mcp<2", "mcp>=1.27,<2", "pkg[extra]==1.2.*", "foo~=1.4".
// The allowlist deliberately excludes quotes, backticks, dollar signs,
// semicolons, parentheses, and backslashes: entries are interpolated into a
// single-quoted shell word inside a Dockerfile RUN instruction, so anything
// that could close the quote or expand in shell context is rejected.
var buildWithPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._+\[\],<>=!~* -]*$`)

// envKeyPattern matches valid environment variable names for RuntimeEnv.
// Must start with an uppercase letter, followed by uppercase letters, numbers, or underscores.
var envKeyPattern = regexp.MustCompile(`^[A-Z][A-Z0-9_]*$`)

// reservedRuntimeEnvKeys lists environment variable names that RuntimeEnv must
// not override, either because the generated Dockerfile sets them itself
// (e.g. PATH) or because overriding them could destabilize the runtime image.
var reservedRuntimeEnvKeys = map[string]bool{
        "PATH": true, "HOME": true, "USER": true, "SHELL": true, "PWD": true,
        "HOSTNAME": true, "TERM": true, "LANG": true, "LC_ALL": true,
        "LD_PRELOAD": true, "LD_LIBRARY_PATH": true,
}

// runtimeEnvDangerousValuePatterns lists substrings that must not appear in a
// RuntimeEnv value. Values are interpolated verbatim into a Dockerfile ENV
// line (ENV KEY="value") with no shell-escaping, so these characters could
// break out of the quoted value and inject arbitrary Dockerfile/shell content.
var runtimeEnvDangerousValuePatterns = []string{
        "`", "$(", "${", "\\", "\n", "\r", "\"", ";", "&&", "||", "|", ">", "<",
}

// RuntimeConfig defines the base images and versions for a specific runtime
type RuntimeConfig struct {
        // BuilderImage is the full image reference for the builder stage.
        // An empty string signals "use the default for this transport type" during config merging.
        // Examples: "golang:1.27-alpine", "node:24-alpine", "python:3.14-slim"
        BuilderImage string `json:"builder_image" yaml:"builder_image"`

        // AdditionalPackages lists extra packages to install in the builder and
        // runtime stages.
        // Examples for Alpine: ["git", "make", "gcc"]
        // Examples for Debian: ["git", "build-essential"]
        AdditionalPackages []string `json:"additional_packages,omitempty" yaml:"additional_packages,omitempty"`

        // BuildWith lists build-time dependency constraints, interpreted per
        // package ecosystem. For uvx:// builds these are PEP 508 requirement
        // specifiers passed to `uv tool install --with`, used to constrain
        // transitive dependencies the package itself leaves unbounded
        // (e.g. "mcp<2"). Ecosystems without constraint support (npx://, go://)
        // reject a non-empty BuildWith at build time.
        BuildWith []string `json:"build_with,omitempty" yaml:"build_with,omitempty"`

        // RuntimeEnv contains environment variables to inject into the Dockerfile's
        // final runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv),
        // which only affects the builder stage, these variables are baked into the
        // shipped image and are present in the running container's process
        // environment at startup. Use this for values a packaged MCP server reads at
        // process start (e.g. feature flags, cache backend selection), not for
        // build-time package manager configuration.
        // Keys must be uppercase with underscores, values are validated for safety.
        RuntimeEnv map[string]string `json:"runtime_env,omitempty" yaml:"runtime_env,omitempty"`
}

// Validate checks that all RuntimeConfig fields contain safe values that cannot
// cause unexpected behavior when interpolated into Dockerfile templates.
// An empty BuilderImage is allowed because it signals "use the default for
// this transport type" during config merging.
// It returns a combined error listing all invalid fields.
func (rc *RuntimeConfig) Validate() error {
        var errs []error

        // Validate BuilderImage using go-containerregistry's ParseReference,
        // which rejects newlines, shell metacharacters, and malformed refs.
        if rc.BuilderImage != "" {
                trimmed := strings.TrimSpace(rc.BuilderImage)
                if trimmed == "" {
                        errs = append(errs, fmt.Errorf("builder_image is blank after trimming whitespace"))
                } else if _, err := nameref.ParseReference(trimmed); err != nil {
                        errs = append(errs, fmt.Errorf("invalid builder_image %q: %w", rc.BuilderImage, err))
                }
        }

        // Validate each AdditionalPackages entry against a strict allowlist regex
        // and a maximum length bound.
        for _, pkg := range rc.AdditionalPackages {
                if len(pkg) > maxPackageNameLength {
                        errs = append(errs, fmt.Errorf(
                                "package name %q exceeds maximum length of %d characters",
                                pkg, maxPackageNameLength,
                        ))
                } else if !packageNamePattern.MatchString(pkg) {
                        errs = append(errs, fmt.Errorf(
                                "invalid package name %q: must match %s",
                                pkg, packageNamePattern.String(),
                        ))
                }
        }

        // Validate each BuildWith entry against a strict allowlist so specifiers
        // cannot escape the single-quoted --with argument in the uvx Dockerfile.
        for _, spec := range rc.BuildWith {
                if len(spec) > maxPackageNameLength {
                        errs = append(errs, fmt.Errorf(
                                "build_with specifier %q exceeds maximum length of %d characters",
                                spec, maxPackageNameLength,
                        ))
                } else if !buildWithPattern.MatchString(spec) {
                        errs = append(errs, fmt.Errorf(
                                "invalid build_with specifier %q: must match %s",
                                spec, buildWithPattern.String(),
                        ))
                }
        }

        // Validate each RuntimeEnv entry to ensure keys and values are safe to
        // interpolate into a Dockerfile ENV instruction.
        for key, value := range rc.RuntimeEnv {
                if !envKeyPattern.MatchString(key) {
                        errs = append(errs, fmt.Errorf(
                                "invalid runtime env key %q: must match %s", key, envKeyPattern.String(),
                        ))
                        continue
                }
                if reservedRuntimeEnvKeys[key] {
                        errs = append(errs, fmt.Errorf("runtime env key %q is reserved and cannot be overridden", key))
                        continue
                }
                for _, pattern := range runtimeEnvDangerousValuePatterns {
                        if strings.Contains(value, pattern) {
                                errs = append(errs, fmt.Errorf(
                                        "runtime env value for key %q contains potentially dangerous characters: %q", key, pattern,
                                ))
                                break
                        }
                }
        }

        return errors.Join(errs...)
}

// Clone returns a deep copy of rc, safe for the caller to mutate without
// affecting the original — including RuntimeDefaults entries, whose slices
// are package-global and would otherwise be aliased by a shallow copy.
// A nil receiver returns nil.
func (rc *RuntimeConfig) Clone() *RuntimeConfig {
        if rc == nil {
                return nil
        }
        clone := *rc
        clone.AdditionalPackages = slices.Clone(rc.AdditionalPackages)
        clone.BuildWith = slices.Clone(rc.BuildWith)
        clone.RuntimeEnv = maps.Clone(rc.RuntimeEnv)
        return &clone
}

// WithOverrides returns a new RuntimeConfig with rc as the base and override
// layered on top:
//   - BuilderImage: override wins if non-empty, else falls back to rc's.
//   - AdditionalPackages: the union, rc's entries first, then any override
//     entries not already present.
//   - RuntimeEnv: merged, with override's value winning on a shared key.
//   - BuildWith: override wins if non-empty, else falls back to rc's — the
//     same "override wins if set" rule as BuilderImage. Fallback rather than
//     union: unioning two constraint sets could hand PEP 508 contradictory
//     specifiers to uv.
//
// WithOverrides(nil) returns rc.Clone() — a distinct object, never rc itself.
// A nil receiver returns override.Clone().
func (rc *RuntimeConfig) WithOverrides(override *RuntimeConfig) *RuntimeConfig {
        if rc == nil {
                return override.Clone()
        }
        if override == nil {
                return rc.Clone()
        }

        // Start from a copy of the base so any future field defaults to
        // base-wins rather than a zero value.
        merged := *rc
        if override.BuilderImage != "" {
                merged.BuilderImage = override.BuilderImage
        }

        seen := make(map[string]bool, len(rc.AdditionalPackages))
        merged.AdditionalPackages = append([]string(nil), rc.AdditionalPackages...)
        for _, pkg := range rc.AdditionalPackages {
                seen[pkg] = true
        }
        for _, pkg := range override.AdditionalPackages {
                if !seen[pkg] {
                        merged.AdditionalPackages = append(merged.AdditionalPackages, pkg)
                        seen[pkg] = true
                }
        }

        merged.RuntimeEnv = mergeEnvMaps(rc.RuntimeEnv, override.RuntimeEnv)
        if len(override.BuildWith) > 0 {
                merged.BuildWith = slices.Clone(override.BuildWith)
        } else {
                merged.BuildWith = slices.Clone(rc.BuildWith)
        }

        return &merged
}

// ValidateFor rejects a non-empty BuildWith for transports whose builder
// doesn't support build-time dependency constraints — only the uvx builder
// currently supports them — and otherwise validates rc via Validate.
// A nil receiver returns nil.
func (rc *RuntimeConfig) ValidateFor(transportType TransportType) error {
        if rc == nil {
                return nil
        }
        if transportType != TransportTypeUVX && len(rc.BuildWith) > 0 {
                return fmt.Errorf(
                        "build_with is not supported for %s:// builds (only uvx://)", transportType,
                )
        }
        return rc.Validate()
}

// IsEmpty reports whether rc has no field set. A nil receiver is empty.
func (rc *RuntimeConfig) IsEmpty() bool {
        if rc == nil {
                return true
        }
        return rc.BuilderImage == "" && len(rc.AdditionalPackages) == 0 &&
                len(rc.BuildWith) == 0 && len(rc.RuntimeEnv) == 0
}

// RuntimeDefaults provides default configurations for each runtime type
var RuntimeDefaults = map[TransportType]RuntimeConfig{
        TransportTypeGO: {
                BuilderImage:       "golang:1.27-alpine",
                AdditionalPackages: []string{"ca-certificates", "git"},
        },
        TransportTypeNPX: {
                BuilderImage:       "node:24-alpine",
                AdditionalPackages: []string{"git", "ca-certificates"},
        },
        TransportTypeUVX: {
                BuilderImage:       "python:3.14-slim",
                AdditionalPackages: []string{"ca-certificates", "git"},
        },
}

// GetDefaultRuntimeConfig returns the default runtime configuration for a given
// transport type. The result is a deep copy detached from RuntimeDefaults, so
// callers may freely mutate it without affecting the package-global defaults.
func GetDefaultRuntimeConfig(transportType TransportType) RuntimeConfig {
        config, ok := RuntimeDefaults[transportType]
        if !ok {
                // Return empty config if transport type not found
                return RuntimeConfig{}
        }
        return *config.Clone()
}

// mergeEnvMaps merges two environment variable maps without mutating either
// input. Entries in override take precedence over entries in base.
func mergeEnvMaps(base, override map[string]string) map[string]string {
        if len(base) == 0 && len(override) == 0 {
                return nil
        }
        merged := make(map[string]string, len(base)+len(override))
        maps.Copy(merged, base)
        maps.Copy(merged, override)
        return merged
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package templates provides utilities for generating Dockerfile templates
// based on different transport types (uvx, npx).
package templates

import (
        "bytes"
        "embed"
        "fmt"
        "regexp"
        "text/template"
)

//go:embed *.tmpl
var templateFS embed.FS

// TemplateData represents the data to be passed to the Dockerfile template.
type TemplateData struct {
        // MCPPackage is the name of the MCP package to run.
        MCPPackage string
        // MCPPackageClean is the package name with version suffix removed.
        // For example: "@org/package@1.2.3" becomes "@org/package", "package@1.0.0" becomes "package"
        // This field is automatically populated by GetDockerfileTemplate.
        MCPPackageClean string
        // CACertContent is the content of the custom CA certificate to include in the image.
        CACertContent string
        // IsLocalPath indicates if the MCPPackage is a local path that should be copied into the container.
        IsLocalPath bool
        // BuildArgs are the arguments to bake into the container's ENTRYPOINT at build time.
        // These are typically required subcommands (e.g., "start") that must always be present.
        // Runtime arguments passed via "-- <args>" will be appended after these build args.
        BuildArgs []string
        // BuildEnv contains environment variables to inject into the Dockerfile builder stage only.
        // These are used for configuring package managers (e.g., custom registry URLs) and do NOT
        // persist into the final runtime image or the running container's environment.
        // For environment variables the running container's process needs at startup
        // (e.g. feature flags, cache backend selection), use RuntimeConfig.RuntimeEnv instead.
        // Keys must be uppercase with underscores, values are validated for safety.
        BuildEnv map[string]string
        // BuildAuthFiles contains auth file contents keyed by file type (npmrc, netrc, etc).
        // These files are injected into the builder stage only for authentication.
        BuildAuthFiles map[string]string
        // RuntimeConfig specifies the base images and packages
        // If nil, defaults for the transport type are used
        RuntimeConfig *RuntimeConfig
}

// TransportType represents the type of transport to use.
type TransportType string

const (
        // TransportTypeUVX represents the uvx transport.
        TransportTypeUVX TransportType = "uvx"
        // TransportTypeNPX represents the npx transport.
        TransportTypeNPX TransportType = "npx"
        // TransportTypeGO represents the go transport.
        TransportTypeGO TransportType = "go"
)

// stripVersionSuffix removes version suffixes from package names.
// It strips @version from the end of package names while preserving scoped package prefixes.
// Examples:
//   - "@org/package@1.2.3" -> "@org/package"
//   - "package@1.0.0" -> "package"
//   - "@org/package" -> "@org/package" (no version, unchanged)
//   - "package" -> "package" (no version, unchanged)
func stripVersionSuffix(pkg string) string {
        // Match @version at the end, where version doesn't contain @ or /
        // This preserves scoped packages like @org/package
        re := regexp.MustCompile(`@[^@/]*$`)
        return re.ReplaceAllString(pkg, "")
}

// GetDockerfileTemplate returns the Dockerfile template for the specified transport type.
func GetDockerfileTemplate(transportType TransportType, data TemplateData) (string, error) {
        // Populate MCPPackageClean with version-stripped package name
        data.MCPPackageClean = stripVersionSuffix(data.MCPPackage)

        // Populate RuntimeConfig with defaults if not provided
        if data.RuntimeConfig == nil {
                defaultConfig := GetDefaultRuntimeConfig(transportType)
                data.RuntimeConfig = &defaultConfig
        }

        var templateName string

        // Determine the template name based on the transport type
        switch transportType {
        case TransportTypeUVX:
                templateName = "uvx.tmpl"
        case TransportTypeNPX:
                templateName = "npx.tmpl"
        case TransportTypeGO:
                templateName = "go.tmpl"
        default:
                return "", fmt.Errorf("unsupported transport type: %s", transportType)
        }

        // Read the template file
        tmplContent, err := templateFS.ReadFile(templateName)
        if err != nil {
                return "", fmt.Errorf("failed to read template file: %w", err)
        }

        // Create template with helper functions
        funcMap := template.FuncMap{
                "contains": func(s, substr string) bool {
                        return bytes.Contains([]byte(s), []byte(substr))
                },
                "isAlpine": func(image string) bool {
                        return bytes.Contains([]byte(image), []byte("alpine"))
                },
                "isDebian": func(image string) bool {
                        img := []byte(image)
                        return bytes.Contains(img, []byte("slim")) ||
                                bytes.Contains(img, []byte("debian")) ||
                                bytes.Contains(img, []byte("ubuntu"))
                },
        }

        // Parse the template with helper functions
        tmpl, err := template.New(templateName).Funcs(funcMap).Parse(string(tmplContent))
        if err != nil {
                return "", fmt.Errorf("failed to parse template: %w", err)
        }

        // Execute the template with the provided data
        var buf bytes.Buffer
        if err := tmpl.Execute(&buf, data); err != nil {
                return "", fmt.Errorf("failed to execute template: %w", err)
        }

        return buf.String(), nil
}

// ParseTransportType parses a string into a transport type.
func ParseTransportType(s string) (TransportType, error) {
        switch s {
        case "uvx":
                return TransportTypeUVX, nil
        case "npx":
                return TransportTypeNPX, nil
        case "go":
                return TransportTypeGO, nil
        default:
                return "", fmt.Errorf("unsupported transport type: %s", s)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package core provides the core domain model for the ToolHive system.
package core

import (
        "sort"
        "time"

        "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// Workload is a domain model representing a workload in the system.
// This is used in our API to hide details of the underlying runtime.
type Workload struct {
        // Name is the name of the workload.
        // It is used as a unique identifier.
        Name string `json:"name"`
        // Package specifies the Workload Package used to create this Workload.
        Package string `json:"package"`
        // URL is the URL of the workload exposed by the ToolHive proxy.
        URL string `json:"url"`
        // Port is the port on which the workload is exposed.
        // This is embedded in the URL.
        Port int `json:"port"`
        // TransportType is the type of transport used for this workload.
        TransportType types.TransportType `json:"transport_type" enums:"stdio,sse,streamable-http,inspector"`
        // ProxyMode is the proxy mode that clients should use to connect.
        // For stdio transports, this will be the proxy mode (sse or streamable-http).
        // For direct transports (sse/streamable-http), this will be the same as TransportType.
        ProxyMode string `json:"proxy_mode,omitempty"`
        // Status is the current status of the workload.
        //nolint:lll // enums tag needed for swagger generation with --parseDependencyLevel
        Status runtime.WorkloadStatus `json:"status" enums:"running,stopped,error,starting,stopping,unhealthy,removing,unknown,unauthenticated,auth_retrying,policy_stopped"`
        // StatusContext provides additional context about the workload's status.
        // The exact meaning is determined by the status and the underlying runtime.
        StatusContext string `json:"status_context,omitempty"`
        // CreatedAt is the timestamp when the workload was created.
        CreatedAt time.Time `json:"created_at"`
        // Labels are the container labels (excluding standard ToolHive labels)
        Labels map[string]string `json:"labels,omitempty"`
        // Group is the name of the group this workload belongs to, if any.
        Group string `json:"group,omitempty"`
        // ToolsFilter is the filter on tools applied to the workload.
        ToolsFilter []string `json:"tools,omitempty"`
        // Remote indicates whether this is a remote workload (true) or a container workload (false).
        Remote bool `json:"remote,omitempty"`
        // StartedAt is when the container was last started (changes on restart)
        StartedAt time.Time `json:"started_at"`
}

// SortWorkloadsByName sorts a slice of Workload by the Name field in ascending alphabetical order.
func SortWorkloadsByName(workloads []Workload) {
        sort.Slice(workloads, func(i, j int) bool {
                return workloads[i].Name < workloads[j].Name
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package desktop

import (
        "encoding/json"
        "errors"
        "os"
        "path/filepath"
)

const (
        // toolhiveDir is the directory name for toolhive files in the user's home.
        toolhiveDir = ".toolhive"
        // markerFileName is the name of the CLI source marker file.
        markerFileName = ".cli-source"
)

// errMarkerNotFound is returned when the marker file does not exist.
var errMarkerNotFound = errors.New("marker file not found")

// errInvalidMarker is returned when the marker file exists but is invalid.
var errInvalidMarker = errors.New("invalid marker file")

// getMarkerFilePath returns the path to the CLI source marker file.
// The marker file is located at ~/.toolhive/.cli-source
func getMarkerFilePath() (string, error) {
        homeDir, err := os.UserHomeDir()
        if err != nil {
                return "", err
        }
        return filepath.Join(homeDir, toolhiveDir, markerFileName), nil
}

// readMarkerFile reads and parses the CLI source marker file.
// Returns errMarkerNotFound if the file doesn't exist.
// Returns errInvalidMarker if the file exists but cannot be parsed or has
// an invalid schema version.
func readMarkerFile() (*cliSourceMarker, error) {
        markerPath, err := getMarkerFilePath()
        if err != nil {
                return nil, err
        }

        return readMarkerFileFromPath(markerPath)
}

// readMarkerFileFromPath reads and parses the CLI source marker file from
// a specific path. This is useful for testing.
func readMarkerFileFromPath(path string) (*cliSourceMarker, error) {
        // #nosec G304 -- path is always the marker file path from getMarkerFilePath or tests
        data, err := os.ReadFile(path)
        if err != nil {
                if os.IsNotExist(err) {
                        return nil, errMarkerNotFound
                }
                return nil, err
        }

        var marker cliSourceMarker
        if err := json.Unmarshal(data, &marker); err != nil {
                return nil, errInvalidMarker
        }

        // Validate schema version
        if marker.SchemaVersion != currentSchemaVersion {
                return nil, errInvalidMarker
        }

        // Validate source field
        if marker.Source != "desktop" {
                return nil, errInvalidMarker
        }

        return &marker, nil
}

// markerFileExists checks if the marker file exists without reading it.
func markerFileExists() (bool, error) {
        markerPath, err := getMarkerFilePath()
        if err != nil {
                return false, err
        }

        _, err = os.Stat(markerPath)
        if err != nil {
                if os.IsNotExist(err) {
                        return false, nil
                }
                return false, err
        }
        return true, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package desktop

import (
        "errors"
        "fmt"
        "os"
        "path/filepath"
        "runtime"
        "strings"
)

// envSkipDesktopCheck is the environment variable name that can be set to
// skip the desktop validation check. Set to "1" or "true" to skip.
const envSkipDesktopCheck = "TOOLHIVE_SKIP_DESKTOP_CHECK"

// ErrDesktopConflict is returned when a conflict is detected between the
// current CLI and the desktop-managed CLI.
var ErrDesktopConflict = errors.New("CLI conflict detected")

// IsDesktopManagedCLI reports whether the current CLI binary is the one
// managed by the ToolHive Desktop application. It returns false on any error
// (fail open: show updates when uncertain).
func IsDesktopManagedCLI() bool {
        result, err := checkDesktopAlignment()
        if err != nil {
                return false
        }
        // No conflict + DesktopCLIPath populated = paths matched, we ARE the desktop binary
        return !result.HasConflict && result.DesktopCLIPath != ""
}

// ValidateDesktopAlignment checks if the current CLI binary conflicts with
// a desktop-managed CLI installation.
//
// Returns nil if:
//   - No marker file exists (no desktop installation)
//   - Marker file is invalid or unreadable (treat as no desktop installation)
//   - The target binary in the marker doesn't exist (desktop was uninstalled)
//   - The current CLI is the desktop-managed CLI (paths match)
//
// Returns an error if:
//   - A valid marker file exists pointing to an existing binary
//   - The current CLI is NOT the desktop-managed binary
func ValidateDesktopAlignment() error {
        // Check for skip override
        if shouldSkipValidation() {
                return nil
        }

        result, err := checkDesktopAlignment()
        if err != nil {
                // Treat errors during validation as non-fatal - don't block the CLI
                return nil
        }

        if result.HasConflict {
                return fmt.Errorf("%w\n\n%s", ErrDesktopConflict, result.Message)
        }

        return nil
}

// checkDesktopAlignment performs the desktop alignment check and returns
// a detailed result.
func checkDesktopAlignment() (*validationResult, error) {
        result := &validationResult{}

        // Read the marker file
        marker, err := readMarkerFile()
        if err != nil {
                if errors.Is(err, errMarkerNotFound) || errors.Is(err, errInvalidMarker) {
                        // No marker or invalid marker - no conflict
                        return result, nil
                }
                return nil, fmt.Errorf("failed to read marker file: %w", err)
        }

        // Get the target path from the marker
        targetPath := getTargetPath(marker)
        if targetPath == "" {
                // No target path available - can't validate
                return result, nil
        }

        // Check if the target binary exists
        if !fileExists(targetPath) {
                // Target doesn't exist - desktop was likely uninstalled but marker
                // wasn't cleaned up. Proceed normally.
                return result, nil
        }

        // Get the current executable path
        currentExePath, err := os.Executable()
        if err != nil {
                return nil, fmt.Errorf("failed to get current executable path: %w", err)
        }

        // Resolve and normalize both paths for comparison
        resolvedCurrent, err := resolvePath(currentExePath)
        if err != nil {
                return nil, fmt.Errorf("failed to resolve current executable path: %w", err)
        }

        resolvedTarget, err := resolvePath(targetPath)
        if err != nil {
                // If we can't resolve the target, we can't compare properly
                // Treat as no conflict to avoid blocking legitimate use
                return result, nil
        }

        result.CurrentCLIPath = resolvedCurrent
        result.DesktopCLIPath = resolvedTarget

        // Compare paths
        if pathsEqual(resolvedCurrent, resolvedTarget) {
                // We ARE the desktop-managed CLI - no conflict
                return result, nil
        }

        // Conflict detected!
        result.HasConflict = true
        result.Message = buildConflictMessage(resolvedTarget, resolvedCurrent, marker)

        return result, nil
}

// shouldSkipValidation checks if the validation should be skipped via
// environment variable.
func shouldSkipValidation() bool {
        val := os.Getenv(envSkipDesktopCheck)
        val = strings.ToLower(strings.TrimSpace(val))
        return val == "1" || val == "true"
}

// getTargetPath extracts the target binary path from the marker based on
// the installation method and platform.
func getTargetPath(marker *cliSourceMarker) string {
        if marker.InstallMethod == "symlink" && marker.SymlinkTarget != "" {
                return marker.SymlinkTarget
        }

        // For Flatpak installations, the target is the host-visible path to the
        // CLI binary inside the Flatpak app directory. The validation logic is
        // the same as symlink: check if the target exists and compare paths.
        if marker.InstallMethod == "flatpak" && marker.FlatpakTarget != "" {
                return marker.FlatpakTarget
        }

        // For Windows/copy method, construct the path to the desktop-managed CLI
        // from the known installation location: %LOCALAPPDATA%\ToolHive\bin\thv.exe
        // Note: copy method is only used on Windows; on other platforms, return empty.
        if marker.InstallMethod == "copy" && runtime.GOOS == "windows" {
                return filepath.Join(getWindowsLocalAppData(), "ToolHive", "bin", "thv.exe")
        }

        return ""
}

// getWindowsLocalAppData returns the LocalAppData path on Windows.
// Falls back to %USERPROFILE%\AppData\Local if LOCALAPPDATA is not set.
func getWindowsLocalAppData() string {
        if localAppData := os.Getenv("LOCALAPPDATA"); localAppData != "" {
                return localAppData
        }
        // Fallback: construct from home directory
        homeDir, err := os.UserHomeDir()
        if err != nil {
                return ""
        }
        return filepath.Join(homeDir, "AppData", "Local")
}

// resolvePath resolves symlinks and normalizes the path for comparison.
func resolvePath(path string) (string, error) {
        // First, evaluate any symlinks
        resolved, err := filepath.EvalSymlinks(path)
        if err != nil {
                return "", err
        }

        // Clean and convert to absolute path
        resolved = filepath.Clean(resolved)
        if !filepath.IsAbs(resolved) {
                resolved, err = filepath.Abs(resolved)
                if err != nil {
                        return "", err
                }
        }

        return resolved, nil
}

// pathsEqual compares two paths accounting for platform-specific
// case sensitivity.
func pathsEqual(path1, path2 string) bool {
        // On Windows and macOS, file systems are typically case-insensitive
        if runtime.GOOS == "windows" || runtime.GOOS == "darwin" {
                return strings.EqualFold(path1, path2)
        }
        // On Linux and other platforms, use case-sensitive comparison
        return path1 == path2
}

// fileExists checks if a file exists at the given path.
func fileExists(path string) bool {
        _, err := os.Stat(path)
        return err == nil
}

// buildConflictMessage creates a user-friendly conflict error message.
func buildConflictMessage(desktopPath, currentPath string, marker *cliSourceMarker) string {
        var sb strings.Builder

        sb.WriteString("The ToolHive Desktop application manages a CLI installation at:\n")
        fmt.Fprintf(&sb, "  %s\n\n", desktopPath)

        sb.WriteString("You are running a different CLI binary at:\n")
        fmt.Fprintf(&sb, "  %s\n\n", currentPath)

        sb.WriteString("To avoid conflicts, please use the desktop-managed CLI or uninstall\n")
        sb.WriteString("the ToolHive Desktop application.\n\n")

        // Provide actionable guidance with platform-specific paths
        binPath, exeName := getDesktopBinPath()

        sb.WriteString("To use the desktop-managed CLI, ensure your PATH includes:\n")
        fmt.Fprintf(&sb, "  %s\n\n", binPath)

        sb.WriteString("Or run the desktop CLI directly:\n")
        fmt.Fprintf(&sb, "  %s [command]\n", filepath.Join(binPath, exeName))

        if marker.DesktopVersion != "" {
                fmt.Fprintf(&sb, "\nDesktop version: %s\n", marker.DesktopVersion)
        }

        return sb.String()
}

// getDesktopBinPath returns the platform-specific path to the desktop-managed
// CLI bin directory and the executable name.
func getDesktopBinPath() (binPath string, exeName string) {
        if runtime.GOOS == "windows" {
                // Windows: %LOCALAPPDATA%\ToolHive\bin\thv.exe
                return filepath.Join(getWindowsLocalAppData(), "ToolHive", "bin"), "thv.exe"
        }
        // macOS/Linux: ~/.toolhive/bin/thv
        homeDir, _ := os.UserHomeDir()
        return filepath.Join(homeDir, toolhiveDir, "bin"), "thv"
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package diagnostics serves operational endpoints on a listener that is
// separate from the application listener.
//
// The reason to separate is enforceability, not middleware coverage. Be precise
// about what this does and does not buy, because the distinction is easy to
// overstate:
//
//   - It does NOT authenticate, rate limit, or audit the endpoints. This
//     listener carries no middleware at all (see the warning below). An endpoint
//     served here is exactly as unauthenticated as it was on the application
//     mux.
//   - It DOES make the endpoints governable by port. Kubernetes NetworkPolicy is
//     L3/L4: it selects pods, ports, and protocols, and cannot filter on HTTP
//     path. While /metrics shares the application port, no policy can permit MCP
//     traffic while denying metrics scraping. On its own port, it can. Route-level
//     controls (Gateway API, Ingress path rules) govern only what reaches the
//     gateway, so they address north-south exposure but leave pod-to-pod traffic
//     untouched.
//   - It DOES keep the endpoints off the port deployments route publicly, so the
//     safe outcome does not depend on every operator getting route rules right.
//
// Serving them on the application mux is what made the port-level control
// impossible: Go's ServeMux resolves the most specific registered pattern first,
// so an explicitly registered "/metrics" always beats the "/" catch-all that
// carries the middleware chain.
//
// This mirrors the ToolHive operator, which already binds its own metrics
// endpoint to a separate address (--metrics-bind-address), and the wider
// convention for diagnostics ports: etcd's --listen-metrics-urls and
// controller-runtime's --metrics-bind-address.
//
// Note that /health deliberately stays on the application listener. Kubernetes
// liveness and readiness probes target the application port, and the proxy
// health response carries no sensitive fields.
//
// Anything registered on this listener is served WITHOUT authentication,
// authorization, rate limiting, or audit — that is the whole point of moving it
// off the application listener, and it is why the listener must not be routed
// publicly. Only add read-only, non-sensitive endpoints here. In particular, do
// not register pprof or any other debug handler that exposes process memory,
// goroutine state, or configuration.
package diagnostics

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "net"
        "net/http"
        "sync"
        "time"

        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/transport/proxy/socket"
)

// MetricsPath is the path the Prometheus metrics handler is served on.
const MetricsPath = "/metrics"

// DefaultPort is the port the diagnostics listener binds when no port is
// configured. 9464 is the OpenTelemetry specification's default for the
// Prometheus exporter (OTEL_EXPORTER_PROMETHEUS_PORT), so scrapers already
// expect metrics there.
//
// A fixed default rather than an arbitrary one matters for deployments: a
// scraper needs a predictable target. When the port is already taken — several
// CLI workloads on one machine, for instance — Start falls back to an available
// port and logs the resolved address.
const DefaultPort = 9464

// Timeouts mirror the proxy and vMCP listeners. They matter more here, not less:
// this listener carries no middleware, so nothing else bounds a slow or
// abandoned client.
const (
        // readHeaderTimeout bounds header reads to prevent Slowloris attacks.
        readHeaderTimeout = 10 * time.Second
        // readTimeout bounds reading the entire request, including a body that a
        // client trickles in. The metrics handler ignores request bodies, so without
        // this a connection could be held open indefinitely past the header phase.
        readTimeout = 30 * time.Second
        // writeTimeout bounds writing the response. There are no long-lived streams
        // on this listener, so it applies unconditionally.
        writeTimeout = 30 * time.Second
        // idleTimeout stops idle keep-alive connections from blocking shutdown.
        idleTimeout = 60 * time.Second
        // maxHeaderBytes caps request header size.
        maxHeaderBytes = 1 << 20 // 1 MiB
)

// startupLogMessage is the message logged when the diagnostics listener binds. It
// is referenced by NotServedHereHandler, which tells the reader to grep for it, so
// the two must not drift apart.
const startupLogMessage = "prometheus metrics are served on a dedicated diagnostics port, not the application port"

// ResolvePort returns the diagnostics port to request for a configured value.
// Zero means "not configured" and selects DefaultPort, so a scraper has a
// predictable target rather than an arbitrary one. Callers that genuinely want
// an OS-assigned port pass 0 to New directly instead.
func ResolvePort(configured int) int {
        if configured == 0 {
                return DefaultPort
        }
        return configured
}

// bindAttempts bounds how many times Start re-resolves and re-binds when the
// port is claimed between the availability check and the bind. See Server.bind.
const bindAttempts = 3

// NotServedHereHandler answers requests for MetricsPath on an application
// listener, where metrics are deliberately not served.
//
// It exists so the response explains itself. A bare 404 is indistinguishable
// from a typo, and the failure is otherwise silent from the server side: an
// upgrade moves the endpoint and the operator sees only a Prometheus target
// going down.
//
// The body deliberately names no port. The diagnostics listener honours a
// configured port and falls back to an available one when that is taken (see
// Server.bind), so any number written here would be wrong for both of those
// supported configurations — and a confidently wrong address is worse than none
// when the whole point is to redirect someone who is already lost. The startup
// log carries the resolved address and is the one source that is always correct.
//
// The status stays 404 rather than 410 Gone: this handler is also registered on
// deployments that never served metrics on this port, where claiming the
// resource was removed would be untrue.
func NotServedHereHandler() http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
                w.Header().Set("Content-Type", "text/plain; charset=utf-8")
                w.Header().Set("X-Content-Type-Options", "nosniff")
                w.WriteHeader(http.StatusNotFound)
                // Best effort: the client has already disconnected if this fails, and
                // there is nothing useful to do about it on a 404.
                _, _ = fmt.Fprintf(w, "%s is not served on this port.\n\n"+
                        "When Prometheus metrics are enabled they are served on a dedicated diagnostics\n"+
                        "listener, so that access can be restricted separately from MCP traffic.\n\n"+
                        "Its address is logged at startup; look for %q.\n"+
                        "See docs/observability.md for how to point a scraper at it.\n",
                        MetricsPath, startupLogMessage)
        })
}

// Server serves diagnostics endpoints on a dedicated listener.
//
// The zero value is not usable; construct one with New.
type Server struct {
        host           string
        port           int
        metricsHandler http.Handler

        // mu guards server and listener, which Start writes and Stop reads.
        mu       sync.Mutex
        server   *http.Server
        listener net.Listener
}

// New creates a diagnostics server that serves metricsHandler at MetricsPath.
//
// host is the bind address. port is the requested port; pass DefaultPort unless
// a deployment needs a different one, since a predictable port is what lets a
// scraper find the endpoint. Passing 0 asks the OS for an arbitrary available
// port, which is useful in tests but leaves nothing for a scraper to target.
// Either way, Start falls back to an available port if the requested one is
// taken, so the resolved port is only known after Start — read it from Addr or
// Port.
//
// The port is not bound until Start is called.
func New(host string, port int, metricsHandler http.Handler) (*Server, error) {
        if host == "" {
                return nil, errors.New("diagnostics: host is required")
        }
        if port < 0 || port > 65535 {
                return nil, fmt.Errorf("diagnostics: port %d out of range", port)
        }
        if metricsHandler == nil {
                return nil, errors.New("diagnostics: metrics handler is required")
        }
        return &Server{
                host:           host,
                port:           port,
                metricsHandler: metricsHandler,
        }, nil
}

// Start binds the diagnostics listener and serves it in the background. It
// returns once the listener is bound, so a caller that observes no error can
// rely on Addr reporting the resolved address.
//
// Calling Start on an already-started Server returns an error.
func (s *Server) Start() error {
        s.mu.Lock()
        defer s.mu.Unlock()

        if s.listener != nil {
                return errors.New("diagnostics: server already started")
        }

        mux := http.NewServeMux()
        mux.Handle(MetricsPath, s.metricsHandler)

        listener, err := s.bind()
        if err != nil {
                return err
        }

        s.listener = listener
        s.server = &http.Server{
                Addr:              listener.Addr().String(),
                Handler:           mux,
                ReadHeaderTimeout: readHeaderTimeout,
                ReadTimeout:       readTimeout,
                WriteTimeout:      writeTimeout,
                IdleTimeout:       idleTimeout,
                MaxHeaderBytes:    maxHeaderBytes,
        }

        // Capture the server locally so the goroutine does not race with Stop
        // clearing the field.
        server := s.server
        go func() {
                if err := server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
                        var opErr *net.OpError
                        if errors.As(err, &opErr) && opErr.Op == "accept" {
                                // Expected once Stop closes the listener.
                                return
                        }
                        slog.Warn("diagnostics server error", "error", err)
                }
        }()

        // Logged at WARN, not INFO, because this is the only server-side signal that
        // the endpoint is not on the application port. Metrics are opt-in, so this
        // fires only for deployments that enabled them — the exact population whose
        // scrape configuration has to point here. Without it, an upgrade moves the
        // endpoint and the operator's only clue is a Prometheus target going down.
        slog.Warn(startupLogMessage+"; update scrape configuration to target this address",
                "address", listener.Addr().String(), "path", MetricsPath)

        return nil
}

// bind resolves a port and binds it, retrying a bounded number of times.
//
// networking.FindOrUsePort only *checks* availability, so another process can
// claim the port in the window before the actual bind. Retrying rather than
// failing keeps a lost race from taking down the whole workload over a
// diagnostics listener; each retry re-resolves, so a genuinely occupied port
// converges on an alternative. Exhausting the attempts is reported as an error,
// because silently serving nothing would hide the metrics endpoint.
//
// Callers must hold s.mu.
func (s *Server) bind() (net.Listener, error) {
        // Use SO_REUSEADDR for parity with the proxy listeners, which allows port
        // reuse after an unclean shutdown left a zombie holding the port.
        lc := socket.ListenConfig()

        var lastErr error
        for attempt := 1; attempt <= bindAttempts; attempt++ {
                port, err := networking.FindOrUsePort(s.port)
                if err != nil {
                        return nil, fmt.Errorf("diagnostics: failed to resolve port: %w", err)
                }

                addr := fmt.Sprintf("%s:%d", s.host, port)
                listener, err := lc.Listen(context.Background(), "tcp", addr)
                if err == nil {
                        return listener, nil
                }

                lastErr = err
                slog.Debug("diagnostics port was claimed between availability check and bind, retrying",
                        "address", addr, "attempt", attempt, "error", err)
        }

        return nil, fmt.Errorf("diagnostics: failed to bind a listener on %s after %d attempts: %w",
                s.host, bindAttempts, lastErr)
}

// Addr returns the resolved listen address, or an empty string before Start
// succeeds or after Stop.
func (s *Server) Addr() string {
        s.mu.Lock()
        defer s.mu.Unlock()

        if s.listener == nil {
                return ""
        }
        return s.listener.Addr().String()
}

// Port returns the resolved port, or 0 before Start succeeds or after Stop.
func (s *Server) Port() int {
        s.mu.Lock()
        defer s.mu.Unlock()

        if s.listener == nil {
                return 0
        }
        tcpAddr, ok := s.listener.Addr().(*net.TCPAddr)
        if !ok {
                return 0
        }
        return tcpAddr.Port
}

// Stop gracefully shuts the diagnostics server down. It is safe to call on a
// Server that was never started, and safe to call more than once.
func (s *Server) Stop(ctx context.Context) error {
        s.mu.Lock()
        server := s.server
        s.server = nil
        s.listener = nil
        s.mu.Unlock()

        if server == nil {
                return nil
        }

        // Shutdown closes the listener, so it is not closed separately here.
        if err := server.Shutdown(ctx); err != nil {
                return fmt.Errorf("diagnostics: failed to shut down: %w", err)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package environment provides utilities for handling environment variables
// and environment-related operations for containers.
package environment

import (
        "context"
        "fmt"
        "strings"

        "github.com/stacklok/toolhive/pkg/secrets"
)

// ParseSecretParameters parses the secret parameters from the command line,
// fetches them from the secrets manager, and returns a map of secrets and
// their environment variable names.
func ParseSecretParameters(ctx context.Context, parameters []string, secretsManager secrets.Provider) (map[string]string, error) {
        secretVariables := make(map[string]string, len(parameters))
        for _, param := range parameters {
                parameter, err := secrets.ParseSecretParameter(param)
                if err != nil {
                        return nil, err
                }

                secret, err := secretsManager.GetSecret(ctx, parameter.Name)
                if err != nil {
                        return nil, err
                }

                secretVariables[parameter.Target] = secret
        }

        return secretVariables, nil
}

// ParseEnvironmentVariables parses environment variables from a slice of strings
// in the format KEY=VALUE
func ParseEnvironmentVariables(envVars []string) (map[string]string, error) {
        result := make(map[string]string)

        for _, env := range envVars {
                parts := strings.SplitN(env, "=", 2)
                if len(parts) != 2 {
                        return nil, fmt.Errorf("invalid environment variable format: %s", env)
                }

                key := parts[0]
                value := parts[1]

                if key == "" {
                        return nil, fmt.Errorf("empty environment variable key")
                }

                result[key] = value
        }

        return result, nil
}

// SetTransportEnvironmentVariables sets transport-specific environment variables
func SetTransportEnvironmentVariables(envVars map[string]string, transportType string, port int) {
        // Set common environment variables
        envVars["MCP_TRANSPORT"] = transportType

        // Set port-related environment variables only if port is greater than 0
        if port > 0 {
                // Set transport-specific environment variables
                switch transportType {
                case "sse", "streamable-http":
                        envVars["MCP_PORT"] = fmt.Sprintf("%d", port)
                        envVars["FASTMCP_PORT"] = fmt.Sprintf("%d", port)
                case "stdio":
                        // No additional environment variables needed for stdio transport
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package fileutils provides file operation utilities including atomic writes.
package fileutils

import (
        "fmt"
        "os"
        "path/filepath"
)

// AtomicWriteFile writes data to a file atomically by writing to a temporary file
// and then renaming it. This ensures that readers either see the complete old file
// or the complete new file, never a partially written file.
func AtomicWriteFile(targetPath string, data []byte, perm os.FileMode) error {
        // Create a temporary file in the same directory as the target file
        // This ensures the temp file is on the same filesystem for atomic rename
        dir := filepath.Dir(targetPath)
        tmpFile, err := os.CreateTemp(dir, ".tmp-*")
        if err != nil {
                return fmt.Errorf("failed to create temp file: %w", err)
        }
        tmpPath := tmpFile.Name()

        // Ensure cleanup of temp file on error
        success := false
        defer func() {
                if !success {
                        tmpFile.Close()    //nolint:errcheck,gosec // best effort cleanup
                        os.Remove(tmpPath) //nolint:errcheck,gosec // best effort cleanup
                }
        }()

        // Write data to temp file
        if _, err := tmpFile.Write(data); err != nil {
                return fmt.Errorf("failed to write to temp file: %w", err)
        }

        // Sync to ensure data is written to disk
        if err := tmpFile.Sync(); err != nil {
                return fmt.Errorf("failed to sync temp file: %w", err)
        }

        // Close the temp file before renaming
        if err := tmpFile.Close(); err != nil {
                return fmt.Errorf("failed to close temp file: %w", err)
        }

        // Set the correct permissions on the temp file
        // #nosec G703 -- tmpPath is from os.CreateTemp in the same directory as targetPath
        if err := os.Chmod(tmpPath, perm); err != nil {
                return fmt.Errorf("failed to set permissions on temp file: %w", err)
        }

        // Atomically rename temp file to target file
        // This is atomic on POSIX systems (Linux, macOS, etc.)
        // #nosec G703 -- tmpPath is from os.CreateTemp, targetPath is caller-controlled
        if err := os.Rename(tmpPath, targetPath); err != nil {
                return fmt.Errorf("failed to rename temp file: %w", err)
        }

        success = true
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package fileutils

import (
        "fmt"
        "os"
        "path/filepath"
        "strings"
)

// WriteContainedFile writes content to filePath (relative) inside targetDir,
// ensuring the resulting path does not escape targetDir. Parent directories
// are created with dirPerm, and the file is written atomically with filePerm.
//
// targetDir must already be filepath.Clean'd by the caller.
func WriteContainedFile(targetDir, filePath string, content []byte, dirPerm, filePerm os.FileMode) error {
        cleanTarget := targetDir + string(os.PathSeparator)
        destPath := filepath.Clean(filepath.Join(targetDir, filepath.FromSlash(filePath)))

        if !strings.HasPrefix(destPath, cleanTarget) {
                return fmt.Errorf("path traversal detected: file %q escapes target directory", filePath)
        }

        parentDir := filepath.Dir(destPath)
        if err := os.MkdirAll(parentDir, dirPerm); err != nil {
                return fmt.Errorf("creating directory %q: %w", parentDir, err)
        }

        if err := AtomicWriteFile(destPath, content, filePerm); err != nil {
                return fmt.Errorf("writing file %q: %w", filePath, err)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package fileutils

import (
        "context"
        "fmt"
        "sync"
        "time"

        "github.com/stacklok/toolhive/pkg/lockfile"
)

const (
        // DefaultLockTimeout is the maximum time to wait for a file lock.
        DefaultLockTimeout = 5 * time.Second

        // defaultLockRetryInterval is the interval between lock acquisition attempts.
        defaultLockRetryInterval = 100 * time.Millisecond
)

// processLocks provides per-path in-process mutual exclusion.
// flock(2) does NOT provide mutual exclusion between different file descriptors
// within the same process — only between different processes. Since each call to
// WithFileLock opens a new file descriptor, concurrent goroutines can all acquire
// the flock simultaneously. This in-process mutex ensures serialization within a
// single process, while the flock continues to protect cross-process access.
//
// This map is never pruned; callers should ensure the number of distinct
// paths remains bounded (e.g. one secrets file per workload).
var processLocks sync.Map

// getProcessLock returns the in-process mutex for the given path,
// creating one if it does not already exist.
func getProcessLock(path string) *sync.Mutex {
        // Fast path: return the existing mutex without allocating.
        if val, ok := processLocks.Load(path); ok {
                return val.(*sync.Mutex)
        }
        val, _ := processLocks.LoadOrStore(path, &sync.Mutex{})
        return val.(*sync.Mutex)
}

// WithFileLock executes fn while holding both an in-process mutex and an
// OS-level advisory file lock on path + ".lock".
// The in-process mutex serializes goroutines within the same process (where
// flock is ineffective), and the file lock serializes across processes.
func WithFileLock(path string, fn func() error) error {
        // Acquire in-process lock first to serialize goroutines.
        mu := getProcessLock(path)
        mu.Lock()
        defer mu.Unlock()

        lockPath := path + ".lock"

        ctx, cancel := context.WithTimeout(context.Background(), DefaultLockTimeout)
        defer cancel()

        fileLock, locked, err := lockfile.AcquireTrackedLock(ctx, lockPath, defaultLockRetryInterval)
        if err != nil {
                return fmt.Errorf("failed to acquire lock: %w", err)
        }
        if !locked {
                return fmt.Errorf("failed to acquire lock: timeout after %v", DefaultLockTimeout)
        }
        defer lockfile.ReleaseTrackedLock(lockPath, fileLock)

        return fn()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package fileutils provides file operation utilities including atomic writes
// and path validation for security.
package fileutils

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/workloads/types"
)

// ValidateWorkloadNameForPath validates a workload name to prevent path traversal attacks.
// It ensures the name is safe for use in file path construction by checking:
// - Path traversal patterns (..)
// - Absolute paths
// - Path separators (/, \)
// - Command injection patterns
// - Null bytes
// - Invalid characters (only alphanumeric, dots, hyphens, underscores allowed)
// - Length limits
//
// This function delegates to types.ValidateWorkloadName which performs comprehensive
// validation including filepath.Clean normalization and filepath.Rel path traversal checks.
//
// Returns nil if the workload name is safe for path construction, or an error describing
// the validation failure.
func ValidateWorkloadNameForPath(workloadName string) error {
        // The types.ValidateWorkloadName function already performs comprehensive validation:
        // - Empty check
        // - Null bytes detection
        // - Path normalization (filepath.Clean)
        // - Path traversal detection (filepath.Rel to check for ".." escapes)
        // - Absolute path rejection
        // - Command injection pattern detection
        // - Character validation (only [a-zA-Z0-9._-] allowed, which excludes / and \)
        // - Length limits (max 100 characters)
        //
        // This provides defense-in-depth against path traversal attacks.
        if err := types.ValidateWorkloadName(workloadName); err != nil {
                return fmt.Errorf("invalid workload name for path construction: %w", err)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package git

import (
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "path/filepath"
        "strings"

        "github.com/go-git/go-billy/v5/memfs"
        "github.com/go-git/go-billy/v5/util"
        "github.com/go-git/go-git/v5"
        "github.com/go-git/go-git/v5/plumbing"
        "github.com/go-git/go-git/v5/plumbing/cache"
        "github.com/go-git/go-git/v5/plumbing/object"
        "github.com/go-git/go-git/v5/plumbing/transport"
        "github.com/go-git/go-git/v5/storage/filesystem"
)

// ErrNilRepository is returned when a nil repository is passed to an operation that requires one.
var ErrNilRepository = errors.New("repository is nil")

// ErrInvalidFilePath is returned when a file path contains traversal or absolute components.
var ErrInvalidFilePath = errors.New("invalid file path")

// ErrInvalidCloneConfig is returned when CloneConfig has conflicting ref specifications.
var ErrInvalidCloneConfig = errors.New("invalid clone config: at most one of Branch, Tag, or Commit may be specified")

// Client defines the interface for Git operations
type Client interface {
        // Clone clones a repository with the given configuration
        Clone(ctx context.Context, config *CloneConfig) (*RepositoryInfo, error)

        // GetFileContent retrieves the content of a file from the repository
        GetFileContent(repoInfo *RepositoryInfo, path string) ([]byte, error)

        // HeadCommit returns the hash and signature of the HEAD commit.
        HeadCommit(repoInfo *RepositoryInfo) (HeadCommit, error)

        // Cleanup removes local repository directory
        Cleanup(ctx context.Context, repoInfo *RepositoryInfo) error
}

// HeadCommit describes the HEAD commit of a cloned repository.
type HeadCommit struct {
        // Hash is the commit hash.
        Hash string
        // Signature is the armored signature attached to the commit, empty
        // when the commit is unsigned. The signature is UNVERIFIED — it is
        // whatever bytes the commit carries; callers must cryptographically
        // verify it before treating it as provenance.
        Signature string
        // Payload is the encoded commit object without its signature — the
        // exact bytes the signature signs. Verifiers check Signature over
        // Payload.
        Payload []byte
}

// DefaultGitClient implements Client using go-git
type DefaultGitClient struct {
        // auth is the optional authentication method for cloning
        auth transport.AuthMethod
}

// ClientOption configures a DefaultGitClient.
type ClientOption func(*DefaultGitClient)

// WithAuth sets the authentication method for git operations.
func WithAuth(auth transport.AuthMethod) ClientOption {
        return func(c *DefaultGitClient) {
                c.auth = auth
        }
}

// NewDefaultGitClient creates a new DefaultGitClient
func NewDefaultGitClient(opts ...ClientOption) *DefaultGitClient {
        c := &DefaultGitClient{}
        for _, o := range opts {
                o(c)
        }
        return c
}

// Clone clones a repository with the given configuration
func (c *DefaultGitClient) Clone(ctx context.Context, config *CloneConfig) (*RepositoryInfo, error) {
        if err := config.validate(); err != nil {
                return nil, err
        }

        // Prepare clone options
        cloneOptions := &git.CloneOptions{
                URL:  config.URL,
                Auth: c.auth,
        }

        // Set reference if specified (but not for commit-based clones)
        if config.Commit == "" {
                cloneOptions.Depth = 1
                if config.Branch != "" {
                        cloneOptions.ReferenceName = plumbing.NewBranchReferenceName(config.Branch)
                        cloneOptions.SingleBranch = true
                } else if config.Tag != "" {
                        cloneOptions.ReferenceName = plumbing.NewTagReferenceName(config.Tag)
                        cloneOptions.SingleBranch = true
                }
        }
        // For commit-based clones, we need the full repository to ensure the commit is available

        // Use in-memory filesystems for the repository and the storer
        // See https://github.com/mindersec/minder/blob/main/internal/providers/git/git.go
        // for more details
        // Clone the repository
        memFS := memfs.New()
        memFS = &LimitedFs{
                Fs:            memFS,
                MaxFiles:      10 * 1000,
                TotalFileSize: 100 * 1024 * 1024,
        }
        // go-git seems to want separate filesystems for the storer and the checked out files
        storerFs := memfs.New()
        storerFs = &LimitedFs{
                Fs:            storerFs,
                MaxFiles:      10 * 1000,
                TotalFileSize: 100 * 1024 * 1024,
        }
        storerCache := cache.NewObjectLRUDefault()
        storer := filesystem.NewStorage(storerFs, storerCache)

        repo, err := git.CloneContext(ctx, storer, memFS, cloneOptions)
        if err != nil {
                return nil, fmt.Errorf("failed to clone repository: %w", err)
        }

        // Get repository information
        repoInfo := &RepositoryInfo{
                Repository:       repo,
                RemoteURL:        config.URL,
                storerFilesystem: storerFs,
                objectCache:      storerCache,
        }

        // If specific commit is requested, checkout that commit
        if config.Commit != "" {
                workTree, err := repo.Worktree()
                if err != nil {
                        return nil, fmt.Errorf("failed to get worktree: %w", err)
                }

                hash := plumbing.NewHash(config.Commit)
                err = workTree.Checkout(&git.CheckoutOptions{
                        Hash: hash,
                })
                if err != nil {
                        return nil, fmt.Errorf("failed to checkout commit %s: %w", config.Commit, err)
                }
        }

        // Update repository info with current state
        if err := c.updateRepositoryInfo(repoInfo); err != nil {
                return nil, fmt.Errorf("failed to update repository info: %w", err)
        }

        return repoInfo, nil
}

// GetFileContent retrieves the content of a file from the repository
func (*DefaultGitClient) GetFileContent(repoInfo *RepositoryInfo, path string) ([]byte, error) {
        if repoInfo == nil || repoInfo.Repository == nil {
                return nil, ErrNilRepository
        }

        // Reject absolute paths, traversal, and null bytes
        if filepath.IsAbs(path) || strings.Contains(path, "..") || strings.ContainsRune(path, 0) {
                return nil, fmt.Errorf("%w: %s", ErrInvalidFilePath, path)
        }

        // Get the HEAD reference
        ref, err := repoInfo.Repository.Head()
        if err != nil {
                return nil, fmt.Errorf("failed to get HEAD reference: %w", err)
        }

        // Get the commit object
        commit, err := repoInfo.Repository.CommitObject(ref.Hash())
        if err != nil {
                return nil, fmt.Errorf("failed to get commit object: %w", err)
        }

        // Get the tree
        tree, err := commit.Tree()
        if err != nil {
                return nil, fmt.Errorf("failed to get tree: %w", err)
        }

        // Get the file
        file, err := tree.File(path)
        if err != nil {
                return nil, fmt.Errorf("failed to get file %s: %w", path, err)
        }

        // Read file contents
        content, err := file.Contents()
        if err != nil {
                return nil, fmt.Errorf("failed to read file contents: %w", err)
        }

        return []byte(content), nil
}

// Cleanup removes local repository directory
func (*DefaultGitClient) Cleanup(_ context.Context, repoInfo *RepositoryInfo) error {
        if repoInfo == nil || repoInfo.Repository == nil {
                return ErrNilRepository
        }

        // 1. Clear object cache explicitly
        if repoInfo.objectCache != nil {
                slog.Debug("Clearing object cache")
                repoInfo.objectCache.Clear()
        }

        // 2. Clear worktree filesystem
        worktree, err := repoInfo.Repository.Worktree()
        if err == nil && worktree.Filesystem != nil {
                slog.Debug("Clearing worktree filesystem")
                if err := util.RemoveAll(worktree.Filesystem, "/"); err != nil {
                        slog.Warn("Failed to clear worktree filesystem", "error", err)
                }
        }

        // 3. Clear storer filesystem (memfs)
        if repoInfo.storerFilesystem != nil {
                slog.Debug("Clearing storer filesystem")
                if err := util.RemoveAll(repoInfo.storerFilesystem, "/"); err != nil {
                        slog.Warn("Failed to clear storer filesystem", "error", err)
                }
        }

        // 4. Nil out all references
        repoInfo.objectCache = nil
        repoInfo.storerFilesystem = nil
        repoInfo.Repository = nil

        return nil
}

// updateRepositoryInfo updates the repository info with current state
func (*DefaultGitClient) updateRepositoryInfo(repoInfo *RepositoryInfo) error {
        if repoInfo == nil || repoInfo.Repository == nil {
                return ErrNilRepository
        }

        // Get current branch name
        ref, err := repoInfo.Repository.Head()
        if err != nil {
                return fmt.Errorf("failed to get HEAD reference: %w", err)
        }

        if ref.Name().IsBranch() {
                repoInfo.Branch = ref.Name().Short()
        }

        return nil
}

// HeadCommit returns the hash and (unverified) signature of the HEAD
// commit in a single lookup, so both describe the same commit.
func (*DefaultGitClient) HeadCommit(repoInfo *RepositoryInfo) (HeadCommit, error) {
        if repoInfo == nil || repoInfo.Repository == nil {
                return HeadCommit{}, ErrNilRepository
        }

        ref, err := repoInfo.Repository.Head()
        if err != nil {
                return HeadCommit{}, fmt.Errorf("failed to get HEAD reference: %w", err)
        }
        commit, err := repoInfo.Repository.CommitObject(ref.Hash())
        if err != nil {
                return HeadCommit{}, fmt.Errorf("failed to read HEAD commit: %w", err)
        }
        payload, err := commitPayload(commit)
        if err != nil {
                return HeadCommit{}, fmt.Errorf("failed to encode HEAD commit payload: %w", err)
        }
        return HeadCommit{
                Hash:      ref.Hash().String(),
                Signature: commit.PGPSignature,
                Payload:   payload,
        }, nil
}

// commitPayload returns the encoded commit object without its signature —
// the bytes a commit signature is computed over.
func commitPayload(commit *object.Commit) ([]byte, error) {
        obj := &plumbing.MemoryObject{}
        if err := commit.EncodeWithoutSignature(obj); err != nil {
                return nil, err
        }
        r, err := obj.Reader()
        if err != nil {
                return nil, err
        }
        defer r.Close() //nolint:errcheck // in-memory reader
        return io.ReadAll(r)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// From https://github.com/mindersec/minder/blob/main/internal/providers/git/memboxfs/fs.go
// Apache License 2.0
// Copyright (c) 2023 MinderSec

package git

import (
        "errors"
        "fmt"
        "io"
        "io/fs"
        "os"

        billy "github.com/go-git/go-billy/v5"
)

// LimitedFs provides a size-limited billy.Filesystem.  This is a struct, there's
// no constructor here. Note that LimitedFs is not thread-safe.
type LimitedFs struct {
        Fs            billy.Filesystem
        MaxFiles      int64
        TotalFileSize int64

        currentFiles int64
        currentSize  int64
}

// ErrNotImplemented is returned when a method is not implemented.
var ErrNotImplemented = errors.New("not implemented")

// ErrTooBig is returned when a file is too big.
var ErrTooBig = errors.New("file too big")

// ErrTooManyFiles is returned when there are too many files.
var ErrTooManyFiles = errors.New("too many files")

var _ billy.Filesystem = (*LimitedFs)(nil)

// Chroot implements billy.Filesystem.
func (*LimitedFs) Chroot(_ string) (billy.Filesystem, error) {
        return nil, ErrNotImplemented
}

// Create implements billy.Filesystem.
func (f *LimitedFs) Create(filename string) (billy.File, error) {
        f.currentFiles++
        if f.currentFiles >= f.MaxFiles {
                return nil, fmt.Errorf("%w: current %d, max %d, %s", ErrTooManyFiles, f.currentFiles, f.MaxFiles, filename)
        }
        file, err := f.Fs.Create(filename)
        return &fileWrapper{f: file, fs: f}, err
}

// Join implements billy.Filesystem.
func (f *LimitedFs) Join(elem ...string) string {
        return f.Fs.Join(elem...)
}

// Lstat implements billy.Filesystem.
func (f *LimitedFs) Lstat(filename string) (fs.FileInfo, error) {
        return f.Fs.Lstat(filename)
}

// MkdirAll implements billy.Filesystem.
func (f *LimitedFs) MkdirAll(filename string, perm fs.FileMode) error {
        // TODO: account for path segments correctly
        f.currentFiles++
        if f.currentFiles >= f.MaxFiles {
                return fmt.Errorf("%w: current %d, max %d, %s", ErrTooManyFiles, f.currentFiles, f.MaxFiles, filename)
        }
        return f.Fs.MkdirAll(filename, perm)
}

// Open implements billy.Filesystem.
func (f *LimitedFs) Open(filename string) (billy.File, error) {
        return f.Fs.Open(filename)
}

// OpenFile implements billy.Filesystem.
func (f *LimitedFs) OpenFile(filename string, flag int, perm fs.FileMode) (billy.File, error) {
        if flag&os.O_CREATE != 0 {
                f.currentFiles++
                if f.currentFiles >= f.MaxFiles {
                        return nil, fmt.Errorf("%w: current %d, max %d, %s", ErrTooManyFiles, f.currentFiles, f.MaxFiles, filename)
                }
        }
        file, err := f.Fs.OpenFile(filename, flag, perm)
        return &fileWrapper{f: file, fs: f}, err
}

// ReadDir implements billy.Filesystem.
func (f *LimitedFs) ReadDir(path string) ([]fs.FileInfo, error) {
        return f.Fs.ReadDir(path)
}

// Readlink implements billy.Filesystem.
func (f *LimitedFs) Readlink(link string) (string, error) {
        return f.Fs.Readlink(link)
}

// Remove implements billy.Filesystem.
func (f *LimitedFs) Remove(filename string) error {
        // TODO: should we decrement currentFiles here?  It's not clear if the underlying
        // fs will reclaim memory on Remove, so we are conservative.
        return f.Fs.Remove(filename)
}

// Rename implements billy.Filesystem.
func (f *LimitedFs) Rename(oldpath string, newpath string) error {
        return f.Fs.Rename(oldpath, newpath)
}

// Root implements billy.Filesystem.
func (f *LimitedFs) Root() string {
        return f.Fs.Root()
}

// Stat implements billy.Filesystem.
func (f *LimitedFs) Stat(filename string) (fs.FileInfo, error) {
        return f.Fs.Stat(filename)
}

// Symlink implements billy.Filesystem.
func (f *LimitedFs) Symlink(target string, link string) error {
        f.currentFiles++
        if f.currentFiles >= f.MaxFiles {
                return fmt.Errorf("%w: current %d, max %d, %s", ErrTooManyFiles, f.currentFiles, f.MaxFiles, link)
        }
        return f.Fs.Symlink(target, link)
}

// TempFile implements billy.Filesystem.
func (f *LimitedFs) TempFile(dir string, prefix string) (billy.File, error) {
        f.currentFiles++
        if f.currentFiles >= f.MaxFiles {
                return nil, fmt.Errorf("%w: current %d, max %d, %s/%s", ErrTooManyFiles, f.currentFiles, f.MaxFiles, dir, prefix)
        }
        file, err := f.Fs.TempFile(dir, prefix)
        return &fileWrapper{f: file, fs: f}, err
}

type fileWrapper struct {
        f billy.File

        fs *LimitedFs
}

var _ billy.File = (*fileWrapper)(nil)

// Close implements billy.File.
func (f *fileWrapper) Close() error {
        return f.f.Close()
}

// Lock implements billy.File.
func (f *fileWrapper) Lock() error {
        return f.f.Lock()
}

// Name implements billy.File.
func (f *fileWrapper) Name() string {
        return f.f.Name()
}

// Read implements billy.File.
func (f *fileWrapper) Read(p []byte) (n int, err error) {
        return f.f.Read(p)
}

// ReadAt implements billy.File.
func (f *fileWrapper) ReadAt(p []byte, off int64) (n int, err error) {
        return f.f.ReadAt(p, off)
}

// Seek implements billy.File.
func (f *fileWrapper) Seek(offset int64, whence int) (int64, error) {
        return f.f.Seek(offset, whence)
}

// Truncate implements billy.File.
func (f *fileWrapper) Truncate(size int64) error {
        existingSize, err := f.fileSize()
        if err != nil {
                return err
        }

        growth := size - existingSize
        if growth+f.fs.currentSize > f.fs.TotalFileSize {
                return fmt.Errorf("%w: current %d, total %d, %s", ErrTooBig, f.fs.currentSize, f.fs.TotalFileSize, f.Name())
        }

        f.fs.currentSize += growth
        return f.f.Truncate(size)
}

// Unlock implements billy.File.
func (f *fileWrapper) Unlock() error {
        return f.f.Unlock()
}

// Write implements billy.File.
func (f *fileWrapper) Write(p []byte) (n int, err error) {
        size, err := f.fileSize()
        if err != nil {
                return 0, err
        }
        offset, err := f.Seek(0, io.SeekCurrent)
        if err != nil {
                return 0, err
        }

        growth := offset + int64(len(p)) - size
        if growth < 0 {
                growth = 0
        }
        if growth+f.fs.currentSize > f.fs.TotalFileSize {
                return 0, fmt.Errorf("%w: %s", ErrTooBig, f.Name())
        }

        f.fs.currentSize += growth
        return f.f.Write(p)
}

func (f *fileWrapper) fileSize() (int64, error) {
        fi, err := f.fs.Stat(f.Name())
        if err != nil {
                return 0, err
        }

        return fi.Size(), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package git

import (
        billy "github.com/go-git/go-billy/v5"
        "github.com/go-git/go-git/v5"
        "github.com/go-git/go-git/v5/plumbing/cache"
)

// CloneConfig contains configuration for cloning a repository
type CloneConfig struct {
        // URL is the repository URL to clone
        URL string

        // Branch is the specific branch to clone (optional)
        Branch string

        // Tag is the specific tag to clone (optional)
        Tag string

        // Commit is the specific commit to clone (optional)
        Commit string
}

// validate checks that the CloneConfig is well-formed.
func (c *CloneConfig) validate() error {
        count := 0
        if c.Branch != "" {
                count++
        }
        if c.Tag != "" {
                count++
        }
        if c.Commit != "" {
                count++
        }
        if count > 1 {
                return ErrInvalidCloneConfig
        }
        return nil
}

// RepositoryInfo contains information about a Git repository
type RepositoryInfo struct {
        // Repository is the go-git repository instance
        Repository *git.Repository

        // Branch is the current branch name
        Branch string

        // RemoteURL is the remote repository URL
        RemoteURL string

        // storerFilesystem holds the in-memory filesystem containing the Git object database (.git/objects).
        // This reference is stored during Clone() and must be explicitly cleared in Cleanup() to release
        // memory, as go-git does not provide automatic cleanup of internal storage structures.
        storerFilesystem billy.Filesystem

        // objectCache holds the LRU cache for decompressed Git objects (commits, trees, blobs).
        // When a repository is cloned, git objects are decompressed and cached here. This cache must
        // be explicitly cleared via Clear() during Cleanup() to release memory, as the garbage collector
        // cannot reclaim cached objects while this reference exists.
        objectCache cache.Object
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package groups provides functionality for managing logical groupings of MCP servers.
// This file contains the CLI/filesystem-based implementation for local environments.
package groups

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "net/http"
        "sort"
        "strings"

        "github.com/stacklok/toolhive-core/httperr"
        groupval "github.com/stacklok/toolhive-core/validation/group"
        "github.com/stacklok/toolhive/pkg/state"
)

// cliManager implements the Manager interface using filesystem-based state storage
type cliManager struct {
        groupStore state.Store
}

// NewCLIManager creates a new CLI-based group manager that uses filesystem storage
func NewCLIManager() (Manager, error) {
        store, err := state.NewGroupConfigStore("toolhive")
        if err != nil {
                return nil, fmt.Errorf("failed to create group state store: %w", err)
        }

        return &cliManager{groupStore: store}, nil
}

// Create creates a new group with the given name
func (m *cliManager) Create(ctx context.Context, name string) error {
        // Validate group name
        if err := groupval.ValidateName(name); err != nil {
                return fmt.Errorf("%w: %s - %w", ErrInvalidGroupName, name, err)
        }

        group := &Group{
                Name:              name,
                RegisteredClients: []string{},
        }

        // Use CreateExclusive for atomic check-and-create to prevent race conditions
        writer, err := m.groupStore.CreateExclusive(ctx, name)
        if err != nil {
                // Check if the error is a conflict (group already exists)
                if httperr.Code(err) == http.StatusConflict {
                        return fmt.Errorf("%w: %s", ErrGroupAlreadyExists, name)
                }
                return fmt.Errorf("failed to create group: %w", err)
        }
        closed := false
        defer func() {
                if !closed {
                        if err := state.AbortWriter(writer); err != nil {
                                slog.Warn("failed to abort group writer", "name", name, "error", err)
                        }
                }
        }()

        // Write the group data
        encoder := json.NewEncoder(writer)
        encoder.SetIndent("", "  ")
        if err := encoder.Encode(group); err != nil {
                return fmt.Errorf("failed to write group: %w", err)
        }

        if err := writer.Close(); err != nil {
                closed = true
                if httperr.Code(err) == http.StatusConflict {
                        return fmt.Errorf("%w: %s", ErrGroupAlreadyExists, name)
                }
                return fmt.Errorf("failed to close group writer: %w", err)
        }
        closed = true

        return nil
}

// Get retrieves a group by name
func (m *cliManager) Get(ctx context.Context, name string) (*Group, error) {
        reader, err := m.groupStore.GetReader(ctx, name)
        if err != nil {
                return nil, fmt.Errorf("failed to get reader for group: %w", err)
        }
        defer func() {
                if err := reader.Close(); err != nil {
                        // Non-fatal: reader cleanup failure
                        slog.Debug("failed to close reader", "error", err)
                }
        }()

        var group Group
        if err := json.NewDecoder(reader).Decode(&group); err != nil {
                return nil, fmt.Errorf("failed to decode group: %w", err)
        }

        return &group, nil
}

// List returns all groups
func (m *cliManager) List(ctx context.Context) ([]*Group, error) {
        names, err := m.groupStore.List(ctx)
        if err != nil {
                return nil, fmt.Errorf("failed to list groups: %w", err)
        }

        groups := make([]*Group, 0, len(names))
        for _, name := range names {
                group, err := m.Get(ctx, name)
                if err != nil {
                        return nil, fmt.Errorf("failed to get group %s: %w", name, err)
                }
                groups = append(groups, group)
        }

        // Sort groups alphanumerically by name (handles mixed characters, numbers, etc.)
        sort.Slice(groups, func(i, j int) bool {
                return strings.Compare(groups[i].Name, groups[j].Name) < 0
        })

        return groups, nil
}

// Delete removes a group by name
func (m *cliManager) Delete(ctx context.Context, name string) error {
        return m.groupStore.Delete(ctx, name)
}

// Exists checks if a group exists
func (m *cliManager) Exists(ctx context.Context, name string) (bool, error) {
        return m.groupStore.Exists(ctx, name)
}

// RegisterClients registers multiple clients with multiple groups
func (m *cliManager) RegisterClients(ctx context.Context, groupNames []string, clientNames []string) error {
        for _, groupName := range groupNames {
                // Get the existing group
                group, err := m.Get(ctx, groupName)
                if err != nil {
                        return fmt.Errorf("failed to get group %s: %w", groupName, err)
                }

                groupModified := false
                for _, clientName := range clientNames {
                        // Check if client is already registered
                        alreadyRegistered := false
                        for _, existingClient := range group.RegisteredClients {
                                if existingClient == clientName {
                                        alreadyRegistered = true
                                        break
                                }
                        }

                        if alreadyRegistered {
                                slog.Debug("client already registered with group, skipping", "client", clientName, "group", groupName)
                                continue
                        }

                        // Add the client to the group
                        group.RegisteredClients = append(group.RegisteredClients, clientName)
                        groupModified = true
                        slog.Debug("successfully registered client with group", "client", clientName, "group", groupName)
                }

                // Only save if the group was actually modified
                if groupModified {
                        err = m.saveGroup(ctx, group)
                        if err != nil {
                                return fmt.Errorf("failed to save group %s: %w", groupName, err)
                        }
                }
        }

        return nil
}

// UnregisterClients removes multiple clients from multiple groups
func (m *cliManager) UnregisterClients(ctx context.Context, groupNames []string, clientNames []string) error {
        for _, groupName := range groupNames {
                // Get the existing group
                group, err := m.Get(ctx, groupName)
                if err != nil {
                        return fmt.Errorf("failed to get group %s: %w", groupName, err)
                }

                groupModified := false
                for _, clientName := range clientNames {
                        // Find and remove the client from the group
                        for i, existingClient := range group.RegisteredClients {
                                if existingClient == clientName {
                                        // Remove client from slice
                                        group.RegisteredClients = append(group.RegisteredClients[:i], group.RegisteredClients[i+1:]...)
                                        groupModified = true
                                        slog.Debug("successfully unregistered client from group", "client", clientName, "group", groupName)
                                        break
                                }
                        }
                }

                // Only save if the group was actually modified
                if groupModified {
                        err = m.saveGroup(ctx, group)
                        if err != nil {
                                return fmt.Errorf("failed to save group %s: %w", groupName, err)
                        }
                }
        }

        return nil
}

// Update persists changes to an existing group.
func (m *cliManager) Update(ctx context.Context, group *Group) error {
        return m.saveGroup(ctx, group)
}

// saveGroup saves the group to the group state store
func (m *cliManager) saveGroup(ctx context.Context, group *Group) error {
        writer, err := m.groupStore.GetWriter(ctx, group.Name)
        if err != nil {
                return fmt.Errorf("failed to get writer for group: %w", err)
        }
        closed := false
        defer func() {
                if !closed {
                        if err := state.AbortWriter(writer); err != nil {
                                slog.Warn("failed to abort group writer", "name", group.Name, "error", err)
                        }
                }
        }()

        encoder := json.NewEncoder(writer)
        encoder.SetIndent("", "  ")
        if err := encoder.Encode(group); err != nil {
                return fmt.Errorf("failed to write group: %w", err)
        }

        if err := writer.Close(); err != nil {
                closed = true
                return fmt.Errorf("failed to close group writer: %w", err)
        }
        closed = true

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package groups provides functionality for managing logical groupings of MCP servers.
// This file contains the CRD-based implementation for Kubernetes environments.
package groups

import (
        "context"
        "fmt"
        "log/slog"
        "sort"
        "strings"

        "k8s.io/apimachinery/pkg/api/errors"
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        "k8s.io/apimachinery/pkg/types"
        "sigs.k8s.io/controller-runtime/pkg/client"

        groupval "github.com/stacklok/toolhive-core/validation/group"
        mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
)

// crdManager implements the Manager interface using Kubernetes CRDs
type crdManager struct {
        k8sClient client.Client
        namespace string
}

// NewCRDManager creates a new CRD-based group manager
func NewCRDManager(k8sClient client.Client, namespace string) Manager {
        return &crdManager{
                k8sClient: k8sClient,
                namespace: namespace,
        }
}

// Create creates a new group with the specified name.
func (m *crdManager) Create(ctx context.Context, name string) error {
        // Validate group name
        if err := groupval.ValidateName(name); err != nil {
                return fmt.Errorf("%w: %w", ErrInvalidGroupName, err)
        }

        // Check if group already exists
        exists, err := m.Exists(ctx, name)
        if err != nil {
                return fmt.Errorf("failed to check if group exists: %w", err)
        }
        if exists {
                return fmt.Errorf("%w: %s", ErrGroupAlreadyExists, name)
        }

        // Create the MCPGroup CRD
        mcpGroup := &mcpv1beta1.MCPGroup{
                ObjectMeta: metav1.ObjectMeta{
                        Name:      name,
                        Namespace: m.namespace,
                },
                Spec: mcpv1beta1.MCPGroupSpec{},
        }

        if err := m.k8sClient.Create(ctx, mcpGroup); err != nil {
                return fmt.Errorf("failed to create MCPGroup: %w", err)
        }

        slog.Info("created mcpgroup", "name", name, "namespace", m.namespace)
        return nil
}

// Get retrieves a group by name.
func (m *crdManager) Get(ctx context.Context, name string) (*Group, error) {
        mcpGroup := &mcpv1beta1.MCPGroup{}
        err := m.k8sClient.Get(ctx, types.NamespacedName{
                Name:      name,
                Namespace: m.namespace,
        }, mcpGroup)

        if err != nil {
                if errors.IsNotFound(err) {
                        return nil, fmt.Errorf("%w: %s - %w", ErrGroupNotFound, name, err)
                }
                return nil, fmt.Errorf("failed to get MCPGroup: %w", err)
        }

        return mcpGroupToGroup(mcpGroup), nil
}

// List returns all groups.
func (m *crdManager) List(ctx context.Context) ([]*Group, error) {
        mcpGroupList := &mcpv1beta1.MCPGroupList{}
        err := m.k8sClient.List(ctx, mcpGroupList, client.InNamespace(m.namespace))
        if err != nil {
                return nil, fmt.Errorf("failed to list MCPGroups: %w", err)
        }

        groups := mcpGroupListToGroups(mcpGroupList)

        // Sort groups alphanumerically by name
        sort.Slice(groups, func(i, j int) bool {
                return strings.Compare(groups[i].Name, groups[j].Name) < 0
        })

        return groups, nil
}

// Delete removes a group by name.
func (m *crdManager) Delete(ctx context.Context, name string) error {
        mcpGroup := &mcpv1beta1.MCPGroup{
                ObjectMeta: metav1.ObjectMeta{
                        Name:      name,
                        Namespace: m.namespace,
                },
        }

        err := m.k8sClient.Delete(ctx, mcpGroup)
        if err != nil {
                if errors.IsNotFound(err) {
                        return fmt.Errorf("%w: %s - %w", ErrGroupNotFound, name, err)
                }
                return fmt.Errorf("failed to delete MCPGroup: %w", err)
        }

        slog.Info("deleted mcpgroup", "name", name, "namespace", m.namespace)
        return nil
}

// Exists checks if a group with the specified name exists.
func (m *crdManager) Exists(ctx context.Context, name string) (bool, error) {
        mcpGroup := &mcpv1beta1.MCPGroup{}
        err := m.k8sClient.Get(ctx, types.NamespacedName{
                Name:      name,
                Namespace: m.namespace,
        }, mcpGroup)

        if err != nil {
                if errors.IsNotFound(err) {
                        return false, nil
                }
                return false, fmt.Errorf("failed to check if MCPGroup exists: %w", err)
        }

        return true, nil
}

func (*crdManager) RegisterClients(context.Context, []string, []string) error {
        // In Kubernetes, client configuration management is not applicable, so this is a no-op.
        return nil
}

func (*crdManager) UnregisterClients(context.Context, []string, []string) error {
        // In Kubernetes, client configuration management is not applicable, so this is a no-op.
        return nil
}

func (*crdManager) Update(context.Context, *Group) error {
        // In Kubernetes, group state is managed via CRDs; filesystem-level updates are not applicable.
        return nil
}

// mcpGroupListToGroups converts an MCPGroupList to a slice of Groups
func mcpGroupListToGroups(mcpGroupList *mcpv1beta1.MCPGroupList) []*Group {
        groups := make([]*Group, 0, len(mcpGroupList.Items))
        for i := range mcpGroupList.Items {
                groups = append(groups, mcpGroupToGroup(&mcpGroupList.Items[i]))
        }
        return groups
}

// mcpGroupToGroup converts an MCPGroup CRD to a Group
func mcpGroupToGroup(mcpGroup *mcpv1beta1.MCPGroup) *Group {
        // In Kubernetes, RegisteredClients is not applicable - always return empty slice
        return &Group{
                Name:              mcpGroup.Name,
                RegisteredClients: []string{},
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package groups provides functionality for managing logical groupings of MCP servers.
// It includes types and interfaces for creating, retrieving, listing, and deleting groups.
package groups

import (
        "context"
        "encoding/json"
        "os"
)

// DefaultGroup is the name of the default group for workloads
const DefaultGroup = "default"

// Group represents a logical grouping of MCP servers.
type Group struct {
        Name              string   `json:"name"`
        RegisteredClients []string `json:"registered_clients"`
        Skills            []string `json:"skills,omitempty"`
        Plugins           []string `json:"plugins,omitempty"`
}

// WriteJSON serializes the Group to JSON and writes it to the provided writer
func (g *Group) WriteJSON(w *os.File) error {
        encoder := json.NewEncoder(w)
        encoder.SetIndent("", "  ")
        return encoder.Encode(g)
}

// Manager defines the interface for managing groups of MCP servers.
// It provides methods for creating, retrieving, listing, and deleting groups.
//
//go:generate mockgen -destination=mocks/mock_manager.go -package=mocks -source=group.go Manager
type Manager interface {
        // Create creates a new group with the specified name.
        // Returns an error if a group with the same name already exists.
        Create(ctx context.Context, name string) error

        // Get retrieves a group by name.
        // Returns an error if the group does not exist.
        Get(ctx context.Context, name string) (*Group, error)

        // List returns all groups.
        List(ctx context.Context) ([]*Group, error)

        // Delete removes a group by name.
        // Returns an error if the group does not exist.
        Delete(ctx context.Context, name string) error

        // Exists checks if a group with the specified name exists.
        Exists(ctx context.Context, name string) (bool, error)

        // RegisterClients registers multiple clients with multiple groups.
        RegisterClients(ctx context.Context, groupNames []string, clientNames []string) error

        // UnregisterClients removes multiple clients from multiple groups.
        UnregisterClients(ctx context.Context, groupNames []string, clientNames []string) error

        // Update persists changes to an existing group.
        // The group must already exist; returns an error if it does not.
        Update(ctx context.Context, group *Group) error
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package groups provides functionality for managing logical groupings of MCP servers.
package groups

import (
        "fmt"

        "k8s.io/apimachinery/pkg/runtime"
        clientgoscheme "k8s.io/client-go/kubernetes/scheme"

        mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/k8s"
)

const (
        // DefaultGroupName is the name of the default group
        DefaultGroupName = "default"
)

// NewManager creates a new group manager based on the runtime environment:
// - In Kubernetes mode: returns a CRD-based manager that uses MCPGroup CRDs
// - In local mode: returns a CLI/filesystem-based manager
func NewManager() (Manager, error) {
        if rt.IsKubernetesRuntime() {
                return newCRDManager()
        }
        return NewCLIManager()
}

// newCRDManager creates a CRD-based group manager for Kubernetes environments
func newCRDManager() (Manager, error) {
        // Create a scheme for controller-runtime client
        scheme := runtime.NewScheme()
        if err := clientgoscheme.AddToScheme(scheme); err != nil {
                return nil, fmt.Errorf("failed to add client-go scheme: %w", err)
        }
        if err := mcpv1beta1.AddToScheme(scheme); err != nil {
                return nil, fmt.Errorf("failed to add MCP v1beta1 scheme: %w", err)
        }

        // Create controller-runtime client with custom scheme
        k8sClient, err := k8s.NewControllerRuntimeClient(scheme)
        if err != nil {
                return nil, fmt.Errorf("failed to create Kubernetes client: %w", err)
        }

        // Detect namespace
        namespace := k8s.GetCurrentNamespace()

        return NewCRDManager(k8sClient, namespace), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package groups

import (
        "context"
        "fmt"
        "slices"
)

// AddPluginToGroup adds pluginName to the Plugins slice of the named group.
// Groups that do not exist return an error. Duplicate plugin names are skipped.
// Empty groupName is a no-op. The bool reports whether this call inserted the
// name (false when it was already a member or groupName is empty), so a later
// rollback can remove it only when this operation added it.
func AddPluginToGroup(ctx context.Context, mgr Manager, groupName string, pluginName string) (added bool, err error) {
        if groupName == "" {
                return false, nil
        }
        group, err := mgr.Get(ctx, groupName)
        if err != nil {
                return false, fmt.Errorf("getting group %q: %w", groupName, err)
        }

        if slices.Contains(group.Plugins, pluginName) {
                return false, nil
        }

        group.Plugins = append(group.Plugins, pluginName)
        if err := mgr.Update(ctx, group); err != nil {
                return false, fmt.Errorf("updating group %q: %w", groupName, err)
        }
        return true, nil
}

// RemovePluginFromGroup removes pluginName from the named group's Plugins
// slice. Missing membership is a no-op. Empty groupName is a no-op.
func RemovePluginFromGroup(ctx context.Context, mgr Manager, groupName string, pluginName string) error {
        if groupName == "" {
                return nil
        }
        group, err := mgr.Get(ctx, groupName)
        if err != nil {
                return fmt.Errorf("getting group %q: %w", groupName, err)
        }

        idx := slices.Index(group.Plugins, pluginName)
        if idx < 0 {
                return nil
        }
        group.Plugins = slices.Delete(group.Plugins, idx, idx+1)
        if err := mgr.Update(ctx, group); err != nil {
                return fmt.Errorf("updating group %q: %w", groupName, err)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package groups

import (
        "context"
        "fmt"
        "slices"
)

// AddSkillToGroup adds skillName to the Skills slice of the named group.
// Groups that do not exist return an error. Duplicate skill names are skipped.
// Empty groupName is a no-op. The bool reports whether this call inserted the
// name (false when it was already a member or groupName is empty), so a later
// rollback can remove it only when this operation added it.
func AddSkillToGroup(ctx context.Context, mgr Manager, groupName string, skillName string) (added bool, err error) {
        if groupName == "" {
                return false, nil
        }
        group, err := mgr.Get(ctx, groupName)
        if err != nil {
                return false, fmt.Errorf("getting group %q: %w", groupName, err)
        }

        if slices.Contains(group.Skills, skillName) {
                return false, nil
        }

        group.Skills = append(group.Skills, skillName)
        if err := mgr.Update(ctx, group); err != nil {
                return false, fmt.Errorf("updating group %q: %w", groupName, err)
        }
        return true, nil
}

// RemoveSkillFromGroup removes skillName from the named group's Skills slice.
// Missing membership is a no-op. Empty groupName is a no-op.
func RemoveSkillFromGroup(ctx context.Context, mgr Manager, groupName string, skillName string) error {
        if groupName == "" {
                return nil
        }
        group, err := mgr.Get(ctx, groupName)
        if err != nil {
                return fmt.Errorf("getting group %q: %w", groupName, err)
        }

        idx := slices.Index(group.Skills, skillName)
        if idx < 0 {
                return nil
        }
        group.Skills = slices.Delete(group.Skills, idx, idx+1)
        if err := mgr.Update(ctx, group); err != nil {
                return fmt.Errorf("updating group %q: %w", group.Name, err)
        }
        return nil
}

// RemoveSkillFromAllGroups removes skillName from every group that references it.
// It is a no-op when the skill is not found in any group.
func RemoveSkillFromAllGroups(ctx context.Context, mgr Manager, skillName string) error {
        allGroups, err := mgr.List(ctx)
        if err != nil {
                return fmt.Errorf("listing groups: %w", err)
        }

        for _, group := range allGroups {
                modified := false
                for i, s := range group.Skills {
                        if s == skillName {
                                group.Skills = append(group.Skills[:i], group.Skills[i+1:]...)
                                modified = true
                                break
                        }
                }
                if modified {
                        if err := mgr.Update(ctx, group); err != nil {
                                return fmt.Errorf("updating group %q: %w", group.Name, err)
                        }
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package healthcheck provides common healthcheck functionality for ToolHive proxies.
// It includes MCP server status information and standardized response formats.
package healthcheck

import (
        "context"
        "encoding/json"
        "log/slog"
        "net/http"
        "time"
)

// HealthStatus represents the overall health status
type HealthStatus string

const (
        // StatusHealthy indicates the service is healthy
        StatusHealthy HealthStatus = "healthy"
        // StatusUnhealthy indicates the service is unhealthy
        StatusUnhealthy HealthStatus = "unhealthy"
        // StatusDegraded indicates the service is partially healthy
        StatusDegraded HealthStatus = "degraded"
)

// MCPStatus represents the status of an MCP server connection
type MCPStatus struct {
        // Available indicates if the MCP server is reachable
        Available bool `json:"available"`
        // ResponseTime is the time taken to ping the MCP server (in milliseconds)
        ResponseTime *int64 `json:"response_time_ms,omitempty"`
        // Error contains any error message if the MCP server is not available
        Error string `json:"error,omitempty"`
        // LastChecked is the timestamp of the last health check
        LastChecked time.Time `json:"last_checked"`
}

// HealthResponse represents the standardized health check response.
//
// This endpoint is unauthenticated (it must stay reachable for Kubernetes
// liveness and readiness probes), so it deliberately exposes no build
// fingerprint — version, commit, build date, Go version, or platform. This
// matches the minimal shape already used by the vMCP server and the v1 API,
// which redact the same information to avoid disclosure on an unauthenticated
// path.
type HealthResponse struct {
        // Status is the overall health status
        Status HealthStatus `json:"status"`
        // Timestamp is when the health check was performed
        Timestamp time.Time `json:"timestamp"`
        // Transport indicates the type of transport used by the MCP server (stdio, sse)
        Transport string `json:"transport"`
        // MCP contains MCP server status information
        MCP *MCPStatus `json:"mcp,omitempty"`
}

// MCPPinger defines the interface for pinging MCP servers
// Implementations should follow the MCP ping specification:
// https://modelcontextprotocol.io/specification/2025-03-26/basic/utilities/ping
type MCPPinger interface {
        // Ping sends a ping request to the MCP server and returns the response time
        // The implementation should send a JSON-RPC request with method "ping" and no parameters,
        // and expect an empty response: {"jsonrpc": "2.0", "id": "123", "result": {}}
        Ping(ctx context.Context) (time.Duration, error)
}

// HealthChecker provides health check functionality for proxies
type HealthChecker struct {
        transport string
        mcpPinger MCPPinger
}

// NewHealthChecker creates a new health checker instance
func NewHealthChecker(transport string, mcpPinger MCPPinger) *HealthChecker {
        return &HealthChecker{
                transport: transport,
                mcpPinger: mcpPinger,
        }
}

// CheckHealth performs a comprehensive health check including MCP server status
func (hc *HealthChecker) CheckHealth(ctx context.Context) *HealthResponse {
        response := &HealthResponse{
                Status:    StatusHealthy,
                Timestamp: time.Now(),
                Transport: hc.transport,
        }

        // Check MCP server status if pinger is available
        if hc.mcpPinger != nil {
                mcpStatus := hc.checkMCPStatus(ctx)
                response.MCP = mcpStatus

                // Adjust overall status based on MCP status
                if !mcpStatus.Available {
                        response.Status = StatusDegraded
                }
        }

        return response
}

// checkMCPStatus checks the status of the MCP server using ping
func (hc *HealthChecker) checkMCPStatus(ctx context.Context) *MCPStatus {
        status := &MCPStatus{
                LastChecked: time.Now(),
        }

        duration, err := hc.mcpPinger.Ping(ctx)

        if err != nil {
                status.Available = false
                status.Error = err.Error()
                slog.Debug("MCP ping failed", "error", err)
        } else {
                status.Available = true
                responseTimeMs := duration.Milliseconds()
                status.ResponseTime = &responseTimeMs
                slog.Debug("MCP ping successful", "duration", duration)
        }

        return status
}

// ServeHTTP implements http.Handler for health check endpoints
func (hc *HealthChecker) ServeHTTP(w http.ResponseWriter, r *http.Request) {
        if r.Method != http.MethodGet {
                http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
                return
        }

        health := hc.CheckHealth(r.Context())

        w.Header().Set("Content-Type", "application/json")

        // Set HTTP status code based on health status
        switch health.Status {
        case StatusHealthy:
                w.WriteHeader(http.StatusOK)
        case StatusDegraded:
                w.WriteHeader(http.StatusOK) // Still return 200 for degraded state
        case StatusUnhealthy:
                w.WriteHeader(http.StatusServiceUnavailable)
        }

        if err := json.NewEncoder(w).Encode(health); err != nil {
                slog.Warn("Failed to encode health response", "error", err)
                http.Error(w, "Internal server error", http.StatusInternalServerError)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package ignore provides functionality for processing .thvignore files
// to filter bind mount contents using tmpfs overlays.
package ignore

import (
        "bufio"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "strings"
        "sync"
        "sync/atomic"

        "github.com/adrg/xdg"
)

// Processor handles loading and processing ignore patterns
type Processor struct {
        GlobalPatterns     []string
        LocalPatterns      []string
        Config             *Config
        sharedEmptyFile    string     // Cached path to a single shared empty file
        overlayArtifacts   []string   // Paths to created overlay artifacts (files and directories)
        overlayArtifactsMu sync.Mutex // Mutex to protect overlayArtifacts
        workloadID         string     // Unique identifier for this workload
        artifactDir        string     // Directory to store overlay artifacts
}

// Config holds configuration for ignore processing
type Config struct {
        LoadGlobal    bool // Whether to load global ignore patterns
        PrintOverlays bool // Whether to print resolved overlay paths for debugging
}

const ignoreFileName = ".thvignore"

// processorSeq distinguishes processors within one process so each owns a
// separate artifact directory and Cleanup never removes another's storage.
var processorSeq atomic.Uint64

// NewProcessor creates a new Processor instance with the given configuration
func NewProcessor(config *Config) *Processor {
        if config == nil {
                config = &Config{
                        LoadGlobal:    true,
                        PrintOverlays: false,
                }
        }

        // Generate a unique workload ID for this processor instance
        workloadID := fmt.Sprintf("thvignore-%d-%d", os.Getpid(), processorSeq.Add(1))

        // Create artifact directory for this workload
        artifactDir := getArtifactDir(workloadID)

        return &Processor{
                GlobalPatterns:   make([]string, 0),
                LocalPatterns:    make([]string, 0),
                Config:           config,
                overlayArtifacts: make([]string, 0),
                workloadID:       workloadID,
                artifactDir:      artifactDir,
        }
}

// getArtifactDir returns the directory path for storing overlay artifacts
func getArtifactDir(workloadID string) string {
        // Use XDG runtime directory if available, otherwise fall back to temp
        runtimeDir := os.Getenv("XDG_RUNTIME_DIR")
        if runtimeDir == "" {
                runtimeDir = os.TempDir()
        }
        return filepath.Join(runtimeDir, "toolhive", "overlays", workloadID)
}

// LoadGlobal loads global ignore patterns from ~/.config/toolhive/thvignore
func (p *Processor) LoadGlobal() error {
        // Skip loading global patterns if disabled in config
        if !p.Config.LoadGlobal {
                slog.Debug("Global ignore patterns disabled by configuration")
                return nil
        }

        globalIgnoreFile, err := xdg.ConfigFile("toolhive/thvignore")
        if err != nil {
                slog.Debug("Failed to get XDG config file path", "error", err)
                return nil // Not a fatal error, continue without global patterns
        }

        patterns, err := p.loadIgnoreFile(globalIgnoreFile)
        if err != nil {
                if os.IsNotExist(err) {
                        slog.Debug("Global ignore file not found", "path", globalIgnoreFile)
                        return nil // Not a fatal error
                }
                return fmt.Errorf("failed to load global ignore file: %w", err)
        }

        p.GlobalPatterns = patterns
        slog.Debug("Loaded global ignore patterns", "count", len(patterns), "path", globalIgnoreFile)
        return nil
}

// LoadLocal loads local ignore patterns from the configured ignore file in the specified directory
func (p *Processor) LoadLocal(sourceDir string) error {
        localIgnoreFile := filepath.Join(sourceDir, ignoreFileName)
        patterns, err := p.loadIgnoreFile(localIgnoreFile)
        if err != nil {
                if os.IsNotExist(err) {
                        slog.Debug("Local ignore file not found", "path", localIgnoreFile)
                        return nil // Not a fatal error
                }
                return fmt.Errorf("failed to load local ignore file: %w", err)
        }

        p.LocalPatterns = append(p.LocalPatterns, patterns...)
        slog.Debug("Loaded local ignore patterns", "count", len(patterns), "path", localIgnoreFile)
        return nil
}

// loadIgnoreFile loads patterns from a .gitignore-style file
func (*Processor) loadIgnoreFile(filePath string) ([]string, error) {
        // #nosec G304 - This is intentional as we're reading user-specified ignore files
        file, err := os.Open(filePath)
        if err != nil {
                return nil, err
        }
        defer func() {
                if err := file.Close(); err != nil {
                        slog.Warn("Failed to close ignore file", "error", err)
                }
        }()

        var patterns []string
        scanner := bufio.NewScanner(file)

        for scanner.Scan() {
                line := strings.TrimSpace(scanner.Text())

                // Skip empty lines and comments
                if line == "" || strings.HasPrefix(line, "#") {
                        continue
                }

                patterns = append(patterns, line)
        }

        if err := scanner.Err(); err != nil {
                return nil, fmt.Errorf("error reading ignore file: %w", err)
        }

        return patterns, nil
}

// OverlayMount represents a mount that should overlay an ignored path
type OverlayMount struct {
        ContainerPath string // Path in the container to overlay
        HostPath      string // Host path (for bind mounts) or empty (for tmpfs)
        Type          string // "tmpfs" for directories, "bind" for files
}

// GetOverlayMounts returns mounts that should overlay ignored paths
// based on the loaded ignore patterns
func (p *Processor) GetOverlayMounts(bindMount, containerPath string) []OverlayMount {
        var overlayMounts []OverlayMount
        overlaySet := make(map[string]bool) // To avoid duplicates

        // Combine global and local patterns
        allPatterns := append(p.GlobalPatterns, p.LocalPatterns...)

        for _, pattern := range allPatterns {
                overlayMounts = append(overlayMounts, p.processPattern(bindMount, containerPath, pattern, overlaySet)...)
        }

        p.printOverlays(overlayMounts, bindMount, containerPath)
        return overlayMounts
}

// processPattern processes a single ignore pattern and returns overlay mounts
func (p *Processor) processPattern(bindMount, containerPath, pattern string, overlaySet map[string]bool) []OverlayMount {
        var overlayMounts []OverlayMount
        matchingPaths := p.getMatchingPaths(bindMount, pattern)

        for _, matchPath := range matchingPaths {
                if overlay := p.createOverlayMount(matchPath, bindMount, containerPath, pattern, overlaySet); overlay != nil {
                        overlayMounts = append(overlayMounts, *overlay)
                }
        }

        return overlayMounts
}

// createOverlayMount creates an overlay mount for a matched path
func (p *Processor) createOverlayMount(
        matchPath, bindMount, containerPath, pattern string, overlaySet map[string]bool,
) *OverlayMount {
        // Calculate relative path from bind mount to matched path
        relPath, err := filepath.Rel(bindMount, matchPath)
        if err != nil {
                slog.Debug("Failed to calculate relative path", "matchPath", matchPath, "error", err)
                return nil
        }

        // Convert to container path
        containerOverlayPath := filepath.Join(containerPath, relPath)

        // Skip if we already have this overlay
        if overlaySet[containerOverlayPath] {
                return nil
        }
        overlaySet[containerOverlayPath] = true

        // Check if the matched path is a directory or file
        info, err := os.Stat(matchPath)
        if err != nil {
                slog.Debug("Failed to stat path", "path", matchPath, "error", err)
                return nil
        }

        if info.IsDir() {
                // For directories, create an empty directory and bind mount it
                emptyDirPath, err := p.createEmptyDirectory()
                if err != nil {
                        slog.Debug("Failed to create empty directory for pattern", "pattern", pattern, "error", err)
                        return nil
                }

                slog.Debug("Adding bind overlay for directory pattern",
                        "pattern", pattern, "containerPath", containerOverlayPath, "hostPath", emptyDirPath)
                return &OverlayMount{
                        ContainerPath: containerOverlayPath,
                        HostPath:      emptyDirPath,
                        Type:          "bind",
                }
        }

        // For files, create empty file and bind mount it
        emptyFilePath, err := p.createEmptyFile()
        if err != nil {
                slog.Debug("Failed to create empty file for pattern", "pattern", pattern, "error", err)
                return nil
        }

        slog.Debug("Adding bind overlay for file pattern",
                "pattern", pattern, "containerPath", containerOverlayPath, "hostPath", emptyFilePath)
        return &OverlayMount{
                ContainerPath: containerOverlayPath,
                HostPath:      emptyFilePath,
                Type:          "bind",
        }
}

// printOverlays prints resolved overlays if requested
func (p *Processor) printOverlays(overlayMounts []OverlayMount, bindMount, containerPath string) {
        if p.Config.PrintOverlays && len(overlayMounts) > 0 {
                slog.Info("Resolved overlays for mount", "bindMount", bindMount, "containerPath", containerPath)
                for _, overlay := range overlayMounts {
                        slog.Info("Overlay mount", "containerPath", overlay.ContainerPath, "hostPath", overlay.HostPath)
                }
        }
}

// createEmptyFile returns a path to a shared empty file for bind mounting
func (p *Processor) createEmptyFile() (string, error) {
        // Return cached shared empty file if it exists
        if p.sharedEmptyFile != "" {
                // Verify the file still exists
                if _, err := os.Stat(p.sharedEmptyFile); err == nil {
                        return p.sharedEmptyFile, nil
                }
                // File was deleted, clear the cache
                p.sharedEmptyFile = ""
        }

        // Create a new shared empty file
        tmpFile, err := os.CreateTemp("", "thvignore-shared-empty-*")
        if err != nil {
                return "", fmt.Errorf("failed to create shared empty file: %w", err)
        }
        if err := tmpFile.Close(); err != nil {
                return "", fmt.Errorf("failed to close shared empty file: %w", err)
        }

        // Cache the path for reuse
        p.sharedEmptyFile = tmpFile.Name()
        slog.Debug("Created shared empty file for bind mounting", "path", p.sharedEmptyFile)

        return p.sharedEmptyFile, nil
}

// createEmptyDirectory creates an empty directory for bind mounting
func (p *Processor) createEmptyDirectory() (string, error) {
        p.overlayArtifactsMu.Lock()
        defer p.overlayArtifactsMu.Unlock()

        // Ensure artifact directory exists
        if err := os.MkdirAll(p.artifactDir, 0750); err != nil {
                return "", fmt.Errorf("failed to create artifact directory: %w", err)
        }

        // Create a unique empty directory
        emptyDir, err := os.MkdirTemp(p.artifactDir, "dir-*")
        if err != nil {
                return "", fmt.Errorf("failed to create empty directory: %w", err)
        }

        // Track this artifact for cleanup
        p.overlayArtifacts = append(p.overlayArtifacts, emptyDir)
        slog.Debug("Created empty directory for bind mounting", "path", emptyDir)

        return emptyDir, nil
}

// Cleanup removes all overlay artifacts (shared empty file and directories)
func (p *Processor) Cleanup() error {
        p.overlayArtifactsMu.Lock()
        defer p.overlayArtifactsMu.Unlock()

        var lastErr error

        // Remove shared empty file
        if p.sharedEmptyFile != "" {
                if err := os.Remove(p.sharedEmptyFile); err != nil && !os.IsNotExist(err) {
                        slog.Debug("Failed to remove shared empty file", "path", p.sharedEmptyFile, "error", err)
                        lastErr = fmt.Errorf("failed to remove shared empty file: %w", err)
                } else {
                        slog.Debug("Cleaned up shared empty file", "path", p.sharedEmptyFile)
                }
                p.sharedEmptyFile = ""
        }

        // Remove all overlay artifacts (empty directories)
        for _, artifact := range p.overlayArtifacts {
                if err := os.RemoveAll(artifact); err != nil && !os.IsNotExist(err) {
                        slog.Debug("Failed to remove overlay artifact", "path", artifact, "error", err)
                        lastErr = fmt.Errorf("failed to remove overlay artifact: %w", err)
                } else {
                        slog.Debug("Cleaned up overlay artifact", "path", artifact)
                }
        }
        p.overlayArtifacts = nil

        // Remove the artifact directory if it's empty
        if p.artifactDir != "" {
                if err := os.Remove(p.artifactDir); err != nil && !os.IsNotExist(err) {
                        // It's okay if the directory is not empty or doesn't exist
                        slog.Debug("Could not remove artifact directory", "path", p.artifactDir, "error", err)
                }
        }

        return lastErr
}

// GetOverlayPaths returns container paths that should be overlaid
// based on the loaded ignore patterns (kept for backward compatibility)
func (p *Processor) GetOverlayPaths(bindMount, containerPath string) []string {
        overlayMounts := p.GetOverlayMounts(bindMount, containerPath)
        var overlayPaths []string

        for _, mount := range overlayMounts {
                overlayPaths = append(overlayPaths, mount.ContainerPath)
        }

        return overlayPaths
}

// getMatchingPaths returns all paths that match the given pattern in the directory
func (*Processor) getMatchingPaths(dir, pattern string) []string {
        var matchingPaths []string

        // Handle directory patterns (ending with /)
        if strings.HasSuffix(pattern, "/") {
                dirPattern := strings.TrimSuffix(pattern, "/")
                targetPath := filepath.Join(dir, dirPattern)
                if info, err := os.Stat(targetPath); err == nil && info.IsDir() {
                        matchingPaths = append(matchingPaths, targetPath)
                }
                return matchingPaths
        }

        // Handle direct file/directory matches
        targetPath := filepath.Join(dir, pattern)
        if _, err := os.Stat(targetPath); err == nil {
                matchingPaths = append(matchingPaths, targetPath)
                return matchingPaths
        }

        // Handle glob patterns
        matches, err := filepath.Glob(filepath.Join(dir, pattern))
        if err != nil {
                slog.Debug("Error matching pattern", "pattern", pattern, "error", err)
                return matchingPaths
        }

        return matches
}

// patternMatchesInDirectory checks if a pattern matches any files/directories in the given directory
func (p *Processor) patternMatchesInDirectory(dir, pattern string) bool {
        return len(p.getMatchingPaths(dir, pattern)) > 0
}

// ShouldIgnore checks if a given path should be ignored based on loaded patterns
func (p *Processor) ShouldIgnore(path string) bool {
        // Combine global and local patterns
        allPatterns := append(p.GlobalPatterns, p.LocalPatterns...)

        for _, pattern := range allPatterns {
                // Simple pattern matching - can be enhanced with more sophisticated gitignore-style matching
                if matched, err := filepath.Match(pattern, filepath.Base(path)); err == nil && matched {
                        return true
                }

                // Check if path contains the pattern
                if strings.Contains(path, pattern) {
                        return true
                }
        }

        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package json provides JSON-related utilities for ToolHive.
//
// This package extends Go's standard json package with types that work
// seamlessly with both Kubernetes CRDs and CLI YAML configurations.
package json

import (
        stdjson "encoding/json"
        "fmt"

        "go.yaml.in/yaml/v3"
)

// Data stores JSON-compatible data of type T. It supports both JSON and YAML
// marshaling/unmarshaling, making it suitable for use in both Kubernetes CRDs
// and CLI YAML configurations.
//
// The Value field stores the Go value directly, which simplifies usage in tests
// and when working with the data programmatically.
//
// Common instantiations:
//   - Data[any] (aliased as Any) for arbitrary JSON values
//   - Data[map[string]any] (aliased as Map) for JSON objects
//
// +kubebuilder:pruning:PreserveUnknownFields
// +kubebuilder:validation:Type=object
type Data[T any] struct {
        // Value holds the typed Go value.
        Value T `json:"-" yaml:"-"`
}

// MarshalJSON implements json.Marshaler.
func (d Data[T]) MarshalJSON() ([]byte, error) {
        if d.IsEmpty() {
                return []byte("null"), nil
        }
        return stdjson.Marshal(d.Value)
}

// UnmarshalJSON implements json.Unmarshaler.
func (d *Data[T]) UnmarshalJSON(data []byte) error {
        if len(data) == 0 || string(data) == "null" {
                var zero T
                d.Value = zero
                return nil
        }
        var v T
        if err := stdjson.Unmarshal(data, &v); err != nil {
                return err
        }
        d.Value = v
        return nil
}

// MarshalYAML implements yaml.Marshaler.
func (d Data[T]) MarshalYAML() (interface{}, error) {
        return d.Value, nil
}

// UnmarshalYAML implements yaml.Unmarshaler.
func (d *Data[T]) UnmarshalYAML(node *yaml.Node) error {
        if node.Kind == yaml.ScalarNode && node.Tag == "!!null" {
                var zero T
                d.Value = zero
                return nil
        }

        var value T
        if err := node.Decode(&value); err != nil {
                return err
        }
        d.Value = value
        return nil
}

// Get returns the stored value.
func (d Data[T]) Get() T {
        return d.Value
}

// IsEmpty returns true if the value is nil or empty.
// For maps and slices, it checks if the length is 0.
func (d Data[T]) IsEmpty() bool {
        v := any(d.Value)
        if v == nil {
                return true
        }
        // Check for empty maps and slices
        switch val := v.(type) {
        case map[string]any:
                return len(val) == 0
        case []any:
                return len(val) == 0
        }
        return false
}

// DeepCopyInto copies the receiver into out. Required for controller-gen.
func (d *Data[T]) DeepCopyInto(out *Data[T]) {
        if any(d.Value) != nil {
                // Deep copy Value by marshaling and unmarshaling
                raw, err := stdjson.Marshal(d.Value)
                if err != nil {
                        panic(fmt.Sprintf("failed to marshal Data[%T]: %v", d.Value, err))
                }

                var v T
                if err := stdjson.Unmarshal(raw, &v); err != nil {
                        panic(fmt.Sprintf("failed to unmarshal Data[%T]: %v", d.Value, err))
                }
                out.Value = v
        } else {
                var zero T
                out.Value = zero
        }
}

// DeepCopy creates a deep copy. Required for controller-gen.
func (d *Data[T]) DeepCopy() *Data[T] {
        if d == nil {
                return nil
        }
        out := new(Data[T])
        d.DeepCopyInto(out)
        return out
}

// Any is a type alias for Data[any], storing arbitrary JSON values.
// This is the most flexible type, suitable when the JSON structure is unknown.
//
// +gendoc
// +kubebuilder:pruning:PreserveUnknownFields
// +kubebuilder:validation:Type=object
type Any = Data[any]

// Map is a type alias for Data[map[string]any], storing JSON objects.
// Use this when you know the data will always be a JSON object (not array, string, etc.).
//
// +gendoc
// +kubebuilder:pruning:PreserveUnknownFields
// +kubebuilder:validation:Type=object
type Map = Data[map[string]any]

// NewAny creates an Any (Data[any]) from a value.
// This is a convenience function for tests and programmatic use.
func NewAny(v any) Any {
        return Any{Value: v}
}

// NewMap creates a Map (Data[map[string]any]) from a map.
func NewMap(m map[string]any) Map {
        return Map{Value: m}
}

// ToMap returns the data as a map[string]any.
// This is a convenience method for Any types.
// Returns nil if there is no data or if the data is not a map.
func (d Data[T]) ToMap() (map[string]any, error) {
        if any(d.Value) == nil {
                return nil, nil
        }
        if m, ok := any(d.Value).(map[string]any); ok {
                return m, nil
        }
        // Data is set but not a map - marshal and unmarshal to convert
        raw, err := stdjson.Marshal(d.Value)
        if err != nil {
                return nil, err
        }
        var m map[string]any
        if err := stdjson.Unmarshal(raw, &m); err != nil {
                return nil, err
        }
        return m, nil
}

// ToAny returns the data as any type.
// This is useful when you need to pass the value to functions expecting any.
func (d Data[T]) ToAny() (any, error) {
        return d.Value, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package k8s

import (
        "fmt"

        "k8s.io/apimachinery/pkg/runtime"
        "k8s.io/client-go/kubernetes"
        "k8s.io/client-go/rest"
        "sigs.k8s.io/controller-runtime/pkg/client"
)

// NewClient creates a new standard Kubernetes clientset using the default config loading.
// It tries in-cluster config first, then falls back to out-of-cluster config.
// Use this when you only need to work with standard Kubernetes resources.
func NewClient() (kubernetes.Interface, *rest.Config, error) {
        config, err := GetConfig()
        if err != nil {
                return nil, nil, fmt.Errorf("failed to create kubernetes config: %w", err)
        }

        clientset, err := kubernetes.NewForConfig(config)
        if err != nil {
                return nil, nil, fmt.Errorf("failed to create kubernetes client: %w", err)
        }

        return clientset, config, nil
}

// NewClientWithConfig creates a new standard Kubernetes clientset from the provided config.
// Use this when you have an existing config and only need standard Kubernetes resources.
func NewClientWithConfig(config *rest.Config) (kubernetes.Interface, error) {
        if config == nil {
                return nil, fmt.Errorf("failed to create kubernetes client: config cannot be nil")
        }

        clientset, err := kubernetes.NewForConfig(config)
        if err != nil {
                return nil, fmt.Errorf("failed to create kubernetes client: %w", err)
        }

        return clientset, nil
}

// NewControllerRuntimeClient creates a new controller-runtime client with a custom scheme.
// This is useful for working with Custom Resource Definitions (CRDs) alongside standard resources.
// The scheme should have all required types registered before calling this function.
//
// Example:
//
//        scheme := runtime.NewScheme()
//        utilruntime.Must(clientgoscheme.AddToScheme(scheme))
//        utilruntime.Must(mycrds.AddToScheme(scheme))
//        k8sClient, err := k8s.NewControllerRuntimeClient(scheme)
func NewControllerRuntimeClient(scheme *runtime.Scheme) (client.Client, error) {
        config, err := GetConfig()
        if err != nil {
                return nil, fmt.Errorf("failed to get kubernetes config: %w", err)
        }

        return newControllerRuntimeClientWithConfig(config, scheme)
}

// newControllerRuntimeClientWithConfig is the internal implementation for creating a controller-runtime client
func newControllerRuntimeClientWithConfig(config *rest.Config, scheme *runtime.Scheme) (client.Client, error) {
        if scheme == nil {
                return nil, fmt.Errorf("failed to create controller-runtime client: scheme cannot be nil")
        }

        k8sClient, err := client.New(config, client.Options{Scheme: scheme})
        if err != nil {
                return nil, fmt.Errorf("failed to create controller-runtime client: %w", err)
        }

        return k8sClient, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package k8s provides common Kubernetes utilities for creating clients,
// configs, and namespace detection that can be shared across packages.
package k8s

import (
        "fmt"

        "k8s.io/client-go/rest"
        "k8s.io/client-go/tools/clientcmd"
)

// configLoader defines the interface for loading Kubernetes configs
type configLoader interface {
        // InClusterConfig returns the in-cluster config
        InClusterConfig() (*rest.Config, error)
        // LoadFromRules loads config using clientcmd loading rules
        LoadFromRules(rules *clientcmd.ClientConfigLoadingRules) (*rest.Config, error)
}

// defaultConfigLoader implements configLoader using real Kubernetes client-go functions
type defaultConfigLoader struct{}

func (*defaultConfigLoader) InClusterConfig() (*rest.Config, error) {
        return rest.InClusterConfig()
}

func (*defaultConfigLoader) LoadFromRules(rules *clientcmd.ClientConfigLoadingRules) (*rest.Config, error) {
        configOverrides := &clientcmd.ConfigOverrides{}
        kubeConfig := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, configOverrides)
        return kubeConfig.ClientConfig()
}

// GetConfig returns a Kubernetes REST config with the following fallback strategy:
//  1. In-cluster config (when running inside a Kubernetes pod)
//  2. Out-of-cluster config using standard kubeconfig loading rules:
//     a. KUBECONFIG environment variable (colon-separated paths)
//     b. ~/.kube/config file
//
// This order prioritizes in-cluster config for security and reliability when
// running as a pod, while supporting local development when running outside the cluster.
//
// The returned config uses secure defaults:
//   - TLS certificate verification is enabled
//   - In-cluster: Service account CA cert is used automatically
//   - Out-of-cluster: certificate-authority-data from kubeconfig is used
//   - Default QPS: 5 requests/second, Burst: 10 (suitable for most use cases)
//
// For high-volume operations (e.g., operators reconciling many resources),
// consider increasing QPS and Burst limits:
//
//        config, err := k8s.GetConfig()
//        if err != nil {
//            return err
//        }
//        config.QPS = 50      // Increase from default 5
//        config.Burst = 100   // Increase from default 10
func GetConfig() (*rest.Config, error) {
        return getConfigWithLoader(&defaultConfigLoader{})
}

// getConfigWithLoader is the internal implementation that accepts a configLoader
func getConfigWithLoader(loader configLoader) (*rest.Config, error) {
        // Try in-cluster config first
        config, err := loader.InClusterConfig()
        if err == nil {
                return config, nil
        }

        // If in-cluster config fails, try out-of-cluster config
        loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()
        config, err = loader.LoadFromRules(loadingRules)
        if err != nil {
                return nil, fmt.Errorf("failed to create kubernetes config (tried both in-cluster and out-of-cluster): %w", err)
        }

        return config, nil
}

// getConfigFromKubeconfigFile loads config from a specific kubeconfig file path
// This is primarily useful for testing
func getConfigFromKubeconfigFile(kubeconfigPath string) (*rest.Config, error) {
        loadingRules := &clientcmd.ClientConfigLoadingRules{
                ExplicitPath: kubeconfigPath,
        }
        configOverrides := &clientcmd.ConfigOverrides{}
        kubeConfig := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, configOverrides)
        config, err := kubeConfig.ClientConfig()
        if err != nil {
                return nil, fmt.Errorf("failed to load config from %s: %w", kubeconfigPath, err)
        }
        return config, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package k8s

import (
        "fmt"
        "os"
        "strings"

        "k8s.io/client-go/tools/clientcmd"
)

const (
        // defaultNamespace is the default Kubernetes namespace
        defaultNamespace = "default"
        // defaultServiceAccountPath is the default path to the service account namespace file
        defaultServiceAccountPath = "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
        // defaultPodNamespaceEnv is the default environment variable for POD_NAMESPACE
        defaultPodNamespaceEnv = "POD_NAMESPACE"
)

// GetCurrentNamespace attempts to determine the current Kubernetes namespace
// using multiple methods, falling back to "default" if none succeed.
func GetCurrentNamespace() string {
        // Method 1: Try to read from the service account namespace file
        if ns, err := getNamespaceFromServiceAccountPath(defaultServiceAccountPath); err == nil {
                return ns
        }

        // Method 2: Try to get the namespace from environment variables
        if ns, err := getNamespaceFromEnvVar(defaultPodNamespaceEnv); err == nil {
                return ns
        }

        // Method 3: Try to get the namespace from the current kubectl context
        if ns, err := getNamespaceFromKubeConfig(); err == nil {
                return ns
        }

        // Method 4: Fall back to default
        return defaultNamespace
}

// getNamespaceFromServiceAccountPath attempts to read the namespace from a service account token file
// This is a thin I/O wrapper - the logic is in parseNamespaceFromFile
func getNamespaceFromServiceAccountPath(path string) (string, error) {
        //nolint:gosec // G304: Reading from configurable path is intentional for testing
        data, err := os.ReadFile(path)
        if err != nil {
                return "", fmt.Errorf("failed to read namespace file: %w", err)
        }
        return parseNamespaceFromFile(data)
}

// parseNamespaceFromFile parses namespace from file data
// This is pure logic, fully testable without I/O
func parseNamespaceFromFile(data []byte) (string, error) {
        // Kubernetes writes the namespace file without trailing newlines, but we trim
        // them for robustness in case the file was manually edited or created incorrectly.
        // We only trim newlines (not all whitespace) to be conservative.
        ns := strings.TrimRight(string(data), "\n\r")
        if ns == "" {
                return "", fmt.Errorf("namespace file is empty")
        }
        return ns, nil
}

// getNamespaceFromEnvVar attempts to get the namespace from a specific environment variable
// This is a thin I/O wrapper - the logic is in validateNamespaceValue
func getNamespaceFromEnvVar(envVar string) (string, error) {
        return validateNamespaceValue(os.Getenv(envVar), envVar)
}

// validateNamespaceValue validates a namespace value from an environment variable
// This is pure logic, fully testable without environment access
func validateNamespaceValue(ns, source string) (string, error) {
        if ns == "" {
                return "", fmt.Errorf("%s environment variable not set", source)
        }
        return ns, nil
}

// getNamespaceFromKubeConfig attempts to get the namespace from the current kubectl context
func getNamespaceFromKubeConfig() (string, error) {
        kubeConfig := loadKubeconfigRaw()
        return extractNamespaceFromKubeconfig(kubeConfig)
}

// loadKubeconfigRaw loads the raw kubeconfig
// This is a thin I/O wrapper
func loadKubeconfigRaw() clientcmd.ClientConfig {
        loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()
        configOverrides := &clientcmd.ConfigOverrides{}
        return clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, configOverrides)
}

// extractNamespaceFromKubeconfig extracts namespace from kubeconfig
// This is pure logic, testable with mock configs
func extractNamespaceFromKubeconfig(kubeConfig clientcmd.ClientConfig) (string, error) {
        rawConfig, err := kubeConfig.RawConfig()
        if err != nil {
                return "", fmt.Errorf("failed to load kubeconfig: %w", err)
        }

        currentContext := rawConfig.CurrentContext
        if currentContext == "" {
                return "", fmt.Errorf("no current context set in kubeconfig")
        }

        contextConfig, exists := rawConfig.Contexts[currentContext]
        if !exists {
                return "", fmt.Errorf("current context %q not found in kubeconfig", currentContext)
        }

        ns := strings.TrimSpace(contextConfig.Namespace)
        if ns == "" {
                return "", fmt.Errorf("no namespace set in current context %q", currentContext)
        }

        return ns, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package labels provides utilities for managing container labels
// used by the toolhive application.
package labels

import (
        "fmt"
        "strconv"
        "strings"
)

const (
        // LabelPrefix is the prefix for all ToolHive labels
        LabelPrefix = "toolhive"

        // LabelToolHive is the label that indicates a container is managed by ToolHive
        LabelToolHive = "toolhive"

        // LabelName is the label that contains the container name
        LabelName = "toolhive-name"

        // LabelBaseName is the label that contains the base container name (without timestamp)
        LabelBaseName = "toolhive-basename"

        // LabelTransport is the label that contains the transport mode
        LabelTransport = "toolhive-transport"

        // LabelPort is the label that contains the port
        LabelPort = "toolhive-port"

        // LabelNetworkIsolation indicates that the network isolation functionality is enabled.
        LabelNetworkIsolation = "toolhive-network-isolation"

        // LabelGroup is the label that contains the group name
        LabelGroup = "toolhive-group"

        // LabelAuxiliary is the label that indicates this is an auxiliary workload (like inspector)
        LabelAuxiliary = "toolhive-auxiliary"

        // LabelToolHiveValue is the value for the LabelToolHive label
        LabelToolHiveValue = "true"
)

// AddStandardLabels adds standard labels to a container
func AddStandardLabels(labels map[string]string, containerName, containerBaseName, transportType string, port int) {
        labels[LabelToolHive] = LabelToolHiveValue
        labels[LabelName] = containerName
        labels[LabelBaseName] = containerBaseName
        labels[LabelTransport] = transportType
        labels[LabelPort] = fmt.Sprintf("%d", port)
}

// AddNetworkLabels adds network-related labels to a network
func AddNetworkLabels(labels map[string]string, networkName string) {
        labels[LabelToolHive] = LabelToolHiveValue
        labels[LabelName] = networkName
}

// AddNetworkIsolationLabel adds the network isolation label to a container
func AddNetworkIsolationLabel(labels map[string]string, networkIsolation bool) {
        labels[LabelNetworkIsolation] = strconv.FormatBool(networkIsolation)
}

// FormatToolHiveFilter formats a filter for ToolHive containers
func FormatToolHiveFilter() string {
        return fmt.Sprintf("%s=%s", LabelToolHive, LabelToolHiveValue)
}

// IsToolHiveContainer checks if a container is managed by ToolHive
func IsToolHiveContainer(labels map[string]string) bool {
        value, ok := labels[LabelToolHive]
        return ok && strings.ToLower(value) == LabelToolHiveValue
}

// HasNetworkIsolation checks if a container has network isolation enabled.
func HasNetworkIsolation(labels map[string]string) bool {
        value, ok := labels[LabelNetworkIsolation]
        // If the label is not present, assume that network isolation is enabled.
        // This is to ensure that workloads created before this label was introduced
        // will be properly cleaned up during stop/rm.
        return !ok || strings.ToLower(value) == "true"
}

// GetContainerName gets the container name from labels
func GetContainerName(labels map[string]string) string {
        return labels[LabelName]
}

// GetContainerBaseName gets the base container name from labels
func GetContainerBaseName(labels map[string]string) string {
        return labels[LabelBaseName]
}

// GetTransportType gets the transport type from labels
func GetTransportType(labels map[string]string) string {
        return labels[LabelTransport]
}

// GetPort gets the port from labels
func GetPort(labels map[string]string) (int, error) {
        portStr, ok := labels[LabelPort]
        if !ok {
                return 0, fmt.Errorf("port label not found")
        }

        var port int
        if _, err := fmt.Sscanf(portStr, "%d", &port); err != nil {
                return 0, fmt.Errorf("invalid port: %s", portStr)
        }

        return port, nil
}

// IsAuxiliaryWorkload checks if a workload is an auxiliary workload (like inspector)
// Auxiliary workloads don't follow standard workload management patterns and don't use proxy processes
func IsAuxiliaryWorkload(labels map[string]string) bool {
        value, ok := labels[LabelAuxiliary]
        return ok && strings.ToLower(value) == LabelToolHiveValue
}

// IsStandardToolHiveLabel checks if a label key is a standard ToolHive label
// that should not be passed through from user input or displayed to users
func IsStandardToolHiveLabel(key string) bool {
        standardLabels := []string{
                LabelToolHive,
                LabelName,
                LabelBaseName,
                LabelTransport,
                LabelPort,
                LabelNetworkIsolation,
        }

        for _, standardLabel := range standardLabels {
                if key == standardLabel {
                        return true
                }
        }

        return false
}

// ParseLabel parses a label string in the format "key=value" and validates it
// according to Kubernetes label naming conventions
func ParseLabel(label string) (string, string, error) {
        parts := strings.SplitN(label, "=", 2)
        if len(parts) != 2 {
                return "", "", fmt.Errorf("invalid label format, expected key=value")
        }

        key := strings.TrimSpace(parts[0])
        value := strings.TrimSpace(parts[1])

        if key == "" {
                return "", "", fmt.Errorf("label key cannot be empty")
        }

        // Validate key according to Kubernetes label naming conventions
        if err := validateLabelKey(key); err != nil {
                return "", "", fmt.Errorf("invalid label key: %w", err)
        }

        // Validate value according to Kubernetes label naming conventions
        if err := validateLabelValue(value); err != nil {
                return "", "", fmt.Errorf("invalid label value: %w", err)
        }

        return key, value, nil
}

// validateLabelKey validates a label key according to Kubernetes naming conventions
func validateLabelKey(key string) error {
        if len(key) == 0 {
                return fmt.Errorf("key cannot be empty")
        }
        if len(key) > 253 {
                return fmt.Errorf("key cannot be longer than 253 characters")
        }

        // Check for valid prefix (optional)
        parts := strings.Split(key, "/")
        if len(parts) > 2 {
                return fmt.Errorf("key can have at most one '/' separator")
        }

        var name string
        if len(parts) == 2 {
                prefix := parts[0]
                name = parts[1]

                // Validate prefix (should be a valid DNS subdomain)
                if len(prefix) > 253 {
                        return fmt.Errorf("prefix cannot be longer than 253 characters")
                }
                if !isValidDNSSubdomain(prefix) {
                        return fmt.Errorf("prefix must be a valid DNS subdomain")
                }
        } else {
                name = parts[0]
        }

        // Validate name part
        if len(name) == 0 {
                return fmt.Errorf("name part cannot be empty")
        }
        if len(name) > 63 {
                return fmt.Errorf("name part cannot be longer than 63 characters")
        }
        if !isValidLabelName(name) {
                return fmt.Errorf("name part must consist of alphanumeric characters, '-', '_' or '.', " +
                        "and must start and end with an alphanumeric character")
        }

        return nil
}

// validateLabelValue validates a label value according to Kubernetes naming conventions
func validateLabelValue(value string) error {
        if len(value) > 63 {
                return fmt.Errorf("value cannot be longer than 63 characters")
        }
        if value != "" && !isValidLabelName(value) {
                return fmt.Errorf("value must consist of alphanumeric characters, '-', '_' or '.', " +
                        "and must start and end with an alphanumeric character")
        }
        return nil
}

// isValidDNSSubdomain checks if a string is a valid DNS subdomain
func isValidDNSSubdomain(s string) bool {
        if len(s) == 0 || len(s) > 253 {
                return false
        }

        parts := strings.Split(s, ".")
        for _, part := range parts {
                if len(part) == 0 || len(part) > 63 {
                        return false
                }
                if !isValidDNSLabel(part) {
                        return false
                }
        }
        return true
}

// isValidDNSLabel checks if a string is a valid DNS label
func isValidDNSLabel(s string) bool {
        if len(s) == 0 || len(s) > 63 {
                return false
        }

        // Must start and end with alphanumeric
        if !isAlphaNumeric(s[0]) || !isAlphaNumeric(s[len(s)-1]) {
                return false
        }

        // Middle characters can be alphanumeric or hyphen
        for i := 1; i < len(s)-1; i++ {
                if !isAlphaNumeric(s[i]) && s[i] != '-' {
                        return false
                }
        }

        return true
}

// isValidLabelName checks if a string is a valid label name
func isValidLabelName(s string) bool {
        if len(s) == 0 {
                return false
        }

        // Must start and end with alphanumeric
        if !isAlphaNumeric(s[0]) || !isAlphaNumeric(s[len(s)-1]) {
                return false
        }

        // Middle characters can be alphanumeric, hyphen, underscore, or dot
        for i := 1; i < len(s)-1; i++ {
                if !isAlphaNumeric(s[i]) && s[i] != '-' && s[i] != '_' && s[i] != '.' {
                        return false
                }
        }

        return true
}

// isAlphaNumeric checks if a character is alphanumeric
func isAlphaNumeric(c byte) bool {
        return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package llm

import (
        "errors"
        "fmt"

        "github.com/stacklok/toolhive/pkg/networking"
        pkgoidc "github.com/stacklok/toolhive/pkg/oidc"
)

const (
        // DefaultProxyListenPort is the default port the localhost proxy listens on.
        DefaultProxyListenPort = 14000
)

// OIDCConfig is a type alias for oidc.ClientConfig, holding OIDC provider
// settings and cached token state for the LLM gateway. Using a type alias
// ensures this type stays in sync with pkg/config.RegistryOAuthConfig, which
// is also an alias for the same underlying type.
type OIDCConfig = pkgoidc.ClientConfig

// Config holds all LLM gateway settings persisted under the llm: key in
// ToolHive's config.yaml.
type Config struct {
        GatewayURL    string `yaml:"gateway_url,omitempty"       json:"gateway_url,omitempty"`
        TLSSkipVerify bool   `yaml:"tls_skip_verify,omitempty"   json:"tls_skip_verify,omitempty"`
        // ShortPromptCache opts Claude Code out of ToolHive's default one-hour
        // prompt-cache lifetime and restores Claude Code's five-minute default.
        ShortPromptCache bool          `yaml:"short_prompt_cache,omitempty" json:"short_prompt_cache,omitempty"`
        OIDC             OIDCConfig    `yaml:"oidc,omitempty"              json:"oidc,omitempty"`
        Proxy            ProxyConfig   `yaml:"proxy,omitempty"             json:"proxy,omitempty"`
        Bedrock          BedrockConfig `yaml:"bedrock,omitempty"           json:"bedrock,omitempty"`
        // Models is the persisted, single source of truth for the model IDs applied
        // during setup. It feeds two consumers: credential-helper clients (Claude
        // Desktop) write it verbatim as inferenceModels, and — when Bedrock compat is
        // on — each entry is also mapped to a Claude Code tier (see BedrockConfig).
        // Persisting it here (rather than passing a transient flag value) keeps both
        // consumers consistent on a later plain "thv llm setup".
        Models          []string     `yaml:"models,omitempty"            json:"models,omitempty"`
        ConfiguredTools []ToolConfig `yaml:"configured_tools,omitempty"  json:"configured_tools,omitempty"`
        // discoveryAccessToken is populated only for the duration of setup. It is
        // deliberately unexported so authentication material can never be persisted.
        discoveryAccessToken string
}

// SetDiscoveryAccessToken makes the token obtained by setup's authentication
// flow available to authenticated model discovery without persisting it.
func (c *Config) SetDiscoveryAccessToken(token string) { c.discoveryAccessToken = token }

// BedrockConfig holds settings for configuring Claude Code to reach an LLM
// gateway that forwards to AWS Bedrock. It is persisted so that a later plain
// "thv llm setup" re-applies these settings rather than silently clearing them.
// That stickiness ends at teardown: reverting the last configured tool resets the
// whole LLM config (see Teardown), so it cannot outlive the tools that used it.
type BedrockConfig struct {
        // Compat, when true, writes CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1 and the
        // per-tier Bedrock model IDs into Claude Code's settings.json. Bedrock rejects
        // Claude Code's experimental anthropic-beta headers, so this is required for a
        // Bedrock-backed gateway.
        Compat bool `yaml:"compat,omitempty" json:"compat,omitempty"`
        // Enable1M appends the "[1m]" suffix to the opus and sonnet model IDs to opt
        // into the 1M-token context window on Bedrock (never haiku, which is 200K).
        // Off by default: 1M-on-Bedrock is a version-dependent Claude Code behavior.
        Enable1M bool `yaml:"enable_1m,omitempty" json:"enable_1m,omitempty"`
}

// ProxyConfig holds configuration for the localhost reverse proxy.
type ProxyConfig struct {
        ListenPort int `yaml:"listen_port,omitempty" json:"listen_port,omitempty"`
}

// ToolConfig records a tool that setup has configured, so teardown knows
// exactly what to reverse.
type ToolConfig struct {
        // Tool is the canonical tool identifier (e.g. "claude-code", "cursor").
        Tool string `yaml:"tool" json:"tool"`
        // Mode is the authentication mode: one of the llmgateway.Mode* values
        // ("direct", "proxy", "credential-helper", "codex-auth", "vscode").
        Mode string `yaml:"mode" json:"mode"`
        // ConfigPath is the absolute path to the tool's config file that was patched.
        ConfigPath string `yaml:"config_path" json:"config_path"`
        // EnvFilePath is the absolute path to the .env file written during setup,
        // or empty if no .env file was managed for this tool.
        EnvFilePath string `yaml:"env_file_path,omitempty" json:"env_file_path,omitempty"`
}

// IsConfigured reports whether the minimum required fields are present for the
// LLM gateway to be usable: gateway URL, OIDC issuer, and OIDC client ID.
func (c *Config) IsConfigured() bool {
        return c.GatewayURL != "" && c.OIDC.Issuer != "" && c.OIDC.ClientID != ""
}

// ValidatePartial validates any fields that are explicitly set, without
// requiring the mandatory trio (gateway_url, oidc.issuer, oidc.client_id).
// Use this to catch URL format or port range errors during incremental
// configuration, before all required fields have been provided.
func (c *Config) ValidatePartial() error {
        var errs []error

        if c.GatewayURL != "" {
                if err := networking.ValidateHTTPSURL(c.GatewayURL); err != nil {
                        errs = append(errs, fmt.Errorf("gateway_url: %w", err))
                }
        }

        if c.OIDC.Issuer != "" {
                if err := networking.ValidateIssuerURL(c.OIDC.Issuer); err != nil {
                        errs = append(errs, fmt.Errorf("oidc.issuer: %w", err))
                }
        }

        if c.Proxy.ListenPort != 0 && (c.Proxy.ListenPort < 1024 || c.Proxy.ListenPort > 65535) {
                errs = append(errs, fmt.Errorf("proxy.listen_port must be between 1024 and 65535, got: %d", c.Proxy.ListenPort))
        }

        // Reuse networking.ValidateCallbackPort for the OIDC callback port — same
        // range check (1024–65535), zero means ephemeral (auto-assigned). Pass the
        // client ID so the validator applies strict availability checking for
        // pre-registered clients (clientID != "").
        if err := networking.ValidateCallbackPort(c.OIDC.CallbackPort, c.OIDC.ClientID); err != nil {
                errs = append(errs, fmt.Errorf("oidc.callback_port: %w", err))
        }

        return errors.Join(errs...)
}

// Validate performs full validation of the LLM config, including HTTPS
// enforcement, port range checks, and OIDC field requirements.
func (c *Config) Validate() error {
        var errs []error

        if c.GatewayURL == "" {
                errs = append(errs, errors.New("gateway_url is required"))
        }

        if c.OIDC.Issuer == "" {
                errs = append(errs, errors.New("oidc.issuer is required"))
        }

        if c.OIDC.ClientID == "" {
                errs = append(errs, errors.New("oidc.client_id is required"))
        }

        return errors.Join(append(errs, c.ValidatePartial())...)
}

// EffectiveProxyPort returns the configured proxy listen port, or
// DefaultProxyListenPort if none is set.
func (c *Config) EffectiveProxyPort() int {
        if c.Proxy.ListenPort > 0 {
                return c.Proxy.ListenPort
        }
        return DefaultProxyListenPort
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package llm

import (
        "context"
        "fmt"
        "io"

        pkgsecrets "github.com/stacklok/toolhive/pkg/secrets"
)

// SetFields applies the non-zero fields from the provided options to the
// config and validates the result. If the mandatory trio (gateway_url,
// oidc.issuer, oidc.client_id) is present after the update, full validation
// runs; otherwise only format/range validation runs to catch bad values early
// while still allowing incremental configuration.
func (c *Config) SetFields(opts SetOptions) error {
        if opts.GatewayURL != "" {
                c.GatewayURL = opts.GatewayURL
        }
        if opts.Issuer != "" {
                c.OIDC.Issuer = opts.Issuer
        }
        if opts.ClientID != "" {
                c.OIDC.ClientID = opts.ClientID
        }
        if opts.Audience != "" {
                c.OIDC.Audience = opts.Audience
        }
        if opts.ProxyPort != 0 {
                c.Proxy.ListenPort = opts.ProxyPort
        }
        if opts.CallbackPort != 0 {
                c.OIDC.CallbackPort = opts.CallbackPort
        }
        if opts.TLSSkipVerify != nil {
                c.TLSSkipVerify = *opts.TLSSkipVerify
        }
        if opts.ShortPromptCache != nil {
                c.ShortPromptCache = *opts.ShortPromptCache
        }
        if opts.BedrockCompat != nil {
                c.Bedrock.Compat = *opts.BedrockCompat
        }
        if opts.Enable1M != nil {
                c.Bedrock.Enable1M = *opts.Enable1M
        }
        if opts.Models != nil {
                c.Models = opts.Models
        }

        if !c.IsConfigured() {
                return c.ValidatePartial()
        }
        return c.Validate()
}

// SetOptions carries the flag values for the "config set" command.
// Zero values are treated as "not provided" and leave the existing config
// field unchanged. TLSSkipVerify uses a pointer so that false can be
// distinguished from "not provided" (enabling explicit clear via config set).
type SetOptions struct {
        GatewayURL       string
        Issuer           string
        ClientID         string
        Audience         string
        ProxyPort        int
        CallbackPort     int
        TLSSkipVerify    *bool // nil = not provided; &false = explicitly disable
        ShortPromptCache *bool // nil = not provided; &false = explicitly restore one-hour caching
        // BedrockCompat and Enable1M use pointers so false can be distinguished from
        // "not provided" (enabling explicit clear via config set). See BedrockConfig.
        BedrockCompat *bool
        Enable1M      *bool
        // Models sets the persisted model IDs (Config.Models). nil = not provided
        // (leave existing config unchanged); an empty non-nil slice clears them.
        Models []string
}

// DeleteCachedTokens removes all cached OIDC tokens stored under the LLM
// scope via the provided secrets provider. It is a no-op if the provider does
// not support listing or deletion (e.g. the environment provider), since such
// providers cannot hold cached tokens.
func DeleteCachedTokens(ctx context.Context, provider pkgsecrets.Provider) error {
        scoped := pkgsecrets.NewScopedProvider(provider, pkgsecrets.ScopeLLM)
        caps := scoped.Capabilities()
        if !caps.CanList || !caps.CanDelete {
                return nil
        }
        descs, err := scoped.ListSecrets(ctx)
        if err != nil {
                return err
        }
        if len(descs) == 0 {
                return nil
        }
        names := make([]string, len(descs))
        for i, d := range descs {
                names[i] = d.Key
        }
        return scoped.DeleteSecrets(ctx, names)
}

// Show writes a human-readable representation of the config to w.
// If the config is not yet configured it prints a hint to run "config set".
func (c *Config) Show(w io.Writer) error {
        if !c.IsConfigured() {
                _, err := fmt.Fprintln(w, "LLM gateway is not configured. Run \"thv llm config set\" to configure it.")
                return err
        }

        var err error
        writef := func(format string, args ...any) {
                if err == nil {
                        _, err = fmt.Fprintf(w, format, args...)
                }
        }

        writef("Gateway URL:     %s\n", c.GatewayURL)
        writef("OIDC Issuer:     %s\n", c.OIDC.Issuer)
        writef("OIDC Client:     %s\n", c.OIDC.ClientID)
        if c.OIDC.Audience != "" {
                writef("Audience:        %s\n", c.OIDC.Audience)
        }
        writef("Proxy Port:      %d\n", c.EffectiveProxyPort())
        writef("Scopes:          %v\n", c.OIDC.EffectiveScopes())
        promptCacheTTL := "1h"
        if c.ShortPromptCache {
                promptCacheTTL = "5m"
        }
        writef("Prompt Cache:    %s\n", promptCacheTTL)
        if c.TLSSkipVerify {
                writef("TLS Skip Verify: true (WARNING: certificate verification disabled)\n")
        }
        if len(c.Models) > 0 {
                writef("Models:          %v\n", c.Models)
        }
        if c.Bedrock.Compat {
                writef("Bedrock compat:  true (CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1)\n")
                writef("  1M context:    %t\n", c.Bedrock.Enable1M)
        }
        if len(c.ConfiguredTools) > 0 {
                writef("Configured tools:\n")
                for _, t := range c.ConfiguredTools {
                        writef("  - %s (%s)  %s\n", t.Tool, t.Mode, t.ConfigPath)
                }
        }
        return err
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package llm

import (
        "context"
        "crypto/tls"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "net/http"
        "net/url"
        "os"
        "regexp"
        "slices"
        "sort"
        "strconv"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/llmgateway"
        "github.com/stacklok/toolhive/pkg/networking"
        pkgsecrets "github.com/stacklok/toolhive/pkg/secrets"
)

// LoginFunc performs the interactive OIDC login during setup. It is a
// parameter so that tests can inject a no-op without touching the keyring.
type LoginFunc func(ctx context.Context, cfg *Config) error

// GatewayManager is the subset of client.ClientManager used by Setup and
// Teardown. Defined here so pkg/llm does not import pkg/client.
type GatewayManager interface {
        // DetectedLLMGatewayClients returns tool names for all installed LLM-gateway-capable tools.
        DetectedLLMGatewayClients() []string
        // ConfigureLLMGateway patches the tool's config file and returns the config path.
        ConfigureLLMGateway(clientType string, cfg llmgateway.ApplyConfig) (string, error)
        // LLMGatewayModeFor returns "direct", "proxy", or "" for the given client.
        LLMGatewayModeFor(clientType string) string
        // PromptCacheConflict names a setting that may override ToolHive's
        // one-hour Claude Code prompt-cache configuration.
        PromptCacheConflict(clientType string) (string, error)
        // IsManaged reports whether a managed configuration overrides the
        // client's local config (so the config setup writes would be ignored).
        IsManaged(clientType string) bool
        // LLMClientDetectionHint explains why clientType was not detected, or "".
        // Used when --client names a CLI tool whose settings dir is leftover after uninstall.
        LLMClientDetectionHint(clientType string) string
        // ConfigureEnvFile writes .env file entries for the client and returns the
        // env file path. Returns ("", nil) when the client has no env-file entries.
        ConfigureEnvFile(clientType string, cfg llmgateway.ApplyConfig) (string, error)
        // RevertEnvFile removes the .env file entries that setup wrote. envFilePath
        // is the value returned by ConfigureEnvFile; a no-op when empty.
        RevertEnvFile(clientType, envFilePath string) error
        // RevertLLMGateway removes the LLM gateway settings from the tool's config file.
        RevertLLMGateway(clientType, configPath string) error
}

// ConfigUpdater is the subset of config.Provider used by Setup and Teardown.
// Defined here so pkg/llm does not import pkg/config.
type ConfigUpdater interface {
        // GetLLMConfig returns the current LLM section of the config.
        GetLLMConfig() Config
        // UpdateLLMConfig atomically reads, applies fn, and persists the LLM config.
        UpdateLLMConfig(fn func(*Config) error) error
}

// SecretsProviderFunc returns the secrets provider used to delete cached LLM
// credentials. Teardown calls it only when the requested operation must purge
// tokens.
type SecretsProviderFunc func() (pkgsecrets.Provider, error)

// Setup configures detected AI tools to use the LLM gateway.
//
// When targetClient is non-empty only that client is configured; an error is
// returned if the client is not installed. Pass an empty string to configure
// all detected clients (the original behaviour).
//
// It applies inlineOpts in-memory before login so a failed login leaves no
// persisted state. Tool config files are patched only after login succeeds;
// on any persistence failure the patches are rolled back.
//
// When lazy is true, the interactive OIDC login (login) is skipped entirely:
// tool detection, config-file patching, and config persistence still run, and a
// message is printed telling the user that login will occur on first gateway
// access. lazy is intended for unattended provisioning (e.g. an MDM profile);
// it is opt-in and does not change behaviour for interactive users.
func Setup(
        ctx context.Context, out, errOut io.Writer,
        gm GatewayManager, provider ConfigUpdater, login LoginFunc,
        inlineOpts SetOptions, anthropicPathPrefix string, anthropicPathPrefixSet bool, targetClient string,
        lazy bool,
) error {
        llmCfg, err := prepareSetupConfig(provider, inlineOpts)
        if err != nil {
                return err
        }

        proxyBaseURL := fmt.Sprintf("http://localhost:%d/v1", llmCfg.EffectiveProxyPort())

        // Detect tools before login so we skip the interactive browser flow when
        // there is nothing to configure. In non-lazy mode login still runs before any
        // files are patched, preserving the guarantee that a failed login leaves no
        // state.
        detected, done, err := setupClients(
                gm, targetClient, lazy, llmCfg.Bedrock.Compat, len(llmCfg.Models) > 0, out, errOut)
        if err != nil {
                return err
        }
        if done {
                return nil
        }

        tokenHelperPath, tokenHelperArgs, err := buildTokenHelperArgv()
        if err != nil {
                return err
        }

        // In lazy mode the interactive login is deferred until a configured tool
        // first accesses the gateway (via "thv llm token" or the proxy). Everything
        // below — tool detection, file patching, config persistence — is
        // non-interactive and still runs, so the on-disk result is identical to a
        // normal setup except that no OIDC token is obtained yet.
        if lazy {
                _, _ = fmt.Fprintln(out, "Lazy mode: skipping OIDC login. You'll be signed in on the first")
                _, _ = fmt.Fprintln(out, "request a configured tool makes to the LLM gateway.")
        } else {
                _, _ = fmt.Fprintln(out, "Ensuring you are logged in to the LLM gateway…")
                if err := login(ctx, &llmCfg); err != nil {
                        if portErr := setupCallbackPortError(err); portErr != nil {
                                return fmt.Errorf("OIDC login failed: %w", portErr)
                        }
                        return fmt.Errorf("OIDC login failed: %s", SanitizeTokenError(err))
                }
                _, _ = fmt.Fprintln(out, "Login successful.")
        }

        // Resolve the effective path prefix for the Anthropic API before discovery
        // so Pi uses the same endpoint to list models and serve requests.
        anthropicPrefix := resolveAnthropicPrefix(ctx, gm, detected, llmCfg, anthropicPathPrefix, anthropicPathPrefixSet)

        detected, discoveredModels, err := discoverSetupModels(
                ctx, errOut, detected, targetClient, llmCfg, anthropicPrefix,
        )
        if err != nil {
                return err
        }
        if len(detected) == 0 {
                return nil
        }

        configured, err := configureDetectedToolsWithDiscovery(
                out, errOut, gm, detected, llmCfg.GatewayURL, proxyBaseURL,
                tokenHelperPath, tokenHelperArgs, llmCfg.TLSSkipVerify, anthropicPrefix,
                llmCfg.Models, discoveredModels, llmCfg.ShortPromptCache, llmCfg.Bedrock,
        )
        if err != nil {
                return err
        }

        // Warn about bedrock flags the user passed THIS run that had no effect:
        // --bedrock-compat when Claude Code was not configured, and --enable-1m when
        // bedrock-compat is off. Keyed off inlineOpts (not the merged config) because
        // bedrock-compat is persisted — gating on llmCfg.Bedrock would nag on every
        // later setup that omits claude-code. llmCfg.Bedrock.Compat is the effective
        // (persisted + inline) compat state used for the --enable-1m check.
        warnBedrockNoEffect(errOut, inlineOpts, llmCfg.Bedrock.Compat, configured)
        warnShortPromptCacheNoEffect(errOut, inlineOpts, configured)
        warnPromptCacheBedrockCompatibility(errOut, inlineOpts, llmCfg.ShortPromptCache, configured)

        warnTLSSkipVerify(errOut, llmCfg.TLSSkipVerify, configured)
        warnCredentialHelperTools(out, errOut, gm, configured)

        if err := provider.UpdateLLMConfig(func(c *Config) error {
                // SetFields applies inline opts to the on-disk config (preserving any
                // concurrent writes to unrelated fields) and merges ConfiguredTools
                // atomically in a single write.
                if err := c.SetFields(inlineOpts); err != nil {
                        return fmt.Errorf("persisting inline flags: %w", err)
                }
                c.ConfiguredTools = mergeToolConfigs(c.ConfiguredTools, configured)
                return nil
        }); err != nil {
                // Roll back every tool we successfully patched so the tool config files
                // are not left in a modified state without a persisted record of what
                // was changed (which would make teardown unable to revert them).
                rollbackConfiguredTools(errOut, gm, configured)
                return fmt.Errorf("persisting tool configuration: %w", err)
        }

        printPiLongCacheHint(out, configured, llmCfg.ShortPromptCache)
        if hasProxyMode(configured) {
                _, _ = fmt.Fprintln(out, "One or more tools use proxy mode. Start the proxy with: thv llm proxy start")
        }

        return nil
}

func prepareSetupConfig(provider ConfigUpdater, inlineOpts SetOptions) (Config, error) {
        llmCfg := provider.GetLLMConfig()
        // Apply inline flags in-memory so login and tool detection use the merged
        // config without touching disk. Persistence happens only after login and
        // tool patching succeed, so a failed login leaves no persisted state.
        if err := llmCfg.SetFields(inlineOpts); err != nil {
                if portErr := setupCallbackPortError(err); portErr != nil {
                        return Config{}, fmt.Errorf("invalid inline flag values: %w", portErr)
                }
                return Config{}, fmt.Errorf("invalid inline flag values: %w", err)
        }
        if !llmCfg.IsConfigured() {
                return Config{}, fmt.Errorf("LLM gateway is not configured — run \"thv llm config set\" first")
        }
        return llmCfg, nil
}

func setupCallbackPortError(err error) error {
        var portErr *networking.CallbackPortInUseError
        if !errors.As(err, &portErr) {
                return nil
        }
        return fmt.Errorf(
                "requested callback port %d is already in use; "+
                        "stop the process using it or choose another provider-registered port with "+
                        `"thv llm setup --callback-port <port>"`,
                portErr.Port,
        )
}

// Teardown removes LLM gateway configuration from all (or one) configured tools.
//
// targetTool selects which tool to revert; pass an empty string to revert all
// configured tools. An error is returned when targetTool is non-empty but not
// found in the configured tool list.
//
// Reverting the last configured tool also resets the persisted LLM config to its
// zero value, so settings that are deliberately sticky across "thv llm setup"
// re-runs (e.g. Bedrock compat) are not silently re-applied by a later setup.
// A targeted teardown that leaves other tools configured keeps the config, which
// those tools still need to reach the gateway.
//
// Cached OIDC tokens are deleted by default when no configured tool remains.
// keepTokens opts out of that full-teardown default. A partial targeted
// teardown retains the shared session unless purgeTokens is true.
func Teardown(
        ctx context.Context,
        out, errOut io.Writer,
        gm GatewayManager,
        targetTool string,
        keepTokens bool,
        purgeTokens bool,
        provider ConfigUpdater,
        getSecretsProvider SecretsProviderFunc,
) error {
        if keepTokens && purgeTokens {
                return fmt.Errorf("keepTokens and purgeTokens cannot both be set")
        }

        llmCfg := provider.GetLLMConfig()
        targets, err := selectTeardownTargets(llmCfg.ConfiguredTools, targetTool)
        if err != nil {
                return err
        }

        // Separate tools into those to revert and those to keep without touching
        // any files yet.
        toRevert, remaining := partitionTeardownTools(llmCfg.ConfiguredTools, targets)
        fullTeardown := len(remaining) == 0
        shouldPurge := purgeTokens || (fullTeardown && !keepTokens)

        // Delete credentials before changing the configured-tool record. If cleanup
        // fails, the same targeted teardown remains fully retryable.
        if err := purgeTokensBeforeTeardown(ctx, shouldPurge, getSecretsProvider); err != nil {
                return err
        }

        // Persist the updated tool list (and clear token metadata if purging) in a
        // single write before mutating any tool config files. Token cleanup may have
        // completed already, but a config-write failure leaves every tool record and
        // settings file intact so the caller can retry teardown safely.
        if err := provider.UpdateLLMConfig(func(c *Config) error {
                applyTeardownToConfig(c, remaining, shouldPurge)
                return nil
        }); err != nil {
                if shouldPurge {
                        return fmt.Errorf("cached LLM tokens were deleted, but persisting tool configuration: %w", err)
                }
                return fmt.Errorf("persisting tool configuration: %w", err)
        }

        if len(targets) == 0 {
                _, _ = fmt.Fprintln(out, "No tools are currently configured.")
        } else if fullTeardown {
                _, _ = fmt.Fprintln(out,
                        "Cleared the LLM gateway configuration: no tools are configured anymore.")
        }

        // Revert tool config files best-effort; warn on failure but do not undo
        // the config update above (the user can re-run setup+teardown to reconcile).
        for _, tc := range toRevert {
                revertToolConfig(out, errOut, gm, tc)
        }

        if shouldPurge {
                _, _ = fmt.Fprintln(out,
                        "Deleted locally cached LLM tokens. This does not revoke tokens at the identity provider or sign out the browser session.")
                _, _ = fmt.Fprintln(out,
                        "Stop any running LLM proxy or token helper before relying on cleanup; an active process may recreate cached credentials.")
                if purgeTokens && !fullTeardown {
                        _, _ = fmt.Fprintln(out,
                                "Remaining configured tools share this session and will need to authenticate again.")
                }
        }

        return nil
}

// applyTeardownToConfig updates the persisted LLM config for a teardown that
// leaves remaining configured. purgeTokens reports whether the teardown drops
// cached OIDC token state.
//
// When no tool remains, the config is reset wholesale rather than pruned field by
// field. Settings such as Bedrock compat are deliberately sticky across "thv llm
// setup" re-runs, so keeping them past the last teardown would let the next setup
// silently re-apply settings the user just removed — possibly against a gateway
// they have since repointed elsewhere. While any tool remains the config is left
// intact, since those tools still need it to reach the gateway.
//
// Cached token state survives a reset only when token retention was explicitly
// requested. The secret lives in the encrypted secrets store, and dropping the
// only reference without deleting it would strand it once the user points at a
// different gateway (the fallback key is derived from the gateway URL and
// issuer).
func applyTeardownToConfig(c *Config, remaining []ToolConfig, purgeTokens bool) {
        if len(remaining) > 0 {
                c.ConfiguredTools = remaining
                if purgeTokens {
                        c.OIDC.CachedRefreshTokenRef = ""
                        c.OIDC.CachedTokenExpiry = time.Time{}
                }
                return
        }

        tokenState := c.OIDC
        *c = Config{}
        if !purgeTokens {
                c.OIDC.CachedRefreshTokenRef = tokenState.CachedRefreshTokenRef
                c.OIDC.CachedTokenExpiry = tokenState.CachedTokenExpiry
        }
}

// PurgeTokens deletes all cached OIDC tokens from the provided secrets provider.
func PurgeTokens(ctx context.Context, provider pkgsecrets.Provider) error {
        if err := DeleteCachedTokens(ctx, provider); err != nil {
                return fmt.Errorf("removing cached LLM tokens: %w", err)
        }
        return nil
}

func selectTeardownTargets(configured []ToolConfig, targetTool string) ([]ToolConfig, error) {
        if targetTool == "" {
                return configured, nil
        }
        for _, tool := range configured {
                if tool.Tool == targetTool {
                        return []ToolConfig{tool}, nil
                }
        }
        return nil, fmt.Errorf("tool %q is not configured", targetTool)
}

func partitionTeardownTools(configured, targets []ToolConfig) ([]ToolConfig, []ToolConfig) {
        var toRevert, remaining []ToolConfig
        for _, tool := range configured {
                if isTarget(targets, tool.Tool) {
                        toRevert = append(toRevert, tool)
                } else {
                        remaining = append(remaining, tool)
                }
        }
        return toRevert, remaining
}

func purgeTokensBeforeTeardown(
        ctx context.Context, shouldPurge bool, getSecretsProvider SecretsProviderFunc,
) error {
        if !shouldPurge {
                return nil
        }
        if getSecretsProvider == nil {
                return fmt.Errorf("cannot remove cached LLM tokens: secrets provider is unavailable")
        }
        provider, err := getSecretsProvider()
        if err != nil {
                return fmt.Errorf("getting secrets provider to remove cached LLM tokens: %w", err)
        }
        if provider == nil {
                return fmt.Errorf("cannot remove cached LLM tokens: secrets provider is unavailable")
        }
        if err := PurgeTokens(ctx, provider); err != nil {
                return fmt.Errorf("cached token cleanup failed; no configuration changes were made: %w", err)
        }
        return nil
}

// isTarget reports whether toolName appears in the targets slice.
func isTarget(targets []ToolConfig, toolName string) bool {
        for _, t := range targets {
                if t.Tool == toolName {
                        return true
                }
        }
        return false
}

// mergeToolConfigs merges newly configured tools into the existing list,
// replacing any entry with the same tool name.
func mergeToolConfigs(existing, incoming []ToolConfig) []ToolConfig {
        index := make(map[string]int, len(existing))
        result := make([]ToolConfig, len(existing))
        copy(result, existing)
        for i, tc := range result {
                index[tc.Tool] = i
        }
        for _, tc := range incoming {
                if i, ok := index[tc.Tool]; ok {
                        result[i] = tc
                } else {
                        index[tc.Tool] = len(result)
                        result = append(result, tc)
                }
        }
        return result
}

// warnTLSSkipVerify prints mode-accurate warnings when TLS verification is
// disabled. The impact differs by tool mode:
//   - direct (Node.js tools like Claude Code, Gemini CLI): NODE_TLS_REJECT_UNAUTHORIZED=0
//     is written to the tool's settings, disabling TLS for ALL of that tool's outbound
//     connections — not just the LLM gateway.
//   - proxy-backed: only the proxy's upstream connection to the gateway has TLS verification
//     disabled; the tool itself is unaffected.
func warnTLSSkipVerify(errOut io.Writer, skip bool, configured []ToolConfig) {
        if !skip {
                return
        }
        for _, tc := range configured {
                switch tc.Mode {
                case llmgateway.ModeDirect:
                        _, _ = fmt.Fprintf(errOut,
                                "Warning: %s uses direct mode — NODE_TLS_REJECT_UNAUTHORIZED=0 has been written to its "+
                                        "settings, disabling TLS certificate verification for ALL of %s's outbound connections "+
                                        "(LLM provider APIs, MCP registry, etc.), not just the LLM gateway. "+
                                        "Use only in isolated local environments.\n", tc.Tool, tc.Tool)
                case llmgateway.ModeProxy, llmgateway.ModeVSCode:
                        if tc.Tool == "gemini-cli" {
                                _, _ = fmt.Fprintf(errOut,
                                        "Note: --tls-skip-verify is not supported for Gemini CLI "+
                                                "(setting NODE_TLS_REJECT_UNAUTHORIZED would affect all HTTPS connections in the process). "+
                                                "Ensure your proxy certificate is trusted by the system store instead.\n")
                        } else {
                                _, _ = fmt.Fprintf(errOut,
                                        "Warning: %s uses proxy mode — TLS certificate verification is disabled for the "+
                                                "proxy's upstream gateway connection only. Use only in isolated local environments.\n", tc.Tool)
                        }
                case llmgateway.ModePi:
                        _, _ = fmt.Fprintf(errOut,
                                "Warning: --tls-skip-verify was NOT applied to %s — trust the gateway certificate in the system store.\n", tc.Tool)
                case llmgateway.ModeCodexAuth:
                        _, _ = fmt.Fprintf(errOut,
                                "Warning: --tls-skip-verify was NOT applied to %s — its config.toml has no "+
                                        "TLS-skip option. Codex will fail against a self-signed gateway until you "+
                                        "trust the certificate in the system store.\n", tc.Tool)
                }
        }
}

// filterDetectedClients narrows the detected client list to a single entry when
// targetClient is non-empty. It returns an error if the named client is not in
// the detected list. When targetClient is empty the list is returned unchanged.
func filterDetectedClients(gm GatewayManager, targetClient string) ([]string, error) {
        detected := gm.DetectedLLMGatewayClients()
        if targetClient == "" {
                return detected, nil
        }
        for _, c := range detected {
                if c == targetClient {
                        return []string{targetClient}, nil
                }
        }
        if hint := gm.LLMClientDetectionHint(targetClient); hint != "" {
                return nil, fmt.Errorf("client %q is not installed or not detected: %s", targetClient, hint)
        }
        return nil, fmt.Errorf("client %q is not installed or not detected", targetClient)
}

func setupClients(
        gm GatewayManager, targetClient string, lazy, bedrockCompat, hasModels bool, out, errOut io.Writer,
) ([]string, bool, error) {
        detected, err := filterDetectedClients(gm, targetClient)
        if err != nil {
                return nil, false, err
        }
        if len(detected) == 0 {
                _, _ = fmt.Fprintln(out, "No supported AI tools detected.")
                return nil, true, nil
        }
        if !lazy {
                return detected, false, nil
        }
        detected, err = filterLazyModelDiscoveryClients(detected, targetClient, bedrockCompat, hasModels, errOut)
        if err != nil {
                return nil, false, err
        }
        return detected, len(detected) == 0, nil
}

const (
        vsCodeClient        = "vscode"
        vsCodeInsiderClient = "vscode-insider"
        // claudeCodeClient is the canonical client identifier for Claude Code.
        // Declared here as a string literal because pkg/llm does not import
        // pkg/client (which owns the ClientApp constant) to avoid an import cycle.
        claudeCodeClient = "claude-code"
        // promptCacheVerificationCommand reports which TTL Claude Code used under
        // usage.cache_creation in the command's JSON output.
        promptCacheVerificationCommand = `claude -p "hello" --output-format json`
)

func isVSCodeClient(clientType string) bool {
        return clientType == vsCodeClient || clientType == vsCodeInsiderClient
}

func hasVSCodeClient(clients []string) bool {
        for _, clientType := range clients {
                if isVSCodeClient(clientType) {
                        return true
                }
        }
        return false
}

func filterLazyModelDiscoveryClients(
        clients []string, targetClient string, bedrockCompat, hasModels bool, errOut io.Writer,
) ([]string, error) {
        if targetClient == "pi" {
                return nil, fmt.Errorf("client %q does not support --lazy: Pi needs authenticated model discovery "+
                        "to select each model's API; rerun without --lazy", targetClient)
        }
        if targetClient != "" && isVSCodeClient(targetClient) {
                return nil, fmt.Errorf(
                        "client %q does not support --lazy: VS Code requires non-lazy setup for authenticated model discovery; "+
                                "rerun without --lazy", targetClient)
        }
        // --bedrock-compat pins each tier to a Bedrock model discovered from the
        // gateway, which lazy mode cannot do (it never logs in). Without pins, Claude
        // Code falls back to Anthropic-form IDs a Bedrock-only gateway rejects — so
        // refuse unless the user pins the tiers explicitly with --models.
        if targetClient == claudeCodeClient && bedrockCompat && !hasModels {
                return nil, fmt.Errorf("client %q does not support --lazy with --bedrock-compat: Bedrock tiers need "+
                        "authenticated model discovery; pass --models to pin them, or rerun without --lazy", targetClient)
        }
        filtered := make([]string, 0, len(clients))
        for _, clientType := range clients {
                if clientType == "pi" {
                        _, _ = fmt.Fprintln(errOut, "Warning: skipping pi in lazy mode; Pi needs authenticated model discovery.")
                        continue
                }
                if isVSCodeClient(clientType) {
                        _, _ = fmt.Fprintf(
                                errOut,
                                "Warning: skipping %s in lazy mode; VS Code requires non-lazy setup for authenticated model discovery.\n",
                                clientType,
                        )
                        continue
                }
                if clientType == claudeCodeClient && bedrockCompat && !hasModels {
                        _, _ = fmt.Fprintln(errOut, "Warning: skipping claude-code in lazy mode; --bedrock-compat needs "+
                                "authenticated model discovery to pin tiers. Pass --models to pin them, or rerun without --lazy.")
                        continue
                }
                filtered = append(filtered, clientType)
        }
        return filtered, nil
}

type discoveredGatewayModels struct {
        openAI                  []string
        openAIContextWindows    map[string]int
        anthropic               []string
        anthropicContextWindows map[string]int
        anthropicReasoning      map[string]bool
        anthropicAdaptive       map[string]bool
}

func discoverSetupModels(
        ctx context.Context, errOut io.Writer, clients []string, targetClient string, cfg Config, anthropicPrefix string,
) ([]string, discoveredGatewayModels, error) {
        hasPi := slices.Contains(clients, "pi")
        hasVSCode := hasVSCodeClient(clients)
        // Bedrock-compat Claude Code pins each tier to the newest Bedrock model the
        // gateway advertises, so it needs the Anthropic catalogue too.
        needsBedrockTiers := cfg.Bedrock.Compat && slices.Contains(clients, claudeCodeClient)
        catalog, openErr, anthropicErr := fetchSetupModelCatalog(
                ctx, cfg, hasPi || needsBedrockTiers, hasVSCode || hasPi, anthropicPrefix)
        piMissing := hasPi && len(catalog.openAI) == 0 && len(catalog.anthropic) == 0
        // Bedrock-compat Claude Code cannot derive its tier pins when the Anthropic
        // catalogue fetch failed, or succeeded without a single Bedrock-form ID. Treat
        // that like the Pi/VS Code discovery failures: abort when the client was named
        // explicitly, skip it with a warning when it was auto-detected (so the user is
        // not left with a Claude Code that routes built-in Anthropic IDs a Bedrock-only
        // gateway rejects). --models pinning every tier needs no discovery. Lazy mode
        // never fetches the catalogue (no access token); filterLazyModelDiscoveryClients
        // already gates it on --models.
        bedrockErr := bedrockDiscoveryError(
                needsBedrockTiers && cfg.discoveryAccessToken != "", cfg.Models, catalog.anthropic, anthropicErr)
        bedrockMissing := bedrockErr != nil
        if err := explicitDiscoveryFailure(targetClient, piMissing, openErr, anthropicErr, bedrockErr); err != nil {
                return nil, catalog, err
        }
        warnModelDiscovery(errOut, hasVSCode, hasPi, piMissing, catalog, openErr, anthropicErr)
        if bedrockMissing {
                _, _ = fmt.Fprintf(errOut, "Warning: skipping auto-detected claude-code because Bedrock model discovery "+
                        "failed: %v\n", bedrockErr)
        }
        return keepDiscoverableClients(clients, openErr, piMissing, bedrockMissing), catalog, nil
}

// bedrockDiscoveryError reports why bedrock-compat Claude Code cannot pin its
// tiers from the gateway catalogue, or nil when it can (or need not, because
// --models pins every tier). A catalogue with no Bedrock-form ID for any tier
// is as unusable as a failed fetch; a partial one is only warned about later.
func bedrockDiscoveryError(needsBedrockTiers bool, models, discovered []string, anthropicErr error) error {
        if !needsBedrockTiers {
                return nil
        }
        if pinned, _ := resolveBedrockModels(models, nil, false); len(pinned.emptyFamilies()) == 0 {
                return nil
        }
        if anthropicErr != nil {
                return anthropicErr
        }
        if newestBedrockTiers(discovered) == (bedrockTiers{}) {
                return errNoBedrockModels
        }
        return nil
}

// explicitDiscoveryFailure returns an error when the explicitly named client
// cannot be configured because its model discovery failed, so setup aborts
// rather than writing a half-configured client.
func explicitDiscoveryFailure(targetClient string, piMissing bool, openErr, anthropicErr, bedrockErr error) error {
        switch {
        case isVSCodeClient(targetClient) && openErr != nil:
                return fmt.Errorf("discovering gateway models: %w", openErr)
        case targetClient == "pi" && piMissing:
                return fmt.Errorf("discovering Pi models: %w", errors.Join(openErr, anthropicErr))
        case targetClient == claudeCodeClient && bedrockErr != nil:
                return fmt.Errorf("discovering Bedrock models for Claude Code: %w", bedrockErr)
        default:
                return nil
        }
}

// keepDiscoverableClients drops the auto-detected clients whose required model
// discovery failed, leaving the rest to be configured.
func keepDiscoverableClients(clients []string, openErr error, piMissing, bedrockMissing bool) []string {
        filtered := make([]string, 0, len(clients))
        for _, clientType := range clients {
                if (isVSCodeClient(clientType) && openErr != nil) || (clientType == "pi" && piMissing) ||
                        (clientType == claudeCodeClient && bedrockMissing) {
                        continue
                }
                filtered = append(filtered, clientType)
        }
        return filtered
}

func fetchSetupModelCatalog(
        ctx context.Context, cfg Config, wantAnthropic, wantOpenAI bool, anthropicPrefix string,
) (discoveredGatewayModels, error, error) {
        var catalog discoveredGatewayModels
        var openErr, anthropicErr error
        if wantOpenAI {
                catalog.openAI, catalog.openAIContextWindows, openErr = discoverGatewayModels(ctx, cfg)
        }
        if wantAnthropic && cfg.discoveryAccessToken != "" {
                catalog.anthropic, catalog.anthropicContextWindows, catalog.anthropicReasoning, catalog.anthropicAdaptive, anthropicErr =
                        discoverAnthropicGatewayModels(ctx, cfg, anthropicPrefix)
        }
        return catalog, openErr, anthropicErr
}

func warnModelDiscovery(
        errOut io.Writer, hasVSCode, hasPi, piMissing bool, catalog discoveredGatewayModels, openErr, anthropicErr error,
) {
        if hasVSCode && openErr != nil {
                _, _ = fmt.Fprintf(errOut,
                        "Warning: skipping auto-detected VS Code clients because gateway model discovery failed: %v\n", openErr)
        }
        if piMissing {
                _, _ = fmt.Fprintf(errOut,
                        "Warning: skipping auto-detected Pi because gateway model discovery failed: %v\n",
                        errors.Join(openErr, anthropicErr))
        }
        if hasPi && len(catalog.openAI) > 0 && anthropicErr != nil {
                _, _ = fmt.Fprintf(errOut, "Warning: skipping Pi Anthropic models: %v\n", anthropicErr)
        }
        if hasPi && len(catalog.anthropic) > 0 && openErr != nil {
                _, _ = fmt.Fprintf(errOut, "Warning: skipping Pi OpenAI models: %v\n", openErr)
        }
}

type gatewayModelsResponse struct {
        Object  string `json:"object"`
        HasMore *bool  `json:"has_more"`
        LastID  string `json:"last_id"`
        Data    []struct {
                ID             string `json:"id"`
                ContextWindow  int    `json:"context_window"`
                MaxInputTokens int    `json:"max_input_tokens"`
                Capabilities   *struct {
                        Thinking struct {
                                Supported bool `json:"supported"`
                                Types     struct {
                                        Adaptive struct {
                                                Supported bool `json:"supported"`
                                        } `json:"adaptive"`
                                        Enabled struct {
                                                Supported bool `json:"supported"`
                                        } `json:"enabled"`
                                } `json:"types"`
                        } `json:"thinking"`
                } `json:"capabilities"`
        } `json:"data"`
}

// errEmptyModelCatalog reports a model list the gateway served successfully but
// with no models in it, which usually means the account has no model access.
var errEmptyModelCatalog = errors.New("gateway returned an empty model list")

// errNoBedrockModels reports an Anthropic catalogue with no Bedrock-form model
// ID for any Claude Code tier, so bedrock-compat setup has nothing to pin.
var errNoBedrockModels = errors.New("the gateway's model list has no Bedrock model IDs " +
        "(expected us.anthropic.claude-…); check that the gateway forwards to Bedrock")

// gatewayStatusError reports a non-200 response to a model-list request, so
// callers can tell a wrong path (404) from missing access (401/403) or a
// gateway fault (5xx).
type gatewayStatusError struct {
        URL        string
        StatusCode int
        Status     string
}

func (e *gatewayStatusError) Error() string {
        return fmt.Sprintf("gateway returned %s for %s", e.Status, e.URL)
}

// describeDiscoveryError rewrites a model-discovery failure into a message that
// names its likely cause and what to check, keeping err wrapped. anthropic
// selects the hints for the Anthropic model list (path prefix, Anthropic access)
// over the OpenAI one.
func describeDiscoveryError(err error, anthropic bool) error {
        list, models := "a model list", "models"
        if anthropic {
                list, models = "an Anthropic model list", "Anthropic models"
        }
        var statusErr *gatewayStatusError
        var urlErr *url.Error
        switch {
        case errors.As(err, &statusErr):
                switch code := statusErr.StatusCode; {
                case code == http.StatusNotFound && anthropic:
                        return fmt.Errorf("the gateway doesn't serve %s at %s; check --anthropic-path-prefix: %w",
                                list, statusErr.URL, err)
                case code == http.StatusNotFound:
                        return fmt.Errorf("the gateway doesn't serve %s at %s; check the gateway URL: %w", list, statusErr.URL, err)
                case code == http.StatusUnauthorized || code == http.StatusForbidden:
                        return fmt.Errorf("the gateway rejected model discovery (%s); check that your account has access to %s: %w",
                                statusErr.Status, models, err)
                case code >= http.StatusInternalServerError:
                        return fmt.Errorf("the gateway returned %s listing models; retry, or contact the gateway operator: %w",
                                statusErr.Status, err)
                }
        case errors.Is(err, errEmptyModelCatalog):
                return fmt.Errorf("the gateway returned no %s for your account; ask the gateway operator for model access: %w",
                        models, err)
        case errors.As(err, &urlErr):
                return fmt.Errorf("could not reach the gateway at %s; check the URL and your network: %w", urlErr.URL, err)
        }
        return err
}

func requestGatewayModels(ctx context.Context, cfg Config, path, afterID string) (gatewayModelsResponse, error) {
        if cfg.discoveryAccessToken == "" {
                return gatewayModelsResponse{}, fmt.Errorf("discovering gateway models: authenticated setup did not provide an access token")
        }
        modelsURL, err := url.JoinPath(cfg.GatewayURL, path)
        if err != nil {
                return gatewayModelsResponse{}, fmt.Errorf("building gateway models URL: %w", err)
        }
        requestCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
        defer cancel()
        req, err := http.NewRequestWithContext(requestCtx, http.MethodGet, modelsURL, nil)
        if err != nil {
                return gatewayModelsResponse{}, fmt.Errorf("creating gateway model-discovery request: %w", err)
        }
        if afterID != "" {
                query := req.URL.Query()
                query.Set("after_id", afterID)
                req.URL.RawQuery = query.Encode()
        }
        req.Header.Set("Authorization", "Bearer "+cfg.discoveryAccessToken)
        resp, err := gatewayHTTPClient(cfg.TLSSkipVerify).Do(req)
        if err != nil {
                return gatewayModelsResponse{}, fmt.Errorf("discovering gateway models: %w", err)
        }
        defer func() {
                _ = resp.Body.Close()
        }()
        if resp.StatusCode != http.StatusOK {
                _, _ = io.Copy(io.Discard, resp.Body)
                // Redacted, because the gateway URL may carry a password in its userinfo.
                return gatewayModelsResponse{}, &gatewayStatusError{
                        URL: req.URL.Redacted(), StatusCode: resp.StatusCode, Status: resp.Status,
                }
        }
        var payload gatewayModelsResponse
        if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
                _, _ = io.Copy(io.Discard, resp.Body)
                return gatewayModelsResponse{}, fmt.Errorf("decoding gateway model response: %w", err)
        }
        return payload, nil
}

func discoverGatewayModels(ctx context.Context, cfg Config) ([]string, map[string]int, error) {
        payload, err := requestGatewayModels(ctx, cfg, "v1/models", "")
        if err != nil {
                return nil, nil, describeDiscoveryError(err, false)
        }
        if payload.Object != "list" {
                return nil, nil, fmt.Errorf("decoding gateway model response: expected a non-empty OpenAI model list")
        }
        if len(payload.Data) == 0 {
                return nil, nil, describeDiscoveryError(errEmptyModelCatalog, false)
        }
        seen := make(map[string]struct{}, len(payload.Data))
        models := make([]string, 0, len(payload.Data))
        contextWindows := make(map[string]int, len(payload.Data))
        for _, model := range payload.Data {
                modelID := strings.TrimSpace(model.ID)
                if modelID == "" {
                        return nil, nil, fmt.Errorf("decoding gateway model response: model ID must not be empty")
                }
                if _, exists := seen[modelID]; exists {
                        continue
                }
                seen[modelID] = struct{}{}
                models = append(models, modelID)
                if model.ContextWindow > 0 {
                        contextWindows[modelID] = model.ContextWindow
                }
        }
        sort.Strings(models)
        return models, contextWindows, nil
}

func discoverAnthropicGatewayModels(
        ctx context.Context, cfg Config, anthropicPrefix string,
) ([]string, map[string]int, map[string]bool, map[string]bool, error) {
        path := strings.Trim(anthropicPrefix, "/")
        if path != "" {
                path += "/"
        }
        path += "v1/models"
        seen := make(map[string]struct{})
        models := make([]string, 0)
        contextWindows := make(map[string]int)
        reasoning := make(map[string]bool)
        adaptive := make(map[string]bool)
        afterID := ""
        for range 100 {
                payload, err := requestGatewayModels(ctx, cfg, path, afterID)
                if err != nil {
                        return nil, nil, nil, nil, describeDiscoveryError(err, true)
                }
                if payload.HasMore == nil {
                        return nil, nil, nil, nil, fmt.Errorf("decoding gateway model response: expected a non-empty Anthropic model page")
                }
                if len(payload.Data) == 0 {
                        return nil, nil, nil, nil, describeDiscoveryError(errEmptyModelCatalog, true)
                }
                for _, model := range payload.Data {
                        id := strings.TrimSpace(model.ID)
                        if id == "" {
                                return nil, nil, nil, nil, fmt.Errorf("decoding gateway model response: model ID must not be empty")
                        }
                        if _, exists := seen[id]; exists {
                                continue
                        }
                        seen[id] = struct{}{}
                        models = append(models, id)
                        if model.MaxInputTokens > 0 {
                                contextWindows[id] = model.MaxInputTokens
                        }
                        if model.Capabilities != nil {
                                thinking := model.Capabilities.Thinking
                                reasoning[id], adaptive[id] = piThinkingModes(
                                        thinking.Supported, thinking.Types.Adaptive.Supported, thinking.Types.Enabled.Supported)
                        }
                }
                if !*payload.HasMore {
                        sort.Strings(models)
                        return models, contextWindows, reasoning, adaptive, nil
                }
                if payload.LastID == "" || payload.LastID == afterID {
                        return nil, nil, nil, nil, fmt.Errorf("decoding gateway model response: invalid Anthropic pagination cursor")
                }
                afterID = payload.LastID
        }
        return nil, nil, nil, nil, fmt.Errorf("decoding gateway model response: too many Anthropic model pages")
}

func piThinkingModes(supported, adaptive, enabled bool) (bool, bool) {
        return supported && (adaptive || enabled), supported && adaptive
}

// oneMSuffix is the model-ID suffix that opts Bedrock opus/sonnet into the
// 1M-token context window.
const oneMSuffix = "[1m]"

// bedrockTiers holds the model ID pinned for each Claude Code model family in
// bedrock-compat mode. An empty tier is left unpinned.
type bedrockTiers struct {
        haiku, opus, sonnet, fable string
}

// emptyFamilies lists the tiers left without a pinned model ID, in a stable
// order, so the caller can tell the user which families it could not pin.
func (t *bedrockTiers) emptyFamilies() []string {
        var empty []string
        for _, family := range bedrockFamilies {
                if *t.slot(family) == "" {
                        empty = append(empty, family)
                }
        }
        return empty
}

// slot returns the field for a family name, or nil for an unknown family.
func (t *bedrockTiers) slot(family string) *string {
        switch family {
        case "haiku":
                return &t.haiku
        case "opus":
                return &t.opus
        case "sonnet":
                return &t.sonnet
        case "fable":
                return &t.fable
        default:
                return nil
        }
}

// bedrockFamilies are the Claude Code model tiers a Bedrock model ID can pin.
var bedrockFamilies = []string{"haiku", "opus", "sonnet", "fable"}

// bedrockModelCorePattern matches a Bedrock-form Anthropic Claude model ID: an
// optional cross-region prefix ("us.", "eu.", "global.", ...) and the required
// "anthropic.claude-" vendor segment, capturing the hyphenated core and dropping
// an optional "-v<rev>[:<n>]" Bedrock revision suffix. Anthropic-form IDs
// ("claude-opus-4-8", which lack the "anthropic." segment) do not match, so a
// gateway that also routes Anthropic directly never has those picked as Bedrock
// pins.
var bedrockModelCorePattern = regexp.MustCompile(
        `^(?:[a-z]+\.)?anthropic\.claude-(.+?)(?:-v\d+(?::\d+)?)?$`)

// parseBedrockModel reports whether id is a Bedrock-form Anthropic Claude model
// ID and, if so, its family and a [major, minor, releaseDate] key for ranking.
// It tokenizes the core on "-" rather than pinning one layout, so it accepts
// both the family-first form ("claude-haiku-4-5-20251001-v1:0") and the
// version-first form used by Claude 3.x on Bedrock
// ("claude-3-5-haiku-20241022-v1:0", "claude-3-haiku-20240307-v1:0"), plus a
// major-only form with a release date ("claude-sonnet-4-20250514-v1:0"). The
// 8-digit release date orders releases within a (major, minor). Any unexpected
// token (a free-form route name, a second family word) rejects the ID so only
// genuine Bedrock model IDs are pinned.
func parseBedrockModel(id string) (family string, version []int, ok bool) {
        m := bedrockModelCorePattern.FindStringSubmatch(id)
        if m == nil {
                return "", nil, false
        }
        var nums []int
        date := 0
        for _, tok := range strings.Split(m[1], "-") {
                switch {
                case slices.Contains(bedrockFamilies, tok):
                        if family != "" {
                                return "", nil, false // two family tokens: not a real model ID
                        }
                        family = tok
                case len(tok) == 8 && isAllDigits(tok):
                        date, _ = strconv.Atoi(tok)
                case isAllDigits(tok):
                        n, _ := strconv.Atoi(tok)
                        nums = append(nums, n)
                default:
                        return "", nil, false // unknown token: free-form route name, not Bedrock
                }
        }
        if family == "" {
                return "", nil, false
        }
        version = []int{0, 0, date} // [major, minor, releaseDate]
        if len(nums) > 0 {
                version[0] = nums[0]
        }
        if len(nums) > 1 {
                version[1] = nums[1]
        }
        return family, version, true
}

func isAllDigits(s string) bool {
        if s == "" {
                return false
        }
        for _, r := range s {
                if r < '0' || r > '9' {
                        return false
                }
        }
        return true
}

// newestBedrockTiers picks, for each family, the Bedrock-form ID in discovered
// with the highest [major, minor, releaseDate] version. Ties are broken by the
// lexicographically smallest ID so the result is deterministic. The gateway's
// created_at is not used: it reflects when a route was configured, not the
// model's release.
func newestBedrockTiers(discovered []string) bedrockTiers {
        var tiers bedrockTiers
        best := make(map[string][]int, len(bedrockFamilies))
        for _, id := range discovered {
                family, version, ok := parseBedrockModel(id)
                if !ok {
                        continue
                }
                dst := tiers.slot(family)
                prev, seen := best[family]
                cmp := slices.Compare(version, prev)
                if !seen || cmp > 0 || (cmp == 0 && id < *dst) {
                        best[family] = version
                        *dst = id
                }
        }
        return tiers
}

// applyBedrockCompat resolves the per-tier Bedrock model IDs for Claude Code and
// writes them onto applyCfg, surfacing override diagnostics and naming any tier
// that could not be pinned. Called only for bedrock-compat Claude Code.
func applyBedrockCompat(
        errOut io.Writer, applyCfg *llmgateway.ApplyConfig, models, discoveredAnthropic []string, enable1M bool,
) {
        tiers, warnings := resolveBedrockModels(models, discoveredAnthropic, enable1M)
        applyCfg.BedrockCompat = true
        applyCfg.BedrockHaikuModel = tiers.haiku
        applyCfg.BedrockOpusModel = tiers.opus
        applyCfg.BedrockSonnetModel = tiers.sonnet
        applyCfg.BedrockFableModel = tiers.fable
        for _, w := range warnings {
                _, _ = fmt.Fprintf(errOut, "Warning: %s\n", w)
        }
        // Discovery found at least one Bedrock-form model (a failed fetch, or a
        // catalogue with none, skips claude-code upstream), but a family may still
        // have none — e.g. a gateway serving only some tiers. Name those tiers so the
        // fallback to built-in IDs is visible.
        if empty := tiers.emptyFamilies(); len(empty) > 0 {
                _, _ = fmt.Fprintf(errOut, "Warning: no Bedrock model found in the gateway catalogue for the %s "+
                        "tier(s); they fall back to Claude Code's built-in model IDs.\n",
                        strings.Join(empty, ", "))
        }
}

// resolveBedrockModels picks the model ID pinned for each Claude Code family.
// Each family starts from the newest Bedrock model in discovered (see
// newestBedrockTiers); override model IDs then replace a family by matching
// "haiku", "opus", "sonnet", or "fable" (case-insensitive) as a substring of
// the ID. A family with neither resolves to "" so no pin is written and Claude
// Code falls back to its built-in model ID. When enable1M is true, the "[1m]"
// suffix is appended to pinned opus and sonnet IDs (never haiku, which has a
// 200K context window) to opt into the 1M context window on Bedrock; an ID
// that already carries the suffix is left as-is so it is not doubled.
//
// warnings collects the human-readable diagnostics for overrides that had no
// (or a surprising) effect — an ID matching no tier, and a second ID clobbering
// a tier an earlier ID already set — so the caller can surface them.
func resolveBedrockModels(models, discovered []string, enable1M bool) (bedrockTiers, []string) {
        tiers := newestBedrockTiers(discovered)
        var warnings []string
        // Track which tiers an override has already claimed so a second override for
        // the same tier is reported rather than silently winning last-write.
        assigned := make(map[string]string, 4)
        for _, m := range models {
                lower := strings.ToLower(m)
                family := ""
                for _, f := range []string{"haiku", "opus", "sonnet", "fable"} {
                        if strings.Contains(lower, f) {
                                family = f
                                break
                        }
                }
                if family == "" {
                        warnings = append(warnings, fmt.Sprintf(
                                "--models entry %q did not match a tier (expected haiku/opus/sonnet/fable in the ID); ignored", m))
                        continue
                }
                if prev, ok := assigned[family]; ok {
                        warnings = append(warnings, fmt.Sprintf(
                                "--models entry %q overrides %q for the %s tier; only the last is used", m, prev, family))
                }
                assigned[family] = m
                *tiers.slot(family) = m
        }
        if enable1M {
                tiers.opus = withOneMSuffix(tiers.opus)
                tiers.sonnet = withOneMSuffix(tiers.sonnet)
        }
        return tiers, warnings
}

// withOneMSuffix appends the "[1m]" suffix unless id is empty (an unpinned tier)
// or already ends with it, so a user-supplied ID that already carries the suffix
// is not doubled into "[1m][1m]".
func withOneMSuffix(id string) string {
        if id == "" || strings.HasSuffix(id, oneMSuffix) {
                return id
        }
        return id + oneMSuffix
}

// warnBedrockNoEffect warns about Bedrock flags the user passed on THIS run that
// had no effect. It keys off opts (the inline flags for this invocation), not the
// persisted config, because bedrock-compat is sticky: warning off the merged state
// would nag on every later setup that omits claude-code even though the flag was
// not passed. effectiveCompat is the merged (persisted + inline) compat state, used
// to decide whether an inline --enable-1m is inert.
//
// Two cases:
//   - --bedrock-compat passed this run but Claude Code was not among the configured
//     tools (the keys only attach to Claude Code), and
//   - --enable-1m passed this run while compat is off (the 1M suffix rides the
//     Bedrock model-ID path, so it is inert without compat) — otherwise a silent
//     no-op.
func warnBedrockNoEffect(errOut io.Writer, opts SetOptions, effectiveCompat bool, configured []ToolConfig) {
        if opts.BedrockCompat != nil && *opts.BedrockCompat && !isTarget(configured, claudeCodeClient) {
                _, _ = fmt.Fprintln(errOut,
                        "Warning: --bedrock-compat was set but Claude Code was not configured; the flag had no effect.")
        }
        if opts.Enable1M != nil && *opts.Enable1M && !effectiveCompat {
                _, _ = fmt.Fprintln(errOut,
                        "Warning: --enable-1m has no effect without --bedrock-compat; it is ignored.")
        }
}

func printPiLongCacheHint(out io.Writer, configured []ToolConfig, shortPromptCache bool) {
        if isTarget(configured, "pi") && !shortPromptCache {
                _, _ = fmt.Fprintln(out, "Pi uses short prompt caching by default. To enable long retention, "+
                        "start Pi with PI_CACHE_RETENTION=long in its environment (Anthropic: 1h; OpenAI: 24h).")
        }
}

func warnShortPromptCacheNoEffect(errOut io.Writer, opts SetOptions, configured []ToolConfig) {
        if opts.ShortPromptCache != nil && *opts.ShortPromptCache && !isTarget(configured, claudeCodeClient) {
                _, _ = fmt.Fprintln(errOut,
                        "Warning: --short-prompt-cache was set but Claude Code was not configured; the flag had no effect on client settings.")
        }
}

func warnPromptCacheBedrockCompatibility(
        errOut io.Writer, opts SetOptions, shortPromptCache bool, configured []ToolConfig,
) {
        if shortPromptCache || opts.BedrockCompat == nil || !*opts.BedrockCompat ||
                !isTarget(configured, claudeCodeClient) {
                return
        }
        _, _ = fmt.Fprintf(errOut,
                "Warning: Claude Code's one-hour prompt-cache lifetime may not take effect with Bedrock compatibility: "+
                        "the gateway may reject or strip the required beta header, and Bedrock support varies by model. "+
                        "Verify the effective lifetime with: %s\n", promptCacheVerificationCommand)
}

// configureDetectedTools patches each detected tool's config file and returns
// the list of successfully configured tools. An error is returned only when no
// tool was configured successfully.
func configureDetectedToolsWithDiscovery(
        out, errOut io.Writer, gm GatewayManager, detected []string,
        gatewayURL, proxyBaseURL, tokenHelperPath string, tokenHelperArgs []string,
        tlsSkipVerify bool, anthropicPathPrefix string,
        models []string, discovered discoveredGatewayModels, shortPromptCache bool, bedrock BedrockConfig,
) ([]ToolConfig, error) {
        var configured []ToolConfig
        for _, clientType := range detected {
                mode := gm.LLMGatewayModeFor(clientType)

                // Apply the Anthropic path prefix for tools that talk the Anthropic API
                // directly — direct-mode (ANTHROPIC_BASE_URL) and credential-helper mode
                // (Claude Desktop's inferenceGatewayBaseUrl). Proxy-mode tools (Cursor,
                // VS Code, Xcode) do not use it.
                anthropicBaseURL := ""
                if usesAnthropicBaseURL(mode) && anthropicPathPrefix != "" {
                        // Trim any leading slash: url.JoinPath docs say elements should not
                        // start with "/", and path.Join already handles the join correctly.
                        if joined, err := url.JoinPath(gatewayURL, strings.TrimLeft(anthropicPathPrefix, "/")); err == nil {
                                anthropicBaseURL = joined
                        }
                }

                applyCfg := llmgateway.ApplyConfig{
                        GatewayURL:                    gatewayURL,
                        AnthropicBaseURL:              anthropicBaseURL,
                        ProxyBaseURL:                  proxyBaseURL,
                        TokenHelperCommand:            tokenHelperShellCommand,
                        TokenHelperPath:               tokenHelperPath,
                        TokenHelperArgs:               tokenHelperArgs,
                        TLSSkipVerify:                 tlsSkipVerify,
                        Models:                        models,
                        DiscoveredModels:              discovered.openAI,
                        DiscoveredModelContextWindows: discovered.openAIContextWindows,
                        AnthropicModels:               discovered.anthropic,
                        AnthropicModelContextWindows:  discovered.anthropicContextWindows,
                        AnthropicModelReasoning:       discovered.anthropicReasoning,
                        AnthropicModelAdaptive:        discovered.anthropicAdaptive,
                }

                if clientType == claudeCodeClient {
                        applyCfg.ShortPromptCache = shortPromptCache
                }

                // Bedrock-compat applies only to Claude Code: it disables the experimental
                // anthropic-beta headers Bedrock rejects and pins each tier to the newest
                // discovered Bedrock model (or the --models override). Resolve tier
                // selection and the optional [1m] suffix here so pkg/client only writes
                // the resolved values.
                if bedrock.Compat && clientType == claudeCodeClient {
                        applyBedrockCompat(errOut, &applyCfg, models, discovered.anthropic, bedrock.Enable1M)
                }

                configPath, err := gm.ConfigureLLMGateway(clientType, applyCfg)
                if err != nil {
                        _, _ = fmt.Fprintf(errOut, "Warning: failed to configure %s: %v\n", clientType, err)
                        continue
                }
                envFilePath, err := gm.ConfigureEnvFile(clientType, applyCfg)
                if err != nil {
                        _, _ = fmt.Fprintf(errOut, "Warning: failed to write .env for %s: %v\n", clientType, err)
                        // Roll back the settings-file patch we just made.
                        if revertErr := gm.RevertLLMGateway(clientType, configPath); revertErr != nil {
                                _, _ = fmt.Fprintf(errOut, "Warning: rollback of %s failed: %v\n", clientType, revertErr)
                        }
                        continue
                }
                configured = append(configured, ToolConfig{
                        Tool:        clientType,
                        Mode:        mode,
                        ConfigPath:  configPath,
                        EnvFilePath: envFilePath,
                })
                _, _ = fmt.Fprintf(out, "Configured %s (%s mode)  →  %s\n", clientType, mode, configPath)
                if clientType == claudeCodeClient && !shortPromptCache {
                        warnPromptCacheConflict(errOut, gm, clientType)
                }
        }
        if len(configured) == 0 {
                return nil, fmt.Errorf("failed to configure any detected tools")
        }
        return configured, nil
}

func warnPromptCacheConflict(errOut io.Writer, gm GatewayManager, clientType string) {
        conflict, err := gm.PromptCacheConflict(clientType)
        if err != nil {
                _, _ = fmt.Fprintf(errOut,
                        "Warning: could not inspect %s for prompt-cache TTL conflicts: %v; "+
                                "ToolHive wrote the one-hour settings anyway.\n", clientType, err)
                return
        }
        if conflict != "" {
                _, _ = fmt.Fprintf(errOut,
                        "Warning: ToolHive wrote the one-hour prompt-cache settings for %s, but %s may override them. "+
                                "Remove that setting to use the one-hour lifetime. Verify the effective lifetime with: %s\n",
                        clientType, conflict, promptCacheVerificationCommand)
        }
}

// resolveAnthropicPrefix returns the effective Anthropic path prefix. When the
// caller explicitly set the flag (anthropicPathPrefixSet), the provided value is
// returned as-is (including empty string, which disables the prefix). Otherwise
// the gateway is auto-probed when at least one client uses its Anthropic API.
func resolveAnthropicPrefix(
        ctx context.Context, gm GatewayManager, detected []string,
        llmCfg Config, anthropicPathPrefix string, anthropicPathPrefixSet bool,
) string {
        if anthropicPathPrefixSet || !hasAnthropicBaseURLClient(gm, detected) {
                return anthropicPathPrefix
        }
        return probeAnthropicPrefix(ctx, llmCfg.GatewayURL, llmCfg.TLSSkipVerify)
}

// probeAnthropicPrefix performs a HEAD request to <gatewayURL>/anthropic/v1/messages.
// If the server responds with HTTP 401 (Unauthorized) — meaning the path exists but
// requires authentication — it returns "/anthropic" as the path prefix for
// clients using the Anthropic API.
// Any other status code (including 404 Not Found) or any network error is treated
// as "no prefix needed" and the function returns "".
func probeAnthropicPrefix(ctx context.Context, gatewayURL string, tlsSkipVerify bool) string {
        if gatewayURL == "" {
                return ""
        }
        probeURL, err := url.JoinPath(gatewayURL, "anthropic/v1/messages")
        if err != nil {
                return ""
        }

        httpClient := gatewayHTTPClient(tlsSkipVerify)

        // Use a short timeout so setup is not significantly slowed by an unreachable gateway.
        probeCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
        defer cancel()

        req, err := http.NewRequestWithContext(probeCtx, http.MethodHead, probeURL, nil)
        if err != nil {
                return ""
        }
        resp, err := httpClient.Do(req)
        if err != nil {
                return ""
        }
        _ = resp.Body.Close()

        if resp.StatusCode == http.StatusUnauthorized {
                return "/anthropic"
        }
        return ""
}

// gatewayHTTPClient clones the default transport so gateway requests retain
// its proxy, timeout, HTTP/2, and connection-pooling behavior.
func gatewayHTTPClient(tlsSkipVerify bool) *http.Client {
        //nolint:forcetypeassert // DefaultTransport is always *http.Transport
        transport := http.DefaultTransport.(*http.Transport).Clone()
        if tlsSkipVerify {
                if transport.TLSClientConfig == nil {
                        transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12}
                }
                transport.TLSClientConfig.InsecureSkipVerify = true //nolint:gosec // G402: intentional for local dev with self-signed certs
        }
        return &http.Client{Transport: transport}
}

// tokenHelperShellCommand is the shell command written into direct-mode tools'
// config as their token helper — e.g. Claude Code's apiKeyHelper, which is run
// through a shell (execa with shell:true; see anthropics/claude-code#42593).
//
// It deliberately names "thv" bare rather than interpolating os.Executable().
// A bare command has nothing to escape, so it is correct in both /bin/sh and
// cmd.exe without a platform branch, and it re-resolves on every invocation —
// so upgrading, reinstalling, or relocating thv keeps working without re-running
// "thv llm setup".
//
// The trade-off is that it resolves via PATH at invocation time. A tool launched
// without the user's shell PATH (e.g. from the macOS Dock, which inherits
// launchd's environment) will not find thv, and a binary earlier on PATH can
// shadow it. Both are accepted here: direct-mode tools are terminal-oriented.
// Claude Desktop cannot accept them — it is only ever GUI-launched — so its
// credential-helper shim uses the absolute TokenHelperPath instead (see
// writeCredentialHelperShim in pkg/client).
const tokenHelperShellCommand = llmgateway.TokenHelperShellCommand

// buildTokenHelperArgv returns the argv-form of the token helper, for config
// formats that invoke an executable directly (no shell) — e.g. Codex's
// [model_providers.<id>.auth] command/args table. Since args are passed as a
// TOML array rather than interpolated into a shell string, no shell-metacharacter
// validation is needed. --skip-browser is always included since Codex has no
// interactive-context signal like Claude Desktop's credential-helper shim.
func buildTokenHelperArgv() (string, []string, error) {
        self, err := os.Executable()
        if err != nil {
                return "", nil, fmt.Errorf("resolving thv executable path: %w", err)
        }
        return self, []string{"llm", "token", "--skip-browser"}, nil
}

// usesAnthropicBaseURL reports whether a client mode consumes the Anthropic base
// URL (gateway plus its probed or configured prefix): direct mode via
// ANTHROPIC_BASE_URL, credential-helper mode via inferenceGatewayBaseUrl,
// and Pi via its Anthropic provider.
func usesAnthropicBaseURL(mode string) bool {
        return mode == llmgateway.ModeDirect || mode == llmgateway.ModeCredentialHelper || mode == llmgateway.ModePi
}

// hasAnthropicBaseURLClient reports whether any detected client needs the
// Anthropic base URL. Proxy-mode tools ignore it entirely.
func hasAnthropicBaseURLClient(gm GatewayManager, detected []string) bool {
        for _, clientType := range detected {
                if usesAnthropicBaseURL(gm.LLMGatewayModeFor(clientType)) {
                        return true
                }
        }
        return false
}

// warnCredentialHelperTools prints the follow-up notes credential-helper clients
// (Claude Desktop) need: they read config only at launch, so the user must fully
// quit and relaunch; and a managed configuration, if present, overrides
// the local config setup just wrote.
func warnCredentialHelperTools(out, errOut io.Writer, gm GatewayManager, configured []ToolConfig) {
        for _, tc := range configured {
                if tc.Mode != llmgateway.ModeCredentialHelper {
                        continue
                }
                _, _ = fmt.Fprintf(out,
                        "%s reads its configuration only at launch — fully quit and reopen it "+
                                "for the gateway change to take effect.\n", tc.Tool)
                if gm.IsManaged(tc.Tool) {
                        _, _ = fmt.Fprintf(errOut,
                                "Warning: a managed configuration for %s is present; it overrides the local "+
                                        "configuration just written, which will be ignored. Remove the managed configuration, "+
                                        "configure the gateway there instead, or fix it if %s rejected it.\n", tc.Tool, tc.Tool)
                }
        }
}

// revertToolConfig reverts the settings-file and .env patches for a single
// tool. Errors are logged as warnings so the caller can continue with others.
func revertToolConfig(out, errOut io.Writer, gm GatewayManager, tc ToolConfig) {
        ok := true
        if err := gm.RevertLLMGateway(tc.Tool, tc.ConfigPath); err != nil {
                _, _ = fmt.Fprintf(errOut, "Warning: failed to revert %s settings: %v\n", tc.Tool, err)
                ok = false
        }
        if tc.EnvFilePath != "" {
                if err := gm.RevertEnvFile(tc.Tool, tc.EnvFilePath); err != nil {
                        _, _ = fmt.Fprintf(errOut, "Warning: failed to revert %s .env: %v\n", tc.Tool, err)
                        ok = false
                }
        }
        if ok {
                _, _ = fmt.Fprintf(out, "Reverted %s  (%s)\n", tc.Tool, tc.ConfigPath)
        }
}

// rollbackConfiguredTools reverts the settings-file and .env patches for every
// entry in configured. Errors are logged as warnings so all tools are attempted.
func rollbackConfiguredTools(errOut io.Writer, gm GatewayManager, configured []ToolConfig) {
        // Use a discard writer for the "Reverted" line — rollback is silent on success.
        for _, tc := range configured {
                revertToolConfig(io.Discard, errOut, gm, tc)
        }
}

// hasProxyMode reports whether any tool requires the localhost proxy.
func hasProxyMode(cfgs []ToolConfig) bool {
        for _, t := range cfgs {
                if t.Mode == llmgateway.ModeProxy || t.Mode == llmgateway.ModeVSCode {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package llm

import (
        "errors"
        "fmt"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth/tokensource"
        "github.com/stacklok/toolhive/pkg/llmgateway"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// ErrTokenRequired is returned when a fresh token is needed but no USABLE
// cached credential exists and the caller is non-interactive (browser flow
// disabled). That covers both an empty cache and a stored refresh token the
// identity provider has since rejected (expired, revoked, or rotated out from
// under us) — the two are one outcome to the user, because both are fixed only
// by an interactive login that persists a new refresh token.
var ErrTokenRequired = errors.New(
        "LLM gateway authentication required: no usable cached credentials; " +
                "run \"thv llm setup\" to log in",
)

// TokenRefUpdater is a callback invoked when the refresh token changes — either
// after a successful browser flow (initial login) or when the OIDC provider
// rotates the refresh token during a refresh. It persists the secret key and
// the new token expiry into the application config so future CLI invocations
// can restore the session. It is NOT called on routine access-token refreshes
// where the refresh token is unchanged.
// Callers typically wire this to config.UpdateConfig.
type TokenRefUpdater = tokensource.ConfigPersister

// TokenSource provides fresh LLM gateway access tokens.
type TokenSource = tokensource.OAuthTokenSource

// NewTokenSource creates a TokenSource for the LLM gateway.
// secretsProvider may be nil if the secrets store is unavailable.
// tokenRefUpdater is called after login/refresh to persist the token reference
// into config — pass nil to skip config persistence (useful in tests).
// interactive controls whether a genuine cache miss may launch the browser
// OIDC flow. thv llm token passes true so a prior "thv llm setup --lazy" signs
// the user in transparently on first use; cached tokens are served without a
// browser prompt either way.
// When skipBrowser is true, an interactive login prints the authorization URL
// instead of opening a browser (headless/SSH/CI use); it has no effect unless
// interactive is also true.
func NewTokenSource(
        cfg *Config, secretsProvider secrets.Provider, interactive, skipBrowser bool, tokenRefUpdater TokenRefUpdater,
) *TokenSource {
        return tokensource.New(tokensource.Options{
                OIDC: tokensource.OIDCParams{
                        Issuer:                   cfg.OIDC.Issuer,
                        ClientID:                 cfg.OIDC.ClientID,
                        Scopes:                   cfg.OIDC.EffectiveScopes(),
                        Audience:                 cfg.OIDC.Audience,
                        CallbackPort:             cfg.OIDC.CallbackPort,
                        RequireExactCallbackPort: true,
                },
                SecretsProvider: secretsProvider,
                Interactive:     interactive,
                SkipBrowser:     skipBrowser,
                KeyProvider: func() string {
                        if cfg.OIDC.CachedRefreshTokenRef != "" {
                                return cfg.OIDC.CachedRefreshTokenRef
                        }
                        return DeriveSecretKey(cfg.GatewayURL, cfg.OIDC.Issuer)
                },
                ConfigPersister: tokenRefUpdater,
                FallbackErr:     ErrTokenRequired,
                // Widen the preemptive refresh window past Claude Code's apiKeyHelper TTL
                // so every helper invocation in the final window forces a refresh and the
                // helper never hands back an about-to-expire token. Kept in sync with the
                // TTL written to settings.json (see pkg/llmgateway constants).
                PreemptiveRefreshWindow: llmgateway.LLMTokenRefreshWindow,
        })
}

// SanitizeTokenError returns a log-safe string for a token-source error.
// If err wraps *oauth2.RetrieveError, only the error code and description are
// included — never the raw response body, which may contain bearer material
// echoed back by the IdP.
func SanitizeTokenError(err error) string {
        var re *oauth2.RetrieveError
        if errors.As(err, &re) {
                if re.ErrorDescription != "" {
                        return fmt.Sprintf("oauth2 error %q: %s", re.ErrorCode, re.ErrorDescription)
                }
                return fmt.Sprintf("oauth2 error %q", re.ErrorCode)
        }
        return err.Error()
}

// DeriveSecretKey computes the secrets-provider key for an LLM gateway refresh
// token. The formula is: LLM_OAUTH_<8 hex chars> where the hex is derived from
// sha256(gatewayURL + "\x00" + issuer)[:4].
func DeriveSecretKey(gatewayURL, issuer string) string {
        return tokensource.DeriveSecretKey("LLM_OAUTH_", gatewayURL, issuer)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package llmgateway contains shared types for the LLM gateway feature,
// imported by both pkg/llm and pkg/client to avoid an import cycle.
package llmgateway

import (
        "net/url"
        "time"
)

// TokenHelperShellCommand is used by shell-invoked LLM gateway token helpers.
const TokenHelperShellCommand = "thv llm token" //nolint:gosec // command line, not a credential

// ClaudeCodeHelperTTL is written to settings.json as
// CLAUDE_CODE_API_KEY_HELPER_TTL_MS: how often Claude Code re-invokes the
// apiKeyHelper ("thv llm token"). Pinned to Claude Code's documented default so
// the cadence is deterministic rather than relying on the implicit default.
const ClaudeCodeHelperTTL = 5 * time.Minute

// LLMTokenRefreshWindow is how far before expiry the LLM gateway token source
// proactively refreshes. It is derived as 2x ClaudeCodeHelperTTL so that every
// helper invocation in the final window forces a refresh, meaning Claude Code
// never receives an about-to-expire token. The invariant
// LLMTokenRefreshWindow > ClaudeCodeHelperTTL is what makes this belt-and-
// suspenders pairing work; keep the two values in sync (guarded by a test).
const LLMTokenRefreshWindow = 2 * ClaudeCodeHelperTTL

// ClaudeDesktopHelperTTL is written to the Claude Desktop config as
// inferenceCredentialHelperTtlSec: how often Claude Desktop re-invokes the
// credential helper (a shim that runs "thv llm token"). Kept equal to
// ClaudeCodeHelperTTL so the same LLMTokenRefreshWindow invariant holds — every
// invocation in the final window forces a refresh, so Claude Desktop never
// receives an about-to-expire token.
const ClaudeDesktopHelperTTL = ClaudeCodeHelperTTL

// ClaudeDesktopHelperTimeout is written as inferenceCredentialHelperTimeoutSec:
// how long Claude Desktop waits for the credential helper to print a token
// before killing it. Sized to allow a silent OIDC refresh (cached refresh token
// → new access token) but not a full interactive re-auth, which "thv llm setup"
// performs up front so the steady-state helper call stays fast.
const ClaudeDesktopHelperTimeout = 30 * time.Second

// CodexHelperTTL is written to config.toml as
// [model_providers.<id>.auth].refresh_interval_ms: how often Codex re-invokes
// the token helper ("thv llm token"). Kept equal to ClaudeCodeHelperTTL so the
// same LLMTokenRefreshWindow invariant holds — every invocation in the final
// window forces a refresh, so Codex never receives an about-to-expire token.
const CodexHelperTTL = ClaudeCodeHelperTTL

// LLM gateway client modes. The single source of truth dispatched on in
// pkg/client (ConfigureLLMGateway/RevertLLMGateway) and pkg/llm
// (usesAnthropicBaseURL/warnCredentialHelperTools). Kept here so both packages
// reference one compiler-checked set of values rather than carrying private
// string copies that can drift.
const (
        // ModeDirect routes traffic through ANTHROPIC_BASE_URL / a token helper.
        ModeDirect = "direct"
        // ModeProxy routes traffic through a localhost reverse proxy.
        ModeProxy = "proxy"
        // ModeCredentialHelper routes traffic through a document + selector +
        // helper-script model (Claude Desktop's third-party inference surface).
        ModeCredentialHelper = "credential-helper" //nolint:gosec // G101: mode identifier, not a credential
        // ModeCodexAuth routes traffic through a custom model_provider written into
        // Codex's TOML config, using Codex's command-backed bearer-token auth
        // ([model_providers.<id>.auth]) to invoke the token helper directly (no
        // shell), rather than the JSON-Pointer LLMGatewayKeys mechanism the other
        // modes share.
        ModeCodexAuth = "codex-auth" //nolint:gosec // G101: mode identifier, not a credential
        // ModePi configures Pi's OpenAI-compatible and Anthropic providers in models.json.
        ModePi = "pi"
        // ModeVSCode configures VS Code's chatLanguageModels.json provider groups.
        ModeVSCode = "vscode"
)

// ProxyOriginOf returns rawURL with its path, query, fragment, and userinfo
// stripped so only the scheme and host remain (the "origin"). Tools like
// Gemini CLI that append their own API path (e.g. /v1beta/...) need the
// origin rather than the full proxy base URL to avoid doubled path segments.
//
// Falls back to rawURL when:
//   - parsing fails, or
//   - the parsed URL has no Host (e.g. scheme-less "localhost:14000/v1" is
//     parsed by url.Parse as Scheme=localhost, Opaque=14000/v1 with no Host),
//     or
//   - the parsed URL has a non-empty Opaque field (opaque URI reference).
func ProxyOriginOf(rawURL string) string {
        u, err := url.Parse(rawURL)
        if err != nil || u == nil || u.Host == "" || u.Opaque != "" {
                return rawURL
        }
        origin := url.URL{
                Scheme: u.Scheme,
                Host:   u.Host,
        }
        return origin.String()
}

// ApplyConfig holds the values needed to configure a single tool's LLM
// gateway settings. Using a struct prevents positional-argument mistakes when
// the caller has multiple similar string values in scope.
type ApplyConfig struct {
        GatewayURL         string // direct-mode: URL of the upstream LLM gateway
        AnthropicBaseURL   string // direct-mode: effective ANTHROPIC_BASE_URL (gateway + optional path prefix)
        ProxyBaseURL       string // proxy-mode: URL of the localhost reverse proxy
        TokenHelperCommand string // direct-mode: shell command that prints a fresh token
        // TokenHelperPath and TokenHelperArgs are the argv-form of the token helper,
        // for config formats where auth invokes an executable directly (no shell) —
        // e.g. Codex's [model_providers.<id>.auth] command/args table.
        TokenHelperPath string
        TokenHelperArgs []string
        TLSSkipVerify   bool // when true, instruct the tool to skip TLS verification
        // Models is the optional explicit model list written to clients that support
        // a model override (e.g. Claude Desktop's inferenceModels). Empty means the
        // client falls back to gateway-side model auto-discovery.
        Models []string
        // DiscoveredModels contains model IDs returned by the authenticated gateway
        // discovery request. It is used only by integrations that require an explicit
        // model catalogue, such as VS Code's customendpoint provider.
        DiscoveredModels []string
        // DiscoveredModelContextWindows contains positive context sizes from gateway
        // discovery for Pi's model catalogue. Missing entries use Pi's default.
        DiscoveredModelContextWindows map[string]int
        // AnthropicModels and AnthropicModelContextWindows contain the separately
        // discovered Anthropic model catalogue for Pi. Sizes are max_input_tokens.
        AnthropicModels              []string
        AnthropicModelContextWindows map[string]int
        // AnthropicModelReasoning is true for Anthropic models with a supported
        // thinking mode. Missing entries use Pi's default.
        AnthropicModelReasoning map[string]bool
        // AnthropicModelAdaptive marks models that accept adaptive thinking; Pi
        // otherwise sends budget-based thinking for reasoning models.
        AnthropicModelAdaptive map[string]bool
        // ShortPromptCache opts supported clients out of ToolHive's one-hour
        // prompt-cache lifetime and restores their shorter default.
        ShortPromptCache bool
        // BedrockCompat and the per-tier Bedrock model IDs configure Claude Code for a
        // gateway that forwards to AWS Bedrock. When BedrockCompat is true, Claude Code
        // is configured with CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1 (Bedrock rejects
        // the experimental anthropic-beta headers) and the per-tier model IDs.
        // The caller (pkg/llm) resolves tier selection and the optional "[1m]"
        // suffix; pkg/client only writes the resolved values. An empty tier is left
        // unpinned. All are empty/false for non-Bedrock setups.
        BedrockCompat      bool
        BedrockHaikuModel  string
        BedrockOpusModel   string
        BedrockSonnetModel string
        BedrockFableModel  string
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package lockfile provides utilities for managing file locks and cleanup.
package lockfile

import (
        "context"
        "log/slog"
        "os"
        "path/filepath"
        "sync"
        "time"

        "github.com/gofrs/flock"
)

var (
        // globalRegistry holds all active lock files for cleanup
        globalRegistry = &lockRegistry{
                locks: make(map[string]*flock.Flock),
        }
)

// lockRegistry manages active file locks for cleanup purposes
type lockRegistry struct {
        mu    sync.RWMutex
        locks map[string]*flock.Flock
}

// RegisterLock adds a lock to the global registry for cleanup
func (lr *lockRegistry) RegisterLock(lockPath string, lock *flock.Flock) {
        lr.mu.Lock()
        defer lr.mu.Unlock()
        lr.locks[lockPath] = lock
}

// UnregisterLock removes a lock from the global registry
func (lr *lockRegistry) UnregisterLock(lockPath string) {
        lr.mu.Lock()
        defer lr.mu.Unlock()
        delete(lr.locks, lockPath)
}

// CleanupAll unlocks and removes all registered lock files
func (lr *lockRegistry) CleanupAll() {
        lr.mu.Lock()
        defer lr.mu.Unlock()

        for lockPath, lock := range lr.locks {
                if err := lock.Unlock(); err != nil && !os.IsNotExist(err) {
                        slog.Warn("failed to unlock file", "path", lockPath, "error", err)
                }

                if err := os.Remove(lockPath); err != nil && !os.IsNotExist(err) {
                        slog.Warn("failed to remove lock file", "path", lockPath, "error", err)
                }
        }

        // Clear the registry
        lr.locks = make(map[string]*flock.Flock)
}

// NewTrackedLock creates a new file lock and registers it for cleanup
func NewTrackedLock(lockPath string) *flock.Flock {
        lock := flock.New(lockPath)
        globalRegistry.RegisterLock(lockPath, lock)
        return lock
}

// ReleaseTrackedLock unlocks, removes, and unregisters a lock file.
//
// The path is unlinked before the flock is released. flock(2) locks attach to
// an inode, not a path, so releasing in the other order leaves a window where
// a waiter can flock this soon-to-be-removed inode and then have the path
// pulled out from under it by this call's Remove, while a third locker
// creates a fresh inode at the same path — two lockers now believe they hold
// "the" lock. Unlinking first means anyone who opens lockPath from here on
// necessarily gets a new inode, which is what AcquireTrackedLock's identity
// check depends on to notice and retry past a stale lock.
func ReleaseTrackedLock(lockPath string, lock *flock.Flock) {
        if err := os.Remove(lockPath); err != nil && !os.IsNotExist(err) {
                slog.Warn("failed to remove lock file", "path", lockPath, "error", err)
        }

        if err := lock.Unlock(); err != nil && !os.IsNotExist(err) {
                slog.Warn("failed to unlock file", "path", lockPath, "error", err)
        }

        globalRegistry.UnregisterLock(lockPath)
}

// AcquireTrackedLock acquires an exclusive lock on lockPath and registers it
// for cleanup, the way NewTrackedLock+TryLockContext do, but additionally
// guards against acquiring a lock on a stale, already-orphaned inode.
//
// Because flock(2) locks an open file description on an inode rather than a
// path, and acquisition (open, then flock) and release (unlink, then unlock —
// see ReleaseTrackedLock) are separate steps, a caller can open lockPath and
// still be blocked in flock() at the moment another holder unlinks and
// releases: that other holder's Remove can land between this caller's open()
// and its successful flock(), or a concurrent third party can recreate
// lockPath under a new inode while this caller's flock succeeds on the old,
// now-unreachable one. Either way this caller would believe it holds "the"
// lock for lockPath while actually holding a lock nothing else will ever
// contend for again. After acquiring, this function checks that the locked
// file's inode is still the one lockPath resolves to (os.SameFile); on
// mismatch it discards the stale lock and retries with a fresh open, bounded
// by ctx.
func AcquireTrackedLock(ctx context.Context, lockPath string, retryDelay time.Duration) (*flock.Flock, bool, error) {
        for {
                lock := flock.New(lockPath)

                locked, err := lock.TryLockContext(ctx, retryDelay)
                if err != nil {
                        return nil, false, err
                }
                if !locked {
                        return nil, false, nil
                }

                if lockMatchesPath(lock, lockPath) {
                        globalRegistry.RegisterLock(lockPath, lock)
                        return lock, true, nil
                }

                if err := lock.Unlock(); err != nil && !os.IsNotExist(err) {
                        slog.Warn("failed to unlock stale lock file", "path", lockPath, "error", err)
                }
                if ctx.Err() != nil {
                        return nil, false, nil
                }
        }
}

// lockMatchesPath reports whether the inode lock currently holds is still the
// one lockPath resolves to, i.e. that the lock has not been orphaned by a
// concurrent unlink (and possibly recreation) of lockPath.
func lockMatchesPath(lock *flock.Flock, lockPath string) bool {
        heldInfo, err := lock.Stat()
        if err != nil {
                return false
        }

        pathInfo, err := os.Stat(lockPath)
        if err != nil {
                return false
        }

        return os.SameFile(heldInfo, pathInfo)
}

// CleanupAllLocks provides global cleanup of all registered lock files
func CleanupAllLocks() {
        globalRegistry.CleanupAll()
}

// CleanupStaleLocks removes stale lock files from the specified directories
// A lock file is considered stale if it's older than the maxAge duration
func CleanupStaleLocks(directories []string, maxAge time.Duration) {
        cutoff := time.Now().Add(-maxAge)

        for _, dir := range directories {
                matches, err := filepath.Glob(filepath.Join(dir, "*.lock"))
                if err != nil {
                        slog.Warn("failed to glob lock files", "dir", dir, "error", err)
                        continue
                }

                for _, lockFile := range matches {
                        info, err := os.Stat(lockFile)
                        if err != nil {
                                continue // File may have been removed already
                        }

                        if info.ModTime().Before(cutoff) {
                                // Try to acquire the lock to check if it's really stale
                                lock := flock.New(lockFile)
                                if locked, err := lock.TryLock(); err == nil && locked {
                                        // Lock was acquired, so it was stale
                                        if err := lock.Unlock(); err != nil && !os.IsNotExist(err) {
                                                slog.Warn("failed to unlock stale lock file", "path", lockFile, "error", err)
                                        }
                                        if err := os.Remove(lockFile); err != nil && !os.IsNotExist(err) {
                                                slog.Warn("failed to remove stale lock file", "path", lockFile, "error", err)
                                        } else {
                                                slog.Debug("removed stale lock file", "path", lockFile)
                                        }
                                }
                        }
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "bytes"
        "net/http"
)

// batchUnsupported is the shared error rendered by both batch-rejection helpers.
var batchUnsupported = &BatchUnsupportedError{}

// BatchUnsupportedError indicates a client sent a JSON-RPC batch (a top-level
// JSON array). Batching was removed from MCP in the 2025-06-18 revision, and
// ToolHive serves only 2025-11-25 and 2026-07-28, so a batch can only originate
// from an unsupported pre-2025-06-18 client.
//
// Batches are rejected outright rather than executed: the parser and the
// downstream authz, audit, and tool-filter middleware each inspect a single
// parsed request per call, so a batch that reached the backend would smuggle
// every nested call past those controls (see #5745).
type BatchUnsupportedError struct{}

func (*BatchUnsupportedError) Error() string {
        return "JSON-RPC batch requests are not supported (batching was removed in MCP revision 2025-06-18)"
}

// Code implements CodedError.
func (*BatchUnsupportedError) Code() int64 { return CodeInvalidRequest }

// Data implements CodedError.
func (*BatchUnsupportedError) Data() map[string]any { return nil }

// IsBatchRequest reports whether body is a JSON-RPC batch, i.e. a top-level JSON
// array. Only the first non-whitespace byte is inspected, so a syntactically
// malformed array (e.g. "[...") still classifies as a batch and is rejected
// rather than silently mis-parsed as a single message.
//
// It trims the same Unicode whitespace as encoding/json (via bytes.TrimSpace) so
// callers that gate execution on this function and callers that later
// json-decode the body agree on what counts as a leading '[' — a narrower trim
// would let a batch prefixed with, say, a vertical-tab slip past detection yet
// still decode as an array. Every path that could forward a batch to a backend
// must gate on this one function so the two never drift (see #5745).
func IsBatchRequest(body []byte) bool {
        trimmed := bytes.TrimSpace(body)
        return len(trimmed) > 0 && trimmed[0] == '['
}

// WriteBatchUnsupportedError writes an HTTP 400 response carrying a JSON-RPC
// "Invalid Request" error for a rejected batch. A batch has no single request
// id to echo, so the "id" key is omitted entirely (schema/2025-11-25 types the
// error response's id as optional, not nullable; see session.HasJSONRPCID).
// Use this where an http.ResponseWriter is available (the streamable proxy,
// ParsingMiddleware).
func WriteBatchUnsupportedError(w http.ResponseWriter) {
        body := classificationErrorBody(nil, batchUnsupported)
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusBadRequest)
        //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client
        _, _ = w.Write(body)
}

// BatchUnsupportedResponse builds an *http.Response carrying the batch-rejection
// JSON-RPC error, for proxies that intercept at the RoundTripper layer (the
// transparent proxy) where no http.ResponseWriter is available. Mirrors
// ClassificationErrorResponse; the JSON-RPC id is omitted, not null.
func BatchUnsupportedResponse(req *http.Request) *http.Response {
        return jsonRPCErrorResponse(req, nil, batchUnsupported)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "bytes"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "net/http"

        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/transport/session"
)

// WriteClassificationError writes an HTTP 400 response with a JSON-RPC error
// body for a coded request-classification or structural-validation failure.
// Use this where an http.ResponseWriter is available.
func WriteClassificationError(w http.ResponseWriter, requestID any, err error) {
        body := classificationErrorBody(requestID, err)
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusBadRequest)
        //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client
        _, _ = w.Write(body)
}

// ClassificationErrorResponse constructs an *http.Response with HTTP 400 and
// a JSON-RPC error body for an mcp.ClassifyRevision failure. Use this in
// httputil.ReverseProxy.ModifyResponse/RoundTrip (transparent proxy) where no
// http.ResponseWriter is available.
//
// req is attached as resp.Request, matching session.NotFoundResponse: a
// RoundTripper-produced *http.Response is generally expected to carry the
// request it answers, and httputil.ReverseProxy relies on that field.
func ClassificationErrorResponse(req *http.Request, requestID any, err error) *http.Response {
        return jsonRPCErrorResponse(req, requestID, err)
}

// jsonRPCErrorResponse builds an HTTP 400 *http.Response carrying the JSON-RPC
// error rendered from a CodedError. It backs both ClassificationErrorResponse
// and BatchUnsupportedResponse so the RoundTripper-layer error shape lives in
// one place.
func jsonRPCErrorResponse(req *http.Request, requestID any, err error) *http.Response {
        body := classificationErrorBody(requestID, err)
        hdr := make(http.Header)
        hdr.Set("Content-Type", "application/json")
        return &http.Response{
                StatusCode:    http.StatusBadRequest,
                Status:        fmt.Sprintf("%d %s", http.StatusBadRequest, http.StatusText(http.StatusBadRequest)),
                Proto:         "HTTP/1.1",
                ProtoMajor:    1,
                ProtoMinor:    1,
                Header:        hdr,
                ContentLength: int64(len(body)),
                Body:          io.NopCloser(bytes.NewReader(body)),
                Request:       req,
        }
}

// classificationErrorBody renders an MCP admission error as a JSON-RPC
// error body, modeled on session.NotFoundBody: the body is marshaled first
// (with a hand-crafted fallback on marshal failure) so callers only write
// headers/status once a valid body is ready.
//
// Parse errors retain the standard -32700 code. CodedError supplies Code(),
// Error(), and optional Data(); other errors fall back to Invalid Params.
//
// When requestID is nil, an id carried by err (RequestID() any) is used; only
// errors raised after the id was validated carry one.
//
// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error
// response as `id?: RequestId` where RequestId = string | number, so an
// absent id is encoded by omitting the "id" key, never as null. See
// session.HasJSONRPCID.
func classificationErrorBody(requestID any, err error) []byte {
        code := CodeInvalidParams
        var coded CodedError
        var data map[string]any
        if errors.Is(err, jsonrpc2.ErrParse) {
                code = -32700
        } else if errors.As(err, &coded) {
                code = coded.Code()
                data = coded.Data()
        }

        errBody := map[string]any{
                "code":    code,
                "message": err.Error(),
        }
        if len(data) > 0 {
                errBody["data"] = data
        }
        resp := map[string]any{
                "jsonrpc": "2.0",
                "error":   errBody,
        }
        if requestID == nil {
                var withID interface{ RequestID() any }
                if errors.As(err, &withID) {
                        requestID = withID.RequestID()
                }
        }
        if session.HasJSONRPCID(requestID) {
                resp["id"] = requestID
        }

        body, marshalErr := json.Marshal(resp)
        if marshalErr != nil {
                // This should never happen with simple map types, but return a
                // hand-crafted fallback to guarantee a valid JSON-RPC error.
                return []byte(`{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid params"}}`)
        }
        return body
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import "net/http"

// JSONRPCCodeDenied is the application-space JSON-RPC error code ToolHive uses
// when a policy denies a call, alongside HTTP 403. Both denial paths reference
// this constant — the single-server authorization middleware
// (pkg/authz.handleUnauthorized) and the vMCP Serve-path call gate — so a denial
// is represented by one code across `thv run` and vMCP, by reference rather than
// by two copies of the literal. It is deliberately outside the reserved
// -32768..-32000 JSON-RPC range so it never collides with an SDK-generated code.
// It is also deliberately untyped: it is consumed both as an int
// (server.Denial.Code, map[string]any envelopes) and as an int64
// (jsonrpc2.NewError, JSONRPCCodeForStatus). Typing it would break
// assert.Equal sites on one side.
const JSONRPCCodeDenied = 403

// JSONRPCCodeForStatus maps an HTTP status code to the JSON-RPC error code
// that belongs in a response body's error.code field. The HTTP status and the
// JSON-RPC code are independent dimensions: ToolHive's error paths pick an
// HTTP status per failure class, and each class has exactly one correct
// JSON-RPC code. This function is the single place that mapping lives, so
// callers never pass an HTTP status straight through as the JSON-RPC code.
//
// http.StatusUnprocessableEntity maps to JSONRPCCodeDenied because ToolHive's
// mutating-webhook path relays a webhook's HTTP 422 always-deny response
// (webhook.IsAlwaysDenyError) as HTTP 422 downstream — the validating path
// relays the same condition as 403. So a 422 arriving here is semantically a
// denial.
//
// Any status with no explicit case below maps to CodeInternalError, so an
// unmapped status fails safe to a standard code instead of leaking a raw HTTP
// status into the JSON-RPC code field.
func JSONRPCCodeForStatus(status int) int64 {
        switch status {
        case http.StatusForbidden, http.StatusUnprocessableEntity:
                return JSONRPCCodeDenied
        case http.StatusRequestEntityTooLarge:
                return CodeInvalidRequest
        default:
                return CodeInternalError
        }
}

// CodedError is implemented by domain errors that should surface a stable error
// code and optional structured data in an MCP tool result.
//
// The mcp-go tool-handler seam maps returned Go errors to a generic JSON-RPC
// INTERNAL_ERROR, so Serve-path handlers convert these errors to IsError tool
// results with StructuredContent instead of returning them as handler errors.
type CodedError interface {
        error
        Code() int64
        Data() map[string]any
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "bytes"
        "encoding/json"
        "io"
        "strings"
        "unicode"
)

var ambiguousRequest = &ambiguousRequestError{}

type ambiguousRequestError struct{}

// Error is deliberately fixed so conflicting request-controlled names and
// values are never reflected to the client.
func (*ambiguousRequestError) Error() string { return "Invalid Request" }

// Code implements CodedError.
func (*ambiguousRequestError) Code() int64 { return CodeInvalidRequest }

// Data implements CodedError.
func (*ambiguousRequestError) Data() map[string]any { return nil }

type jsonContainer struct {
        delim        json.Delim
        seen         map[string]struct{}
        expectingKey bool
}

type jsonAmbiguityScanner struct {
        containers   []jsonContainer
        ambiguous    bool
        rootComplete bool
}

// hasAmbiguousJSONMembers reports whether a syntactically valid JSON value has
// duplicate or case-fold-equivalent member names in any object. JSON permits
// case-distinct names, but ToolHive rejects them deliberately because downstream
// decoders may merge them and act on a value different from the one authorized.
// Syntax errors and trailing values are left to the request parser.
func hasAmbiguousJSONMembers(body []byte) bool {
        body = bytes.TrimPrefix(body, UTF8BOM)
        decoder := json.NewDecoder(bytes.NewReader(body))
        decoder.UseNumber()
        scanner := jsonAmbiguityScanner{
                containers: make([]jsonContainer, 0, 8),
        }

        for {
                token, err := decoder.Token()
                if err == io.EOF {
                        return scanner.isAmbiguous()
                }
                if err != nil || scanner.rootComplete {
                        return false
                }
                if !scanner.consumeToken(token) {
                        return false
                }
        }
}

func (s *jsonAmbiguityScanner) isAmbiguous() bool {
        return s.rootComplete && len(s.containers) == 0 && s.ambiguous
}

func (s *jsonAmbiguityScanner) consumeToken(token json.Token) bool {
        if len(s.containers) == 0 {
                return s.consumeRootToken(token)
        }

        current := &s.containers[len(s.containers)-1]
        if current.delim == '{' && current.expectingKey {
                return s.consumeObjectKey(current, token)
        }
        if current.delim == '{' {
                current.expectingKey = true
        }

        delim, ok := token.(json.Delim)
        if !ok {
                return true
        }
        return s.consumeValueDelimiter(current, delim)
}

func (s *jsonAmbiguityScanner) consumeRootToken(token json.Token) bool {
        delim, ok := token.(json.Delim)
        if !ok {
                s.rootComplete = true
                return true
        }
        if delim != '{' && delim != '[' {
                return false
        }
        s.containers = append(s.containers, newJSONContainer(delim))
        return true
}

func (s *jsonAmbiguityScanner) consumeObjectKey(current *jsonContainer, token json.Token) bool {
        if delim, ok := token.(json.Delim); ok && delim == '}' {
                s.closeContainer()
                return true
        }

        key, ok := token.(string)
        if !ok {
                return false
        }
        folded := foldJSONMemberName(key)
        if _, exists := current.seen[folded]; exists {
                s.ambiguous = true
        } else {
                current.seen[folded] = struct{}{}
        }
        current.expectingKey = false
        return true
}

func (s *jsonAmbiguityScanner) consumeValueDelimiter(current *jsonContainer, delim json.Delim) bool {
        switch delim {
        case '{', '[':
                s.containers = append(s.containers, newJSONContainer(delim))
        case ']':
                if current.delim != '[' {
                        return false
                }
                s.closeContainer()
        default:
                return false
        }
        return true
}

func (s *jsonAmbiguityScanner) closeContainer() {
        s.containers = s.containers[:len(s.containers)-1]
        if len(s.containers) == 0 {
                s.rootComplete = true
        }
}

func newJSONContainer(delim json.Delim) jsonContainer {
        container := jsonContainer{delim: delim}
        if delim == '{' {
                container.seen = make(map[string]struct{})
                container.expectingKey = true
        }
        return container
}

// foldJSONMemberName returns one stable representative for each Unicode simple
// case-folding equivalence class, matching strings.EqualFold semantics.
func foldJSONMemberName(name string) string {
        var folded strings.Builder
        folded.Grow(len(name))
        for _, r := range name {
                canonical := r
                for next := unicode.SimpleFold(r); next != r; next = unicode.SimpleFold(next) {
                        if next < canonical {
                                canonical = next
                        }
                }
                folded.WriteRune(canonical)
        }
        return folded.String()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "log/slog"
        "net/http"

        "golang.org/x/exp/jsonrpc2"
)

// EncodeJSONRPCError returns the wire encoding of a jsonrpc2 error response.
//
// This is the single sanctioned way to serialize a jsonrpc2 envelope in this
// codebase. NEVER json.Marshal (or json.NewEncoder().Encode) a jsonrpc2 type:
// jsonrpc2.Response carries no json tags and jsonrpc2.ID's only field is
// unexported, so reflection produces Go-capitalized keys, drops the mandatory
// "jsonrpc":"2.0" tag, and renders every id — valid or not — as "{}" (#5950).
// jsonrpc2.EncodeMessage marshals through the library's tagged wire struct:
// lowercase keys, the version tag stamped unconditionally, and an id that is
// omitted when absent (never null, which MCP cannot express and the reference
// TypeScript SDK client rejects, #6038) while a legitimate id of 0 survives.
//
// EncodeMessage cannot fail for a Response built from an ID and a
// jsonrpc2.NewError; if it somehow does, this logs and falls back to a
// hardcoded, conformant Internal Error body rather than returning nothing —
// the original code and id are not trustworthy at that point.
func EncodeJSONRPCError(resp *jsonrpc2.Response) []byte {
        body, err := jsonrpc2.EncodeMessage(resp)
        if err != nil {
                slog.Error("failed to encode JSON-RPC error response", "error", err)
                return []byte(`{"jsonrpc":"2.0","error":{"code":-32603,"message":"Internal error"}}`)
        }
        return body
}

// WriteJSONRPCError writes resp as an HTTP response: the JSON-RPC error body
// under Content-Type application/json with the given HTTP status. The body is
// encoded before any header is written, so an encode failure never leaves a
// half-written response (see EncodeJSONRPCError for the fallback). Returns
// the body write error; callers with no recovery path may discard it.
//
// Do not use this to write into an SSE stream — a bare JSON object in a
// text/event-stream parses as an unrecognized field and is silently
// discarded (#6037); SSE paths must frame the body as an event instead.
func WriteJSONRPCError(w http.ResponseWriter, httpStatus int, resp *jsonrpc2.Response) error {
        body := EncodeJSONRPCError(resp)
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(httpStatus)
        _, err := w.Write(body)
        return err
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "bytes"
        "encoding/json"
        "errors"
        "log/slog"
        "net/http"
        "strings"

        "golang.org/x/exp/jsonrpc2"
)

// DecodeMessage admits and decodes a single, unambiguous MCP JSON-RPC envelope,
// independently of HTTP headers. IDs must be strings or signed int64 JSON integers
// and are preserved exactly. Duplicate or case-fold-colliding object keys are
// rejected recursively, including in params, results, errors and extensions.
// Unknown extension members are allowed; aliases of reserved members are not.
// Request and notification params that do not decode as an object (e.g. a number
// outside float64 range) are rejected.
// Result values retain all accepted JSON shapes; no result-object schema is imposed.
// Only requests and notifications should be published as ParsedMCPRequest.
//
// Invalid JSON returns jsonrpc2.ErrParse; valid batches return *BatchUnsupportedError;
// invalid envelopes, ambiguous members and unsupported IDs return a CodedError
// with CodeInvalidRequest. A leading UTF-8 BOM and JSON whitespace are accepted.
func DecodeMessage(body []byte) (msg jsonrpc2.Message, err error) {
        defer func() {
                if err != nil {
                        slog.Warn("rejected invalid MCP JSON-RPC message", "reason", messageRejectionReason(err))
                }
        }()
        body = bytes.Trim(bytes.TrimPrefix(body, UTF8BOM), " \t\r\n")
        if !json.Valid(body) {
                return nil, jsonrpc2.ErrParse
        }
        if IsBatchRequest(body) {
                return nil, &BatchUnsupportedError{}
        }
        var envelope map[string]json.RawMessage
        if body[0] != '{' || json.Unmarshal(body, &envelope) != nil ||
                hasAmbiguousJSONMembers(body) || reservedMemberAlias(envelope, "jsonrpc", "id", "method", "params", "result", "error") {
                return nil, ambiguousRequest
        }
        if !validMessageEnvelope(envelope) {
                return nil, ambiguousRequest
        }
        id, err := decodeMessageID(envelope["id"])
        if err != nil {
                return nil, err
        }
        if _, isRequest := envelope["method"]; isRequest && !decodableParams(envelope["params"]) {
                return nil, newUndecodableParams(id)
        }
        return decodeWithExactID(body, envelope, id)
}

// decodeWithExactID decodes an admitted envelope, preserving id exactly.
func decodeWithExactID(body []byte, envelope map[string]json.RawMessage, id jsonrpc2.ID) (jsonrpc2.Message, error) {
        // MCP permits an error response without an ID. The decoder requires one;
        // supply it only for decoding and restore the absent ID below.
        if envelope["error"] != nil && !id.IsValid() {
                body = append([]byte(`{"id":0,`), body[1:]...)
        }
        msg, err := jsonrpc2.DecodeMessage(body)
        if err != nil {
                return nil, ambiguousRequest
        }
        // DecodeMessage rounds numeric IDs through float64. Restore the exact ID.
        switch message := msg.(type) {
        case *jsonrpc2.Request:
                message.ID = id
        case *jsonrpc2.Response:
                message.ID = id
        }
        return msg, nil
}

func messageRejectionReason(err error) string {
        var batch *BatchUnsupportedError
        var undecodable *undecodableParamsError
        switch {
        case errors.Is(err, jsonrpc2.ErrParse):
                return "parse"
        case errors.As(err, &batch):
                return "batch"
        case errors.As(err, &undecodable):
                return "undecodable_params"
        default:
                return "invalid_request"
        }
}

// undecodableParamsError rejects a request whose envelope is otherwise valid
// but whose params do not decode. Only then is the request ID unambiguous, so
// it is carried to let the client correlate the error. Error is fixed so no
// request content is reflected.
type undecodableParamsError struct{ id any }

// newUndecodableParams carries id only when it is valid (a notification has none).
func newUndecodableParams(id jsonrpc2.ID) *undecodableParamsError {
        if !id.IsValid() {
                return &undecodableParamsError{}
        }
        return &undecodableParamsError{id: id.Raw()}
}

func (*undecodableParamsError) Error() string { return "Invalid Request" }

// Code implements CodedError.
func (*undecodableParamsError) Code() int64 { return CodeInvalidRequest }

// Data implements CodedError.
func (*undecodableParamsError) Data() map[string]any { return nil }

// RequestID returns the decoded request ID (string or int64), or nil.
func (e *undecodableParamsError) RequestID() any { return e.id }

// writeBodyReadError never exposes reader-provided errors to clients or logs.
func writeBodyReadError(w http.ResponseWriter, err error) {
        var tooLarge *http.MaxBytesError
        if errors.As(err, &tooLarge) {
                slog.Warn("rejected unreadable MCP request body", "reason", "body_too_large")
                http.Error(w, "Request Entity Too Large", http.StatusRequestEntityTooLarge)
                return
        }
        slog.Warn("rejected unreadable MCP request body", "reason", "read_error")
        WriteClassificationError(w, nil, ambiguousRequest)
}

func decodeMessageID(raw json.RawMessage) (jsonrpc2.ID, error) {
        if raw == nil {
                return jsonrpc2.ID{}, nil
        }
        if raw[0] == '"' {
                var id string
                if json.Unmarshal(raw, &id) == nil {
                        return jsonrpc2.StringID(id), nil
                }
        } else {
                var id int64
                if json.Unmarshal(raw, &id) == nil && string(raw) != "null" {
                        return jsonrpc2.Int64ID(id), nil
                }
        }
        return jsonrpc2.ID{}, ambiguousRequest
}

func validMessageEnvelope(envelope map[string]json.RawMessage) bool {
        var version string
        if json.Unmarshal(envelope["jsonrpc"], &version) != nil || version != "2.0" {
                return false
        }
        if method, request := envelope["method"]; request {
                var name string
                params := envelope["params"]
                return json.Unmarshal(method, &name) == nil && name != "" &&
                        envelope["result"] == nil && envelope["error"] == nil && (params == nil || params[0] == '{')
        }
        return validResponseEnvelope(envelope)
}

// decodableParams reports whether request params decode the way the parser
// decodes them. Syntactically valid JSON can still fail to decode: a number
// outside float64 range (e.g. 1e1000) is accepted by lenient backend parsers
// but makes the decode fail, which would leave authorization, rate limiting and
// audit acting on an empty resource ID while the backend executes the real one.
// Such requests are rejected rather than decoded with json.Number, because
// authorizers only consume float64 numeric arguments. Responses are not checked:
// their results are opaque to ToolHive and may legitimately carry any number.
func decodableParams(params json.RawMessage) bool {
        _, err := decodeParams(params)
        return err == nil
}

// decodeParams decodes request params into an object map. It is the single
// decode shared by admission (decodableParams) and the parser, so the two
// cannot disagree about which params are usable. Nil params decode to a nil map.
func decodeParams(params json.RawMessage) (map[string]any, error) {
        if params == nil {
                return nil, nil
        }
        var m map[string]any
        if err := json.Unmarshal(params, &m); err != nil {
                return nil, err
        }
        return m, nil
}

func validResponseEnvelope(envelope map[string]json.RawMessage) bool {
        result, rpcError := envelope["result"], envelope["error"]
        if envelope["params"] != nil || (result == nil) == (rpcError == nil) {
                return false
        }
        if rpcError != nil {
                return validMessageError(rpcError)
        }
        return envelope["id"] != nil
}

func validMessageError(raw json.RawMessage) bool {
        var fields map[string]json.RawMessage
        if raw[0] != '{' || json.Unmarshal(raw, &fields) != nil || reservedMemberAlias(fields, "code", "message", "data") {
                return false
        }
        var code int64
        var message string
        return fields["code"] != nil && string(fields["code"]) != "null" && json.Unmarshal(fields["code"], &code) == nil &&
                fields["message"] != nil && string(fields["message"]) != "null" && json.Unmarshal(fields["message"], &message) == nil
}

func reservedMemberAlias(fields map[string]json.RawMessage, reserved ...string) bool {
        for key := range fields {
                for _, name := range reserved {
                        if key != name && strings.EqualFold(key, name) {
                                return true
                        }
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "encoding/json"
        "fmt"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// Middleware type constants
const (
        ParserMiddlewareType         = "mcp-parser"
        ToolFilterMiddlewareType     = "tool-filter"
        ToolCallFilterMiddlewareType = "tool-call-filter"
)

// ParserMiddlewareParams represents the parameters for MCP parser middleware
type ParserMiddlewareParams struct {
        // No parameters needed for MCP parser
}

// ToolOverride represents a tool override entry.
type ToolOverride struct {
        Name        string `json:"name"`
        Description string `json:"description"`
}

// ToolFilterMiddlewareParams represents the parameters for tool filter middleware
type ToolFilterMiddlewareParams struct {
        FilterTools   []string                `json:"filter_tools"`
        ToolsOverride map[string]ToolOverride `json:"tools_override"`
}

// ParserMiddleware wraps MCP parser middleware functionality
type ParserMiddleware struct{}

// Handler returns the middleware function used by the proxy.
func (*ParserMiddleware) Handler() types.MiddlewareFunction {
        return ParsingMiddleware
}

// Close cleans up any resources used by the middleware.
func (*ParserMiddleware) Close() error {
        // MCP parser middleware doesn't need cleanup
        return nil
}

// ToolFilterMiddleware wraps tool filter middleware functionality
type ToolFilterMiddleware struct {
        middleware types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *ToolFilterMiddleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (*ToolFilterMiddleware) Close() error {
        // Tool filter middleware doesn't need cleanup
        return nil
}

// CreateParserMiddleware factory function for MCP parser middleware
func CreateParserMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {

        mcpMw := &ParserMiddleware{}
        runner.AddMiddleware(config.Type, mcpMw)
        return nil
}

// CreateToolFilterMiddleware factory function for tool filter middleware
func CreateToolFilterMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {

        var params ToolFilterMiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal tool filter middleware parameters: %w", err)
        }

        opts := []ToolMiddlewareOption{}
        opts = append(opts, WithToolsFilter(params.FilterTools...))
        for actualName, tool := range params.ToolsOverride {
                opts = append(opts, WithToolsOverride(actualName, tool.Name, tool.Description))
        }

        middleware, err := NewListToolsMappingMiddleware(opts...)
        if err != nil {
                return fmt.Errorf("failed to create tool filter middleware: %w", err)
        }

        toolFilterMw := &ToolFilterMiddleware{middleware: middleware}
        runner.AddMiddleware(config.Type, toolFilterMw)
        return nil
}

// CreateToolCallFilterMiddleware factory function for tool call filter middleware
func CreateToolCallFilterMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {

        var params ToolFilterMiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal tool call filter middleware parameters: %w", err)
        }

        opts := []ToolMiddlewareOption{}
        opts = append(opts, WithToolsFilter(params.FilterTools...))
        for actualName, tool := range params.ToolsOverride {
                opts = append(opts, WithToolsOverride(actualName, tool.Name, tool.Description))
        }

        middleware, err := NewToolCallMappingMiddleware(opts...)
        if err != nil {
                return fmt.Errorf("failed to create tool call filter middleware: %w", err)
        }

        toolCallFilterMw := &ToolFilterMiddleware{middleware: middleware}
        runner.AddMiddleware(config.Type, toolCallFilterMw)
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package mcp provides MCP (Model Context Protocol) parsing utilities and middleware.
package mcp

import (
        "bytes"
        "context"
        "encoding/json"
        "errors"
        "io"
        "log/slog"
        "mime"
        "net/http"
        "strconv"
        "strings"

        "golang.org/x/exp/jsonrpc2"
)

// contextKey is a type for context keys to avoid collisions.
type contextKey string

const (
        // MCPRequestContextKey is the context key for storing parsed MCP request data.
        MCPRequestContextKey contextKey = "mcp_request"
)

// ParsedMCPRequest contains the parsed MCP request information.
type ParsedMCPRequest struct {
        // Method is the MCP method name (e.g., "tools/call", "resources/read")
        Method string
        // ID is the JSON-RPC request ID
        ID interface{}
        // Params contains the raw JSON parameters
        Params json.RawMessage
        // ResourceID is the extracted resource identifier (tool name, resource URI, etc.)
        ResourceID string
        // Arguments contains the extracted arguments for the operation
        Arguments map[string]interface{}
        // Meta contains the _meta field from the request params for protocol-level metadata
        // such as progress tokens, trace IDs, or custom namespaced metadata
        Meta map[string]interface{}
        // MCPMethodHeader is the value of the Modern (stateless MCP) "Mcp-Method"
        // request header, if present. Mandatory on every Modern POST per the draft spec.
        MCPMethodHeader string
        // MCPNameHeader is the raw, as-received value of the Modern (stateless MCP)
        // "Mcp-Name" request header, if present. Required for tools/call,
        // resources/read, prompts/get. Stored undecoded: the spec allows the header
        // value to be sentinel-encoded (=?base64?...?=), and a caller comparing it to
        // the plain body name/uri during validation must decode the header first.
        MCPNameHeader string
        // ClientInfo is the client implementation info surfaced via _meta for Modern
        // (stateless) requests, sourced from _meta["io.modelcontextprotocol/clientInfo"].
        ClientInfo map[string]interface{}
        // ProtocolVersion is the per-request protocol version surfaced via _meta for
        // Modern (stateless) requests, sourced from _meta["io.modelcontextprotocol/protocolVersion"].
        ProtocolVersion string
        // IsRequest indicates if this is a JSON-RPC request (vs response or notification)
        IsRequest bool
        // IsBatch indicates if this is a batch request. It is always false: JSON-RPC
        // batches are rejected by ParsingMiddleware before a ParsedMCPRequest is ever
        // constructed (batching was removed in MCP 2025-06-18; see batch.go). The
        // field is retained for wire/telemetry compatibility.
        IsBatch bool
}

// ParsingMiddleware creates an HTTP middleware that parses MCP JSON-RPC requests
// and stores the parsed information in the request context for use by downstream
// middleware (authorization, audit, etc.).
//
// The middleware:
// 1. Checks if the request should be parsed (every POST in an MCP chain, on any path)
// 2. Reads and parses the JSON-RPC message
// 3. Extracts method, parameters, and resource information
// 4. Stores the parsed data in request context
// 5. Restores the request body for downstream handlers
//
// Example usage:
//
//        middlewares := []types.Middleware{
//            auditMiddleware,       // Audit wraps the chain; parsed data flows back via ParsedRequestHolder
//            authMiddleware,        // Authentication
//            mcp.ParsingMiddleware, // MCP parsing after auth
//            authzMiddleware,       // Authorization uses parsed data
//        }
func ParsingMiddleware(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                // Skip if already parsed by an outer middleware (e.g. auth composes
                // ParsingMiddleware and server.go applies it again for the no-auth case).
                if GetParsedMCPRequest(r.Context()) != nil {
                        next.ServeHTTP(w, r)
                        return
                }

                // Check if we should parse this request
                if !shouldParseMCPRequest(r) {
                        // Only a POST body is parsed as an MCP message. Rather than forward a
                        // body the chain has not inspected, reject any non-POST request that
                        // carries one. A bodyless GET (SSE stream), DELETE (session
                        // termination) or OPTIONS (CORS preflight) passes through unchanged;
                        // the method itself may be legitimate, so this is a 400 for the
                        // unexpected body, not a 405. Method-specific 405/Allow responses are
                        // left to the transport handlers, which know the methods their
                        // endpoint serves. Method names are case-sensitive (RFC 9110 §9.1).
                        if requestCarriesBody(r) {
                                slog.WarnContext(r.Context(), "rejected request body on a non-POST method",
                                        "method", r.Method)
                                http.Error(w, "request body is only accepted on POST", http.StatusBadRequest)
                                return
                        }
                        next.ServeHTTP(w, r)
                        return
                }

                // Read the request body
                bodyBytes, err := io.ReadAll(r.Body)
                if err != nil {
                        writeBodyReadError(w, err)
                        return
                }

                msg, err := DecodeMessage(bodyBytes)
                if err != nil {
                        WriteClassificationError(w, nil, err)
                        return
                }

                // Consumers must see the same JSON accepted by DecodeMessage. In
                // particular, a BOM cannot be embedded in a webhook's json.RawMessage.
                bodyBytes = bytes.TrimPrefix(bodyBytes, UTF8BOM)
                r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
                r.ContentLength = int64(len(bodyBytes))
                if r.GetBody != nil {
                        r.GetBody = func() (io.ReadCloser, error) {
                                return io.NopCloser(bytes.NewReader(bodyBytes)), nil
                        }
                }

                // Responses remain transport-owned and do not populate request context.
                parsedRequest, err := parsedMCPMessage(msg)
                if err != nil {
                        WriteClassificationError(w, nil, err)
                        return
                }
                if parsedRequest != nil {
                        parsedRequest.MCPMethodHeader = r.Header.Get("Mcp-Method")
                        parsedRequest.MCPNameHeader = r.Header.Get("Mcp-Name")
                        // Publish to a ParsedRequestHolder if one is present, so middleware
                        // wrapping the parser (e.g. audit) can observe the parsed request
                        // even though the derived context only flows downstream.
                        if holder, ok := ParsedRequestHolderFromContext(r.Context()); ok {
                                holder.Parsed = parsedRequest
                        }
                        ctx := context.WithValue(r.Context(), MCPRequestContextKey, parsedRequest)
                        r = r.WithContext(ctx)
                }

                // Call the next handler
                next.ServeHTTP(w, r)
        })
}

// RepublishParsedMCPRequest refreshes the cached parse after middleware has
// rewritten the request body, returning the request to pass downstream.
// ParsingMiddleware deliberately parses only once, so any middleware that
// replaces r.Body MUST call this or downstream consumers (authorization,
// audit, telemetry) will decide on the bytes that arrived rather than the
// bytes the backend executes. Replacement bodies must satisfy the same message
// admission checks enforced by ParsingMiddleware and must be requests.
//
// On error, the returned request is nil and must not be passed downstream:
// the caller is responsible for terminating the request (e.g. writing an
// error response) instead of proceeding with a stale or absent parse.
//
// This refreshes downstream context consumers and outer wrappers that inspect
// ParsedRequestHolder after the inner chain returns. It cannot change decisions
// already made before the rewrite: tool-call filtering and rate limiting still
// evaluate the request as received.
//
// The caller must also refresh r.ContentLength when it replaces r.Body, or the
// reverse proxy will reject the forwarded request.
func RepublishParsedMCPRequest(r *http.Request, body []byte) (*http.Request, error) {
        msg, err := DecodeMessage(body)
        if err != nil {
                return nil, err
        }
        parsed, err := parsedMCPMessage(msg)
        if err != nil {
                return nil, err
        }
        if parsed == nil {
                return nil, errors.New("republished body is not a valid JSON-RPC request")
        }
        parsed.MCPMethodHeader = r.Header.Get("Mcp-Method")
        parsed.MCPNameHeader = r.Header.Get("Mcp-Name")

        if holder, ok := ParsedRequestHolderFromContext(r.Context()); ok {
                holder.Parsed = parsed
        }

        return r.WithContext(context.WithValue(r.Context(), MCPRequestContextKey, parsed)), nil
}

// parsedRequestHolderContextKey is the context key for ParsedRequestHolder.
type parsedRequestHolderContextKey struct{}

// ParsedRequestHolder is a mutable carrier that lets middleware running
// OUTSIDE the parsing middleware observe the parsed MCP request. Context
// values only flow downstream, so a wrapper such as the audit middleware
// cannot read the parsed request that the parser attaches for inner handlers.
// The wrapper injects an empty holder via WithParsedRequestHolder before
// calling the inner chain; ParsingMiddleware fills it; the wrapper reads it
// after the inner chain returns.
//
// The holder is written and read by the single request goroutine (writes
// happen-before the wrapper's post-ServeHTTP read), so no synchronization is
// needed.
type ParsedRequestHolder struct {
        Parsed *ParsedMCPRequest
}

// WithParsedRequestHolder returns a new context carrying the given holder.
func WithParsedRequestHolder(ctx context.Context, holder *ParsedRequestHolder) context.Context {
        return context.WithValue(ctx, parsedRequestHolderContextKey{}, holder)
}

// ParsedRequestHolderFromContext retrieves the ParsedRequestHolder from the
// context. Returns (nil, false) if no holder is present.
func ParsedRequestHolderFromContext(ctx context.Context) (*ParsedRequestHolder, bool) {
        holder, ok := ctx.Value(parsedRequestHolderContextKey{}).(*ParsedRequestHolder)
        return holder, ok && holder != nil
}

// authzDenialMarkerContextKey is the context key for AuthzDenialMarker.
type authzDenialMarkerContextKey struct{}

// AuthzDenialMarker is a mutable carrier that lets the authorization
// middleware (pkg/authz), which runs INSIDE the audit middleware, flag a
// request it refused before message-level authorization could run, such as a
// client response or a POST reaching authz without parsing middleware. It follows the same
// propagation pattern as ParsedRequestHolder: the audit wrapper injects an
// empty marker via WithAuthzDenialMarker, the inner middleware fills it, and
// the wrapper reads it back after the inner chain returns so the refusal is
// audited as a denial rather than a generic 400 failure. The type lives in
// this package because pkg/authz already depends on pkg/audit transitively,
// so the carrier cannot live in either of those two packages.
//
// Like ParsedRequestHolder, the marker is written and read by the single
// request goroutine, so no synchronization is needed.
type AuthzDenialMarker struct {
        Denied bool
}

// WithAuthzDenialMarker returns a new context carrying the given marker.
func WithAuthzDenialMarker(ctx context.Context, marker *AuthzDenialMarker) context.Context {
        return context.WithValue(ctx, authzDenialMarkerContextKey{}, marker)
}

// AuthzDenialMarkerFromContext retrieves the AuthzDenialMarker from the
// context. Returns (nil, false) if no marker is present.
func AuthzDenialMarkerFromContext(ctx context.Context) (*AuthzDenialMarker, bool) {
        marker, ok := ctx.Value(authzDenialMarkerContextKey{}).(*AuthzDenialMarker)
        return marker, ok && marker != nil
}

// RequestHasJSONContentType reports whether r declares an application/json
// Content-Type. Media types are case-insensitive per RFC 9110 8.3.1, so the
// header is parsed with mime.ParseMediaType and compared with EqualFold; the
// case-sensitive prefix match this replaces also wrongly accepted longer
// types such as "application/jsonx".
//
// Deprecated: Content-Type must not gate MCP admission. Use DecodeMessage or
// ParsingMiddleware to validate every POST in an MCP chain instead.
func RequestHasJSONContentType(r *http.Request) bool {
        mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
        if err != nil {
                return false
        }
        return strings.EqualFold(mediaType, "application/json")
}

// GetParsedMCPRequest retrieves the parsed MCP request from the request context.
// Returns nil if no parsed request is available.
func GetParsedMCPRequest(ctx context.Context) *ParsedMCPRequest {
        if parsed, ok := ctx.Value(MCPRequestContextKey).(*ParsedMCPRequest); ok {
                return parsed
        }
        return nil
}

// shouldParseMCPRequest determines if the request should be parsed as an MCP request.
//
// Deliberately path-agnostic. There is no fixed MCP endpoint path: the
// streamable transport's path is server-chosen, and in 2024-11-05 HTTP+SSE the
// server names its own message endpoint in the `endpoint` event, which
// ToolHive's proxies forward as given. Any POST may carry an MCP message,
// regardless of Content-Type (not mandated by all supported MCP revisions).
// Path or media-type gates would bypass downstream request-level controls.
// ParsingMiddleware rejects a body sent with any other method instead of
// forwarding it unparsed (see requestCarriesBody).
func shouldParseMCPRequest(r *http.Request) bool {
        return r.Method == http.MethodPost
}

// requestCarriesBody reports whether r has a non-empty body. When the length
// is unknown (chunked) it reads one byte and puts it back, so the body
// forwarded downstream is unchanged. A read error other than EOF counts as a
// body: the request is rejected rather than forwarded unexamined.
func requestCarriesBody(r *http.Request) bool {
        if r.Body == nil || r.Body == http.NoBody || r.ContentLength == 0 {
                return false
        }
        if r.ContentLength > 0 {
                return true
        }
        var first [1]byte
        n, err := io.ReadFull(r.Body, first[:])
        if n > 0 {
                r.Body = struct {
                        io.Reader
                        io.Closer
                }{io.MultiReader(bytes.NewReader(first[:n]), r.Body), r.Body}
                return true
        }
        return !errors.Is(err, io.EOF)
}

// parsedMCPMessage returns nil, nil for responses. Request params that do not
// decode fail closed with a CodedError rather than yielding an empty resource ID.
func parsedMCPMessage(msg jsonrpc2.Message) (*ParsedMCPRequest, error) {
        // Handle only request messages (both calls with ID and notifications without ID)
        req, ok := msg.(*jsonrpc2.Request)
        if !ok {
                // Response or error messages are not parsed here
                return nil, nil
        }

        // Extract resource ID, arguments, and meta based on the method
        resourceID, arguments, meta, err := extractResourceAndArguments(req.Method, req.Params)
        if err != nil {
                // The decode error text can quote the offending literal, so log a fixed reason only.
                slog.Warn("rejected invalid MCP JSON-RPC message", "reason", "undecodable_params")
                return nil, newUndecodableParams(req.ID)
        }
        clientInfo, protocolVersion := extractModernMeta(meta)

        // Determine the ID - will be nil for notifications
        var id interface{}
        if req.ID.IsValid() {
                id = req.ID.Raw()
        }

        return &ParsedMCPRequest{
                Method:          req.Method,
                ID:              id,
                Params:          req.Params,
                ResourceID:      resourceID,
                Arguments:       arguments,
                Meta:            meta,
                ClientInfo:      clientInfo,
                ProtocolVersion: protocolVersion,
                IsRequest:       true,
                // Batches are rejected in ParsingMiddleware before reaching here, so a
                // parsed request is always a single, non-batch message.
                IsBatch: false,
        }, nil
}

// methodHandler defines a function type for handling specific MCP methods
type methodHandler func(map[string]interface{}) (string, map[string]interface{})

// methodHandlers maps MCP methods to their respective handlers
var methodHandlers = map[string]methodHandler{
        "initialize":                         handleInitializeMethod,
        "tools/call":                         handleNamedResourceMethod,
        "prompts/get":                        handleNamedResourceMethod,
        "resources/read":                     handleResourceReadMethod,
        "skills/get":                         handleResourceReadMethod,
        "resources/list":                     handleListMethod,
        "skills/list":                        handleListMethod,
        "tools/list":                         handleListMethod,
        "prompts/list":                       handleListMethod,
        "notifications/message":              handleNotificationMethod,
        "logging/setLevel":                   handleLoggingMethod,
        "completion/complete":                handleCompletionMethod,
        "elicitation/create":                 handleElicitationMethod,
        "sampling/createMessage":             handleSamplingMethod,
        "resources/subscribe":                handleResourceSubscribeMethod,
        "resources/unsubscribe":              handleResourceUnsubscribeMethod,
        "resources/templates/list":           handleListMethod,
        "roots/list":                         handleListMethod,
        "notifications/progress":             handleProgressNotificationMethod,
        "notifications/cancelled":            handleCancelledNotificationMethod,
        "tasks/list":                         handleListMethod,
        "tasks/get":                          handleTaskIDMethod,
        "tasks/cancel":                       handleTaskIDMethod,
        "tasks/result":                       handleTaskIDMethod,
        "notifications/tasks/status":         handleTaskStatusNotificationMethod,
        "notifications/elicitation/complete": handleElicitationCompleteNotificationMethod,
}

// staticResourceIDs maps methods to their static resource IDs
var staticResourceIDs = map[string]string{
        "ping":                                 "ping",
        "notifications/roots/list_changed":     "roots",
        "notifications/initialized":            "initialized",
        "notifications/prompts/list_changed":   "prompts",
        "notifications/resources/list_changed": "resources",
        "notifications/resources/updated":      "resources",
        "notifications/tools/list_changed":     "tools",
        "server/discover":                      "discover",
}

// extractResourceAndArguments extracts the resource ID, arguments, and _meta field from the JSON-RPC params
// based on the MCP method type. It returns an error if params do not decode as an object.
func extractResourceAndArguments(
        method string, params json.RawMessage,
) (string, map[string]interface{}, map[string]interface{}, error) {
        if params == nil {
                return getStaticResourceID(method), nil, nil, nil
        }

        paramsMap, err := decodeParams(params)
        if err != nil {
                return "", nil, nil, err
        }

        meta := metaFromParamsMap(paramsMap)

        resourceID, arguments := processMethodWithHandler(method, paramsMap)
        return resourceID, arguments, meta, nil
}

// extractModernMeta surfaces the Modern (stateless MCP) clientInfo and
// protocolVersion fields from a parsed _meta map, if present. It delegates to
// the reserved-key helpers in revision.go so the guarded type assertions live
// in one place; a wrong-shaped value is treated as absent rather than causing
// an error.
func extractModernMeta(meta map[string]interface{}) (clientInfo map[string]interface{}, protocolVersion string) {
        clientInfo, _ = objectMetaValue(meta, metaKeyClientInfo)
        protocolVersion, _ = stringMetaValue(meta, metaKeyProtocolVersion)
        return clientInfo, protocolVersion
}

// getStaticResourceID returns the static resource ID for methods that don't need parameter parsing
func getStaticResourceID(method string) string {
        if resourceID, exists := staticResourceIDs[method]; exists {
                return resourceID
        }
        return ""
}

// processMethodWithHandler processes the method using the appropriate handler
func processMethodWithHandler(method string, paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if handler, exists := methodHandlers[method]; exists {
                return handler(paramsMap)
        }
        return getStaticResourceID(method), nil
}

// handleInitializeMethod extracts resource ID and arguments for initialize method
func handleInitializeMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        var resourceID string
        if clientInfo, ok := paramsMap["clientInfo"].(map[string]interface{}); ok {
                if name, ok := clientInfo["name"].(string); ok {
                        resourceID = name
                }
        }
        return resourceID, paramsMap
}

// handleNamedResourceMethod extracts resource ID and arguments for methods with name parameter
func handleNamedResourceMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        var resourceID string
        var arguments map[string]interface{}

        if name, ok := paramsMap["name"].(string); ok {
                resourceID = name
        }
        if args, ok := paramsMap["arguments"].(map[string]interface{}); ok {
                arguments = args
        }

        return resourceID, arguments
}

// handleResourceReadMethod extracts resource ID for resource read operations
func handleResourceReadMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if uri, ok := paramsMap["uri"].(string); ok {
                return uri, nil
        }
        return "", nil
}

// handleListMethod extracts resource ID for list operations
func handleListMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if cursor, ok := paramsMap["cursor"].(string); ok && cursor != "" {
                return cursor, nil
        }
        return "", nil
}

// handleNotificationMethod extracts resource ID for notification messages
func handleNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if notifMethod, ok := paramsMap["method"].(string); ok {
                return notifMethod, nil
        }
        return "", nil
}

// handleLoggingMethod extracts resource ID for logging operations
func handleLoggingMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if level, ok := paramsMap["level"].(string); ok {
                return level, nil
        }
        return "", nil
}

// handleCompletionMethod extracts resource ID for completion requests.
// For PromptReference: extracts the prompt name
// For ResourceTemplateReference: extracts the template URI
// For legacy string ref: returns the string value
// Always returns paramsMap as arguments since completion requests need the full context
// including the argument being completed and any context from previous completions.
func handleCompletionMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        // Check if ref is a map (PromptReference or ResourceTemplateReference)
        if ref, ok := paramsMap["ref"].(map[string]interface{}); ok {
                // Try to extract name for PromptReference
                if name, ok := ref["name"].(string); ok {
                        return name, paramsMap
                }
                // Try to extract uri for ResourceTemplateReference
                if uri, ok := ref["uri"].(string); ok {
                        return uri, paramsMap
                }
        }
        // Fallback to string ref (legacy support)
        if ref, ok := paramsMap["ref"].(string); ok {
                return ref, paramsMap
        }
        return "", paramsMap
}

// handleElicitationMethod extracts resource ID for elicitation requests
func handleElicitationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        // The message field could be used as a resource identifier
        if message, ok := paramsMap["message"].(string); ok {
                return message, paramsMap
        }
        return "", paramsMap
}

// handleElicitationCompleteNotificationMethod extracts resource ID for elicitation complete notifications.
// This notification is sent by the server when an out-of-band URL-mode elicitation is completed.
// Returns the elicitationId as the resource identifier.
func handleElicitationCompleteNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if elicitationId, ok := paramsMap["elicitationId"].(string); ok {
                return elicitationId, paramsMap
        }
        return "", paramsMap
}

// handleSamplingMethod extracts resource ID for sampling/createMessage requests.
// Returns the model name from modelPreferences if available, otherwise returns a
// truncated version of the systemPrompt. The 50-character truncation provides a
// reasonable balance between uniqueness and readability for authorization and audit logs.
func handleSamplingMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        // Use model preferences or system prompt as identifier if available
        if modelPrefs, ok := paramsMap["modelPreferences"].(map[string]interface{}); ok && modelPrefs != nil {
                // Try direct name field first (simplified structure)
                if name, ok := modelPrefs["name"].(string); ok && name != "" {
                        return name, paramsMap
                }
                // Try to get model name from hints array (full spec structure)
                if hints, ok := modelPrefs["hints"].([]interface{}); ok && len(hints) > 0 {
                        if hint, ok := hints[0].(map[string]interface{}); ok {
                                if name, ok := hint["name"].(string); ok && name != "" {
                                        return name, paramsMap
                                }
                        }
                }
        }
        if systemPrompt, ok := paramsMap["systemPrompt"].(string); ok && systemPrompt != "" {
                // Use first 50 chars of system prompt as identifier
                // This provides a reasonable balance between uniqueness and readability
                if len(systemPrompt) > 50 {
                        return systemPrompt[:50], paramsMap
                }
                return systemPrompt, paramsMap
        }
        return "", paramsMap
}

// handleResourceSubscribeMethod extracts resource ID for resource subscribe operations
func handleResourceSubscribeMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if uri, ok := paramsMap["uri"].(string); ok {
                return uri, nil
        }
        return "", nil
}

// handleResourceUnsubscribeMethod extracts resource ID for resource unsubscribe operations
func handleResourceUnsubscribeMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if uri, ok := paramsMap["uri"].(string); ok {
                return uri, nil
        }
        return "", nil
}

// handleProgressNotificationMethod extracts resource ID for progress notifications.
// Extracts the progressToken which can be either a string or numeric value.
func handleProgressNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if token, ok := paramsMap["progressToken"].(string); ok {
                return token, paramsMap
        }
        // Also handle numeric progress tokens
        if token, ok := paramsMap["progressToken"].(float64); ok {
                return strconv.FormatFloat(token, 'f', 0, 64), paramsMap
        }
        return "", paramsMap
}

// handleCancelledNotificationMethod extracts resource ID for cancelled notifications.
// Extracts the requestId which can be either a string or numeric value.
func handleCancelledNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        // Extract request ID as the resource identifier
        if requestId, ok := paramsMap["requestId"].(string); ok {
                return requestId, paramsMap
        }
        // Handle numeric request IDs
        if requestId, ok := paramsMap["requestId"].(float64); ok {
                return strconv.FormatFloat(requestId, 'f', 0, 64), paramsMap
        }
        return "", paramsMap
}

// handleTaskIDMethod extracts resource ID for task operations (tasks/get, tasks/cancel, tasks/result).
// Returns the taskId parameter as the resource identifier, or empty string if not present.
// Handles both string and numeric taskId values.
func handleTaskIDMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if taskId, ok := paramsMap["taskId"].(string); ok {
                return taskId, nil
        }
        // Handle numeric task IDs
        if taskId, ok := paramsMap["taskId"].(float64); ok {
                return strconv.FormatFloat(taskId, 'f', 0, 64), nil
        }
        return "", nil
}

// handleTaskStatusNotificationMethod extracts resource ID for task status notifications.
// Returns the taskId parameter as the resource identifier while preserving all notification parameters.
// Handles both string and numeric taskId values.
func handleTaskStatusNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
        if taskId, ok := paramsMap["taskId"].(string); ok {
                return taskId, paramsMap
        }
        // Handle numeric task IDs
        if taskId, ok := paramsMap["taskId"].(float64); ok {
                return strconv.FormatFloat(taskId, 'f', 0, 64), paramsMap
        }
        return "", paramsMap
}

// GetMCPMethod is a convenience function to get the MCP method from the context.
func GetMCPMethod(ctx context.Context) string {
        if parsed := GetParsedMCPRequest(ctx); parsed != nil {
                return parsed.Method
        }
        return ""
}

// GetMCPResourceID is a convenience function to get the MCP resource ID from the context.
func GetMCPResourceID(ctx context.Context) string {
        if parsed := GetParsedMCPRequest(ctx); parsed != nil {
                return parsed.ResourceID
        }
        return ""
}

// GetMCPArguments is a convenience function to get the MCP arguments from the context.
func GetMCPArguments(ctx context.Context) map[string]interface{} {
        if parsed := GetParsedMCPRequest(ctx); parsed != nil {
                return parsed.Arguments
        }
        return nil
}

// GetMCPMeta is a convenience function to get the MCP _meta field from the context.
// Returns nil if no parsed request is available or if _meta field is not present.
func GetMCPMeta(ctx context.Context) map[string]interface{} {
        if parsed := GetParsedMCPRequest(ctx); parsed != nil {
                return parsed.Meta
        }
        return nil
}

// GetMCPClientInfo is a convenience function to get the Modern (stateless MCP)
// per-request clientInfo from the context.
// Returns nil if no parsed request is available or clientInfo is not present.
func GetMCPClientInfo(ctx context.Context) map[string]interface{} {
        if parsed := GetParsedMCPRequest(ctx); parsed != nil {
                return parsed.ClientInfo
        }
        return nil
}

// GetMCPProtocolVersion is a convenience function to get the Modern (stateless
// MCP) per-request protocol version from the context.
// Returns "" if no parsed request is available or protocolVersion is not present.
func GetMCPProtocolVersion(ctx context.Context) string {
        if parsed := GetParsedMCPRequest(ctx); parsed != nil {
                return parsed.ProtocolVersion
        }
        return ""
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "encoding/json"
)

// ParsedMCPResponse contains the result of inspecting a JSON-RPC response
// body for application-level errors. Only the error-related fields are
// extracted; the full result payload is intentionally not captured to avoid
// duplicating the privacy-sensitive IncludeResponseData path.
type ParsedMCPResponse struct {
        // HasError is true when the response body contains a top-level "error" field.
        HasError bool
        // ErrorCode is the JSON-RPC error code (e.g., -32603 for internal error).
        ErrorCode int
        // ErrorMessage is the raw error message from the JSON-RPC response.
        ErrorMessage string
}

// jsonRPCError mirrors the JSON-RPC 2.0 error object for decoding purposes.
// We use a minimal custom struct rather than jsonrpc2.DecodeMessage because
// the library's wireError type is unexported, making it impossible to extract
// the numeric error code from outside the package.
type jsonRPCError struct {
        Code    int    `json:"code"`
        Message string `json:"message"`
}

// jsonRPCResponseEnvelope is the minimal structure needed to detect a
// JSON-RPC error in a response body. We intentionally omit "result" to
// keep the parse lightweight.
type jsonRPCResponseEnvelope struct {
        Error *jsonRPCError `json:"error,omitempty"`
}

// ParseMCPResponse inspects a response body and returns a ParsedMCPResponse
// indicating whether it contains a JSON-RPC error. The function is
// intentionally lenient: if the body is not valid JSON or does not contain
// an "error" field, it returns HasError=false rather than an error.
func ParseMCPResponse(body []byte) *ParsedMCPResponse {
        if len(body) == 0 {
                return &ParsedMCPResponse{}
        }

        var envelope jsonRPCResponseEnvelope
        if err := json.Unmarshal(body, &envelope); err != nil {
                return &ParsedMCPResponse{}
        }

        if envelope.Error == nil {
                return &ParsedMCPResponse{}
        }

        return &ParsedMCPResponse{
                HasError:     true,
                ErrorCode:    envelope.Error.Code,
                ErrorMessage: envelope.Error.Message,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "encoding/base64"
        "encoding/json"
        "fmt"
        "maps"
        "strings"
)

// Revision identifies which MCP protocol era a request belongs to.
type Revision int

const (
        // RevisionLegacy is the current 2025-11-25 MCP revision: session-based,
        // initialize handshake, Mcp-Session-Id.
        RevisionLegacy Revision = iota
        // RevisionModern is the 2026-07-28 MCP revision: stateless, no initialize,
        // protocol metadata carried per-request in _meta.
        RevisionModern
)

// MCPVersionModern is the single Modern (stateless) protocol version this
// build understands.
const MCPVersionModern = "2026-07-28"

// MCPVersionLegacy is the single Legacy (session-based) protocol version this
// build understands. It is what RevisionLegacy names on the wire and the one
// version mcpcompat's initialize handshake negotiates.
const MCPVersionLegacy = "2025-11-25"

// String returns the wire protocol-version string for the revision:
// MCPVersionModern for RevisionModern, MCPVersionLegacy otherwise.
func (r Revision) String() string {
        if r == RevisionModern {
                return MCPVersionModern
        }
        return MCPVersionLegacy
}

// metaKeyProtocolVersion is the reserved _meta key that carries the per-request
// protocol version on Modern (stateless) MCP requests, per the draft MCP schema's
// RequestMetaObject.
const metaKeyProtocolVersion = "io.modelcontextprotocol/protocolVersion"

// metaKeyClientInfo is the reserved _meta key that carries the per-request client
// implementation info on Modern (stateless) MCP requests, per the draft MCP
// schema's RequestMetaObject.
const metaKeyClientInfo = "io.modelcontextprotocol/clientInfo"

// metaKeyClientCapabilities is the reserved _meta key that carries the per-request
// client capabilities on Modern (stateless) MCP requests, per the draft MCP
// schema's RequestMetaObject.
const metaKeyClientCapabilities = "io.modelcontextprotocol/clientCapabilities"

// MetaKeyLogLevel carries the per-request minimum log level on Modern
// (2026-07-28) requests (draft schema RequestMetaObject; go-sdk protocol.go).
// It is a reserved per-hop key that must be stripped before a Legacy backend
// hop, but — unlike the other reserved keys — its mere presence is NOT a claim
// of the Modern revision: go-sdk's validateRequestMeta gates Modern-ness purely
// on protocolVersion, and SEP-2577 already deprecates logLevel. It therefore
// belongs in the strip set (ReservedMetaPrefix, minus passthroughMetaKeys) but
// not the ingress signal set (modernSignalMetaKeys).
//
// Exported so the vMCP Modern client can overlay it onto the request _meta it
// mints (the Modern replacement for the removed logging/setLevel RPC).
const MetaKeyLogLevel = "io.modelcontextprotocol/logLevel"

// ReservedMetaPrefix is the _meta key namespace the MCP spec reserves for the
// protocol's own use. StripReservedMeta removes every key carrying it except
// the passthroughMetaKeys below.
//
// The trailing slash is load-bearing. The registry's
// "io.modelcontextprotocol.registry/publisher-provided" (docs/registry/schema.md)
// is ALSO spec-reserved — per the 2025-11-25 _meta key-name rules any prefix
// whose second label is "modelcontextprotocol" or "mcp" is reserved — but it is
// server-descriptor provenance, not per-hop control, and never rides a
// tools/call result. The slash is what keeps this predicate off it.
//
// Matching one literal prefix deliberately does not implement that full
// second-label rule: "dev.mcp/", "org.modelcontextprotocol.api/" and friends are
// equally reserved and sail through. Nothing in this repo emits them and no
// threat model calls for parsing the general rule.
const ReservedMetaPrefix = "io.modelcontextprotocol/"

// passthroughMetaKeys are the reserved keys StripReservedMeta must NOT remove.
//
// The reserved namespace holds two different kinds of key, and only one of them
// is a gateway's to terminate:
//
//   - per-hop protocol control (protocolVersion, clientInfo, clientCapabilities,
//     logLevel, serverInfo, subscriptionId) — scoped to a single MCP connection.
//     vMCP is its client's peer and its backends' peer on two different hops, so
//     it must strip these and mint its own.
//   - end-to-end semantic payload — meaningful to the ORIGINAL endpoints and
//     merely carried by intermediaries. Stripping these breaks the feature.
//
// Both entries below are the second kind, from the 2025-11-25 task facility:
// "related-task" is a MUST on every task-related request AND response
// (tasks/result in particular carries the task id nowhere else), and
// "model-immediate-response" rides CreateTaskResult._meta. vMCP does not route
// tasks/* today, so these are inert — but this strip is what would silently
// break them later, and silent broken correlation is expensive to debug.
//
// Extend this by category, not by string: ask whether the key names something
// about THIS hop (strip) or about the two endpoints (pass through).
// A set, not a lookup table: struct{} values make membership the only question
// this map can answer. A map[string]bool would let a future "key": false entry
// read as "in the set" while silently stripping.
var passthroughMetaKeys = map[string]struct{}{
        "io.modelcontextprotocol/related-task":             {},
        "io.modelcontextprotocol/model-immediate-response": {},
}

// modernSignalMetaKeys is the INGRESS/detection set consumed by hasModernSignal:
// the reserved keys whose presence — independent of whether the value is
// well-formed — is a claim of the Modern revision when nothing else contradicts
// it, and must not then be silently downgraded to Legacy.
//
// "Nothing else contradicts it" is load-bearing. Real Legacy clients DO set
// these keys (the ChatGPT connector sets them while negotiating 2025-11-25),
// so an explicit non-Modern MCP-Protocol-Version header outranks them — see
// ClassifyRevision. The claim stands where the key is the only signal present,
// and there a malformed/absent protocolVersion alongside one of them turns into
// a rejection, never a downgrade.
//
// It is deliberately narrower than what StripReservedMeta removes: logLevel is
// excluded so a request carrying only logLevel — which go-sdk's
// validateRequestMeta accepts (gating purely on protocolVersion) and which
// SEP-2577 deprecates — is classified Legacy rather than misdetected as Modern
// and then rejected. "Strip this key on the way out" and "this key means the
// request is Modern" are separate decisions; conflating them would make a
// request rejectable merely for carrying a strippable-but-not-signalling key.
var modernSignalMetaKeys = []string{
        metaKeyProtocolVersion,
        metaKeyClientInfo,
        metaKeyClientCapabilities,
}

// StripReservedMeta returns a copy of meta with every ReservedMetaPrefix key
// removed except the passthroughMetaKeys, leaving all other caller-supplied keys
// (progressToken, trace context, backend-custom fields) untouched. The input is
// never mutated (maps.Clone).
//
// Use it on BOTH hops, in both directions, wherever a _meta map crosses vMCP:
//
//   - Legacy backend egress (request): a downstream Modern caller's
//     _meta.protocolVersion is only valid on a stateless Modern hop. go-sdk v1.7
//     rejects the request outright (HTTP 400: "protocol version ... is only
//     supported on stateless HTTP servers") for ANY _meta.protocolVersion on a
//     stateful streamable-HTTP server, regardless of its value.
//   - Modern backend egress (request): vMCP overlays its own authoritative
//     values afterwards — see ModernRequestMeta / mergeModernMeta.
//   - client egress (response, and server->client requests): a backend must not
//     speak for vMCP. Only serverInfo is even schema-legal on a result
//     (ResultMetaObject); the request-only keys arriving on one are a backend
//     fabricating the client's own identity. The spec says nothing about what a
//     gateway should forward, so this is derived rather than quoted: serverInfo
//     identifies "the server software producing the response", vMCP is that
//     software, therefore vMCP's value is the correct one.
//
// TRIPWIRE: if vMCP ever relays resource subscriptions it must MINT its own
// io.modelcontextprotocol/subscriptionId from its own downstream listen request
// — never forward a backend's, which names a request id on the vMCP<->backend
// connection and collides across backends. That belongs in the notification
// relay (pkg/vmcp/client/forwarding.go), not here; this helper only sees
// requests and results. See pkg/transport/proxy/streamable/dispatcher_streams.go
// for the same concern in the streamable proxy.
//
// Returns nil whenever the result would be empty -- for nil or empty input, and
// also for a map whose keys were ALL reserved. Callers can therefore treat nil
// as "no _meta to send" without a second length check (mergeModernMeta's
// caller-tolerant convention, and what lets conversion.ToMCPMeta omit _meta
// entirely rather than emit an empty object). A map with keys left over is
// returned as a copy (maps.Clone), never the original.
func StripReservedMeta(meta map[string]any) map[string]any {
        if len(meta) == 0 {
                return nil
        }
        stripped := maps.Clone(meta)
        maps.DeleteFunc(stripped, func(k string, _ any) bool {
                _, passthrough := passthroughMetaKeys[k]
                return strings.HasPrefix(k, ReservedMetaPrefix) && !passthrough
        })
        if len(stripped) == 0 {
                return nil
        }
        return stripped
}

// ModernRequestMeta builds the reserved _meta object every Modern (2026-07-28)
// request must carry: protocolVersion, clientInfo, and (empty) clientCapabilities.
// It is the single source of truth for the client side of the reserved
// io.modelcontextprotocol/* keys, kept consistent with what ClassifyRevision and
// ValidateHeaderConsistency require of the server side — protocolVersion equal to
// MCPVersionModern and clientCapabilities present as a JSON object.
func ModernRequestMeta(clientName, clientVersion string) map[string]any {
        return map[string]any{
                metaKeyProtocolVersion:    MCPVersionModern,
                metaKeyClientInfo:         map[string]any{"name": clientName, "version": clientVersion},
                metaKeyClientCapabilities: map[string]any{},
        }
}

// The following JSON-RPC error codes are defined by the draft MCP spec
// (schema/draft/schema.ts) for the stateless "Modern" revision, except for the
// last three, which are standard JSON-RPC 2.0 codes rather than draft-spec-
// specific ones. They are declared as local literals, matching this repo's
// existing convention for JSON-RPC codes (e.g. streamable_proxy.go's -32603).
// The only SDK with equivalents, modelcontextprotocol/go-sdk, is reachable
// solely as a transitive dependency (via toolhive-core's mcpcompat) and is
// pinned to an older draft snapshot (2026-06-30, CodeHeaderMismatch = -32001)
// with no equivalents at all for the other two codes below, so importing it
// would risk wiring in stale values. Revisit once the go-sdk dependency is
// bumped to a revision that matches MCPVersionModern.
//
// These are exported so other packages (e.g. the HTTP layer) can reference
// the same wire values instead of hardcoding or redeclaring them.
const (
        // CodeHeaderMismatch signals a mismatch between the MCP-Protocol-Version
        // header and the _meta protocol version (schema.ts HeaderMismatchError).
        // Also reused by RequestHeaderMismatchError for the Mcp-Method/Mcp-Name
        // headers: keep both error types' Code() in sync with this constant.
        CodeHeaderMismatch int64 = -32020
        // CodeMissingClientCapability signals that _meta is missing a client
        // capability required for the request (schema.ts MissingRequiredClientCapabilityError).
        CodeMissingClientCapability int64 = -32021
        // CodeUnsupportedProtocolVersion signals that the _meta protocol version
        // is not one this server supports (schema.ts UnsupportedProtocolVersionError).
        CodeUnsupportedProtocolVersion int64 = -32022
        // CodeInvalidParams is the standard JSON-RPC Invalid Params code, used
        // as a fallback when the draft spec defines no dedicated code for the failure.
        CodeInvalidParams int64 = -32602
        // CodeInvalidRequest is the standard JSON-RPC Invalid Request code. ToolHive
        // uses it to reject a message whose shape is not a valid single request for
        // the negotiated protocol version — currently a JSON-RPC batch, which was
        // removed from MCP in the 2025-06-18 revision (see batch.go) — and also for
        // a request body over the configured size limit (HTTP 413, see
        // JSONRPCCodeForStatus).
        CodeInvalidRequest int64 = -32600
        // CodeInternalError is the standard JSON-RPC Internal Error code, used for
        // server-side failures that are not attributable to the request's shape.
        CodeInternalError int64 = -32603
)

// HeaderMismatchError indicates the MCP-Protocol-Version HTTP header did not match
// the io.modelcontextprotocol/protocolVersion carried in the request's _meta field.
// Neither value wins: this is a hard rejection of the request, not a signal to
// proceed with either value.
type HeaderMismatchError struct {
        // Header is the MCP-Protocol-Version header value.
        Header string
        // Body is the _meta protocol version value.
        Body string
}

func (e *HeaderMismatchError) Error() string {
        return fmt.Sprintf("MCP-Protocol-Version header %q does not match _meta protocol version %q", e.Header, e.Body)
}

// Code implements CodedError.
func (*HeaderMismatchError) Code() int64 { return CodeHeaderMismatch }

// Data implements CodedError.
func (e *HeaderMismatchError) Data() map[string]any {
        return map[string]any{"header": e.Header, "body": e.Body}
}

// requestHeaderMismatchReason distinguishes why a Mcp-Method/Mcp-Name header
// check failed, so Error() can describe the actual problem instead of always
// reporting a value "mismatch" — the wire code (CodeHeaderMismatch, -32020)
// is the same in every case; only the human-readable message differs.
type requestHeaderMismatchReason int

const (
        // headerMismatchReasonValue is the default: the header and body both
        // carry a value, but they disagree.
        headerMismatchReasonValue requestHeaderMismatchReason = iota
        // headerMismatchReasonMissing means the header was required but absent.
        headerMismatchReasonMissing
        // headerMismatchReasonMalformed means the header value itself could not
        // be decoded (e.g. invalid base64 in a Mcp-Name sentinel).
        headerMismatchReasonMalformed
)

// RequestHeaderMismatchError indicates a Modern (2026-07-28) request header
// contradicted, was missing from, or was malformed relative to the value
// carried in the request body. This is the generalized sibling of
// HeaderMismatchError for the Mcp-Method/Mcp-Name headers: HeaderMismatchError
// is specific to MCP-Protocol-Version, while this type covers any other
// header/body consistency check and names which header failed.
type RequestHeaderMismatchError struct {
        // Header is the name of the header that failed the check (e.g. "Mcp-Method", "Mcp-Name").
        Header string
        // HeaderValue is the (decoded, where applicable) value carried by the header.
        HeaderValue string
        // BodyValue is the value the header was compared against in the request body.
        BodyValue string

        // reason is unexported: it only shapes Error()'s message, not Data()'s
        // three-key shape or Code()'s wire value.
        reason requestHeaderMismatchReason
}

func (e *RequestHeaderMismatchError) Error() string {
        switch e.reason {
        case headerMismatchReasonValue:
                return fmt.Sprintf("%s header %q does not match request body value %q", e.Header, e.HeaderValue, e.BodyValue)
        case headerMismatchReasonMissing:
                return fmt.Sprintf("%s header is missing (required for this request)", e.Header)
        case headerMismatchReasonMalformed:
                return fmt.Sprintf("%s header value %q is malformed", e.Header, e.HeaderValue)
        default:
                return fmt.Sprintf("%s header %q does not match request body value %q", e.Header, e.HeaderValue, e.BodyValue)
        }
}

// Code implements CodedError.
func (*RequestHeaderMismatchError) Code() int64 { return CodeHeaderMismatch }

// Data implements CodedError.
func (e *RequestHeaderMismatchError) Data() map[string]any {
        return map[string]any{"header": e.Header, "headerValue": e.HeaderValue, "bodyValue": e.BodyValue}
}

// UnsupportedVersionError indicates the _meta protocol version named a Modern
// revision this server does not support.
type UnsupportedVersionError struct {
        // Requested is the _meta protocol version the client asked for.
        Requested string
        // Supported lists the protocol versions this server supports.
        Supported []string
}

func (e *UnsupportedVersionError) Error() string {
        return fmt.Sprintf("unsupported MCP protocol version %q (supported: %v)", e.Requested, e.Supported)
}

// Code implements CodedError.
func (*UnsupportedVersionError) Code() int64 { return CodeUnsupportedProtocolVersion }

// Data implements CodedError.
func (e *UnsupportedVersionError) Data() map[string]any {
        return map[string]any{"supported": e.Supported, "requested": e.Requested}
}

// MissingClientCapabilityError indicates a Modern request's _meta is missing
// clientCapabilities (clientInfo is optional per the draft schema and is not
// checked here).
//
// The draft types MissingRequiredClientCapabilityError.data.requiredCapabilities
// as a ClientCapabilities object, not a list of names. The classifier cannot
// compute per-method required capabilities (that check is deferred to the
// caller/handler layer), so RequiredCapabilities is populated best-effort and
// may be an empty object.
type MissingClientCapabilityError struct {
        // RequiredCapabilities is the ClientCapabilities object the request was
        // missing, if known.
        RequiredCapabilities map[string]any
}

func (e *MissingClientCapabilityError) Error() string {
        return fmt.Sprintf("request _meta is missing required client capabilities: %v", e.RequiredCapabilities)
}

// Code implements CodedError.
func (*MissingClientCapabilityError) Code() int64 { return CodeMissingClientCapability }

// Data implements CodedError.
func (e *MissingClientCapabilityError) Data() map[string]any {
        return map[string]any{"requiredCapabilities": e.RequiredCapabilities}
}

// MissingModernMetadataError indicates the request carried a Modern signal
// via one of the reserved io.modelcontextprotocol/* _meta keys — with no
// MCP-Protocol-Version header present at all — but _meta carried no valid
// protocolVersion. (When a header IS present, a missing or invalid body
// protocolVersion is a header/body mismatch instead; see HeaderMismatchError.)
// The draft spec defines no dedicated error code for this reserved-key-only
// case, so it falls back to the standard JSON-RPC Invalid Params code.
type MissingModernMetadataError struct{}

func (*MissingModernMetadataError) Error() string {
        return "request carries a Modern signal but no valid io.modelcontextprotocol/protocolVersion in _meta"
}

// Code implements CodedError.
func (*MissingModernMetadataError) Code() int64 { return CodeInvalidParams }

// Data implements CodedError.
func (*MissingModernMetadataError) Data() map[string]any { return map[string]any{} }

// ClassifyRevision determines whether a single MCP request is Legacy or Modern,
// and whether it is valid.
//
// When err != nil, the caller MUST reject the request; the returned Revision is
// then purely informational — it records that the request claimed Modern, not
// that classification succeeded.
//
// method == "initialize" always classifies Legacy immediately, unconditionally —
// Modern never sends initialize (it is the Legacy session-start marker by
// definition) — which also guards against a spoofed Modern _meta on a Legacy
// call.
//
// Otherwise, a request "signals" Modern if the MCP-Protocol-Version header is
// exactly MCPVersionModern, OR _meta carries ANY of the reserved
// io.modelcontextprotocol/* keys (protocolVersion, clientInfo,
// clientCapabilities) — presence of the key is the signal, independent of
// whether its value is well-formed, since a Legacy client never sets these
// keys at all. A request with no signal anywhere classifies Legacy, the safe
// default. A request with a signal is never silently downgraded: it either
// classifies Modern with a nil error, or Modern with an error the caller must
// reject on.
//
// A non-empty MCP-Protocol-Version header that names some OTHER version (not
// MCPVersionModern) and carries no reserved _meta key is, by design, not a
// Modern signal: the request body's _meta is authoritative for the protocol
// version, and an unrecognized header value alone classifies Legacy rather
// than erroring.
//
// Given a Modern signal, checks run in this order:
//  1. meta[metaKeyProtocolVersion] must be a non-empty string. If it is not,
//     and protoHeader is non-empty, the body has nothing valid to match against
//     the header: this is a header/body mismatch (*HeaderMismatchError). If
//     protoHeader is empty — the signal came only from a reserved _meta key,
//     so there is no header to mismatch against — the request is malformed
//     (*MissingModernMetadataError).
//  2. that string must equal MCPVersionModern, or it names an unsupported
//     version (*UnsupportedVersionError).
//  3. if protoHeader is non-empty it must equal the body version, or the two
//     conflict (*HeaderMismatchError, a hard rejection — neither value wins).
//  4. _meta must carry clientCapabilities (clientInfo is optional per the
//     draft schema), or the client capabilities are missing
//     (*MissingClientCapabilityError, with per-method requirements deferred to
//     the caller).
//
// A request that passes all four classifies Modern with a nil error.
func ClassifyRevision(method string, meta map[string]any, protoHeader string) (Revision, error) {
        if method == "initialize" {
                return RevisionLegacy, nil
        }

        if !hasModernSignal(meta, protoHeader) {
                return RevisionLegacy, nil
        }

        bodyVersion, hasBodyVersion := stringMetaValue(meta, metaKeyProtocolVersion)
        if !hasBodyVersion {
                if protoHeader != "" && protoHeader != MCPVersionModern {
                        // An explicit non-Modern header is the client's negotiated
                        // declaration and outranks a stray reserved _meta key that carries
                        // no version of its own. Without this, merely including e.g.
                        // clientInfo flipped an otherwise-fine Legacy request from accepted
                        // to -32020 -- the same request with no reserved key at all is
                        // classified Legacy by hasModernSignal above, so rejecting this one
                        // was incoherent rather than strict. It broke every tools/call from
                        // the ChatGPT connector, which negotiates 2025-11-25 and sets
                        // reserved keys without a _meta protocolVersion (#6188).
                        //
                        // Modern enforcement is untouched: a Modern header with no _meta
                        // version still falls through to the mismatch below, and a reserved
                        // key with no header at all still yields MissingModernMetadata --
                        // which is where "a reserved key is a claim of Modern that must not
                        // be silently downgraded" (see modernSignalMetaKeys) actually
                        // applies, because there the key is the only signal present.
                        return RevisionLegacy, nil
                }
                if protoHeader != "" {
                        return RevisionModern, &HeaderMismatchError{Header: protoHeader, Body: ""}
                }
                // TODO: this always returns -32602, but ClassifyRevision is transport-agnostic
                // and cannot tell "stdio, no header concept" (where -32602 is correct) apart from
                // "HTTP, header omitted" (where the draft's Server Validation rules make a missing
                // MCP-Protocol-Version header itself a -32020 HeaderMismatch condition). Revisit
                // once the classifier is wired into the HTTP request path and can be given
                // transport context.
                return RevisionModern, &MissingModernMetadataError{}
        }

        if bodyVersion != MCPVersionModern {
                return RevisionModern, &UnsupportedVersionError{Requested: bodyVersion, Supported: []string{MCPVersionModern}}
        }

        if protoHeader != "" && protoHeader != bodyVersion {
                return RevisionModern, &HeaderMismatchError{Header: protoHeader, Body: bodyVersion}
        }

        if !hasObjectMetaValue(meta, metaKeyClientCapabilities) {
                return RevisionModern, &MissingClientCapabilityError{RequiredCapabilities: map[string]any{}}
        }

        return RevisionModern, nil
}

// nameRequiredMethods are the Modern (2026-07-28) methods for which the draft
// spec's "Server Validation" section requires a Mcp-Name header naming the
// body's tool/resource/prompt identifier. Methods outside this set (e.g.
// tools/list) have no per-request name to check, so Mcp-Name is never
// required for them — though if a client sends one anyway, it is still
// validated for consistency.
var nameRequiredMethods = map[string]bool{
        "tools/call":     true,
        "resources/read": true,
        "prompts/get":    true,
}

// IsNameRequiredMethod reports whether the Modern (2026-07-28) method requires
// an Mcp-Name request header naming the target tool/resource/prompt. It shares
// nameRequiredMethods with ValidateHeaderConsistency so a Modern client sets the
// header for exactly the methods the server validates it on.
func IsNameRequiredMethod(method string) bool {
        return nameRequiredMethods[method]
}

// ValidateHeaderConsistency enforces the Modern (2026-07-28) Mcp-Method and
// Mcp-Name request headers against the corresponding parsed request body
// fields (Method and ResourceID).
//
// Caller contract: only call this for a request ClassifyRevision has already
// classified Modern. Mcp-Method/Mcp-Name are HTTP-only fields — populated
// solely by ParsingMiddleware reading real HTTP headers (see parser.go) —
// so there is no stdio/HTTP transport ambiguity here to excuse a Legacy
// request from this check; the caller must simply not invoke this function
// for Legacy traffic in the first place.
//
// Mcp-Method is required on every Modern request: a missing or empty header
// is itself a rejection (*RequestHeaderMismatchError), not a silent no-op,
// and a present-but-different value is a mismatch.
//
// Mcp-Name is required only for the methods in nameRequiredMethods
// (tools/call, resources/read, prompts/get) — for any other method, an
// absent Mcp-Name header is fine. When present, it is decoded via
// decodeSentinelName before comparison against ResourceID, since the draft
// spec allows it to be sentinel-encoded; a decode failure or mismatch is a
// rejection.
func ValidateHeaderConsistency(parsed *ParsedMCPRequest) error {
        if parsed.MCPMethodHeader == "" {
                return &RequestHeaderMismatchError{Header: "Mcp-Method", reason: headerMismatchReasonMissing}
        }
        if parsed.MCPMethodHeader != parsed.Method {
                return &RequestHeaderMismatchError{Header: "Mcp-Method", HeaderValue: parsed.MCPMethodHeader, BodyValue: parsed.Method}
        }

        if parsed.MCPNameHeader == "" {
                if nameRequiredMethods[parsed.Method] {
                        return &RequestHeaderMismatchError{Header: "Mcp-Name", reason: headerMismatchReasonMissing}
                }
                return nil
        }

        decoded, err := decodeSentinelName(parsed.MCPNameHeader)
        if err != nil {
                return &RequestHeaderMismatchError{
                        Header:      "Mcp-Name",
                        HeaderValue: parsed.MCPNameHeader,
                        BodyValue:   parsed.ResourceID,
                        reason:      headerMismatchReasonMalformed,
                }
        }
        if decoded != parsed.ResourceID {
                return &RequestHeaderMismatchError{Header: "Mcp-Name", HeaderValue: decoded, BodyValue: parsed.ResourceID}
        }

        return nil
}

// ExtractMeta pulls the "_meta" object out of raw JSON-RPC request params, for
// use with ClassifyRevision. It is deliberately tolerant: absent params, params
// that don't decode as a JSON object, or a "_meta" value that isn't itself an
// object all yield a nil map rather than an error. Only a well-formed object
// "_meta" is returned.
func ExtractMeta(params json.RawMessage) map[string]any {
        if len(params) == 0 {
                return nil
        }
        var paramsMap map[string]any
        if err := json.Unmarshal(params, &paramsMap); err != nil {
                return nil
        }
        return metaFromParamsMap(paramsMap)
}

// metaFromParamsMap reports the "_meta" value of an already-decoded JSON-RPC
// params map, if it decodes as a JSON object. A missing key or a wrong-typed
// value (e.g. a string or number) both yield nil.
func metaFromParamsMap(paramsMap map[string]any) map[string]any {
        meta, ok := paramsMap["_meta"].(map[string]any)
        if !ok {
                return nil
        }
        return meta
}

// hasModernSignal reports whether the request signals the Modern revision:
// either the header exactly names MCPVersionModern, or _meta carries any of the
// modernSignalMetaKeys (regardless of whether their values are well-formed).
// It reads modernSignalMetaKeys, NOT everything StripReservedMeta removes — a
// key can be reserved-for-stripping without signalling Modern (logLevel).
func hasModernSignal(meta map[string]any, protoHeader string) bool {
        if protoHeader == MCPVersionModern {
                return true
        }
        for _, key := range modernSignalMetaKeys {
                if _, ok := meta[key]; ok {
                        return true
                }
        }
        return false
}

// stringMetaValue reports the non-empty string value of meta[key], if present.
func stringMetaValue(meta map[string]any, key string) (string, bool) {
        raw, ok := meta[key]
        if !ok {
                return "", false
        }
        s, ok := raw.(string)
        if !ok || s == "" {
                return "", false
        }
        return s, true
}

// hasObjectMetaValue reports whether meta[key] is present and decodes as a JSON
// object; see objectMetaValue for the shape this covers.
func hasObjectMetaValue(meta map[string]any, key string) bool {
        _, ok := objectMetaValue(meta, key)
        return ok
}

// objectMetaValue reports the value of meta[key] if it decodes as a JSON
// object (map[string]any) — the shape clientInfo and clientCapabilities take
// in _meta. A missing key or a wrong-typed value (e.g. a string or number)
// both count as "not present".
func objectMetaValue(meta map[string]any, key string) (map[string]any, bool) {
        raw, ok := meta[key]
        if !ok {
                return nil, false
        }
        obj, ok := raw.(map[string]any)
        if !ok {
                return nil, false
        }
        return obj, true
}

// sentinelPrefix and sentinelSuffix wrap a base64-encoded Mcp-Name header
// value per the draft MCP spec. This is NOT RFC 2047 encoded-word syntax
// (there is no encoding-letter field); both markers are literal and
// case-sensitive.
const (
        sentinelPrefix = "=?base64?"
        sentinelSuffix = "?="
)

// decodeSentinelName decodes a Mcp-Name header value that may be wrapped in
// the draft spec's base64 sentinel format (=?base64?<payload>?=). A value
// that isn't wrapped in the sentinel markers is returned unchanged, since it
// is already the plain name/uri. A wrapped value whose payload fails to
// base64-decode is reported as an error.
func decodeSentinelName(v string) (string, error) {
        if !strings.HasPrefix(v, sentinelPrefix) || !strings.HasSuffix(v, sentinelSuffix) {
                return v, nil
        }

        payload := strings.TrimSuffix(strings.TrimPrefix(v, sentinelPrefix), sentinelSuffix)
        decoded, err := base64.StdEncoding.DecodeString(payload)
        if err != nil {
                return "", fmt.Errorf("decoding base64 sentinel Mcp-Name payload: %w", err)
        }
        return string(decoded), nil
}

// EncodeSentinelName encodes v into the draft spec's base64 sentinel format
// (=?base64?<payload>?=) when required — the mirror of decodeSentinelName.
// Encoding is required when EITHER:
//   - v is empty: a plain empty HTTP header value is indistinguishable from an
//     omitted header to receivers, so its sentinel form preserves that the
//     parameter was present;
//   - v is not safely representable as a plain ASCII header value: any byte
//     falls outside printable ASCII 0x21-0x7E, which also covers leading/
//     trailing whitespace and CR/LF (both fall below 0x21 and would
//     otherwise make net/http reject the request outright), or
//   - v already matches the sentinel pattern (has both sentinelPrefix and
//     sentinelSuffix), which must be escaped so the server doesn't mistake a
//     literal name for an encoded payload.
//
// Otherwise v is returned unchanged. The result always round-trips through
// decodeSentinelName back to v.
func EncodeSentinelName(v string) string {
        if !needsSentinelEncoding(v) {
                return v
        }
        return sentinelPrefix + base64.StdEncoding.EncodeToString([]byte(v)) + sentinelSuffix
}

// needsSentinelEncoding reports whether v requires sentinel encoding; see
// EncodeSentinelName for the three conditions checked.
func needsSentinelEncoding(v string) bool {
        if v == "" || (strings.HasPrefix(v, sentinelPrefix) && strings.HasSuffix(v, sentinelSuffix)) {
                return true
        }
        for i := 0; i < len(v); i++ {
                if v[i] < 0x21 || v[i] > 0x7E {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "bytes"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "strings"

        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// filteredToolNotFoundMessage is the generic message returned to a client
// that called a tool blocked by the tool filter. It deliberately matches what
// a client would see for a tool that doesn't exist at all: see the NOTE above
// the *toolCallFilter case in NewToolCallMappingMiddleware for why a filtered
// tool must not be distinguishable from a nonexistent one.
const filteredToolNotFoundMessage = "tool not found"

var errToolNameNotFound = errors.New("tool name not found")
var errBug = errors.New("there's a bug")
var errKeepBuffering = errors.New("keep buffering")
var errUnsupportedMimeType = errors.New("unsupported mime type")

// toolOverrideEntry is a struct that represents a tool override entry.
type toolOverrideEntry struct {
        ActualName          string
        OverrideName        string
        OverrideDescription string
}

// toolMiddlewareConfig is a helper struct used to configure the tool middleware,
// and it's meant to map from a tool's actual name to a config entry.
//
// The two separate structs are necessary because it must be possible to specify
// tool overrides without tool filtering.
//
// Assume a User only specified an override for a single tool out of a list of
// n tools; in such a case, it would become unclear whether the tool is the only
// one allowed or is the only one overridden.
//
// Sufficient information could be represented in a more complex structure, but
// this gets the job and is easy enough to understand.
type toolMiddlewareConfig struct {
        filterTools          map[string]struct{}
        actualToUserOverride map[string]toolOverrideEntry
        userToActualOverride map[string]toolOverrideEntry
}

func (c *toolMiddlewareConfig) isToolInFilter(toolName string) bool {
        if len(c.filterTools) == 0 {
                return true
        }

        _, ok := c.filterTools[toolName]
        return ok
}

func (c *toolMiddlewareConfig) getToolCallActualName(toolName string) (string, bool) {
        if len(c.userToActualOverride) == 0 {
                return "", false
        }

        entry, ok := c.userToActualOverride[toolName]
        return entry.ActualName, ok
}

func (c *toolMiddlewareConfig) getToolListOverride(toolName string) (*toolOverrideEntry, bool) {
        if len(c.actualToUserOverride) == 0 {
                return nil, false
        }

        entry, ok := c.actualToUserOverride[toolName]
        return &entry, ok
}

// ToolMiddlewareOption is a function that can be used to configure the tool
// middleware.
type ToolMiddlewareOption func(*toolMiddlewareConfig) error

// SimpleTool represents a minimal tool with name and description.
// This is used by ApplyToolFiltering to work with tools in a generic way.
type SimpleTool struct {
        Name        string
        Description string
}

// ApplyToolFiltering applies filtering and overriding to a list of tools.
// This is the core logic used by both the HTTP middleware and other components
// that need to apply the same filtering/overriding behavior.
//
// Returns the filtered and overridden tools.
func ApplyToolFiltering(opts []ToolMiddlewareOption, tools []SimpleTool) ([]SimpleTool, error) {
        config := &toolMiddlewareConfig{
                filterTools:          make(map[string]struct{}),
                actualToUserOverride: make(map[string]toolOverrideEntry),
                userToActualOverride: make(map[string]toolOverrideEntry),
        }

        // Apply options to build config
        for _, opt := range opts {
                if err := opt(config); err != nil {
                        return nil, err
                }
        }

        // Use the shared core logic
        return applyFilteringAndOverrides(config, tools), nil
}

// WithToolsFilter is a function that can be used to configure the tool
// middleware to use a filter list of tools.
func WithToolsFilter(toolsFilter ...string) ToolMiddlewareOption {
        return func(mw *toolMiddlewareConfig) error {
                for _, tf := range toolsFilter {
                        if tf == "" {
                                return fmt.Errorf("tool name cannot be empty")
                        }

                        mw.filterTools[tf] = struct{}{}
                }

                return nil
        }
}

// WithToolsOverride is a function that can be used to configure the tool
// middleware to use a map of tools to override the actual list of tools.
//
// If an empty string is provided for either overrideName or overrideDescription,
// that field will be left unchanged. An error is returned if actualName is empty.
func WithToolsOverride(actualName string, overrideName string, overrideDescription string) ToolMiddlewareOption {
        return func(mw *toolMiddlewareConfig) error {
                if actualName == "" {
                        return fmt.Errorf("tool name cannot be empty")
                }

                if overrideName == "" && overrideDescription == "" {
                        return fmt.Errorf("override name and description cannot both be empty")
                }

                entry := toolOverrideEntry{
                        ActualName:          actualName,
                        OverrideName:        overrideName,        // empty string means no override
                        OverrideDescription: overrideDescription, // empty string means no override
                }
                mw.actualToUserOverride[actualName] = entry
                mw.userToActualOverride[overrideName] = entry

                return nil
        }
}

// NewListToolsMappingMiddleware creates an HTTP middleware that parses SSE responses
// and plain JSON objects to extract tool names from JSON-RPC messages containing
// tool lists or tool calls.
//
// The middleware looks for SSE events with:
// - event: message
// - data: {"jsonrpc":"2.0","id":X,"result":{"tools":[...]}}
//
// This middleware is designed to be used ONLY when tool filtering or
// override are enabled, and expects the list of tools to be "correct"
// (i.e. not empty and not containing nonexisting tools).
func NewListToolsMappingMiddleware(opts ...ToolMiddlewareOption) (types.MiddlewareFunction, error) {
        config := &toolMiddlewareConfig{
                filterTools:          make(map[string]struct{}),
                actualToUserOverride: make(map[string]toolOverrideEntry),
                userToActualOverride: make(map[string]toolOverrideEntry),
        }
        for _, opt := range opts {
                if err := opt(config); err != nil {
                        return nil, err
                }
        }

        if len(config.filterTools) == 0 && len(config.actualToUserOverride) == 0 {
                return nil, fmt.Errorf("tools list for filtering or overriding is empty")
        }

        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // NOTE: this middleware only checks the response body, whose
                        // format at this point is not yet known and might be either a
                        // JSON payload or an SSE stream.
                        //
                        // The way this is implemented is that we wrap the response writer
                        // in order to buffer the response body. Once Flush() is called, we
                        // process the buffer according to its content type and possibly
                        // modify it before returning it to the client.
                        rw := &toolFilterWriter{
                                ResponseWriter: w,
                                config:         config,
                                statusCode:     http.StatusOK, // matches net/http's implicit status when WriteHeader is never called
                        }

                        // Call the next handler
                        next.ServeHTTP(rw, r)

                        // Explicitly drain our buffered writer after the handler returns. An inner
                        // buffering middleware (e.g. authz's response filter) may write the final
                        // body into our buffer during its own post-ServeHTTP flush; without this
                        // call that body is stranded and the client gets a 0-byte response. See #5797.
                        //
                        // This is the terminal drain, not a streaming Flush: there is no more data
                        // coming, so a body we can't process (wrong/missing content type, malformed
                        // tools list) is written through unchanged rather than dropped. See #5809 (review).
                        rw.finish()
                })
        }, nil
}

// NewToolCallMappingMiddleware creates an HTTP middleware that parses tool call
// requests and filters out tools that are not in the filter list.
//
// The middleware looks for JSON-RPC messages with:
// - method: tool/call
// - params: {"name": "tool_name"}
//
// This middleware is designed to be used ONLY when tool filtering or override
// is enabled, and expects the list of tools to be "correct" (i.e. not empty
// and not containing nonexisting tools).
func NewToolCallMappingMiddleware(opts ...ToolMiddlewareOption) (types.MiddlewareFunction, error) {
        config := &toolMiddlewareConfig{
                filterTools:          make(map[string]struct{}),
                actualToUserOverride: make(map[string]toolOverrideEntry),
                userToActualOverride: make(map[string]toolOverrideEntry),
        }
        for _, opt := range opts {
                if err := opt(config); err != nil {
                        return nil, err
                }
        }

        if len(config.filterTools) == 0 && len(config.actualToUserOverride) == 0 {
                return nil, fmt.Errorf("tools list for filtering or overriding is empty")
        }

        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        bodyBytes, ok := readToolCallBody(w, r)
                        if !ok {
                                return
                        }

                        // Restore the normalized request body for downstream handlers.
                        r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
                        r.ContentLength = int64(len(bodyBytes))

                        // Try to parse the request as a tool call request. If it succeeds,
                        // check if the tool is in the filter. If it is not a tool call request,
                        // just pass it through.
                        var toolCallRequest toolCallRequest
                        decoder := json.NewDecoder(bytes.NewReader(bodyBytes))
                        decoder.UseNumber()
                        err := decoder.Decode(&toolCallRequest)
                        if err == nil && json.Valid(bodyBytes) && toolCallRequest.Method == "tools/call" {
                                fix := processToolCallRequest(config, toolCallRequest)

                                switch fix := fix.(type) {

                                // If the tool call request is allowed, and the tool name is not overridden,
                                // we just pass it through unmodified.
                                case *toolCallNoAction:
                                        next.ServeHTTP(w, r)
                                        return

                                // NOTE: ideally, trying to call that was filtered out by config should be
                                // equivalent to calling a nonexisting tool; in such cases and when the SSE
                                // transport is used, the behaviour of the official Python SDK is to return
                                // a 202 Accepted to THIS call and return an success message in the SSE
                                // stream saying that the tool does not exist.
                                //
                                // It basically fails successfully.
                                //
                                // Unfortunately, implementing this behaviour is not trivial and requires
                                // session management, as the SSE stream is managed by the proxy in an entirely
                                // different thread of execution. As a consequence, the best thing we can
                                // do that is still compliant with the spec is to return a JSON-RPC error
                                // for this call, over HTTP 200, that looks the same as calling a tool that
                                // doesn't exist (see filteredToolNotFoundMessage above).
                                //
                                // A client that only accepts an event stream (the legacy HTTP+SSE
                                // transport, whose real response is delivered on a separate stream we
                                // don't control from this middleware) can't be answered this way: the
                                // body written here IS the only response this call gets, and it must be
                                // something other than a stray JSON blob on a connection the client
                                // expects to be text/event-stream. So that path keeps the 400, same as
                                // the session-management limitation described above.
                                case *toolCallFilter:
                                        if clientAcceptsJSON(r) {
                                                writeFilteredToolCallError(w, toolCallRequest.ID)
                                                return
                                        }
                                        w.WriteHeader(http.StatusBadRequest)
                                        return

                                // In case of a tool name override, we need to fix the tool call request
                                // and then forward it to the next handler.
                                case *toolCallOverride:
                                        (*toolCallRequest.Params)["name"] = fix.Name()
                                        bodyBytes, err = json.Marshal(toolCallRequest)
                                        if err != nil {
                                                slog.Error("error marshalling tool call request",
                                                        "error", err)
                                                next.ServeHTTP(w, r)
                                                return
                                        }

                                        r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
                                        r.ContentLength = int64(len(bodyBytes))

                                // According to the current version of the MCP spec at
                                // https://modelcontextprotocol.io/specification/2025-11-25/schema#calltoolrequest
                                // this case can only happen if the request is malformed. The proxied MCP
                                // server should be able to process the request, but since we detect it here
                                // we short-circuit returning an error.
                                case *toolCallBogus:
                                        w.WriteHeader(http.StatusBadRequest)
                                        return

                                // This should never happen, but we handle it just in case.
                                default:
                                        slog.Error("error processing tool call of a filtered tool",
                                                "error", err)
                                        next.ServeHTTP(w, r)
                                        return
                                }
                        }

                        next.ServeHTTP(w, r)
                })
        }, nil
}

func readToolCallBody(w http.ResponseWriter, r *http.Request) ([]byte, bool) {
        body, err := io.ReadAll(r.Body)
        if err != nil {
                writeBodyReadError(w, err)
                return nil, false
        }
        // Validate before an override can erase invalid envelope evidence.
        if r.Method == http.MethodPost {
                if _, err := DecodeMessage(body); err != nil {
                        WriteClassificationError(w, nil, err)
                        return nil, false
                }
        }
        return bytes.TrimPrefix(body, UTF8BOM), true
}

// clientAcceptsJSON reports whether r's Accept header allows an
// application/json response. An empty/absent header is treated as
// accepting JSON, matching the common case of a client that doesn't bother
// to negotiate. Each comma-separated entry's media-type token (i.e.
// everything before any ";param=..." qualifier) is compared case-
// insensitively against "application/json", the type wildcard
// "application/*", and the full wildcard "*/*". Quality values (";q=") are
// not honored: a token present with any q-value counts as accepted.
func clientAcceptsJSON(r *http.Request) bool {
        accept := r.Header.Get("Accept")
        if accept == "" {
                return true
        }

        for _, entry := range strings.Split(accept, ",") {
                mediaType := strings.TrimSpace(strings.Split(entry, ";")[0])
                if strings.EqualFold(mediaType, "application/json") ||
                        strings.EqualFold(mediaType, "application/*") ||
                        mediaType == "*/*" {
                        return true
                }
        }

        return false
}

// writeFilteredToolCallError writes a JSON-RPC error response, over HTTP 200,
// for a tools/call request blocked by the tool filter. The error message is
// the generic filteredToolNotFoundMessage: see the NOTE in NewToolCallMappingMiddleware
// for why a filtered tool must look the same as one that doesn't exist.
//
// HTTP 200 (not 400) is deliberate: under MCP streamable HTTP, a validly
// received JSON-RPC request that fails at the application level rides back
// in a 200 response carrying a JSON-RPC error object, mirroring how
// WriteClassificationError's peers (e.g. writeRateLimited) shape their body,
// modulo status code.
func writeFilteredToolCallError(w http.ResponseWriter, id any) {
        body := filteredToolCallErrorBody(id)
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusOK)
        //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client
        _, _ = w.Write(body)
}

// filteredToolCallErrorBody renders the JSON-RPC error body for a filtered
// tool call, modeled on classificationErrorBody: the body is marshaled first
// (with a hand-crafted fallback on marshal failure) so the caller only writes
// headers/status once a valid body is ready.
//
// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error
// response as `id?: RequestId` where RequestId = string | number, so an
// absent id is encoded by omitting the "id" key, never as null. See
// session.HasJSONRPCID.
func filteredToolCallErrorBody(id any) []byte {
        resp := map[string]any{
                "jsonrpc": "2.0",
                "error": map[string]any{
                        "code":    CodeInvalidParams,
                        "message": filteredToolNotFoundMessage,
                },
        }
        if session.HasJSONRPCID(id) {
                resp["id"] = id
        }

        body, err := json.Marshal(resp)
        if err != nil {
                // This should never happen with simple map types, but return a
                // hand-crafted fallback to guarantee a valid JSON-RPC error.
                return []byte(`{"jsonrpc":"2.0","error":{"code":-32602,"message":"tool not found"}}`)
        }
        return body
}

// toolFilterWriter wraps http.ResponseWriter to capture and process SSE responses
type toolFilterWriter struct {
        http.ResponseWriter
        buffer     []byte
        config     *toolMiddlewareConfig
        statusCode int
}

// WriteHeader captures the status code.
//
// Content-Length is stripped because applying tool filters or overrides may
// change the body size (e.g. a longer override description). Without this,
// net/http rejects the rewritten body with "http: wrote more than the declared
// Content-Length" and the client receives only the headers. Removing the
// header lets Go fall back to chunked transfer encoding.
func (rw *toolFilterWriter) WriteHeader(statusCode int) {
        rw.statusCode = statusCode
        rw.ResponseWriter.Header().Del("Content-Length")
        rw.ResponseWriter.WriteHeader(statusCode)
}

// Write captures the response body and processes SSE events
func (rw *toolFilterWriter) Write(data []byte) (int, error) {
        rw.buffer = append(rw.buffer, data...)
        return len(data), nil
}

// Flush processes any remaining buffered data, writes it to the underlying
// ResponseWriter, and forwards the flush downstream. This is the interim/
// streaming path: a downstream handler or reverse proxy may call this
// mid-response, so an incomplete SSE event or partial JSON body is left
// buffered for a later call to complete, and the transport flush is not
// forwarded in that case -- there's nothing new to push over the wire, and
// forwarding it could send a still-incomplete frame. Use finish() for the
// one-time terminal drain after the wrapped handler has returned.
func (rw *toolFilterWriter) Flush() {
        if !rw.drainBuffer(false) {
                return
        }

        if flusher, ok := rw.ResponseWriter.(http.Flusher); ok {
                flusher.Flush()
        }
}

// finish performs the terminal drain of any remaining buffered data after the
// wrapped handler has returned. Unlike Flush, there is no more data coming:
// an incomplete body (a truncated stream, a connection cut short) is written
// through unchanged rather than held forever, and an unrecognized/missing
// Content-Type on a response with nothing to filter is likewise passed
// through as-is. It does not forward a transport flush, since returning from
// ServeHTTP already completes the response.
//
// A body that DOES look like a tools list but fails to filter is the one
// case finish() will not pass through -- see drainBuffer.
func (rw *toolFilterWriter) finish() {
        rw.drainBuffer(true)
}

// drainBuffer processes the buffered response body according to its MIME
// type and writes the result to the underlying ResponseWriter. It reports
// whether the caller may safely forward a transport-level flush: false means
// data is being deliberately withheld, either because it's incomplete and
// more is expected (Flush, non-terminal) or because it looked like a tools
// list we could not safely filter (fail closed, in both Flush and finish).
//
// When terminal is false (Flush, mid-stream), a body processBuffer reports as
// incomplete is left buffered for a later call. When terminal is true
// (finish, called once ServeHTTP has returned), there won't be a later call,
// so an incomplete body is written through unchanged instead of being
// stranded.
//
// A hard filtering failure -- a response that resolved to a tools list
// (whether correctly labeled or sniffed, see processUnrecognizedMimeType) but
// couldn't be filtered, e.g. a tool missing its required name -- is never
// passed through, in either mode: doing so would leak the unfiltered list.
func (rw *toolFilterWriter) drainBuffer(terminal bool) bool {
        if len(rw.buffer) == 0 {
                return true
        }

        // Normalize once before MIME dispatch so every response path applies the
        // same parsing and forwarding semantics as clients that strip a leading BOM.
        rw.buffer = bytes.TrimPrefix(rw.buffer, UTF8BOM)

        mimeType := strings.Split(rw.ResponseWriter.Header().Get("Content-Type"), ";")[0]
        successResponse := rw.statusCode == http.StatusOK || rw.statusCode == http.StatusAccepted

        var b bytes.Buffer
        err := processBuffer(rw.config, rw.buffer, mimeType, successResponse, &b)

        switch {
        case errors.Is(err, errKeepBuffering):
                if !terminal {
                        slog.Debug("keep buffering", "buffered_bytes", len(rw.buffer))
                        return false
                }
                slog.Warn("response ended with an incomplete buffered body; writing it through unchanged",
                        "buffered_bytes", len(rw.buffer))
                rw.writeBuffer(rw.buffer)
                return true

        case errors.Is(err, errUnsupportedMimeType):
                // processBuffer only returns this once it has ruled out the buffer being
                // a tools list: either the response has nothing to filter (a non-2xx
                // error body), or it's a 2xx body that doesn't look like a tools list
                // under either supported wire format. Either way, passing it through
                // unchanged cannot leak an unfiltered list.
                rw.writeBuffer(rw.buffer)
                return true

        case err != nil:
                // A recognized (or sniffed) tools list that failed to filter -- e.g. a
                // malformed tool entry. Fail closed: writing the raw, unfiltered body
                // through would defeat the filter entirely.
                slog.Error("error filtering buffered response; refusing to pass through unfiltered", "error", err)
                rw.buffer = rw.buffer[:0]
                return false

        default:
                slog.Debug("flushing buffer", "bytes", len(b.Bytes()))
                rw.writeBuffer(b.Bytes())
                return true
        }
}

// writeBuffer writes data to the underlying ResponseWriter and resets rw.buffer.
// data is expected to be either rw.buffer itself or a value derived from it, so
// resetting after the write is always correct.
func (rw *toolFilterWriter) writeBuffer(data []byte) {
        _, err := rw.ResponseWriter.Write(data)
        if err != nil {
                slog.Error("error writing buffer", "error", err)
        }
        rw.buffer = rw.buffer[:0] // Reset buffer
}

type toolsListResponse struct {
        JSONRPC string `json:"jsonrpc"`
        ID      any    `json:"id"`
        Result  struct {
                Tools *[]map[string]any `json:"tools"`
        } `json:"result"`
}

type toolCallRequest struct {
        JSONRPC string          `json:"jsonrpc"`
        ID      any             `json:"id,omitempty"`
        Method  string          `json:"method"`
        Params  *map[string]any `json:"params,omitempty"`
}

// processBuffer processes the buffered response body according to its
// declared MIME type. successResponse indicates whether the response's
// status code was 200/202: it's only consulted for an unrecognized MIME type,
// to decide whether the body needs sniffing before it can safely be passed
// through -- see processUnrecognizedMimeType.
func processBuffer(
        config *toolMiddlewareConfig,
        buffer []byte,
        mimeType string,
        successResponse bool,
        w io.Writer,
) error {
        if len(buffer) == 0 {
                return nil
        }

        switch mimeType {
        case "application/json":
                return processJSONBuffer(config, buffer, w)
        case "text/event-stream":
                return processEventStream(config, buffer, w)
        default:
                return processUnrecognizedMimeType(config, buffer, mimeType, successResponse, w)
        }
}

// processJSONBuffer filters a JSON-RPC response body if it carries a tools
// list, otherwise passes it through unchanged (e.g. a tools/call result, or
// any other JSON-RPC message this middleware isn't meant to touch).
func processJSONBuffer(config *toolMiddlewareConfig, buffer []byte, w io.Writer) error {
        var toolsListResponse toolsListResponse
        var syntaxError *json.SyntaxError
        err := json.Unmarshal(buffer, &toolsListResponse)
        if errors.As(err, &syntaxError) {
                return fmt.Errorf("%w: %w", errKeepBuffering, err)
        }
        if err == nil && toolsListResponse.Result.Tools != nil {
                return processToolsListResponse(config, toolsListResponse, w)
        }

        _, err = w.Write(buffer)
        return err
}

// processUnrecognizedMimeType handles a body whose Content-Type is missing or
// isn't one of the two MCP-supported types (application/json,
// text/event-stream).
//
// A non-2xx response has nothing to filter -- the tool filter only ever acts
// on a tools/list result, which a non-2xx status can't carry -- so it's
// always safe to pass through.
//
// A 2xx response, on the other hand, is spec-required to declare one of the
// two supported Content-Types. A missing or wrong header there could be an
// accident, or a backend deliberately omitting/mislabeling it to smuggle an
// unfiltered tools list past the filter (the same disguised-result concern
// handled in pkg/authz for #5257). So the raw body is sniffed the same way a
// correctly-labeled response would be processed: if it resolves to a tools
// list under either supported shape, it's filtered (or, on failure, rejected)
// exactly as it would be under the correct Content-Type. Only once neither
// shape resolves to a tools list is it safe to pass through unchanged.
//
// Known limitation: unlike the correctly-labeled paths, this sniff has no
// notion of "incomplete, wait for more" -- a tools list that both omits/
// mislabels its Content-Type AND arrives across multiple Flush() calls may
// be sniffed as "not a tools list" on an early, truncated call and passed
// through piecemeal before the full body is available. A compliant backend
// that streams a tools list always declares text/event-stream, so this only
// affects a backend that is already violating the spec on two axes at once.
func processUnrecognizedMimeType(
        config *toolMiddlewareConfig,
        buffer []byte,
        mimeType string,
        successResponse bool,
        w io.Writer,
) error {
        if !successResponse {
                return fmt.Errorf("%w: %s", errUnsupportedMimeType, mimeType)
        }

        var candidate toolsListResponse
        if err := json.Unmarshal(buffer, &candidate); err == nil && candidate.Result.Tools != nil {
                return processToolsListResponse(config, candidate, w)
        }

        if sniffSSEToolsList(buffer) {
                return processEventStream(config, buffer, w)
        }

        return fmt.Errorf("%w: %s", errUnsupportedMimeType, mimeType)
}

// sniffSSEToolsList reports whether buffer contains an SSE "data:" line whose
// payload decodes to a tools list. It's a lightweight detector only, used to
// decide whether a mislabeled/untyped 2xx body needs the full SSE processing
// path (which enforces its own completeness and separator rules).
func sniffSSEToolsList(buffer []byte) bool {
        normalized := bytes.ReplaceAll(buffer, []byte("\r\n"), []byte("\n"))
        normalized = bytes.ReplaceAll(normalized, []byte("\r"), []byte("\n"))

        for _, line := range bytes.Split(normalized, []byte("\n")) {
                data, ok := bytes.CutPrefix(line, []byte("data:"))
                if !ok {
                        continue
                }

                var candidate toolsListResponse
                if json.Unmarshal(bytes.TrimSpace(data), &candidate) == nil && candidate.Result.Tools != nil {
                        return true
                }
        }

        return false
}

//nolint:gocyclo
func processEventStream(
        config *toolMiddlewareConfig,
        buffer []byte,
        w io.Writer,
) error {
        if len(buffer) > 1 && buffer[len(buffer)-1] != '\n' && buffer[len(buffer)-1] != '\r' {
                return fmt.Errorf("%w: %v", errKeepBuffering, "event separator not found")
        }

        // NOTE: this looks uglier, but is more efficient than scanning the whole buffer
        var linesep []byte
        if len(buffer) >= 2 && bytes.Equal(buffer[len(buffer)-2:], []byte("\r\n")) {
                linesep = []byte("\r\n")
        } else if len(buffer) >= 1 && buffer[len(buffer)-1] == '\n' {
                linesep = []byte("\n")
        } else if len(buffer) >= 1 && buffer[len(buffer)-1] == '\r' {
                linesep = []byte("\r")
        } else {
                // Length only, not the buffer: buffer is the full pre-filter payload
                // (e.g. the unfiltered tools/list), and this error is logged upstream.
                return fmt.Errorf("unsupported SSE line separator in %d-byte buffer", len(buffer))
        }

        var linesepTotal, linesepCount int
        linesepTotal = bytes.Count(buffer, linesep)
        lines := bytes.Split(buffer, linesep)
        for _, line := range lines {
                if len(line) == 0 {
                        continue
                }

                var written bool
                if data, ok := bytes.CutPrefix(line, []byte("data:")); ok {
                        var toolsListResponse toolsListResponse
                        if err := json.Unmarshal(data, &toolsListResponse); err == nil && toolsListResponse.Result.Tools != nil {
                                // We got to the point of processing a real tools list response,
                                // so we need to write the "data: " prefix first.
                                _, err := w.Write([]byte("data: "))
                                if err != nil {
                                        return fmt.Errorf("%w: %w", errBug, err)
                                }

                                if err := processToolsListResponse(config, toolsListResponse, w); err != nil {
                                        return err
                                }
                                written = true
                        }
                }

                if !written {
                        _, err := w.Write(line)
                        if err != nil {
                                return fmt.Errorf("%w: %w", errBug, err)
                        }
                }

                _, err := w.Write(linesep)
                if err != nil {
                        return fmt.Errorf("%w: %w", errBug, err)
                }
                linesepCount++
        }

        // This ensures we don't send too few line separators, which might break
        // SSE parsing.
        if linesepCount < linesepTotal {
                _, err := w.Write(linesep)
                if err != nil {
                        return fmt.Errorf("%w: %w", errBug, err)
                }
        }

        return nil
}

// processToolsListResponse processes a tools list response filtering out
// tools that are not in the filter list.
func processToolsListResponse(
        config *toolMiddlewareConfig,
        toolsListResponse toolsListResponse,
        w io.Writer,
) error {
        // Convert to SimpleTool format for shared processing
        simpleTools := make([]SimpleTool, 0, len(*toolsListResponse.Result.Tools))
        toolMaps := make([]map[string]any, 0, len(*toolsListResponse.Result.Tools))

        for _, tool := range *toolsListResponse.Result.Tools {
                // NOTE: the spec does not allow for name to be missing.
                toolName, ok := tool["name"].(string)
                if !ok {
                        return errToolNameNotFound
                }

                // NOTE: the spec does not allow for empty tool names.
                if toolName == "" {
                        return errToolNameNotFound
                }

                // Get description if present (optional in MCP spec)
                description, _ := tool["description"].(string)

                simpleTools = append(simpleTools, SimpleTool{
                        Name:        toolName,
                        Description: description,
                })
                toolMaps = append(toolMaps, tool)
        }

        // Apply the shared filtering/override logic
        processedTools := applyFilteringAndOverrides(config, simpleTools)

        // Build the filtered response by matching processed tools with their original maps
        // Note: This is O(n²) complexity, but acceptable because:
        // - Tool lists are typically small (< 100 tools per backend)
        // - Only runs once during tool list retrieval (not in hot path)
        // - Inner loop breaks early on match
        filteredTools := make([]map[string]any, 0, len(processedTools))
        for _, processed := range processedTools {
                // Find the original tool map by matching names
                for i, simple := range simpleTools {
                        if simple.Name == processed.Name || simple.Name == findOriginalName(config, processed.Name) {
                                // Clone the original map and update name/description
                                toolCopy := make(map[string]any, len(toolMaps[i]))
                                for k, v := range toolMaps[i] {
                                        toolCopy[k] = v
                                }
                                toolCopy["name"] = processed.Name
                                if processed.Description != "" {
                                        toolCopy["description"] = processed.Description
                                }
                                filteredTools = append(filteredTools, toolCopy)
                                break
                        }
                }
        }

        toolsListResponse.Result.Tools = &filteredTools
        if err := json.NewEncoder(w).Encode(toolsListResponse); err != nil {
                return fmt.Errorf("%w: %w", errBug, err)
        }

        return nil
}

// applyFilteringAndOverrides is the core logic for filtering and overriding tools.
// This implements the exact same logic as before but is now extracted for reuse.
func applyFilteringAndOverrides(config *toolMiddlewareConfig, tools []SimpleTool) []SimpleTool {
        result := make([]SimpleTool, 0, len(tools))
        for _, tool := range tools {
                description := tool.Description

                // If the tool is overridden, we need to use the override name and description.
                if entry, ok := config.getToolListOverride(tool.Name); ok {
                        if entry.OverrideName != "" {
                                tool.Name = entry.OverrideName
                        }
                        if entry.OverrideDescription != "" {
                                description = entry.OverrideDescription
                        }
                }

                // If the tool is in the filter, we add it to the filtered tools list.
                // Note that lookup is done using the user-known name (tool.Name after override).
                if config.isToolInFilter(tool.Name) {
                        result = append(result, SimpleTool{
                                Name:        tool.Name,
                                Description: description,
                        })
                }
        }
        return result
}

// findOriginalName attempts to find the original tool name before override.
func findOriginalName(config *toolMiddlewareConfig, overriddenName string) string {
        // Iterate through overrides to find reverse mapping
        for actualName, entry := range config.actualToUserOverride {
                if entry.OverrideName == overriddenName {
                        return actualName
                }
        }
        return overriddenName
}

// toolCallFix mimics a sum type in Go. The actual types represent the
// possible manipulations to perform on the tool call request, namely:
// - filter the tool call request
// - override the tool call request
// - return a bogus tool call request
// - do nothing
//
// The actual types are not exported, and the only way to get a value of a specific type
// is to use a type assertion.
//
// Technical note: it might be tempting to build this into toolMiddlewareConfig, but this
// would leave out the case in which the request is malformed, scenario that does not
// belong to the logic implementing config.
type toolCallFixAction interface{}

// toolCallFilter is a struct that represents a tool call filter, i.e.
// the tool call request is not allowed.
type toolCallFilter struct{}

// toolCallOverride is a struct that represents a tool call override, i.e.
// the tool call request is allowed, but the tool name is overridden.
type toolCallOverride struct {
        actualName string
}

// Name returns the actual name of the tool.
func (t *toolCallOverride) Name() string {
        return t.actualName
}

// toolCallBogus is a struct that represents a bogus tool call request, i.e.
// the tool call request is not allowed and the tool name is not overridden.
type toolCallBogus struct{}

// toolCallNoAction is a struct that represents a tool call no action, i.e.
// the tool call request is allowed and the tool name is not overridden.
type toolCallNoAction struct{}

// processToolCallRequest processes a tool call request checking if the tool
// is in the filter list. Note that the tool name received in the toolCallRequest
// is going to be the user-provided name, which might be different from the actual
// tool name.
func processToolCallRequest(
        config *toolMiddlewareConfig,
        toolCallRequest toolCallRequest,
) toolCallFixAction {
        // NOTE: the spec does not allow for nil params.
        if toolCallRequest.Params == nil {
                return &toolCallBogus{}
        }

        // NOTE: the spec does not allow for name to be missing.
        toolName, ok := (*toolCallRequest.Params)["name"].(string)
        if !ok {
                return &toolCallBogus{}
        }

        // NOTE: the spec does not allow for empty tool names.
        if toolName == "" {
                return &toolCallBogus{}
        }

        // If the tool is not in the filter list, return an error.
        // Note that the tool name we use here is the user-provided name, which
        // might be different from the actual tool name, but filters are expressed
        // in terms of tool names as known to the user, so this is correct.
        if !config.isToolInFilter(toolName) {
                return &toolCallFilter{}
        }

        // If the tool is allowed by the filter, and has an override, return the
        // actual name to fix the tool call request.
        if actualName, ok := config.getToolCallActualName(toolName); ok {
                return &toolCallOverride{actualName: actualName}
        }

        // If the tool is allowed by the filter, and does not have an override,
        // return an empty string and no error, signaling the fact that the tool
        // call request is ok as is.
        return &toolCallNoAction{}
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "fmt"

        "golang.org/x/exp/jsonrpc2"
)

// UTF8BOM is the three-byte UTF-8 byte order mark. Clients following the
// WHATWG UTF-8 decode algorithm strip one leading BOM before parsing, so code
// decoding or sniffing JSON-RPC messages from the wire must do the same.
var UTF8BOM = []byte("\xEF\xBB\xBF")

// ConvertToJSONRPC2ID converts an interface{} ID to jsonrpc2.ID
func ConvertToJSONRPC2ID(id interface{}) (jsonrpc2.ID, error) {
        if id == nil {
                return jsonrpc2.ID{}, nil
        }

        switch v := id.(type) {
        case string:
                return jsonrpc2.StringID(v), nil
        case int:
                return jsonrpc2.Int64ID(int64(v)), nil
        case int64:
                return jsonrpc2.Int64ID(v), nil
        case float64:
                // JSON numbers are often unmarshaled as float64
                return jsonrpc2.Int64ID(int64(v)), nil
        default:
                return jsonrpc2.ID{}, fmt.Errorf("unsupported ID type: %T", id)
        }
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mcp

import (
        "encoding/json"
        "errors"
        "fmt"
        "math"
        "slices"
        "strconv"
        "strings"
)

// SEP-2243 ("HTTP standardization") lets an MCP server mark individual tool
// parameters for mirroring into HTTP request headers, via an x-mcp-header
// annotation inside the parameter's schema in the tool's inputSchema. A server
// MAY use the annotation; a client MUST honour it, sending each designated
// parameter's value as the header Mcp-Param-{name} on the tools/call request. A
// server that designated a parameter and did not receive its header rejects the
// call with -32020.
//
// This file owns the annotation's vocabulary and its validation. It is
// deliberately the single place that reads x-mcp-header, so the two consumers --
// ingestion-time validation of a backend's advertised tools, and call-time
// derivation of the outgoing headers -- cannot drift in their reading of the
// constraints. (Compare the "third independent copy" problem called out for the
// reserved _meta keys in #5986.)
const (
        // XMCPHeaderAnnotation is the JSON Schema extension key a server sets on a
        // tool parameter to designate it for header mirroring. Its value is the
        // header's name suffix, NOT the full header name.
        XMCPHeaderAnnotation = "x-mcp-header"

        // ParamHeaderPrefix prefixes every mirrored parameter header: an annotation
        // of "Region" is sent as "Mcp-Param-Region".
        ParamHeaderPrefix = "Mcp-Param-"
)

// maxSchemaDepth bounds the inputSchema traversal. Annotations are legal at any
// nesting depth, so the walk has to recurse, and the schema is attacker-supplied
// from vMCP's perspective (a backend can advertise any tool list it likes). A
// depth cap keeps a hostile or accidentally-recursive schema from exhausting the
// stack. 64 is far past any hand-written tool schema.
const maxSchemaDepth = 64

// ErrSchemaTooDeep is returned when an inputSchema nests past maxSchemaDepth.
// It is deliberately distinguishable from a constraint violation: the schema may
// be perfectly valid and merely beyond what this walk will inspect, so a caller
// that wants to fail open on depth alone can single it out.
var ErrSchemaTooDeep = errors.New("tool inputSchema exceeds maximum inspection depth")

// ParamHeader is one x-mcp-header annotation found in a tool's inputSchema.
type ParamHeader struct {
        // Path is the property path from the root of inputSchema to the annotated
        // parameter, e.g. ["filter", "region"] for a nested property. Length is
        // always >= 1.
        Path []string

        // Name is the annotation's value -- the suffix appended to
        // ParamHeaderPrefix, not the full header name. Use HeaderName for that.
        Name string

        // Type is the annotated parameter's declared JSON Schema type, guaranteed by
        // validation to be one of "string", "integer", or "boolean".
        Type string
}

// HeaderName is the full HTTP header name this annotation mirrors into.
func (p ParamHeader) HeaderName() string {
        return ParamHeaderPrefix + p.Name
}

// ParamHeaders walks a tool's inputSchema and returns every x-mcp-header
// annotation it declares, in a deterministic order (by path). A schema with no
// annotations -- the overwhelmingly common case -- returns nil, nil.
//
// It returns an error when the schema violates any of SEP-2243's constraints on
// the annotation:
//
//   - the value must be a non-empty string;
//   - it must be a valid HTTP field-name token (RFC 9110 tchar), which also
//     excludes control characters and whitespace;
//   - it must be unique, case-insensitively, within the whole inputSchema
//     (HTTP field names are case-insensitive, so two spellings would collide on
//     the wire);
//   - it may only annotate a parameter whose declared type is "string",
//     "integer", or "boolean" -- notably NOT "number", which SEP-2243 excludes
//     because a float has no canonical wire spelling.
//
// A missing or non-string "type" on an annotated parameter is treated as a
// violation. SEP-2243 permits the annotation only on primitive parameters, and
// an undeclared type is not a declared primitive; equally, mirroring cannot
// serialize a value whose type it does not know. This is the strict reading:
// it rejects the tool rather than guessing a spelling for the header value.
//
// Traversal covers "properties", array "items", and the "oneOf"/"anyOf"/"allOf"
// combinators, so an annotation is found wherever SEP-2243 allows one. Nesting
// past maxSchemaDepth yields ErrSchemaTooDeep.
func ParamHeaders(schema map[string]any) ([]ParamHeader, error) {
        if len(schema) == 0 {
                return nil, nil
        }
        // seen maps the case-folded annotation value to the path that first claimed
        // it, so a duplicate can name both colliding parameters in its error.
        seen := map[string]string{}
        var found []ParamHeader
        if err := walkSchema(schema, nil, 0, seen, &found); err != nil {
                return nil, err
        }
        if len(found) == 0 {
                return nil, nil
        }
        // Deterministic order so callers (and tests) see a stable sequence
        // regardless of Go's map iteration order.
        slices.SortFunc(found, func(a, b ParamHeader) int {
                return strings.Compare(strings.Join(a.Path, "."), strings.Join(b.Path, "."))
        })
        return found, nil
}

// ValidateParamHeaders reports whether a tool's inputSchema declares only
// SEP-2243-conformant x-mcp-header annotations. It is ParamHeaders with the
// annotations discarded, for the ingestion-time check where only the verdict
// matters.
func ValidateParamHeaders(schema map[string]any) error {
        _, err := ParamHeaders(schema)
        return err
}

// walkSchema recurses one schema node, appending any annotation it finds to
// found. path is the property path to this node ("" at the root, which cannot
// itself be annotated -- an annotation designates a parameter, and the root is
// the parameter object).
func walkSchema(node map[string]any, path []string, depth int, seen map[string]string, found *[]ParamHeader) error {
        if depth > maxSchemaDepth {
                return fmt.Errorf("%w (%d)", ErrSchemaTooDeep, maxSchemaDepth)
        }

        // An annotation on the root node designates no parameter, so it is ignored
        // rather than treated as an error: len(path) == 0 only at the root.
        if len(path) > 0 {
                if raw, ok := node[XMCPHeaderAnnotation]; ok {
                        hdr, err := parseAnnotation(raw, node, path, seen)
                        if err != nil {
                                return err
                        }
                        *found = append(*found, hdr)
                }
        }

        return walkChildren(node, path, depth, seen, found)
}

// walkChildren recurses into every sub-schema of node that SEP-2243 allows an
// annotation to appear in: object properties, array element schemas, and the
// oneOf/anyOf/allOf combinator branches. Split from walkSchema to keep each
// within the cyclomatic limit.
func walkChildren(node map[string]any, path []string, depth int, seen map[string]string, found *[]ParamHeader) error {
        if props, ok := node["properties"].(map[string]any); ok {
                for name, sub := range props {
                        subSchema, ok := sub.(map[string]any)
                        if !ok {
                                continue // not a schema object; nothing to inspect
                        }
                        if err := walkSchema(subSchema, childPath(path, name), depth+1, seen, found); err != nil {
                                return err
                        }
                }
        }

        // Array element schemas: "items" is a schema object in the JSON Schema
        // dialect MCP uses. The element carries no property name of its own, so it
        // inherits the array's path with an "[]" marker for legibility in errors.
        if items, ok := node["items"].(map[string]any); ok {
                if err := walkSchema(items, childPath(path, "[]"), depth+1, seen, found); err != nil {
                        return err
                }
        }

        return walkCombinators(node, path, depth, seen, found)
}

// walkCombinators recurses into the oneOf/anyOf/allOf branches of node. These
// carry real schemas in this repo's backend tool sets (#5976 fixed ingestion
// dropping them), so an annotation inside a branch must be validated like any
// other.
func walkCombinators(
        node map[string]any, path []string, depth int, seen map[string]string, found *[]ParamHeader,
) error {
        for _, combinator := range []string{"oneOf", "anyOf", "allOf"} {
                branches, ok := node[combinator].([]any)
                if !ok {
                        continue
                }
                for i, branch := range branches {
                        branchSchema, ok := branch.(map[string]any)
                        if !ok {
                                continue
                        }
                        marker := fmt.Sprintf("%s[%d]", combinator, i)
                        if err := walkSchema(branchSchema, childPath(path, marker), depth+1, seen, found); err != nil {
                                return err
                        }
                }
        }
        return nil
}

// maxSafeInteger is JavaScript's Number.MAX_SAFE_INTEGER (2^53 - 1). SEP-2243
// requires a mirrored integer to sit within this range, since a peer that parses
// the header with JSON-number semantics could not round-trip a larger value.
const maxSafeInteger = 1<<53 - 1

// ErrUnmirrorableValue is returned when a designated parameter's value cannot be
// mirrored into a header: a control character (header injection), a non-integral
// or out-of-safe-range integer, or a value whose type contradicts the schema.
var ErrUnmirrorableValue = errors.New("parameter value cannot be mirrored into an HTTP header")

// MirrorParamHeaders derives the Mcp-Param-* headers to send with a tools/call,
// given the tool's x-mcp-header annotations (from ParamHeaders) and the call's
// arguments. The result maps full header names to values, ready to set on the
// request; it is nil when nothing is to be mirrored.
//
// A designated parameter that is absent from args contributes no header. That is
// deliberate rather than an error: an optional parameter the caller did not
// supply has no value to mirror, and SEP-2243's -32020 covers the server's view
// of a genuinely missing designated value. A provided empty string is instead
// sent as the base64 sentinel so it remains distinguishable from an absent
// header.
//
// Annotations reached through an array element ("[]" in the path) are skipped: an
// array holds many elements and a header holds one value, so there is no
// well-defined single value to send. Combinator segments (oneOf/anyOf/allOf) are
// dropped when resolving, because they are structural to the schema and absent
// from the arguments the schema describes.
//
// It returns an error wrapping ErrUnmirrorableValue when a present value cannot
// be safely rendered. Values originate from the caller (ultimately a model), so
// they are untrusted: a CR, LF, or NUL in a string would let a caller forge
// additional headers on vMCP's outgoing request, and is refused rather than
// silently stripped.
func MirrorParamHeaders(headers []ParamHeader, args map[string]any) (map[string]string, error) {
        if len(headers) == 0 || len(args) == 0 {
                return nil, nil
        }
        var out map[string]string
        for _, h := range headers {
                if slices.Contains(h.Path, "[]") {
                        continue
                }
                value, ok := resolveArg(args, h.Path)
                if !ok {
                        continue
                }
                rendered, err := renderHeaderValue(h, value)
                if err != nil {
                        return nil, err
                }
                if out == nil {
                        out = map[string]string{}
                }
                out[h.HeaderName()] = rendered
        }
        return out, nil
}

// ParamHeadersForSchema is ParamHeaders followed by MirrorParamHeaders: it takes
// a tool's inputSchema and a call's arguments and returns the Mcp-Param-* headers
// to send. It exists so the several call sites that mirror headers share one
// reading of the two-step dance rather than each re-deriving it.
//
// A schema error is returned unwrapped from ParamHeaders; a value error wraps
// ErrUnmirrorableValue. Callers distinguish them because the first indicts the
// backend's tool definition and the second the caller's arguments.
func ParamHeadersForSchema(schema map[string]any, args map[string]any) (map[string]string, error) {
        annotations, err := ParamHeaders(schema)
        if err != nil {
                return nil, err
        }
        if len(annotations) == 0 {
                return nil, nil
        }
        return MirrorParamHeaders(annotations, args)
}

// resolveArg walks path through args and returns the value at its end. Segments
// naming a combinator are skipped: they exist in the schema's structure, not in
// the data it describes. A missing or non-object intermediate reports false.
func resolveArg(args map[string]any, path []string) (any, bool) {
        current := any(args)
        for _, seg := range path {
                if isCombinatorSegment(seg) {
                        continue
                }
                obj, ok := current.(map[string]any)
                if !ok {
                        return nil, false
                }
                current, ok = obj[seg]
                if !ok {
                        return nil, false
                }
        }
        return current, true
}

// isCombinatorSegment reports whether a path segment is a combinator marker
// emitted by walkCombinators (e.g. "oneOf[0]") rather than a property name.
func isCombinatorSegment(seg string) bool {
        for _, combinator := range []string{"oneOf", "anyOf", "allOf"} {
                if strings.HasPrefix(seg, combinator+"[") && strings.HasSuffix(seg, "]") {
                        return true
                }
        }
        return false
}

// renderHeaderValue converts a designated parameter's value to its header
// spelling, enforcing the schema's declared type and SEP-2243's integer range.
// A present empty string is sentinel-encoded so it remains distinguishable from
// an absent mirrored parameter; absent values are omitted by MirrorParamHeaders.
func renderHeaderValue(h ParamHeader, value any) (string, error) {
        where := strings.Join(h.Path, ".")
        switch h.Type {
        case "string":
                s, ok := value.(string)
                if !ok {
                        return "", fmt.Errorf("%w: parameter %q is declared string but got %T", ErrUnmirrorableValue, where, value)
                }
                if bad, invalid := firstControlChar(s); invalid {
                        return "", fmt.Errorf(
                                "%w: parameter %q contains a control character (%q)", ErrUnmirrorableValue, where, bad)
                }
                if s == "" {
                        return EncodeSentinelName(s), nil
                }
                return s, nil
        case "boolean":
                b, ok := value.(bool)
                if !ok {
                        return "", fmt.Errorf("%w: parameter %q is declared boolean but got %T", ErrUnmirrorableValue, where, value)
                }
                return strconv.FormatBool(b), nil
        case "integer":
                return renderIntegerHeaderValue(where, value)
        default:
                // Unreachable: ParamHeaders admits only the three types above.
                return "", fmt.Errorf("%w: parameter %q has unsupported type %q", ErrUnmirrorableValue, where, h.Type)
        }
}

// renderIntegerHeaderValue renders an integer-declared parameter. JSON decoding
// yields float64 for every number, so an integer arrives as a float that must be
// checked for integrality and for SEP-2243's safe-integer range; the int/int64
// cases cover arguments built in Go rather than decoded from JSON.
func renderIntegerHeaderValue(where string, value any) (string, error) {
        switch n := value.(type) {
        case float64:
                if n != math.Trunc(n) {
                        return "", fmt.Errorf("%w: parameter %q is declared integer but got %v", ErrUnmirrorableValue, where, n)
                }
                if n > maxSafeInteger || n < -maxSafeInteger {
                        return "", fmt.Errorf(
                                "%w: parameter %q value %v is outside the safe integer range", ErrUnmirrorableValue, where, n)
                }
                return strconv.FormatInt(int64(n), 10), nil
        case int:
                return renderIntegerHeaderValue(where, float64(n))
        case int64:
                return renderIntegerHeaderValue(where, float64(n))
        case json.Number:
                i, err := n.Int64()
                if err != nil {
                        return "", fmt.Errorf("%w: parameter %q is declared integer but got %q", ErrUnmirrorableValue, where, n)
                }
                return renderIntegerHeaderValue(where, float64(i))
        default:
                return "", fmt.Errorf(
                        "%w: parameter %q is declared integer but got %T", ErrUnmirrorableValue, where, value)
        }
}

// firstControlChar returns the first control character in s, reporting true when
// one exists. CR and LF are the header-injection vectors; NUL and the other C0
// controls are equally illegal in a header value.
func firstControlChar(s string) (string, bool) {
        for _, r := range s {
                if r < 0x20 || r == 0x7f {
                        return string(r), true
                }
        }
        return "", false
}

// childPath returns path with seg appended, always in freshly allocated storage.
// Appending to path directly would let sibling recursions share (and overwrite)
// one backing array, so a path captured deeper in the walk could be rewritten by
// the next sibling. Allocating per child keeps every captured path independent.
func childPath(path []string, seg string) []string {
        child := make([]string, len(path)+1)
        copy(child, path)
        child[len(path)] = seg
        return child
}

// parseAnnotation validates a single x-mcp-header annotation against SEP-2243
// and records it in seen for the uniqueness check.
func parseAnnotation(raw any, node map[string]any, path []string, seen map[string]string) (ParamHeader, error) {
        where := strings.Join(path, ".")

        name, ok := raw.(string)
        if !ok {
                return ParamHeader{}, fmt.Errorf(
                        "parameter %q: %s must be a string, got %T", where, XMCPHeaderAnnotation, raw)
        }
        if name == "" {
                return ParamHeader{}, fmt.Errorf("parameter %q: %s must not be empty", where, XMCPHeaderAnnotation)
        }
        if bad, invalid := firstNonTokenChar(name); invalid {
                return ParamHeader{}, fmt.Errorf(
                        "parameter %q: %s value %q is not a valid HTTP field-name token (offending character %q)",
                        where, XMCPHeaderAnnotation, name, bad)
        }

        // HTTP field names are case-insensitive, so two annotations differing only in
        // case would mirror onto the same header and one would silently win.
        folded := strings.ToLower(name)
        if first, dup := seen[folded]; dup {
                return ParamHeader{}, fmt.Errorf(
                        "%s value %q on parameter %q collides case-insensitively with parameter %q",
                        XMCPHeaderAnnotation, name, where, first)
        }
        seen[folded] = where

        declared, ok := node["type"].(string)
        if !ok {
                return ParamHeader{}, fmt.Errorf(
                        "parameter %q: %s requires a declared primitive type (string, integer, or boolean)",
                        where, XMCPHeaderAnnotation)
        }
        switch declared {
        case "string", "integer", "boolean":
        default:
                return ParamHeader{}, fmt.Errorf(
                        "parameter %q: %s is not permitted on type %q (only string, integer, or boolean)",
                        where, XMCPHeaderAnnotation, declared)
        }

        // childPath already gives each branch its own storage, so cloning is belt and
        // braces -- it keeps the returned Path independent of the walk even if the
        // traversal's path handling is later changed.
        return ParamHeader{Path: slices.Clone(path), Name: name, Type: declared}, nil
}

// firstNonTokenChar returns the first byte of s that is not an RFC 9110 tchar,
// reporting true when one exists. Operating on bytes rather than runes is
// correct here: every tchar is ASCII, so any multi-byte rune is invalid and its
// leading byte is a faithful thing to name in the error.
func firstNonTokenChar(s string) (string, bool) {
        for i := 0; i < len(s); i++ {
                if !isTokenChar(s[i]) {
                        return string(s[i]), true
                }
        }
        return "", false
}

// isTokenChar reports whether c is an RFC 9110 tchar, the character set HTTP
// field names are drawn from.
func isTokenChar(c byte) bool {
        switch {
        case c >= 'a' && c <= 'z',
                c >= 'A' && c <= 'Z',
                c >= '0' && c <= '9':
                return true
        }
        switch c {
        case '!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', '^', '_', '`', '|', '~':
                return true
        }
        return false
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package networking

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "mime"
        "net/http"
        "net/url"
        "strings"
)

const (
        // maxResponseSize is the maximum response body size (1MB).
        maxResponseSize = 1024 * 1024

        // contentTypeJSON is the JSON content type.
        contentTypeJSON = "application/json"

        // contentTypeFormURLEncoded is the form-urlencoded content type.
        contentTypeFormURLEncoded = "application/x-www-form-urlencoded"
)

// FetchResult contains the result of a successful JSON fetch operation.
type FetchResult[T any] struct {
        // Data is the parsed JSON response body.
        Data T

        // Headers are the response headers.
        Headers http.Header
}

// FetchOption configures a fetch request.
type FetchOption func(*fetchOptions)

// fetchOptions holds the configuration for a fetch request.
type fetchOptions struct {
        method          string
        headers         http.Header
        body            io.Reader
        errorHandler    func(*http.Response, []byte) error
        maxResponseSize int64
}

// newFetchOptions creates default fetch options.
func newFetchOptions() *fetchOptions {
        return &fetchOptions{
                method:          http.MethodGet,
                headers:         make(http.Header),
                maxResponseSize: maxResponseSize,
        }
}

// WithMethod sets the HTTP method for the request.
func WithMethod(method string) FetchOption {
        return func(opts *fetchOptions) {
                opts.method = method
        }
}

// WithHeader adds a single header to the request.
func WithHeader(key, value string) FetchOption {
        return func(opts *fetchOptions) {
                opts.headers.Set(key, value)
        }
}

// WithBody sets the request body.
func WithBody(body io.Reader) FetchOption {
        return func(opts *fetchOptions) {
                opts.body = body
        }
}

// WithMaxResponseSize sets a custom maximum response body size in bytes.
// The default is 1 MB. Use this to enforce tighter limits for endpoints that
// are expected to return small documents (e.g. OAuth metadata, CIMD documents).
func WithMaxResponseSize(size int64) FetchOption {
        return func(opts *fetchOptions) {
                opts.maxResponseSize = size
        }
}

// WithErrorHandler sets a custom error handler for non-200 responses.
// The handler receives the response and body, and should return an error.
// If the handler returns nil, the default HTTPError will be returned.
// This is useful for parsing structured error responses (e.g., OAuth error responses).
func WithErrorHandler(handler func(*http.Response, []byte) error) FetchOption {
        return func(opts *fetchOptions) {
                opts.errorHandler = handler
        }
}

// FetchJSON performs an HTTP request and parses the JSON response body.
// It sets the Accept header to application/json by default.
// For non-200 responses, it returns an HTTPError or the result of a custom error handler.
func FetchJSON[T any](
        ctx context.Context,
        client HTTPClient,
        requestURL string,
        opts ...FetchOption,
) (*FetchResult[T], error) {
        options := newFetchOptions()
        for _, opt := range opts {
                opt(options)
        }

        // Set default Accept header if not already set
        if options.headers.Get("Accept") == "" {
                options.headers.Set("Accept", contentTypeJSON)
        }

        req, err := http.NewRequestWithContext(ctx, options.method, requestURL, options.body)
        if err != nil {
                return nil, fmt.Errorf("failed to create request: %w", err)
        }

        // Apply headers
        for key, values := range options.headers {
                for _, value := range values {
                        req.Header.Add(key, value)
                }
        }

        resp, err := client.Do(req)
        if err != nil {
                return nil, fmt.Errorf("request failed: %w", err)
        }
        defer func() { _ = resp.Body.Close() }()

        // Read body with size limit
        body, err := io.ReadAll(io.LimitReader(resp.Body, options.maxResponseSize))
        if err != nil {
                return nil, fmt.Errorf("failed to read response body: %w", err)
        }

        // Handle non-200 responses
        if resp.StatusCode != http.StatusOK {
                // Try custom error handler first
                if options.errorHandler != nil {
                        if customErr := options.errorHandler(resp, body); customErr != nil {
                                return nil, customErr
                        }
                }

                // Fall back to default HTTPError using status text to avoid leaking sensitive body content
                return nil, NewHTTPError(resp.StatusCode, requestURL, resp.Status)
        }

        // Validate Content-Type for successful responses. Accept application/json
        // and application/*+json subtypes (RFC 6839) — the old strings.Contains check
        // incorrectly rejected valid subtypes like application/ld+json.
        contentType := resp.Header.Get("Content-Type")
        mediaType, _, _ := mime.ParseMediaType(contentType)
        if mediaType != contentTypeJSON && (!strings.HasPrefix(mediaType, "application/") || !strings.HasSuffix(mediaType, "+json")) {
                return nil, fmt.Errorf("unexpected content type: %s", contentType)
        }

        // Parse JSON response
        var data T
        if err := json.Unmarshal(body, &data); err != nil {
                return nil, fmt.Errorf("failed to parse JSON response: %w", err)
        }

        return &FetchResult[T]{
                Data:    data,
                Headers: resp.Header,
        }, nil
}

// FetchJSONWithForm performs a POST request with form-urlencoded body and parses JSON response.
// This is a convenience wrapper around FetchJSON for token endpoints and similar APIs.
// It sets Content-Type to application/x-www-form-urlencoded and Accept to application/json.
func FetchJSONWithForm[T any](
        ctx context.Context,
        client HTTPClient,
        requestURL string,
        formData url.Values,
        opts ...FetchOption,
) (*FetchResult[T], error) {
        // Prepend form-specific options
        formOpts := []FetchOption{
                WithMethod(http.MethodPost),
                WithHeader("Content-Type", contentTypeFormURLEncoded),
                WithBody(strings.NewReader(formData.Encode())),
        }

        // Append user options (they can override form options if needed)
        allOpts := append(formOpts, opts...)

        return FetchJSON[T](ctx, client, requestURL, allOpts...)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package networking

import (
        "context"
        "crypto/tls"
        "crypto/x509"
        "errors"
        "fmt"
        "net"
        "net/http"
        "net/url"
        "os"
        "strings"
        "syscall"
        "time"

        "golang.org/x/oauth2"
)

// HTTPClient is an interface for making HTTP requests.
// This interface is satisfied by *http.Client and allows for dependency injection in testing.
type HTTPClient interface {
        Do(req *http.Request) (*http.Response, error)
}

var privateIPBlocks []*net.IPNet

// HttpTimeout is the timeout for outgoing HTTP requests
const HttpTimeout = 30 * time.Second

// HttpsScheme is the HTTPS scheme
const HttpsScheme = "https"

// HttpScheme is the HTTP scheme
const HttpScheme = "http"

// MaxRedirects bounds how many HTTP redirects an SSRF-guarded client follows
// before giving up. Matches the cap used by the transparent proxy data path.
const MaxRedirects = 10

// Connection-pool bounds applied by Build.
const (
        // idleConnTimeout bounds how long a pooled keep-alive connection is
        // retained after its last use. This is the leak: http.Transport's zero
        // value means "never expire", so a client that performs one request and is
        // then dropped pins a socket and its readLoop/writeLoop goroutine pair for
        // the lifetime of the process. 90s matches http.DefaultTransport.
        idleConnTimeout = 90 * time.Second
        // maxIdleConns caps pooled idle connections across all hosts. The zero
        // value is unlimited. 100 matches http.DefaultTransport.
        maxIdleConns = 100
        // maxIdleConnsPerHost caps them per host. The zero value here is already
        // bounded (http.DefaultMaxIdleConnsPerHost is 2), so this is a deliberate
        // increase rather than a new bound: these clients are host-scoped, and a
        // slightly larger pool serves concurrent requests without dialing now that
        // the idle timeout is finite.
        maxIdleConnsPerHost = 4
)

// ErrRedirectRefused is wrapped by SameHostRedirectPolicy when it declines to
// follow a redirect, so callers can match it with errors.Is.
var ErrRedirectRefused = errors.New("redirect refused")

// SameHostRedirectPolicy returns a value for http.Client.CheckRedirect that
// follows only same-host redirects, refuses HTTPS-to-HTTP downgrades, and caps
// the chain at MaxRedirects.
//
// Any client that fetches a URL derived from an untrusted remote server — auth
// discovery probes, RFC 9728 resource-metadata fetches, OIDC issuer discovery —
// must install this. Validating only the originally-supplied URL is not enough:
// a malicious server can return a 30x that points the request at an internal
// address (cloud IMDS, RFC1918 services), and the host-side client would follow
// it (CWE-918). Restricting redirects to the same host as the original request
// keeps the request on the endpoint the operator actually configured.
//
// This mirrors the data-plane guard in pkg/transport/proxy/transparent
// (followRedirects); keep the two policies in sync.
func SameHostRedirectPolicy() func(req *http.Request, via []*http.Request) error {
        return func(req *http.Request, via []*http.Request) error {
                if len(via) >= MaxRedirects {
                        return fmt.Errorf("stopped after %d redirects: %w", MaxRedirects, ErrRedirectRefused)
                }
                // via[0] is the original request; CheckRedirect is only invoked once at
                // least one redirect has occurred, so via is never empty.
                original := via[0]
                // Compare host:port (not just hostname): a redirect to a different port
                // on the same host can reach a different internal service (a co-located
                // metadata/admin port), so it must be treated as cross-host.
                if !strings.EqualFold(req.URL.Host, original.URL.Host) {
                        return fmt.Errorf("refusing cross-host redirect to %q (original host %q): %w",
                                req.URL.Host, original.URL.Host, ErrRedirectRefused)
                }
                if original.URL.Scheme == HttpsScheme && req.URL.Scheme != HttpsScheme {
                        return fmt.Errorf("refusing redirect that downgrades from HTTPS to %q: %w",
                                req.URL.Scheme, ErrRedirectRefused)
                }
                return nil
        }
}

// Dialer control function for validating addresses prior to connection
func protectedDialerControl(_, address string, _ syscall.RawConn) error {
        err := AddressReferencesPrivateIp(address)
        if err != nil {
                return err
        }

        return nil
}

// NewPrivateIPBlockingDialContext returns a DialContext that refuses to connect
// to private, loopback, or link-local addresses. The check runs after DNS
// resolution on the address actually being dialed, so it also defends against
// DNS rebinding and is re-applied on every redirect hop. Pair it with
// Transport.DisableKeepAlives so a pooled connection cannot skip the check on a
// later request.
//
// Use this on clients that fetch a URL derived from untrusted input when the
// operator-configured target is public; SameHostRedirectPolicy is the
// redirect-following counterpart.
func NewPrivateIPBlockingDialContext() func(ctx context.Context, network, addr string) (net.Conn, error) {
        return (&net.Dialer{Control: protectedDialerControl}).DialContext
}

// Dial timeouts applied to backend connections. Both match the Go standard
// library's http.DefaultTransport dialer.
const (
        backendDialTimeout   = 30 * time.Second
        backendDialKeepAlive = 30 * time.Second
)

// CloneDefaultTransportWithDialControl returns an *http.Transport equivalent to
// http.DefaultTransport, optionally carrying a per-connection dial Control hook.
//
// When http.DefaultTransport is the standard *http.Transport it is cloned
// (preserving proxy, HTTP/2, and idle-connection settings). If it has been
// replaced (e.g. in tests) a transport with the Go standard-library defaults is
// reconstructed instead, so proxy/timeout/HTTP2 settings are not silently
// dropped.
//
// When control is non-nil, a net.Dialer carrying it — with the standard backend
// dial timeout and keep-alive — is installed as DialContext. The hook fires on
// the resolved peer IP before the TCP handshake, which is what lets callers
// defeat DNS-rebinding: a name-based check cannot, because the name can resolve
// to a blocked IP after the check passes. A nil control leaves the cloned
// dialer untouched.
//
// This is the single construction point for the backend transport used by every
// caller that needs a dial-control hook, so callers building similar transports
// cannot drift apart. Currently used by pkg/vmcp/client and
// pkg/vmcp/session/internal/backend.
func CloneDefaultTransportWithDialControl(
        control func(network, address string, c syscall.RawConn) error,
) *http.Transport {
        var t *http.Transport
        if dt, ok := http.DefaultTransport.(*http.Transport); ok {
                t = dt.Clone()
        } else {
                t = &http.Transport{
                        Proxy:                 http.ProxyFromEnvironment,
                        DialContext:           (&net.Dialer{Timeout: backendDialTimeout, KeepAlive: backendDialKeepAlive}).DialContext,
                        ForceAttemptHTTP2:     true,
                        MaxIdleConns:          maxIdleConns,
                        IdleConnTimeout:       idleConnTimeout,
                        TLSHandshakeTimeout:   10 * time.Second,
                        ExpectContinueTimeout: 1 * time.Second,
                }
        }
        if control != nil {
                t.DialContext = (&net.Dialer{
                        Timeout:   backendDialTimeout,
                        KeepAlive: backendDialKeepAlive,
                        Control:   control,
                }).DialContext
        }
        return t
}

// IdleConnectionCloser is the capability http.Client discovers on its outermost
// transport to drain pooled connections (it asserts an identical unexported
// interface). Any RoundTripper in this repo that wraps another must implement it
// and forward the call, or every CloseIdleConnections on the resulting client is
// a silent no-op. Assert against this named type rather than redeclaring an
// anonymous `interface{ CloseIdleConnections() }`.
//
// upstream.IdleConnectionCloser is the provider-level analogue one layer up.
type IdleConnectionCloser interface {
        CloseIdleConnections()
}

// SetIdleConnBounds applies Build's idle-connection-pool bounds to a
// hand-constructed *http.Transport — the transports this package cannot build
// through Build because they need their own DialContext, CheckRedirect, or TLS
// config. A zero IdleConnTimeout never expires a pooled connection, so a client
// that performs one request and is dropped pins a socket and its goroutine pair
// for the process lifetime; this is the leak Build fixes for its own clients
// (see the constant block above). Setting all three fields together keeps a
// call site from bounding two of the three by mistake.
//
// pkg/networking imports pkg/oauthproto, so leaf packages that oauthproto's
// import graph reaches cannot call this without a cycle; those mirror the
// values with a local const block instead.
func SetIdleConnBounds(t *http.Transport) {
        t.IdleConnTimeout = idleConnTimeout
        t.MaxIdleConns = maxIdleConns
        t.MaxIdleConnsPerHost = maxIdleConnsPerHost
}

// ForwardCloseIdle drains the idle-connection pool reachable through rt when rt
// implements IdleConnectionCloser — as *http.Transport, and every wrapping
// RoundTripper in this repo that forwards the call, do. A wrapping RoundTripper
// must call this from its own CloseIdleConnections; otherwise
// http.Client.CloseIdleConnections stops at the wrapper and the pool underneath
// is never drained (see IdleConnectionCloser). Using this helper keeps the
// wrapper half from silently regressing into an anonymous
// `interface{ CloseIdleConnections() }` assertion that drifts out of sync.
func ForwardCloseIdle(rt http.RoundTripper) {
        if closer, ok := rt.(IdleConnectionCloser); ok {
                closer.CloseIdleConnections()
        }
}

// ValidatingTransport is for validating URLs prior to request
type ValidatingTransport struct {
        Transport         http.RoundTripper
        InsecureAllowHTTP bool
}

// Compile-time assertion: a rename or typo of CloseIdleConnections would
// otherwise silently re-hide the pool it wraps (see IdleConnectionCloser).
var _ IdleConnectionCloser = (*ValidatingTransport)(nil)

// RoundTrip validates the request URL prior to forwarding
func (t *ValidatingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
        // Skip validation if INSECURE_DISABLE_URL_VALIDATION is set or if InsecureAllowHTTP is true
        if strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true") || t.InsecureAllowHTTP {
                return t.Transport.RoundTrip(req)
        }

        // Check for valid URL specification
        parsedUrl, err := url.Parse(req.URL.String())
        if err != nil {
                return nil, fmt.Errorf("the supplied URL %s is malformed", req.URL.String())
        }

        // Check for HTTPS scheme
        if parsedUrl.Scheme != HttpsScheme {
                return nil, fmt.Errorf("the supplied URL %s is not HTTPS scheme", req.URL.String())
        }

        return t.Transport.RoundTrip(req)
}

// CloseIdleConnections forwards to the wrapped transport so that
// http.Client.CloseIdleConnections reaches the real connection pool.
//
// Because Build always wraps the pool in a ValidatingTransport, omitting this
// method would make every client's CloseIdleConnections a silent no-op — see
// IdleConnectionCloser. The assertion is required because the field is an
// http.RoundTripper; it holds an *http.Transport for every client Build produces.
func (t *ValidatingTransport) CloseIdleConnections() {
        ForwardCloseIdle(t.Transport)
}

// closeIdlerTransport wraps a RoundTripper that does not implement
// CloseIdleConnections (oauth2.Transport) and forwards the call to the
// *http.Transport owning the connection pool underneath it.
type closeIdlerTransport struct {
        http.RoundTripper
        pool *http.Transport
}

var _ IdleConnectionCloser = (*closeIdlerTransport)(nil)

// CloseIdleConnections closes the idle connections held by the underlying pool.
func (t *closeIdlerTransport) CloseIdleConnections() {
        t.pool.CloseIdleConnections()
}

// createTokenSourceFromFile creates an oauth2.TokenSource from a token file
func createTokenSourceFromFile(tokenFile string) (oauth2.TokenSource, error) {
        tokenBytes, err := os.ReadFile(tokenFile) // #nosec G304 - tokenFile path is provided by user via CLI flag
        if err != nil {
                return nil, fmt.Errorf("failed to read auth token file: %w", err)
        }

        // Remove any trailing newlines/whitespace
        tokenStr := strings.TrimSpace(string(tokenBytes))
        if tokenStr == "" {
                return nil, fmt.Errorf("auth token file is empty")
        }

        // Create a static token source
        token := &oauth2.Token{
                AccessToken: tokenStr,
                TokenType:   "Bearer",
        }

        return oauth2.StaticTokenSource(token), nil
}

// HttpClientBuilder provides a fluent interface for building HTTP clients
type HttpClientBuilder struct {
        clientTimeout         time.Duration
        tlsHandshakeTimeout   time.Duration
        responseHeaderTimeout time.Duration
        caCertPath            string
        caCertUsesSystemRoots bool
        authTokenFile         string
        allowPrivate          bool
        insecureAllowHTTP     bool
        disableKeepAlives     bool
}

// NewHttpClientBuilder returns a new HttpClientBuilder
func NewHttpClientBuilder() *HttpClientBuilder {
        return &HttpClientBuilder{
                clientTimeout:         HttpTimeout,
                tlsHandshakeTimeout:   10 * time.Second,
                responseHeaderTimeout: 10 * time.Second,
        }
}

// NewHostScopedClientBuilder returns an HttpClientBuilder pre-configured with
// the SSRF-guard policy (CWE-918) appropriate for a host explicitly configured
// by an operator. By default the returned builder blocks plain HTTP and
// connections to private/loopback/link-local IP ranges. Both gates are relaxed
// automatically for loopback hosts (development/testing) and when
// INSECURE_DISABLE_URL_VALIDATION is set.
//
// allowPrivateIPs widens only the private-IP gate — for example an in-cluster
// provider reachable solely over an RFC-1918 address — without enabling plain
// HTTP for non-loopback hosts. insecureAllowHTTP additionally permits
// plain-HTTP for non-loopback hosts and must never be set in production.
//
// The returned builder is not yet built: callers may chain further options
// (e.g. WithTimeout, WithDisableKeepAlives) before calling Build. This builder
// carries the policy for operator-configured hosts only; a host named by a
// remote server or by a document it serves must use
// NewServerSuppliedHostClientBuilder instead.
func NewHostScopedClientBuilder(host string, allowPrivateIPs, insecureAllowHTTP bool) *HttpClientBuilder {
        allowInsecure := IsLocalhost(host) ||
                insecureAllowHTTP ||
                strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true")
        return NewHttpClientBuilder().
                WithInsecureAllowHTTP(allowInsecure).
                WithPrivateIPs(allowInsecure || allowPrivateIPs)
}

// NewServerSuppliedHostClientBuilder returns an HttpClientBuilder for a host
// supplied by a remote server or its metadata. It permits HTTP for localhost,
// explicit insecure configuration, or INSECURE_DISABLE_URL_VALIDATION to
// preserve compatibility, but private/loopback/link-local dialing is enabled
// only by allowPrivateIPs. In particular, localhost and the environment
// override never widen the private-IP gate for server-supplied endpoints:
// INSECURE_DISABLE_URL_VALIDATION deliberately relaxes only the HTTPS
// requirement here, because a remote server that can name an endpoint must not
// be able to reach an internal address just because a developer set that
// variable.
func NewServerSuppliedHostClientBuilder(host string, allowPrivateIPs, insecureAllowHTTP bool) *HttpClientBuilder {
        allowInsecure := IsLocalhost(host) ||
                insecureAllowHTTP ||
                strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true")
        return NewHttpClientBuilder().
                WithInsecureAllowHTTP(allowInsecure).
                WithPrivateIPs(allowPrivateIPs)
}

// WithCABundle sets a pinned CA certificate bundle path. When a bundle is
// configured, only certificates from that bundle are trusted (system roots are
// not included).
func (b *HttpClientBuilder) WithCABundle(path string) *HttpClientBuilder {
        b.caCertPath = path
        b.caCertUsesSystemRoots = false
        return b
}

// WithSystemRootsPlusCABundle sets a CA certificate bundle path and preserves
// trust in the system root pool. Use this when an upstream may use either a
// publicly trusted certificate or a private CA.
func (b *HttpClientBuilder) WithSystemRootsPlusCABundle(path string) *HttpClientBuilder {
        b.caCertPath = path
        b.caCertUsesSystemRoots = true
        return b
}

// WithTokenFromFile sets the auth token file path
func (b *HttpClientBuilder) WithTokenFromFile(path string) *HttpClientBuilder {
        b.authTokenFile = path
        return b
}

// WithPrivateIPs allows connections to private IP addresses
func (b *HttpClientBuilder) WithPrivateIPs(allow bool) *HttpClientBuilder {
        b.allowPrivate = allow
        return b
}

// WithInsecureAllowHTTP allows HTTP (non-HTTPS) URLs
// WARNING: This is insecure and should NEVER be used in production
func (b *HttpClientBuilder) WithInsecureAllowHTTP(allow bool) *HttpClientBuilder {
        b.insecureAllowHTTP = allow
        return b
}

// WithDisableKeepAlives disables HTTP keep-alive on the transport. When true,
// each request uses a fresh connection, ensuring the per-dial SSRF check fires
// on every request rather than being bypassed by a reused connection.
func (b *HttpClientBuilder) WithDisableKeepAlives(disable bool) *HttpClientBuilder {
        b.disableKeepAlives = disable
        return b
}

// WithTimeout sets the HTTP client timeout
func (b *HttpClientBuilder) WithTimeout(timeout time.Duration) *HttpClientBuilder {
        b.clientTimeout = timeout
        return b
}

// WithResponseHeaderTimeout sets how long the transport waits for a server's
// response headers after writing the request. Zero means no limit, as for
// http.Transport.ResponseHeaderTimeout; the client timeout set by WithTimeout
// still bounds the whole request.
func (b *HttpClientBuilder) WithResponseHeaderTimeout(timeout time.Duration) *HttpClientBuilder {
        b.responseHeaderTimeout = timeout
        return b
}

// Build creates the configured HTTP client
func (b *HttpClientBuilder) Build() (*http.Client, error) {
        transport := &http.Transport{
                TLSHandshakeTimeout:   b.tlsHandshakeTimeout,
                ResponseHeaderTimeout: b.responseHeaderTimeout,
        }
        SetIdleConnBounds(transport)
        transport.DisableKeepAlives = b.disableKeepAlives

        if !b.allowPrivate {
                transport.DialContext = (&net.Dialer{
                        Control: protectedDialerControl,
                }).DialContext
        }

        if b.caCertPath != "" {
                caCert, err := os.ReadFile(b.caCertPath)
                if err != nil {
                        return nil, fmt.Errorf("failed to read CA certificate bundle: %w", err)
                }

                caCertPool := x509.NewCertPool()
                if b.caCertUsesSystemRoots {
                        caCertPool, err = x509.SystemCertPool()
                        if err != nil {
                                return nil, fmt.Errorf("failed to load system CA certificate pool: %w", err)
                        }
                        if caCertPool == nil {
                                return nil, fmt.Errorf("failed to load system CA certificate pool: pool is nil")
                        }
                }

                if !caCertPool.AppendCertsFromPEM(caCert) {
                        return nil, fmt.Errorf("failed to parse CA certificate bundle")
                }

                if transport.TLSClientConfig == nil {
                        transport.TLSClientConfig = &tls.Config{
                                MinVersion: tls.VersionTLS12,
                        }
                }
                transport.TLSClientConfig.RootCAs = caCertPool
        }

        // Start with validation transport
        var clientTransport http.RoundTripper = &ValidatingTransport{
                Transport:         transport,
                InsecureAllowHTTP: b.insecureAllowHTTP,
        }

        // Add auth transport if token file is provided using oauth2.Transport
        if b.authTokenFile != "" {
                tokenSource, err := createTokenSourceFromFile(b.authTokenFile)
                if err != nil {
                        return nil, fmt.Errorf("failed to create token source: %w", err)
                }

                // oauth2.Transport wraps our existing transport and adds Bearer token
                // authentication. It exposes only RoundTrip/CancelRequest, so wrap it
                // again to keep http.Client.CloseIdleConnections reaching the pool.
                clientTransport = &closeIdlerTransport{
                        RoundTripper: &oauth2.Transport{
                                Source: tokenSource,
                                Base:   clientTransport, // Preserves our ValidatingTransport
                        },
                        pool: transport,
                }
        }

        client := &http.Client{
                Transport: clientTransport,
                Timeout:   b.clientTimeout,
        }

        return client, nil
}

// Copyright 2025 Stacklok, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package networking

import (
        "errors"
        "fmt"
)

// HTTPError represents an HTTP error response with status code, URL, and message.
type HTTPError struct {
        // StatusCode is the HTTP status code.
        StatusCode int

        // Message is a description of the error (may be a preview of the response body).
        Message string

        // URL is the requested URL.
        URL string
}

// Error implements the error interface.
func (e *HTTPError) Error() string {
        return fmt.Sprintf("HTTP %d for URL %s: %s", e.StatusCode, e.URL, e.Message)
}

// NewHTTPError creates a new HTTP error.
func NewHTTPError(statusCode int, url, message string) error {
        return &HTTPError{
                StatusCode: statusCode,
                URL:        url,
                Message:    message,
        }
}

// IsHTTPError checks if an error is an HTTPError with the specified status code.
// If statusCode is 0, it matches any HTTPError.
func IsHTTPError(err error, statusCode int) bool {
        var httpErr *HTTPError
        if !errors.As(err, &httpErr) {
                return false
        }
        if statusCode == 0 {
                return true
        }
        return httpErr.StatusCode == statusCode
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package networking provides utilities for network operations,
// such as finding available ports and checking network connectivity.
package networking

import (
        "crypto/rand"
        "fmt"
        "log/slog"
        "math/big"
        "net"
        "strconv"
        "strings"

        gopsutilnet "github.com/shirou/gopsutil/v4/net"
)

const (
        // MinPort is the minimum port number to use
        MinPort = 10000
        // MaxPort is the maximum port number to use
        MaxPort = 65535
        // MaxAttempts is the maximum number of attempts to find an available port
        MaxAttempts = 10
)

// CallbackPortInUseError reports that a requested OAuth callback port is unavailable.
type CallbackPortInUseError struct {
        Port int
}

func (e *CallbackPortInUseError) Error() string {
        return fmt.Sprintf("OAuth callback port %d is already in use", e.Port)
}

// IsAvailable checks if a port is available
func IsAvailable(port int) bool {
        // Check TCP
        tcpAddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("127.0.0.1:%d", port))
        if err != nil {
                return false
        }

        tcpListener, err := net.ListenTCP("tcp", tcpAddr)
        if err != nil {
                return false
        }
        if err := tcpListener.Close(); err != nil {
                // Log the error but continue, as we're just checking if the port is available
                slog.Warn("Failed to close TCP listener", "error", err)
        }

        // Check UDP
        udpAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("127.0.0.1:%d", port))
        if err != nil {
                return false
        }

        udpConn, err := net.ListenUDP("udp", udpAddr)
        if err != nil {
                return false
        }
        if err := udpConn.Close(); err != nil {
                // Log the error but continue, as we're just checking if the port is available
                slog.Warn("Failed to close UDP connection", "error", err)
        }

        return true
}

// FindAvailable finds an available port
func FindAvailable() int {
        for i := 0; i < MaxAttempts; i++ {
                // Generate a cryptographically secure random number
                n, err := rand.Int(rand.Reader, big.NewInt(int64(MaxPort-MinPort)))
                if err != nil {
                        // Fall back to sequential search if random generation fails
                        break
                }
                port := int(n.Int64()) + MinPort
                if IsAvailable(port) {
                        return port
                }
        }

        // If we can't find a random port, try sequential ports
        for port := MinPort; port <= MaxPort; port++ {
                if IsAvailable(port) {
                        return port
                }
        }

        // If we still can't find a port, return 0
        return 0
}

// FindOrUsePort checks if the provided port is available or finds an available port if none is provided.
// If port is 0, it will find an available port.
// If port is not 0, it will check if the port is available.
// Returns the selected port and an error if any.
func FindOrUsePort(port int) (int, error) {
        if port == 0 {
                // Find an available port
                port = FindAvailable()
                if port == 0 {
                        return 0, fmt.Errorf("could not find an available port")
                }
                return port, nil
        }

        if IsAvailable(port) {
                return port, nil
        }

        // Requested port is busy — find an alternative
        alt := FindAvailable()
        if alt == 0 {
                return 0, fmt.Errorf("failed to find an alternative port after requested port %d was unavailable", port)
        }
        return alt, nil
}

// ValidateCallbackPort validates that the specified callback port is valid and available.
// It checks that the port is within the valid range (1-65535) and, for pre-registered
// clients (with clientID), it returns an error if the port is not available.
func ValidateCallbackPort(callbackPort int, clientID string) error {
        // If port is 0, we'll find an available port later, so no need to validate
        if callbackPort == 0 {
                return nil
        }

        // Validate port range
        if callbackPort < 1024 || callbackPort > 65535 {
                return fmt.Errorf("OAuth callback port must be between 1024 and 65535, got: %d", callbackPort)
        }

        // Check if this is a pre-registered client (has client credentials)
        // For pre-registered clients, we need strict port checking
        isPreRegisteredClient := IsPreRegisteredClient(clientID)

        if isPreRegisteredClient {
                // For pre-registered clients, the port must be available
                // The user likely configured this port in their IdP/app
                if !IsAvailable(callbackPort) {
                        return &CallbackPortInUseError{Port: callbackPort}
                }
        }

        return nil
}

// IsPreRegisteredClient determines if the OAuth client is pre-registered (has client ID)
func IsPreRegisteredClient(clientID string) bool {
        return clientID != ""
}

// GetProcessOnPort returns the PID of the process listening on the given TCP port.
// Returns 0 if the port is free or if the holder cannot be determined.
// Uses gopsutil which provides cross-platform support (Linux: /proc, Windows: GetExtendedTcpTable,
// Darwin/FreeBSD: lsof).
func GetProcessOnPort(port int) (int, error) {
        if port <= 0 || port > MaxPort {
                return 0, fmt.Errorf("invalid port %d", port)
        }

        conns, err := gopsutilnet.Connections("tcp")
        if err != nil {
                return 0, fmt.Errorf("failed to get TCP connections: %w", err)
        }

        for _, c := range conns {
                if c.Laddr.Port == uint32(port) && c.Status == "LISTEN" && c.Pid > 0 { //nolint:gosec // G115 - port validated in [1, 65535]
                        return int(c.Pid), nil
                }
        }
        return 0, nil
}

// ParsePortSpec parses a port specification string in the format "hostPort:containerPort" or just "containerPort".
// Returns the host port string and container port integer.
// If only a container port is provided, a random available host port is selected.
func ParsePortSpec(portSpec string) (string, int, error) {
        slog.Debug("Parsing port spec", "spec", portSpec)
        // Check if it's in host:container format
        if strings.Contains(portSpec, ":") {
                parts := strings.Split(portSpec, ":")
                if len(parts) != 2 {
                        return "", 0, fmt.Errorf("invalid port specification: %s (expected 'hostPort:containerPort')", portSpec)
                }

                hostPortStr := parts[0]
                containerPortStr := parts[1]

                // Verify host port is a valid integer (or empty string if we supported random host port with :, but here we expect explicit)
                if _, err := strconv.Atoi(hostPortStr); err != nil {
                        return "", 0, fmt.Errorf("invalid host port in spec '%s': %w", portSpec, err)
                }

                containerPort, err := strconv.Atoi(containerPortStr)
                if err != nil {
                        return "", 0, fmt.Errorf("invalid container port in spec '%s': %w", portSpec, err)
                }

                return hostPortStr, containerPort, nil
        }

        // Try parsing as just container port
        containerPort, err := strconv.Atoi(portSpec)
        if err == nil {
                // Find a random available host port
                hostPort := FindAvailable()
                if hostPort == 0 {
                        return "", 0, fmt.Errorf("could not find an available port for container port %d", containerPort)
                }
                return fmt.Sprintf("%d", hostPort), containerPort, nil
        }

        return "", 0, fmt.Errorf("invalid port specification: %s (expected 'hostPort:containerPort' or 'containerPort')", portSpec)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package networking

import (
        "context"
        "errors"
        "fmt"
        "net"
        "net/url"
        "os"
        "strings"

        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// AuthorityMatchesAny reports whether endpointURL's authority (host and port)
// equals the authority of any of operatorURLs. Comparison is case-insensitive
// and normalizes the implicit scheme default port, so "https://idp.example.com"
// and "idp.example.com:443" match. Empty and malformed entries never match.
//
// This is the trust test for endpoints read out of an authorization-server
// metadata document: an AS naming its own authority is still covered by the
// operator's decision to configure it, while an AS naming some other host is
// not (CWE-918).
func AuthorityMatchesAny(endpointURL string, operatorURLs ...string) bool {
        target := canonicalAuthority(endpointURL)
        if target == "" {
                return false
        }
        for _, operatorURL := range operatorURLs {
                if authority := canonicalAuthority(operatorURL); authority != "" && authority == target {
                        return true
                }
        }
        return false
}

// canonicalAuthority reduces a URL — or a bare "host:port" authority — to a
// lowercase host:port form suitable for equality comparison. It returns the
// empty string whenever the authority cannot be determined, including when a
// bare host carries no port and no scheme to imply one: callers treat that as
// "does not match", so an undecidable input fails closed.
func canonicalAuthority(rawURL string) string {
        if rawURL == "" {
                return ""
        }
        u, err := url.Parse(rawURL)
        if err != nil || u.Host == "" {
                // Not a full URL. Retry as a bare authority ("idp.example.com:443",
                // "[::1]:8443"), which url.Parse otherwise reads as scheme:opaque or
                // rejects outright.
                u, err = url.Parse("//" + rawURL)
                if err != nil || u.Host == "" {
                        return ""
                }
        }

        host := strings.ToLower(u.Hostname())
        if host == "" {
                return ""
        }

        port := u.Port()
        if port == "" {
                switch strings.ToLower(u.Scheme) {
                case HttpScheme:
                        port = "80"
                case HttpsScheme:
                        port = "443"
                default:
                        return ""
                }
        }
        return net.JoinHostPort(host, port)
}

// TargetIsPrivate reports whether the host in rawURL refers to — or resolves to —
// a private, loopback, or link-local address. It is used to detect when an
// operator has deliberately pointed ToolHive at an internal target, so that
// discovery fetches derived from untrusted server input may also be allowed to
// reach internal addresses for that deployment.
//
// IP literals and "localhost" are classified without DNS. Hostnames are resolved
// and reported private if ANY resolved address is private. Unparsable input or
// resolution failure returns false (treat as public — the SSRF guard then stays
// engaged, failing secure).
func TargetIsPrivate(ctx context.Context, rawURL string) bool {
        u, err := url.Parse(rawURL)
        if err != nil {
                return false
        }
        host := u.Hostname()
        if host == "" {
                return false
        }
        if IsLocalhost(host) {
                return true
        }
        if ip := net.ParseIP(host); ip != nil {
                return IsPrivateIP(ip)
        }
        addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host)
        if err != nil {
                return false
        }
        for _, a := range addrs {
                if IsPrivateIP(a.IP) {
                        return true
                }
        }
        return false
}

const (
        // ErrPrivateIpAddress is the error returned when the provided URL redirects to a private IP address
        ErrPrivateIpAddress = "the provided URL redirects to a private IP address, which is not allowed"
)

// low32EmbeddingPrefixes are the IPv6 prefixes whose addresses carry a full
// IPv4 address in the low 32 bits. An address inside one of these reaches that
// IPv4 through a translator or a dual-stack host, so its true reachability is
// determined by the embedded IPv4, not by the IPv6 form.
//   - 64:ff9b::/96     NAT64 well-known prefix (RFC 6052 §2.2), always a /96
//   - 64:ff9b:1::/96   the /96 sub-prefix of the RFC 8215 local-use 64:ff9b:1::/48
//   - ::/96            IPv4-compatible addresses (RFC 4291 §2.5.5.1, deprecated)
//   - ::ffff:0:0:0/96  IPv4-translated addresses (SIIT, RFC 2765 §2.1)
//
// IPv4-mapped addresses (::ffff:0:0/96) are not listed: net.IP.To4 already
// unwraps them to IPv4.
//
// The rest of the local-use 64:ff9b:1::/48 uses a shorter NAT64 prefix, where
// the embedded IPv4 is NOT in the low 32 bits and cannot be located from the
// address alone. Decoding the low 32 bits there could read attacker-chosen
// suffix bytes and misclassify an internal target as public, so that remainder
// is blocked wholesale via privateIPBlocks to avoid a false-negative bypass.
var low32EmbeddingPrefixes []*net.IPNet

// sixToFourPrefix is the 6to4 prefix 2002::/16 (RFC 3056, deprecated by
// RFC 7526). A 6to4 address embeds a full IPv4 address in bits 16-47
// (2002:aabb:ccdd::/48 encodes aa.bb.cc.dd), so — like NAT64 — its true
// reachability is determined by the embedded IPv4, not by the (global-unicast)
// IPv6 form. Unlike the NAT64 remainder above, this embedding is exact and
// reversible for the whole prefix, so it is decoded rather than blocked
// wholesale (blocking 2002::/16 outright would also reject legitimate public
// 6to4 addresses).
var sixToFourPrefix *net.IPNet

// embeddedIPv4 returns the IPv4 address embedded in ip if ip falls inside one
// of low32EmbeddingPrefixes or the 6to4 prefix, or nil otherwise.
func embeddedIPv4(ip net.IP) net.IP {
        v6 := ip.To16()
        if v6 == nil || ip.To4() != nil {
                return nil
        }
        for _, block := range low32EmbeddingPrefixes {
                if block.Contains(ip) {
                        return net.IPv4(v6[12], v6[13], v6[14], v6[15])
                }
        }
        if sixToFourPrefix.Contains(ip) {
                return net.IPv4(v6[2], v6[3], v6[4], v6[5])
        }
        return nil
}

func init() {
        for _, cidr := range []string{
                "0.0.0.0/8",       // RFC1122 "this host" (incl. 0.0.0.0 unspecified)
                "127.0.0.0/8",     // IPv4 loopback
                "10.0.0.0/8",      // RFC1918
                "172.16.0.0/12",   // RFC1918
                "192.168.0.0/16",  // RFC1918
                "169.254.0.0/16",  // RFC3927 link-local
                "::1/128",         // IPv6 loopback
                "fe80::/10",       // IPv6 link-local
                "fc00::/7",        // IPv6 unique local addr
                "100.64.0.0/10",   // RFC6598 shared address space (CGN)
                "192.0.2.0/24",    // RFC5737 documentation (TEST-NET-1)
                "198.51.100.0/24", // RFC5737 documentation (TEST-NET-2)
                "203.0.113.0/24",  // RFC5737 documentation (TEST-NET-3)
                "224.0.0.0/4",     // IPv4 multicast
                "240.0.0.0/4",     // RFC1112 reserved (Class E), incl. 255.255.255.255 broadcast
                "ff00::/8",        // IPv6 multicast
                // NAT64 local-use range (RFC 8215). The 64:ff9b:1::/96 subset is decoded
                // to its embedded IPv4 below; this catch-all blocks the remaining
                // non-/96 embeddings, which cannot be decoded from the address alone.
                "64:ff9b:1::/48",
                // Teredo (RFC 4380). The embedded client IPv4 lives in the low 32 bits
                // but is obfuscated via bitwise NOT (RFC 4380 §4), so — unlike NAT64 and
                // 6to4 — it cannot be reliably decoded from the address alone. Blocked
                // wholesale to avoid a false-negative bypass.
                "2001::/32",
        } {
                _, block, err := net.ParseCIDR(cidr)
                if err != nil {
                        panic(fmt.Errorf("parse error on %q: %w", cidr, err))
                }
                privateIPBlocks = append(privateIPBlocks, block)
        }
        for _, cidr := range []string{
                "64:ff9b::/96",    // NAT64 well-known prefix, /96 embedding (RFC 6052)
                "64:ff9b:1::/96",  // /96 sub-prefix of the local-use range (RFC 8215)
                "::/96",           // IPv4-compatible (RFC 4291 §2.5.5.1)
                "::ffff:0:0:0/96", // IPv4-translated (RFC 2765 §2.1)
        } {
                _, block, err := net.ParseCIDR(cidr)
                if err != nil {
                        panic(fmt.Errorf("parse error on %q: %w", cidr, err))
                }
                low32EmbeddingPrefixes = append(low32EmbeddingPrefixes, block)
        }
        const sixToFourCIDR = "2002::/16"
        _, sixToFourBlock, err := net.ParseCIDR(sixToFourCIDR)
        if err != nil {
                panic(fmt.Errorf("parse error on %q: %w", sixToFourCIDR, err))
        }
        sixToFourPrefix = sixToFourBlock
}

// IsPrivateIP reports whether ip is a private, loopback, link-local,
// unspecified, or otherwise reserved/non-public address.
//
// NAT64- and 6to4-translated addresses are evaluated by the IPv4 address they
// embed: a NAT64 address whose low 32 bits map to a private/link-local IPv4
// (e.g. 64:ff9b:1::a9fe:a9fe -> 169.254.169.254, the cloud metadata endpoint)
// is treated as private, because behind a NAT64 gateway it reaches exactly
// that internal IPv4, while NAT64 addresses embedding a genuinely public IPv4
// remain allowed. This /96 decoding covers the well-known 64:ff9b::/96
// (RFC 6052) and the 64:ff9b:1::/96 sub-prefix of the RFC 8215 local-use
// range; the rest of 64:ff9b:1::/48 uses a non-/96 embedding that cannot be
// decoded from the address alone and is blocked wholesale (see
// privateIPBlocks). IPv4-compatible (::/96) and IPv4-translated
// (::ffff:0:0:0/96) addresses use the same low-32-bit embedding and are decoded
// the same way. 6to4 addresses (2002::/16, RFC 3056) are decoded too,
// extracting the IPv4 from bits 16-47 (e.g. 2002:a9fe:a9fe:: ->
// 169.254.169.254). Teredo addresses (2001::/32, RFC 4380) obfuscate their
// embedded IPv4 via bitwise NOT and so cannot be decoded from the address
// alone; they are blocked wholesale (see privateIPBlocks).
func IsPrivateIP(ip net.IP) bool {
        if ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() {
                return true
        }
        if v4 := embeddedIPv4(ip); v4 != nil {
                return IsPrivateIP(v4)
        }
        for _, block := range privateIPBlocks {
                if block.Contains(ip) {
                        return true
                }
        }
        return false
}

// AddressReferencesPrivateIp returns an error if the address references a private IP address
func AddressReferencesPrivateIp(address string) error {
        host, _, err := net.SplitHostPort(address)
        if err != nil {
                return err
        }
        // Check for a private IP address or loopback
        ip := net.ParseIP(host)
        if IsPrivateIP(ip) {
                return errors.New(ErrPrivateIpAddress)
        }

        return nil
}

// ValidateEndpointURL validates that an endpoint URL is secure
func ValidateEndpointURL(endpoint string) error {
        skipValidation := strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true")
        return validateEndpointURLWithSkip(endpoint, skipValidation)
}

// ValidateEndpointURLWithInsecure validates that an endpoint URL is secure, allowing HTTP if insecureAllowHTTP is true
// WARNING: This is insecure and should NEVER be used in production
func ValidateEndpointURLWithInsecure(endpoint string, insecureAllowHTTP bool) error {
        skipValidation := strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true")
        return validateEndpointURLWithSkip(endpoint, skipValidation || insecureAllowHTTP)
}

// validateEndpointURLWithSkip validates that an endpoint URL is secure, with an option to skip validation
func validateEndpointURLWithSkip(endpoint string, skipValidation bool) error {
        if skipValidation {
                return nil // Skip validation
        }
        u, err := url.Parse(endpoint)
        if err != nil {
                return fmt.Errorf("invalid URL: %w", err)
        }

        // Ensure HTTPS for security (except localhost for development)
        if u.Scheme != HttpsScheme && !IsLocalhost(u.Host) {
                return fmt.Errorf("endpoint must use HTTPS: %s", endpoint)
        }

        return nil
}

// ValidateHTTPSURL checks that rawURL is a valid URL using the https scheme.
// Unlike ValidateEndpointURL, no localhost exception is made — HTTPS is always
// required (suitable for gateway URLs and other production endpoints).
func ValidateHTTPSURL(rawURL string) error {
        parsed, err := url.Parse(rawURL)
        if err != nil {
                return fmt.Errorf("invalid URL: %w", err)
        }
        if parsed.Host == "" {
                return fmt.Errorf("URL must include a host: %s", rawURL)
        }
        if parsed.Scheme != HttpsScheme {
                return fmt.Errorf("must use HTTPS, got scheme %q", parsed.Scheme)
        }
        return nil
}

// ValidateIssuerURL validates that an OIDC issuer URL is well-formed and uses
// HTTPS. HTTP is permitted only for localhost (development). Per OIDC Core
// Section 3.1.2.1 and RFC 8414 Section 2, the issuer MUST use the "https"
// scheme.
func ValidateIssuerURL(rawURL string) error {
        u, err := url.Parse(rawURL)
        if err != nil {
                return fmt.Errorf("invalid issuer URL %q: %w", rawURL, err)
        }
        if u.Host == "" {
                return fmt.Errorf("issuer URL must include a host: %s", rawURL)
        }
        if u.Scheme != HttpsScheme && !IsLocalhost(u.Host) {
                return fmt.Errorf("issuer URL must use HTTPS (except localhost for development): %s", rawURL)
        }
        return nil
}

// ValidateLoopbackAddress returns an error if addr (a host:port string) does
// not contain a literal loopback IP address. Both IPv4 (127.x.x.x) and IPv6
// (::1) loopback addresses are accepted. Hostnames (including "localhost") are
// not resolved and will be rejected.
func ValidateLoopbackAddress(addr string) error {
        host, _, err := net.SplitHostPort(addr)
        if err != nil {
                return fmt.Errorf("invalid listen address %q: %w", addr, err)
        }
        ip := net.ParseIP(host)
        if ip == nil || !ip.IsLoopback() {
                return fmt.Errorf("listen address %q must be a loopback interface (127.x.x.x or ::1)", addr)
        }
        return nil
}

// IsLocalhost checks if a host is a loopback address (for development).
// The canonical implementation lives in pkg/oauthproto.IsLoopbackHost;
// this function is a thin wrapper that preserves backward compatibility for
// all callers in this package and beyond.
func IsLocalhost(host string) bool {
        return oauthproto.IsLoopbackHost(host)
}

// IsLoopbackHost reports whether the Host header value refers to a loopback
// address. It is intended for DNS-rebinding guards on loopback-only listeners.
// It accepts the hostname "localhost" (case-insensitive), any 127.x.x.x
// address, and the IPv6 loopback ::1. Both plain-host and host:port forms are
// accepted. Hostnames other than "localhost" are NOT resolved.
func IsLoopbackHost(host string) bool {
        h, _, err := net.SplitHostPort(host)
        if err != nil {
                // No port present — treat the whole value as the host.
                h = host
                // Strip brackets from bare IPv6 literals like "[::1]".
                if len(h) > 2 && h[0] == '[' && h[len(h)-1] == ']' {
                        h = h[1 : len(h)-1]
                }
        }
        if strings.EqualFold(h, "localhost") {
                return true
        }
        ip := net.ParseIP(h)
        return ip != nil && ip.IsLoopback()
}

// Network isolation log messages. Network isolation is only functional in bridge
// mode; host/none/custom network modes bypass the isolation sidecars entirely, so
// isolation must be dropped for them. These constants are shared across the config
// build, config load, and deploy paths so the wording cannot drift.
const (
        // NetworkIsolationHostDroppedMsg is logged at WARN when isolation was requested
        // but is dropped because the network mode (host/custom) cannot enforce it.
        // Dropping isolation here is a real reduction in confinement.
        NetworkIsolationHostDroppedMsg = "network isolation is not enforced with host/custom networking; " +
                "disabling isolation for this workload (pass --isolate-network=false to opt out explicitly and silence this warning)"
        // NetworkIsolationNoneRedundantMsg is logged at DEBUG when isolation is dropped
        // for the "none" network mode, where it is merely redundant (already maximally confined).
        NetworkIsolationNoneRedundantMsg = "network isolation is redundant with none networking; " +
                "disabling isolation for this workload"
)

// IsBridgeMode reports whether the given network mode is a bridge-style mode
// (""/"bridge"/"default"). Network isolation is only functional in bridge mode:
// it builds a per-workload internal network plus DNS/egress/ingress sidecars, and
// host/none/custom modes bypass those sidecars entirely.
func IsBridgeMode(mode string) bool {
        return mode == "" || mode == "bridge" || mode == "default"
}

// IsURL checks if the input is a valid HTTP or HTTPS URL
func IsURL(input string) bool {
        parsedURL, err := url.Parse(input)
        if err != nil {
                return false
        }
        // Must have HTTP or HTTPS scheme and a valid host
        return (parsedURL.Scheme == HttpScheme || parsedURL.Scheme == HttpsScheme) &&
                parsedURL.Host != "" &&
                parsedURL.Host != "//"
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauthproto

import (
        "net/url"
        "strings"
)

// ToolHiveClientMetadataDocumentURL is the stable HTTPS URL where ToolHive's
// client metadata document is hosted. ToolHive presents this URL as its
// client_id to remote authorization servers that support CIMD. The URL must
// be live and serving the client metadata document before this feature can
// be used in production.
const ToolHiveClientMetadataDocumentURL = "https://toolhive.dev/oauth/client-metadata.json"

// IsClientIDMetadataDocumentURL returns true if clientID looks like a CIMD URL.
// In production this means any HTTPS URL; in local development http://localhost,
// http://127.0.0.1, and http://[::1] URLs are also accepted so that integration
// tests can use plain httptest.Server instances without TLS setup. DCR-issued IDs
// are always opaque strings that never begin with a URL scheme, so there is no
// risk of false positives. Do not tighten this to an exact match against
// ToolHiveClientMetadataDocumentURL — the embedded AS must accept CIMD URLs
// from third-party clients too.
//
// This predicate is also used by pkg/auth/remote/handler.go to decide whether to
// persist a client_id as DCR credentials or as a CIMD URL. The loopback HTTP
// acceptance is safe in that context: real DCR-issued client IDs are always opaque
// strings that never start with "http://", so no live deployment is affected.
func IsClientIDMetadataDocumentURL(clientID string) bool {
        if strings.HasPrefix(clientID, "https://") {
                return true
        }
        // Allow loopback HTTP URLs (localhost, 127.0.0.1, [::1]) for local
        // development and integration testing. These are the only HTTP URLs that
        // FetchClientMetadataDocument / validateCIMDClientURL also accept.
        if strings.HasPrefix(clientID, "http://") {
                parsed, err := url.Parse(clientID)
                if err != nil {
                        return false
                }
                return IsLoopbackHost(parsed.Host)
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package cimd implements fetching and validation of OAuth 2.0 Client ID
// Metadata Documents (CIMD) per draft-ietf-oauth-client-id-metadata-document.
//
// This package is a sub-package of pkg/oauthproto and is allowed to import
// pkg/networking (for SSRF utilities) without violating the leaf-package
// invariant of the parent package.
package cimd

import (
        "context"
        "fmt"
        "net"
        "net/http"
        "net/url"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// ClientMetadataDocument represents an OAuth 2.0 Client ID Metadata Document
// per draft-ietf-oauth-client-id-metadata-document.
type ClientMetadataDocument struct {
        // Required
        ClientID     string   `json:"client_id"`
        RedirectURIs []string `json:"redirect_uris"`

        // Recommended
        ClientName string `json:"client_name,omitempty"`
        LogoURI    string `json:"logo_uri,omitempty"`
        ClientURI  string `json:"client_uri,omitempty"`

        // Optional
        TosURI                  string   `json:"tos_uri,omitempty"`
        PolicyURI               string   `json:"policy_uri,omitempty"`
        GrantTypes              []string `json:"grant_types,omitempty"`
        ResponseTypes           []string `json:"response_types,omitempty"`
        Scope                   string   `json:"scope,omitempty"`
        TokenEndpointAuthMethod string   `json:"token_endpoint_auth_method,omitempty"`
        // TokenEndpointAuthMethodsSupported is not part of the CIMD draft itself;
        // it is defined by OpenID Connect RP Metadata Choices 1.0 as the plural
        // counterpart of TokenEndpointAuthMethod, listing every method the client
        // supports while TokenEndpointAuthMethod names its preferred one. Live
        // clients (e.g. ChatGPT) publish it alongside a stricter singular
        // preference the client_id-issuing party does not control per-AS; the
        // decorator in pkg/authserver/storage uses it to negotiate a mutually
        // supported method instead of rejecting the document outright (#6278).
        TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported,omitempty"`
        ApplicationType                   string   `json:"application_type,omitempty"`
        PostLogoutRedirectURIs            []string `json:"post_logout_redirect_uris,omitempty"`
}

// forbiddenAuthMethods lists token_endpoint_auth_method values that MUST NOT
// appear in a CIMD document. Per the spec, CIMD clients may not use symmetric
// shared-secret methods because no pre-shared secret is ever established.
var forbiddenAuthMethods = []string{
        "client_secret_post",
        "client_secret_basic",
        "client_secret_jwt",
}

// FetchClientMetadataDocument fetches and validates a Client ID Metadata Document
// from the given URL. The URL must use the HTTPS scheme (http://localhost is
// accepted in development). The document is fetched with a 5-second timeout,
// redirects disabled, a 10 KB body cap, and SSRF protection via a per-dial IP
// check. After fetching, ValidateClientMetadataDocument is called and any
// validation error is returned.
func FetchClientMetadataDocument(ctx context.Context, rawURL string) (*ClientMetadataDocument, error) {
        allowHTTPLoopback, err := validateCIMDClientURL(rawURL)
        if err != nil {
                return nil, err
        }

        result, err := networking.FetchJSON[ClientMetadataDocument](
                ctx,
                newCIMDHTTPClient(allowHTTPLoopback),
                rawURL,
                networking.WithMaxResponseSize(10*1024),
        )
        if err != nil {
                return nil, fmt.Errorf("failed to fetch client metadata document: %w", err)
        }

        doc := &result.Data
        if err := ValidateClientMetadataDocument(doc, rawURL); err != nil {
                return nil, err
        }

        return doc, nil
}

// newCIMDHTTPClient builds an *http.Client for fetching CIMD documents.
// Keep-alive connections are disabled so the per-dial SSRF check fires on
// every request. A custom DialContext permits loopback connections only when
// allowHTTPLoopback was granted by validated HTTP-loopback URL input; it rejects
// all other private/special-use ranges via networking.IsPrivateIP and dials by
// IP literal to prevent DNS-rebinding. A CheckRedirect hook prevents automatic
// redirect following.
//
// HttpClientBuilder.WithPrivateIPs(false) is not used here because it also
// blocks loopback addresses (127.0.0.0/8), which would block the intentional,
// per-request validated HTTP-loopback exception.
func newCIMDHTTPClient(allowHTTPLoopback bool) *http.Client {
        transport := &http.Transport{
                TLSHandshakeTimeout:   5 * time.Second,
                ResponseHeaderTimeout: 5 * time.Second,
                DisableKeepAlives:     true,
                DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
                        host, port, err := net.SplitHostPort(address)
                        if err != nil {
                                return nil, fmt.Errorf("cimd: invalid dial address %q: %w", address, err)
                        }
                        ips, err := net.DefaultResolver.LookupHost(ctx, host)
                        if err != nil {
                                return nil, fmt.Errorf("cimd: failed to resolve %q: %w", host, err)
                        }
                        for _, ipStr := range ips {
                                ip := net.ParseIP(ipStr)
                                if ip == nil {
                                        continue
                                }
                                // Private/special-use addresses are rejected unless this request
                                // received the validated HTTP-loopback permission. Dial by IP literal
                                // to prevent DNS-rebinding TOCTOU.
                                if networking.IsPrivateIP(ip) && (!allowHTTPLoopback || !ip.IsLoopback()) {
                                        return nil, fmt.Errorf("cimd: refusing connection to private address %s", ipStr)
                                }
                                dialer := &net.Dialer{Timeout: 5 * time.Second}
                                return dialer.DialContext(ctx, network, net.JoinHostPort(ipStr, port))
                        }
                        return nil, fmt.Errorf("cimd: no valid address found for %q", host)
                },
        }

        client := &http.Client{
                Timeout:   5 * time.Second,
                Transport: transport,
                // CIMD §4: "The authorization server MUST NOT automatically follow
                // HTTP redirects when retrieving the client registration information."
                // Return the redirect response unchanged; FetchJSON will surface it
                // as a non-200 error.
                CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
                        return http.ErrUseLastResponse
                },
        }

        return client
}

// validateCIMDClientURL enforces the client_id URL requirements from
// draft-ietf-oauth-client-id-metadata-document §3:
//   - MUST use https scheme (http://localhost permitted for development only)
//   - MUST contain a non-empty path component (not just "/")
//   - MUST NOT contain a fragment component
//   - MUST NOT contain userinfo (username or password)
//   - MUST NOT contain single-dot or double-dot path segments
//
// It returns true only when the fully validated URL is the intentional HTTP
// recognized-loopback-host exception; valid HTTPS URLs and validation failures
// return false.
func validateCIMDClientURL(rawURL string) (bool, error) {
        parsed, err := url.Parse(rawURL)
        if err != nil {
                return false, fmt.Errorf("client_id URL is not a valid URL: %w", err)
        }
        allowHTTPLoopback := parsed.Scheme == "http" && oauthproto.IsLoopbackHost(parsed.Hostname())
        if parsed.Scheme != "https" && !allowHTTPLoopback {
                return false, fmt.Errorf("client_id URL must use the https scheme: %s", rawURL)
        }
        if parsed.Host == "" {
                return false, fmt.Errorf("client_id URL must contain a host")
        }
        if parsed.Fragment != "" {
                return false, fmt.Errorf("client_id URL must not contain a fragment component")
        }
        if parsed.User != nil {
                return false, fmt.Errorf("client_id URL must not contain userinfo (username or password)")
        }
        if parsed.Path == "" || parsed.Path == "/" {
                return false, fmt.Errorf("client_id URL must contain a non-empty path component")
        }
        for _, segment := range strings.Split(parsed.Path, "/") {
                if segment == "." || segment == ".." {
                        return false, fmt.Errorf("client_id URL must not contain dot-segment path components")
                }
        }
        return allowHTTPLoopback, nil
}

// ValidateClientMetadataDocument validates a parsed ClientMetadataDocument
// against the URL it was fetched from. Per the CIMD draft spec, the
// client_id field must exactly equal the URL — no normalization is applied,
// because allowing normalization would permit subtle spoofing attacks where a
// document at URL A claims the identity of URL B.
func ValidateClientMetadataDocument(doc *ClientMetadataDocument, fetchedFrom string) error {
        if doc.ClientID == "" {
                return fmt.Errorf("client metadata document missing required field: client_id")
        }

        if len(doc.RedirectURIs) == 0 {
                return fmt.Errorf("client metadata document missing required field: redirect_uris")
        }

        if doc.ClientID != fetchedFrom {
                return fmt.Errorf("client_id %q does not match the URL it was fetched from %q", doc.ClientID, fetchedFrom)
        }

        for _, uri := range doc.RedirectURIs {
                if err := oauthproto.ValidateRedirectURI(uri, oauthproto.RedirectURIPolicyStrict); err != nil {
                        return err
                }
        }

        // Only symmetric shared-secret auth methods are forbidden; asymmetric methods
        // (e.g. private_key_jwt, tls_client_auth, none) are permitted by the spec.
        if doc.TokenEndpointAuthMethod != "" {
                for _, forbidden := range forbiddenAuthMethods {
                        if doc.TokenEndpointAuthMethod == forbidden {
                                return fmt.Errorf("token_endpoint_auth_method %q is not permitted in a CIMD document: "+
                                        "symmetric shared-secret methods cannot be used without pre-registration", forbidden)
                        }
                }
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package claims provides strict decoders for OAuth protocol claim values.
package claims

import (
        "errors"
        "fmt"
)

// DecodeStringArray decodes a claim containing []string or []any whose elements
// are all strings. Nil values (including typed nil slices such as []string(nil)
// or []any(nil)), other types, and non-string elements return an error. A
// non-nil empty array returns (nil, nil). The result does not alias the input.
//
// Typed nil slices are rejected rather than treated as empty arrays so a
// security-relevant caller can rely on a consistent nil-is-error contract: an
// absent or malformed claim must never be silently equivalent to a legitimate
// empty-array claim.
func DecodeStringArray(raw any) ([]string, error) {
        if raw == nil {
                return nil, errors.New("claim value is nil")
        }
        switch v := raw.(type) {
        case []string:
                if v == nil {
                        return nil, errors.New("claim value is nil")
                }
                return append([]string(nil), v...), nil
        case []any:
                if v == nil {
                        return nil, errors.New("claim value is nil")
                }
                if len(v) == 0 {
                        return nil, nil
                }
                out := make([]string, len(v))
                for i, item := range v {
                        s, ok := item.(string)
                        if !ok {
                                return nil, fmt.Errorf("array element at index %d is %T, want string", i, item)
                        }
                        out[i] = s
                }
                return out, nil
        default:
                return nil, fmt.Errorf("value is %T, want JSON string array", raw)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauthproto

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

        "github.com/go-jose/go-jose/v3"
)

// ToolHiveMCPClientName is the name advertised in dynamic client registration requests.
const ToolHiveMCPClientName = "ToolHive MCP Client"

// DynamicClientRegistrationRequest represents the request for dynamic client registration (RFC 7591).
type DynamicClientRegistrationRequest struct {
        // Required field according to RFC 7591
        RedirectURIs []string `json:"redirect_uris"`

        // Essential fields for OAuth flow
        ClientName              string    `json:"client_name,omitempty"`
        TokenEndpointAuthMethod string    `json:"token_endpoint_auth_method,omitempty"`
        GrantTypes              []string  `json:"grant_types,omitempty"`
        ResponseTypes           []string  `json:"response_types,omitempty"`
        Scopes                  ScopeList `json:"scope,omitempty"`
        // JWKS is the inline public key set used by private_key_jwt clients.
        JWKS                        *jose.JSONWebKeySet `json:"jwks,omitempty"`
        JWKSURI                     string              `json:"jwks_uri,omitempty"`
        TokenEndpointAuthSigningAlg string              `json:"token_endpoint_auth_signing_alg,omitempty"`

        // SoftwareID is the RFC 7591 Section 2 "software_id": a unique identifier
        // for the client software. Optional; servers may capture it for audit
        // or telemetry but it is not required to register.
        SoftwareID string `json:"software_id,omitempty"`
}

// ScopeList represents the "scope" field in both dynamic client registration requests and responses.
//
// Marshaling (requests): Per RFC 7591 Section 2, scopes are serialized as a space-delimited string.
// Examples:
//   - []string{"openid", "profile", "email"} → "openid profile email"
//   - []string{"openid"}                     → "openid"
//   - nil or []string{}                      → omitted (via omitempty)
//
// Unmarshaling (responses): Some servers return scopes as a space-delimited string per RFC 7591,
// while others return a JSON array. This type normalizes both formats into []string.
// Examples:
//   - "openid profile email"       → []string{"openid", "profile", "email"}
//   - ["openid","profile","email"] → []string{"openid", "profile", "email"}
//   - null                         → nil
//   - "" or ["", "  "]             → nil
type ScopeList []string

// MarshalJSON implements custom encoding for ScopeList. It converts the slice
// of scopes into a space-delimited string as required by RFC 7591 Section 2.
//
// Important: This method does NOT handle empty slices. Go's encoding/json package
// evaluates omitempty by checking if the Go value is "empty" (len(slice) == 0)
// BEFORE calling MarshalJSON. Empty slices are omitted at the struct level, so this
// method is never invoked for empty slices. This means we don't need to return null
// or handle the empty case - omitempty does it for us automatically.
//
// See: https://pkg.go.dev/encoding/json (omitempty checks zero values before marshaling)
func (s ScopeList) MarshalJSON() ([]byte, error) {
        // Join scopes with spaces and marshal as a string (RFC 7591 Section 2)
        scopeString := strings.Join(s, " ")
        result, err := json.Marshal(scopeString)
        if err == nil {
                slog.Debug("RFC 7591: Marshaled ScopeList", "scopes", []string(s), "result", scopeString)
        }
        return result, err
}

// UnmarshalJSON implements custom decoding for ScopeList. It supports both
// string and array encodings of the "scope" field, trimming whitespace and
// normalizing empty values to nil for consistent semantics.
func (s *ScopeList) UnmarshalJSON(data []byte) error {
        // Handle explicit null
        if strings.TrimSpace(string(data)) == "null" {
                *s = nil
                return nil
        }

        // Case 1: space-delimited string
        var str string
        if err := json.Unmarshal(data, &str); err == nil {
                if strings.TrimSpace(str) == "" {
                        *s = nil
                        return nil
                }
                *s = strings.Fields(str)
                return nil
        }

        // Case 2: JSON array
        var arr []string
        if err := json.Unmarshal(data, &arr); err == nil {
                cleaned := make([]string, 0, len(arr))
                for _, v := range arr {
                        if v = strings.TrimSpace(v); v != "" {
                                cleaned = append(cleaned, v)
                        }
                }
                // Normalize: treat all-empty/whitespace arrays the same as ""
                if len(cleaned) == 0 {
                        *s = nil
                } else {
                        *s = cleaned
                }
                return nil
        }

        return fmt.Errorf("invalid scope format: %s", string(data))
}

// DynamicClientRegistrationResponse represents the response from dynamic client registration (RFC 7591).
type DynamicClientRegistrationResponse struct {
        // Required fields
        ClientID     string `json:"client_id"`
        ClientSecret string `json:"client_secret,omitempty"` //nolint:gosec // G117: field legitimately holds sensitive data

        // Optional fields that may be returned
        ClientIDIssuedAt int64 `json:"client_id_issued_at,omitempty"`
        // ClientSecretExpiresAt is a pointer so the key is present (with value 0,
        // "does not expire") exactly when client_secret is issued and absent
        // otherwise: RFC 7591 §3.2.1 requires the key on any response carrying a
        // secret, and 0 is the value this server legitimately returns for a
        // confidential client — an int64 with omitempty would silently drop
        // exactly that case, while a plain int64 would emit the key on public
        // registrations that carry no secret.
        ClientSecretExpiresAt   *int64 `json:"client_secret_expires_at,omitempty"`
        RegistrationAccessToken string `json:"registration_access_token,omitempty"`
        RegistrationClientURI   string `json:"registration_client_uri,omitempty"`

        // Echo back the essential request fields. response_types is omitted for
        // private_key_jwt registrations because they support no response types;
        // RFC 7591 has no wire encoding for that zero-element set.
        ClientName                  string              `json:"client_name,omitempty"`
        RedirectURIs                []string            `json:"redirect_uris"`
        TokenEndpointAuthMethod     string              `json:"token_endpoint_auth_method"`
        GrantTypes                  []string            `json:"grant_types"`
        ResponseTypes               []string            `json:"response_types,omitempty"`
        Scopes                      ScopeList           `json:"scope,omitempty"`
        JWKS                        *jose.JSONWebKeySet `json:"jwks,omitempty"`
        JWKSURI                     string              `json:"jwks_uri,omitempty"`
        TokenEndpointAuthSigningAlg string              `json:"token_endpoint_auth_signing_alg,omitempty"`
}

// RegisterClientDynamically performs RFC 7591 Dynamic Client Registration against
// the given registrationEndpoint.
//
// If client is nil, a default *http.Client with a 30 s timeout, 10 s TLS handshake
// timeout, and 10 s response-header timeout is used. Pass a non-nil client to supply
// custom transport settings (e.g., in tests using httptest.NewServer).
func RegisterClientDynamically(
        ctx context.Context,
        registrationEndpoint string,
        request *DynamicClientRegistrationRequest,
        client *http.Client,
) (*DynamicClientRegistrationResponse, error) {
        // Validate registration endpoint URL
        if _, err := validateRegistrationEndpoint(registrationEndpoint); err != nil {
                return nil, err
        }

        // Reject a nil request before the dereference below; the nil check previously
        // lived inside validateAndSetDefaults, but the shallow copy must come first.
        if request == nil {
                return nil, fmt.Errorf("registration request cannot be nil")
        }

        // Shallow-copy the request before passing it to validateAndSetDefaults so
        // that the caller's original struct is never mutated. Slice fields (RedirectURIs,
        // GrantTypes, ResponseTypes, Scopes) share the same backing arrays, but
        // validateAndSetDefaults only assigns new slices to nil/zero fields — it never
        // appends to or modifies existing ones — so a shallow copy is safe here.
        reqCopy := *request
        if err := validateAndSetDefaults(&reqCopy); err != nil {
                return nil, err
        }

        // Create HTTP request
        req, err := createHTTPRequest(ctx, registrationEndpoint, &reqCopy)
        if err != nil {
                return nil, err
        }

        // Use caller-supplied client or build a default one
        httpClient := buildHTTPClient(client)

        // Make the request
        resp, err := httpClient.Do(req)
        if err != nil {
                return nil, fmt.Errorf("failed to perform dynamic client registration: %w", err)
        }

        // Handle response
        response, err := handleHTTPResponse(resp)
        if err != nil {
                return nil, err
        }

        //nolint:gosec // G706: client_id is public metadata from DCR response
        slog.Debug("Successfully registered OAuth client dynamically",
                "client_id", response.ClientID)
        return response, nil
}

// validateRegistrationEndpoint validates the registration endpoint URL.
func validateRegistrationEndpoint(registrationEndpoint string) (*url.URL, error) {
        registrationURL, err := url.Parse(registrationEndpoint)
        if err != nil {
                return nil, fmt.Errorf("invalid registration endpoint URL: %w", err)
        }

        // Ensure HTTPS for security (except loopback addresses for development)
        if registrationURL.Scheme != schemeHTTPS && !IsLoopbackHost(registrationURL.Host) {
                return nil, fmt.Errorf("registration endpoint must use HTTPS: %s", registrationEndpoint)
        }

        return registrationURL, nil
}

// validateAndSetDefaults validates the request and sets default values.
func validateAndSetDefaults(request *DynamicClientRegistrationRequest) error {
        if len(request.RedirectURIs) == 0 {
                return fmt.Errorf("at least one redirect URI is required")
        }

        // Validate that individual scope values don't contain spaces (RFC 6749 Section 3.3)
        // Scopes must be space-separated tokens, so spaces within a scope value are invalid
        for _, scope := range request.Scopes {
                if strings.Contains(scope, " ") {
                        return fmt.Errorf("invalid scope value %q: scope values cannot contain spaces (RFC 6749)", scope)
                }
        }

        // Set default values if not provided
        if request.ClientName == "" {
                request.ClientName = ToolHiveMCPClientName
        }
        if len(request.GrantTypes) == 0 {
                request.GrantTypes = []string{GrantTypeAuthorizationCode, GrantTypeRefreshToken}
        }
        if len(request.ResponseTypes) == 0 {
                request.ResponseTypes = []string{ResponseTypeCode}
        }
        if request.TokenEndpointAuthMethod == "" {
                request.TokenEndpointAuthMethod = TokenEndpointAuthMethodNone // For PKCE flow
        }

        return nil
}

// createHTTPRequest creates the HTTP request for dynamic client registration.
func createHTTPRequest(
        ctx context.Context,
        registrationEndpoint string,
        request *DynamicClientRegistrationRequest,
) (*http.Request, error) {
        // Serialize request to JSON
        requestBody, err := json.Marshal(request)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal registration request: %w", err)
        }

        // Create HTTP request
        req, err := http.NewRequestWithContext(ctx, http.MethodPost, registrationEndpoint, strings.NewReader(string(requestBody)))
        if err != nil {
                return nil, fmt.Errorf("failed to create registration request: %w", err)
        }

        // Set headers
        req.Header.Set("Content-Type", "application/json")
        req.Header.Set("Accept", "application/json")
        req.Header.Set("User-Agent", UserAgent)

        return req, nil
}

// NewDefaultDCRClient returns the canonical bounded *http.Client used by
// RegisterClientDynamically when its caller does not supply one. It is
// exported so callers that need to wrap the transport (for example, to
// inject an RFC 7591 initial access token as an Authorization header) can
// reuse the same timeout policy and benefit automatically from any future
// tightening of these bounds.
//
// Timeouts:
//
//   - Overall request timeout: 30 s
//   - TLS handshake timeout: 10 s
//   - Response-header timeout: 10 s
func NewDefaultDCRClient() *http.Client {
        return &http.Client{
                Timeout: 30 * time.Second,
                Transport: &http.Transport{
                        TLSHandshakeTimeout:   10 * time.Second,
                        ResponseHeaderTimeout: 10 * time.Second,
                        IdleConnTimeout:       idleConnTimeout,
                        MaxIdleConns:          maxIdleConns,
                        MaxIdleConnsPerHost:   maxIdleConnsPerHost,
                },
        }
}

// buildHTTPClient returns the caller-supplied client, or a default client if nil.
func buildHTTPClient(client *http.Client) *http.Client {
        if client != nil {
                return client
        }
        return NewDefaultDCRClient()
}

// handleHTTPResponse handles the HTTP response and validates it.
func handleHTTPResponse(resp *http.Response) (*DynamicClientRegistrationResponse, error) {
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("Failed to close response body", "error", err)
                }
        }()

        // Check response status
        if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
                // Cap the upstream's error body so a misbehaving or malicious AS
                // cannot inflate operator logs by echoing arbitrary content. 8 KiB
                // is far larger than any conformant RFC 7591 error response and
                // still small enough that the surrounding error chain stays
                // readable.
                const maxErrorBodySize = 8 * 1024
                errorBody, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBodySize))
                // Drain whatever exceeded the cap so net/http can reuse the TCP
                // connection — without this the deferred Close terminates the
                // connection mid-stream when the upstream returned more than 8 KiB.
                _, _ = io.Copy(io.Discard, resp.Body)

                // Detect if DCR is not supported by the provider.
                // Common HTTP status codes when DCR is unsupported:
                //   - 404 Not Found: endpoint doesn't exist
                //   - 405 Method Not Allowed: endpoint exists but POST not allowed
                //   - 501 Not Implemented: DCR feature not implemented
                if resp.StatusCode == http.StatusNotFound ||
                        resp.StatusCode == http.StatusMethodNotAllowed ||
                        resp.StatusCode == http.StatusNotImplemented {
                        return nil, fmt.Errorf(
                                "the provider does not support RFC 7591 Dynamic Client Registration (HTTP %d); "+
                                        "configure client credentials out of band. Error details: %s",
                                resp.StatusCode, string(errorBody))
                }

                return nil, fmt.Errorf("dynamic client registration failed with status %d: %s", resp.StatusCode, string(errorBody))
        }

        // Check content type; drain before returning to allow TCP connection reuse.
        contentType := resp.Header.Get("Content-Type")
        if !strings.Contains(contentType, "application/json") {
                _, _ = io.Copy(io.Discard, resp.Body)
                return nil, fmt.Errorf("unexpected content type: %s", contentType)
        }

        // Limit response size to prevent DoS
        const maxResponseSize = 1024 * 1024 // 1MB
        limitedReader := io.LimitReader(resp.Body, maxResponseSize)

        // Parse the response
        var response DynamicClientRegistrationResponse
        decoder := json.NewDecoder(limitedReader)
        if err := decoder.Decode(&response); err != nil {
                return nil, fmt.Errorf("failed to decode registration response: %w", err)
        }

        // Validate required response fields
        if response.ClientID == "" {
                return nil, fmt.Errorf("registration response missing client_id")
        }

        return &response, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauthproto

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

// discoveryTimeout is the bounded per-call timeout applied when fetching
// authorization server metadata, regardless of the caller's context deadline.
//
// Declared as a var rather than a const so tests that need to exercise the
// timeout path can shorten it without waiting 10 s per run. Production code
// MUST NOT mutate this value — it is effectively constant outside tests.
var discoveryTimeout = 10 * time.Second

// discoveryMaxResponseSize caps the response body size accepted from a
// discovery endpoint to prevent resource exhaustion from hostile servers.
const discoveryMaxResponseSize = 1024 * 1024 // 1MB

// FetchAuthorizationServerMetadata fetches RFC 8414 authorization server
// metadata for the given issuer, using a three-path fallback that mirrors
// the guidance in RFC 8414 §3.1 and OpenID Connect Discovery 1.0.
//
// The paths are tried in this order:
//  1. RFC 8414 path-insertion: {origin}/.well-known/oauth-authorization-server{path}
//  2. OIDC Discovery:          {issuer}/.well-known/openid-configuration
//  3. Bare RFC 8414:           {origin}/.well-known/oauth-authorization-server
//
// A 200 OK response with a clean application/json body that passes the
// RFC 8414 §3.3 issuer check is a candidate. The first candidate with a
// non-empty registration_endpoint wins immediately. If no candidate has a
// registration_endpoint but at least one was otherwise valid, the first such
// partial document is returned with ErrRegistrationEndpointMissing — never
// short-circuiting on a partial doc, since a tenant-aware IdP may publish
// path-insertion metadata without DCR while the OIDC document advertises it.
//
// If client is nil, a default *http.Client with a bounded 10 s timeout is
// used. When client is non-nil, the same bounded 10 s per-call timeout is
// applied via context.WithTimeout on top of the caller's context to protect
// against callers that pass context.Background or a long deadline.
//
// Return contract:
//
//   - On full success, returns (metadata, nil) with a non-empty
//     RegistrationEndpoint.
//
//   - When at least one candidate document was issuer-validated but every such
//     document omits registration_endpoint, returns the first partial document
//     paired with ErrRegistrationEndpointMissing. The metadata is non-nil so
//     callers can reuse the other fields (TokenEndpoint, JWKSURI, etc.). Note
//     this deviates from the usual Go "err != nil ⇒ value is invalid" idiom;
//     callers must check with errors.Is and explicitly decide whether to use
//     the partial doc:
//
//     md, err := FetchAuthorizationServerMetadata(ctx, issuer, nil)
//     switch {
//     case errors.Is(err, ErrRegistrationEndpointMissing):
//     // md is non-nil; DCR unavailable but token/JWKS endpoints usable.
//     case err != nil:
//     // md is nil; discovery failed.
//     default:
//     // md fully populated.
//     }
//
//   - On any other failure (no URL served a valid doc, transport/decode error,
//     issuer mismatch), returns (nil, err). The returned error wraps the
//     per-attempt errors via errors.Join so callers can use errors.Is /
//     errors.As to inspect underlying causes (e.g. context.DeadlineExceeded).
func FetchAuthorizationServerMetadata(
        ctx context.Context,
        issuer string,
        client *http.Client,
) (*AuthorizationServerMetadata, error) {
        urls, err := buildDiscoveryURLs(issuer)
        if err != nil {
                return nil, err
        }

        httpClient := buildDiscoveryHTTPClient(client)

        // Bound per-call timeout regardless of caller context.
        fetchCtx, cancel := context.WithTimeout(ctx, discoveryTimeout)
        defer cancel()

        var attemptErrs []error
        var partialMetadata *AuthorizationServerMetadata
        for _, u := range urls {
                metadata, err := fetchDiscoveryDocument(fetchCtx, httpClient, u, issuer)
                if err != nil {
                        attemptErrs = append(attemptErrs, fmt.Errorf("%s: %w", u, err))
                        continue
                }
                if metadata.RegistrationEndpoint == "" {
                        // Stash the first partial doc as a fallback, but keep iterating:
                        // a later URL (e.g. OIDC discovery) may advertise DCR even when
                        // the first hit (e.g. tenant-aware path-insertion) does not.
                        if partialMetadata == nil {
                                partialMetadata = metadata
                        }
                        continue
                }
                return metadata, nil
        }

        if partialMetadata != nil {
                return partialMetadata, ErrRegistrationEndpointMissing
        }

        return nil, fmt.Errorf(
                "failed to discover authorization server metadata for issuer %q: %w",
                issuer, errors.Join(attemptErrs...),
        )
}

// buildDiscoveryURLs constructs the discovery URLs to try per RFC 8414 §3.1
// and OpenID Connect Discovery 1.0. The issuer must have an "https" scheme
// (or "http" with a loopback host, for development).
//
// URLs are returned in priority order with exact duplicates removed: for an
// issuer with no tenant path, path-insertion (1) and bare RFC 8414 (3)
// collapse to the same URL, so only two distinct URLs are returned. Callers
// must not rely on a fixed slice length.
func buildDiscoveryURLs(issuer string) ([]string, error) {
        if issuer == "" {
                return nil, fmt.Errorf("issuer is required")
        }

        parsed, err := url.Parse(issuer)
        if err != nil {
                return nil, fmt.Errorf("invalid issuer URL: %w", err)
        }
        if parsed.Scheme == "" || parsed.Host == "" {
                return nil, fmt.Errorf("invalid issuer URL: scheme and host are required")
        }
        if parsed.Scheme != schemeHTTPS && !IsLoopbackHost(parsed.Host) {
                return nil, fmt.Errorf("issuer must use https (got %q)", parsed.Scheme)
        }

        // Strip trailing slash from issuer path so URL joins stay predictable.
        // Use the unescaped Path: assigning back into url.URL.Path means String()
        // will escape exactly once. Using EscapedPath() here would re-escape
        // percent-encoded tenants and produce e.g. "%252F" from "%2F".
        tenant := strings.Trim(parsed.Path, "/")

        origin := &url.URL{Scheme: parsed.Scheme, Host: parsed.Host}

        // (1) RFC 8414 §3.1 path-insertion: /.well-known/oauth-authorization-server/{tenant}
        pathInsertion := *origin
        pathInsertion.Path = path.Join(WellKnownOAuthServerPath, tenant)

        // (2) OIDC Discovery: {issuer}/.well-known/openid-configuration
        oidc := *origin
        oidc.Path = path.Join("/", tenant, WellKnownOIDCPath)

        // (3) Bare RFC 8414: {origin}/.well-known/oauth-authorization-server
        bare := *origin
        bare.Path = WellKnownOAuthServerPath

        // Deduplicate while preserving order. A tenant-less issuer causes (1) and
        // (3) to collapse to the same URL; emitting both would waste a round-trip
        // and produce confusing duplicate entries in the aggregated error message.
        candidates := []string{pathInsertion.String(), oidc.String(), bare.String()}
        seen := make(map[string]struct{}, len(candidates))
        urls := make([]string, 0, len(candidates))
        for _, u := range candidates {
                if _, ok := seen[u]; ok {
                        continue
                }
                seen[u] = struct{}{}
                urls = append(urls, u)
        }
        return urls, nil
}

// buildDiscoveryHTTPClient returns the caller-supplied client, or a default
// bounded client. The per-call timeout enforced via context.WithTimeout in
// FetchAuthorizationServerMetadata guarantees the bound is honored even when
// the caller supplies a client with a larger or missing Timeout.
func buildDiscoveryHTTPClient(client *http.Client) *http.Client {
        if client != nil {
                return client
        }
        return &http.Client{
                Timeout: discoveryTimeout,
                Transport: &http.Transport{
                        TLSHandshakeTimeout:   5 * time.Second,
                        ResponseHeaderTimeout: 5 * time.Second,
                        IdleConnTimeout:       idleConnTimeout,
                        MaxIdleConns:          maxIdleConns,
                        MaxIdleConnsPerHost:   maxIdleConnsPerHost,
                },
        }
}

// FetchAuthorizationServerMetadataFromURL fetches RFC 8414 authorization
// server metadata from a single caller-supplied URL, bypassing the
// well-known-path fallback used by FetchAuthorizationServerMetadata.
//
// It is intended for cases where the operator has configured an explicit
// discovery document URL (for example a multi-tenant IdP that does not
// advertise the tenant-aware path at {issuer}/.well-known/...). The same
// RFC 8414 §3.3 issuer-equality check is enforced: the returned metadata's
// issuer field must exactly match the caller-supplied expectedIssuer.
//
// If client is nil, the same bounded default client used by
// FetchAuthorizationServerMetadata is constructed. A 10 s per-call timeout
// is applied via context.WithTimeout regardless of the caller's context
// deadline.
//
// Return contract mirrors FetchAuthorizationServerMetadata:
//
//   - On full success, returns (metadata, nil) with a non-empty
//     RegistrationEndpoint.
//   - When the document is otherwise valid but omits
//     registration_endpoint, returns (metadata, ErrRegistrationEndpointMissing).
//     The metadata is non-nil so callers can reuse the other fields.
//   - On any other failure (transport/decode error, issuer mismatch),
//     returns (nil, err).
func FetchAuthorizationServerMetadataFromURL(
        ctx context.Context,
        discoveryURL string,
        expectedIssuer string,
        client *http.Client,
) (*AuthorizationServerMetadata, error) {
        if discoveryURL == "" {
                return nil, fmt.Errorf("discovery URL is required")
        }
        if expectedIssuer == "" {
                return nil, fmt.Errorf("expected issuer is required")
        }

        parsed, err := url.Parse(discoveryURL)
        if err != nil {
                return nil, fmt.Errorf("invalid discovery URL: %w", err)
        }
        if parsed.Scheme == "" || parsed.Host == "" {
                return nil, fmt.Errorf("invalid discovery URL: scheme and host are required")
        }
        if parsed.Scheme != schemeHTTPS && !IsLoopbackHost(parsed.Host) {
                return nil, fmt.Errorf("discovery URL must use https (got %q)", parsed.Scheme)
        }

        httpClient := buildDiscoveryHTTPClient(client)

        fetchCtx, cancel := context.WithTimeout(ctx, discoveryTimeout)
        defer cancel()

        metadata, err := fetchDiscoveryDocument(fetchCtx, httpClient, discoveryURL, expectedIssuer)
        if err != nil {
                return nil, fmt.Errorf("fetch discovery document from %q: %w", discoveryURL, err)
        }
        if metadata.RegistrationEndpoint == "" {
                return metadata, ErrRegistrationEndpointMissing
        }
        return metadata, nil
}

// fetchDiscoveryDocument performs a single GET against a discovery URL and
// returns the parsed AuthorizationServerMetadata, enforcing RFC 8414 §3.3
// issuer equality.
func fetchDiscoveryDocument(
        ctx context.Context,
        client *http.Client,
        urlStr string,
        expectedIssuer string,
) (*AuthorizationServerMetadata, error) {
        req, err := http.NewRequestWithContext(ctx, http.MethodGet, urlStr, nil)
        if err != nil {
                return nil, fmt.Errorf("build request: %w", err)
        }
        req.Header.Set("Accept", "application/json")
        req.Header.Set("User-Agent", UserAgent)

        resp, err := client.Do(req) //nolint:gosec // G107: URL is constructed from configured issuer
        if err != nil {
                return nil, fmt.Errorf("GET failed: %w", err)
        }
        defer func() {
                // Drain fully before closing so the underlying TCP connection is
                // eligible for reuse. The JSON decode path below caps the accepted
                // body via io.LimitReader; we intentionally do NOT limit the drain
                // here because a bounded drain leaves bytes in the kernel buffer
                // and defeats connection reuse.
                _, _ = io.Copy(io.Discard, resp.Body)
                if cerr := resp.Body.Close(); cerr != nil {
                        slog.Debug("failed to close discovery response body", "error", cerr)
                }
        }()

        if resp.StatusCode != http.StatusOK {
                return nil, fmt.Errorf("HTTP %d", resp.StatusCode)
        }

        contentType := strings.ToLower(resp.Header.Get("Content-Type"))
        if !strings.Contains(contentType, "application/json") {
                return nil, fmt.Errorf("unexpected content-type %q", contentType)
        }

        var metadata AuthorizationServerMetadata
        if err := json.NewDecoder(io.LimitReader(resp.Body, discoveryMaxResponseSize)).Decode(&metadata); err != nil {
                return nil, fmt.Errorf("decode response: %w", err)
        }

        // RFC 8414 §3.3: the issuer in the metadata MUST exactly match the caller's expected issuer.
        if metadata.Issuer != expectedIssuer {
                return nil, fmt.Errorf("issuer mismatch (RFC 8414 §3.3): expected %q, got %q", expectedIssuer, metadata.Issuer)
        }

        return &metadata, nil
}

// AuthorizationServerMetadata represents the OAuth 2.0 Authorization Server Metadata
// per RFC 8414. This is the base structure that OIDC Discovery extends.
type AuthorizationServerMetadata struct {
        // Issuer is the authorization server's issuer identifier (REQUIRED per RFC 8414).
        Issuer string `json:"issuer"`

        // AuthorizationEndpoint is the URL of the authorization endpoint (RECOMMENDED).
        // It is omitted for token-only authorization servers.
        AuthorizationEndpoint string `json:"authorization_endpoint,omitempty"`

        // TokenEndpoint is the URL of the token endpoint (RECOMMENDED).
        // Note: No omitempty to maintain backward compatibility with existing JSON serialization.
        TokenEndpoint string `json:"token_endpoint"`

        // JWKSURI is the URL of the JSON Web Key Set document (RECOMMENDED).
        // Note: No omitempty to maintain backward compatibility with existing JSON serialization.
        JWKSURI string `json:"jwks_uri"`

        // RegistrationEndpoint is the URL of the Dynamic Client Registration endpoint (OPTIONAL).
        RegistrationEndpoint string `json:"registration_endpoint,omitempty"`

        // DeviceAuthorizationEndpoint is the RFC 8628 device authorization endpoint (OPTIONAL).
        DeviceAuthorizationEndpoint string `json:"device_authorization_endpoint,omitempty"`

        // IntrospectionEndpoint is the URL of the token introspection endpoint (OPTIONAL, RFC 7662).
        IntrospectionEndpoint string `json:"introspection_endpoint,omitempty"`

        // UserinfoEndpoint is the URL of the UserInfo endpoint (RECOMMENDED per OIDC Discovery, not in RFC 8414).
        // Omitted from JSON when empty to avoid serializing an invalid URL value.
        UserinfoEndpoint string `json:"userinfo_endpoint,omitempty"`

        // ResponseTypesSupported lists the response types supported (RECOMMENDED).
        ResponseTypesSupported []string `json:"response_types_supported,omitempty"`

        // GrantTypesSupported lists the grant types supported (OPTIONAL).
        GrantTypesSupported []string `json:"grant_types_supported,omitempty"`

        // CodeChallengeMethodsSupported lists the PKCE code challenge methods supported (OPTIONAL).
        CodeChallengeMethodsSupported []string `json:"code_challenge_methods_supported,omitempty"`

        // TokenEndpointAuthMethodsSupported lists the authentication methods supported at the token endpoint (OPTIONAL).
        TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported,omitempty"`

        // TokenEndpointAuthSigningAlgValuesSupported lists the JWS algorithms supported
        // for private_key_jwt client assertions (RFC 8414).
        TokenEndpointAuthSigningAlgValuesSupported []string `json:"token_endpoint_auth_signing_alg_values_supported,omitempty"`

        // ScopesSupported lists the OAuth 2.0 scope values supported (RECOMMENDED per RFC 8414).
        // For MCP authorization servers, this typically includes "openid" and "offline_access".
        ScopesSupported []string `json:"scopes_supported,omitempty"`

        // ClientIDMetadataDocumentSupported indicates the server accepts HTTPS URLs as client_id
        // values per draft-ietf-oauth-client-id-metadata-document.
        ClientIDMetadataDocumentSupported bool `json:"client_id_metadata_document_supported,omitempty"`
}

// OIDCDiscoveryDocument represents the OpenID Connect Discovery 1.0 document.
// It extends OAuth 2.0 Authorization Server Metadata (RFC 8414) with OIDC-specific fields.
// This unified type supports both producer (server) and consumer (client) use cases.
type OIDCDiscoveryDocument struct {
        // Embed OAuth 2.0 AS Metadata (RFC 8414) as the base
        AuthorizationServerMetadata

        // SubjectTypesSupported lists the subject identifier types supported (REQUIRED for OIDC).
        SubjectTypesSupported []string `json:"subject_types_supported,omitempty"`

        // IDTokenSigningAlgValuesSupported lists the JWS algorithms supported for ID tokens (REQUIRED for OIDC).
        IDTokenSigningAlgValuesSupported []string `json:"id_token_signing_alg_values_supported,omitempty"`

        // ClaimsSupported lists the claims that can be returned (RECOMMENDED for OIDC).
        ClaimsSupported []string `json:"claims_supported,omitempty"`
}

// Validate performs basic validation on the discovery document.
// It checks for required fields based on whether this is an OIDC or pure OAuth document.
func (d *OIDCDiscoveryDocument) Validate(isOIDC bool) error {
        if d.Issuer == "" {
                return ErrMissingIssuer
        }
        tokenOnly := !isOIDC && d.AuthorizationEndpoint == "" && d.SupportsGrantType(GrantTypeTokenExchange)
        if d.AuthorizationEndpoint == "" && !tokenOnly {
                return ErrMissingAuthorizationEndpoint
        }
        if d.TokenEndpoint == "" {
                return ErrMissingTokenEndpoint
        }
        if isOIDC && d.JWKSURI == "" {
                return ErrMissingJWKSURI
        }
        if isOIDC && len(d.ResponseTypesSupported) == 0 {
                return ErrMissingResponseTypesSupported
        }
        return nil
}

// SupportsPKCE returns true if the authorization server supports PKCE with S256.
func (d *OIDCDiscoveryDocument) SupportsPKCE() bool {
        for _, method := range d.CodeChallengeMethodsSupported {
                if method == PKCEMethodS256 {
                        return true
                }
        }
        return false
}

// SupportsGrantType returns true if the authorization server supports the given grant type.
func (d *OIDCDiscoveryDocument) SupportsGrantType(grantType string) bool {
        for _, gt := range d.GrantTypesSupported {
                if gt == grantType {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauthproto

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

        "golang.org/x/oauth2"
)

// The tokenJSON struct shape and expirationTime.UnmarshalJSON below are
// adapted from golang.org/x/oauth2/internal/token.go (BSD-3-Clause,
// compatible with Apache-2.0). The upstream file is the authoritative
// reference for handling the PayPal-style JSON-string expires_in field and
// for the RFC 6749 Section 5.2 "2xx body carries an error" rule. The shape
// here is extended with the RFC 8693 issued_token_type field and the scope
// field so a single entry point serves both the authorization-code and
// token-exchange grants.

// Redact returns "<empty>" for an empty input and "[REDACTED]" otherwise.
// Grant-subpackage Config.String() methods use it to keep secrets (client
// secrets, JWT assertions, refresh tokens) out of logs and error output
// without each Config reimplementing the empty-vs-nonempty branch.
func Redact(value string) string {
        if value == "" {
                return "<empty>"
        }
        return "[REDACTED]"
}

// TokenResponse is the public result of a successful token endpoint exchange.
// It composes *oauth2.Token rather than embedding it, so the Token's
// Valid() / Extra() / Type() helpers are not promoted onto this type and
// future oauth2.Token additions cannot leak into this package's API.
//
// For the fields that RFC 8693 Section 2.2.1 requires beyond the standard
// oauth2 token (issued_token_type and scope), callers read them off the
// response directly.
type TokenResponse struct {
        // Token carries AccessToken, TokenType, RefreshToken, and Expiry populated
        // from the response body. The Raw and other unexported fields on
        // oauth2.Token are not set — use the sibling fields on TokenResponse.
        Token *oauth2.Token

        // IssuedTokenType is the RFC 8693 issued_token_type URN when the server
        // performed a token exchange. Empty for plain RFC 6749 responses.
        IssuedTokenType string

        // Scope is the raw (space-separated) scope string returned by the server.
        // Callers that need the list form can strings.Fields(scope).
        Scope string
}

// ParseTokenResponse is the single entry point for decoding a token endpoint
// response. It enforces RFC 6749 Sections 5.1 and 5.2 in a single place:
//
//   - The body is decoded as JSON first, independent of the HTTP status.
//   - If the status is non-2xx, OR the body carries an "error" field (RFC
//     6749 §5.2 — some providers return HTTP 200 with an error payload),
//     an *oauth2.RetrieveError is returned and *TokenResponse is nil.
//   - On success, access_token must be non-empty (RFC 6749 §5.1). token_type
//     is intentionally NOT required here: the x/oauth2 library treats it as
//     optional (Token.Type() defaults to "Bearer") and Google historically
//     omits it. Per-grant callers are responsible for any stricter validation
//     their specification demands.
//
// The caller is responsible for grant-specific validation (for example,
// RFC 8693 Section 2.2.1 requires the issued_token_type field to be present;
// that check belongs in the token-exchange grant's call site, not here).
//
// Malformed JSON on a failure status still yields an *oauth2.RetrieveError
// with the raw body preserved, so callers can surface the server's reply
// verbatim. Malformed JSON on a 2xx status is returned as a wrapped
// json.SyntaxError / json.UnmarshalTypeError via fmt.Errorf("%w", ...).
func ParseTokenResponse(resp *http.Response, body []byte) (*TokenResponse, error) {
        failureStatus := resp.StatusCode < 200 || resp.StatusCode > 299

        var tj tokenJSON
        if err := json.Unmarshal(body, &tj); err != nil {
                if failureStatus {
                        return nil, parseRetrieveError(resp, body)
                }
                return nil, fmt.Errorf("oauth: cannot parse token response: %w", err)
        }

        if failureStatus || tj.ErrorCode != "" {
                return nil, parseRetrieveError(resp, body)
        }

        if tj.AccessToken == "" {
                return nil, fmt.Errorf("oauth: token response missing access_token (RFC 6749 Section 5.1)")
        }

        token := &oauth2.Token{
                AccessToken:  tj.AccessToken,
                TokenType:    tj.TokenType,
                RefreshToken: tj.RefreshToken,
        }
        if tj.ExpiresIn > 0 {
                token.Expiry = time.Now().Add(time.Duration(tj.ExpiresIn) * time.Second)
        }

        return &TokenResponse{
                Token:           token,
                IssuedTokenType: tj.IssuedTokenType,
                Scope:           tj.Scope,
        }, nil
}

// NewFormRequest builds an HTTP POST request for an OAuth 2.0 token endpoint
// call. The body is the URL-encoded form in data, with Content-Type and
// Content-Length set. When both clientID and clientSecret are non-empty,
// HTTP Basic authentication is attached per RFC 6749 Section 2.3.1; the
// credentials are URL-encoded before being passed to SetBasicAuth, which is
// what Go's SetBasicAuth and the OAuth 2.0 spec both require.
//
// Callers own the request's Context — pass the deadline or cancellation
// signal they want DoTokenRequest to honour.
func NewFormRequest(
        ctx context.Context,
        endpoint string,
        data url.Values,
        clientID, clientSecret string,
) (*http.Request, error) {
        encoded := data.Encode()
        req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(encoded))
        if err != nil {
                return nil, fmt.Errorf("oauth: build token request: %w", err)
        }

        req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
        req.Header.Set("Content-Length", strconv.Itoa(len(encoded)))

        // RFC 6749 Section 2.3.1 requires URL-encoding of client credentials when
        // sent via Basic auth, and Go's SetBasicAuth docs mirror that requirement
        // for OAuth2 compatibility.
        if clientID != "" && clientSecret != "" {
                req.SetBasicAuth(url.QueryEscape(clientID), url.QueryEscape(clientSecret))
        }

        return req, nil
}

// DoTokenRequest executes a prepared token endpoint request and returns the
// parsed response. It is the high-level counterpart to NewFormRequest: most
// grants call NewFormRequest followed by DoTokenRequest.
//
// Passing a nil client is explicitly supported and selects the package-level
// shared default (see DefaultHTTPClient). Callers that need custom timeouts,
// a custom transport, or a test double MUST supply their own *http.Client —
// the nil shortcut does not take any caller-visible options.
//
// Behavior:
//
//   - If client is nil, DefaultHTTPClient is used so callers automatically
//     get the shared transport (connection reuse, consistent timeouts).
//   - The response body is initially read with io.LimitReader capped at
//     maxResponseBodySize (1 MiB, matching x/oauth2) before any parsing, so
//     a pathological server cannot exhaust memory. Before closing, at most one
//     additional maxResponseBodySize is drained to enable connection reuse
//     without unbounded reads from oversized or never-ending bodies.
//   - On every exit path — success, JSON decode failure, and RetrieveError
//     — the body is bounded-drained and closed.
//   - RFC 6749 Section 5.2 routing (a 2xx body with an "error" field) is
//     handled inside ParseTokenResponse; DoTokenRequest surfaces the
//     resulting *oauth2.RetrieveError unchanged.
//
// The request's own Context is authoritative: NewFormRequest builds the
// request with the caller's context attached, and client.Do observes
// req.Context() for cancellation and deadlines. A cancelled context fails
// fast without reaching the server.
func DoTokenRequest(client *http.Client, req *http.Request) (*TokenResponse, error) {
        if client == nil {
                client = DefaultHTTPClient()
        }

        resp, err := client.Do(req)
        if err != nil {
                return nil, fmt.Errorf("oauth: token request failed: %w", err)
        }
        defer func() {
                // A bounded drain permits connection reuse for ordinary response bodies
                // without indefinitely reading oversized or never-ending IdP responses.
                _, _ = io.CopyN(io.Discard, resp.Body, maxResponseBodySize)
                if closeErr := resp.Body.Close(); closeErr != nil {
                        slog.Debug("oauth: close token response body", "error", closeErr)
                }
        }()

        body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBodySize))
        if err != nil {
                return nil, fmt.Errorf("oauth: read token response body: %w", err)
        }

        return ParseTokenResponse(resp, body)
}

// DefaultHTTPClient returns the process-wide *http.Client used by OAuth
// grant helpers when no explicit client is injected. Callers that build
// their own requests and call client.Do directly (without going through
// DoTokenRequest) use this to pick up the shared transport and inherit
// the same connection-reuse and timeout behavior as the helper path.
//
// The returned client is a process-wide singleton — callers MUST NOT mutate
// its Timeout or Transport fields. Code wanting custom timeouts must
// construct its own *http.Client and pass it to DoTokenRequest.
//
// http.Client is documented as safe for concurrent use by multiple goroutines
// (see https://pkg.go.dev/net/http#Client), so the returned value can be
// shared across goroutines without additional synchronization.
//
// TODO: consider a future opt-in SSRF-protected variant backed by
// pkg/networking.NewHttpClientBuilder. The builder blocks loopback and RFC
// 1918 ranges, which would break localhost IdPs (dex, Keycloak-in-Docker)
// and the httptest.NewServer-based tests that bind to 127.0.0.1. Not a
// default today for behavior-compatibility with pkg/oauthproto/tokenexchange.
func DefaultHTTPClient() *http.Client {
        return sharedHTTPClient
}

// sharedHTTPClient is the process-wide default client used by OAuth grant
// helpers (see DoTokenRequest, DefaultHTTPClient). Initialized once in init
// so callers share the underlying transport and benefit from connection
// reuse across grants.
var sharedHTTPClient *http.Client

func init() {
        // Base on http.DefaultTransport.Clone() so we inherit stdlib defaults:
        // DialContext, Proxy: ProxyFromEnvironment (honors HTTP_PROXY/NO_PROXY
        // for corporate deployments), MaxIdleConns=100, IdleConnTimeout=90s,
        // ForceAttemptHTTP2=true. Overriding only the handshake timeout keeps
        // pool tuning and HTTP/2 auto-upgrade consistent with the rest of the
        // ecosystem.
        transport := http.DefaultTransport.(*http.Transport).Clone()
        // Tighter than stdlib's 10s default is not needed here, but an explicit
        // cap prevents a hanging TLS handshake from silently consuming the
        // outer 30s request budget.
        transport.TLSHandshakeTimeout = 10 * time.Second
        // Deliberately NOT setting ResponseHeaderTimeout: some corporate IdP
        // chains (Entra OBO, federated Okta) legitimately take >10s to produce
        // the first byte. The outer Client.Timeout (defaultHTTPTimeout) is the
        // authoritative budget for the whole exchange.
        sharedHTTPClient = &http.Client{
                Timeout:   defaultHTTPTimeout,
                Transport: transport,
        }
}

// tokenJSON is the wire shape decoded from a token endpoint response body.
// It intentionally covers both success and failure paths: RFC 6749 §5.2
// allows an "error" code to appear in a body returned with a 2xx status, so
// ParseTokenResponse decodes into this struct unconditionally before
// deciding which branch to return.
type tokenJSON struct {
        AccessToken     string         `json:"access_token"`
        TokenType       string         `json:"token_type"`
        RefreshToken    string         `json:"refresh_token"`
        ExpiresIn       expirationTime `json:"expires_in"` // number OR decimal string (PayPal, etc.)
        IssuedTokenType string         `json:"issued_token_type"`
        Scope           string         `json:"scope"`

        // RFC 6749 Section 5.2 error fields; may legitimately appear in 2xx bodies.
        ErrorCode        string `json:"error"`
        ErrorDescription string `json:"error_description"`
        ErrorURI         string `json:"error_uri"`
}

// expirationTime accepts either a JSON number or a JSON string containing a
// decimal integer. At least PayPal returns expires_in as a string; a naive
// int field would fail to decode there. Negative values are treated as zero
// (no expiry), values larger than math.MaxInt32 are clamped. Copied with
// attribution from golang.org/x/oauth2/internal/token.go.
type expirationTime int32

// UnmarshalJSON implements json.Unmarshaler.
func (e *expirationTime) UnmarshalJSON(b []byte) error {
        if len(b) == 0 || string(b) == "null" {
                return nil
        }
        var n json.Number
        if err := json.Unmarshal(b, &n); err != nil {
                return err
        }
        i, err := n.Int64()
        if err != nil {
                return err
        }
        if i < 0 {
                i = 0
        }
        if i > math.MaxInt32 {
                i = math.MaxInt32
        }
        *e = expirationTime(i)
        return nil
}

// retrieveErrorBody mirrors the RFC 6749 Section 5.2 fields that may appear
// in a token endpoint error response (both 2xx and non-2xx bodies — some
// providers return 200 with "error" set).
type retrieveErrorBody struct {
        ErrorCode        string `json:"error"`
        ErrorDescription string `json:"error_description"`
        ErrorURI         string `json:"error_uri"`
}

// parseRetrieveError builds an *oauth2.RetrieveError from a token endpoint
// response. The resulting error always has Response and Body populated; the
// three OAuth fields are best-effort — a malformed or non-JSON body yields
// empty strings rather than an error. This helper is the single funnel
// ParseTokenResponse uses to report both non-2xx responses and 2xx bodies
// that carry an "error" field (RFC 6749 §5.2).
//
// Unexported: callers route through ParseTokenResponse.
func parseRetrieveError(resp *http.Response, body []byte) *oauth2.RetrieveError {
        retrieveErr := &oauth2.RetrieveError{
                Response: resp,
                Body:     body,
        }

        // Best-effort decode; malformed JSON leaves the OAuth fields empty.
        var parsed retrieveErrorBody
        if err := json.Unmarshal(body, &parsed); err == nil {
                retrieveErr.ErrorCode = parsed.ErrorCode
                retrieveErr.ErrorDescription = parsed.ErrorDescription
                retrieveErr.ErrorURI = parsed.ErrorURI
        }

        return retrieveErr
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauthproto

import (
        "net"
        "strings"
)

// IsLoopbackHost reports whether host is a loopback hostname or IP address.
// pkg/networking wraps this function in its own IsLocalhost to avoid a
// reverse import dependency from this leaf package into networking.
//
// Recognised forms: "localhost" (case-insensitive) and any loopback IP
// (127.0.0.0/8, ::1), each with or without a port, and IPv6 literals with or
// without brackets. url.Hostname() strips brackets, so callers passing a
// parsed hostname hand us the bare "::1" form.
func IsLoopbackHost(host string) bool {
        h, _, err := net.SplitHostPort(host)
        if err != nil {
                // No port (or a bare IPv6 literal, which SplitHostPort rejects for
                // missing brackets) — treat the whole value as the host.
                h = host
        }
        h = strings.TrimPrefix(strings.TrimSuffix(h, "]"), "[")
        if strings.EqualFold(h, "localhost") {
                return true
        }
        ip := net.ParseIP(h)
        return ip != nil && ip.IsLoopback()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauthproto

import (
        "context"
        "fmt"
        "net/url"

        "github.com/ory/fosite"
)

// MaxRedirectURILength is the maximum allowed length for a single redirect URI.
// This limit provides DoS protection during URI parsing per RFC 3986 practical constraints.
const MaxRedirectURILength = 2048

// RedirectURIPolicy controls which URI schemes are accepted during redirect URI validation.
type RedirectURIPolicy int

const (
        // RedirectURIPolicyStrict allows only https and http-loopback schemes.
        // This follows RFC 8252 Section 8.4 strict security recommendations and
        // is appropriate for dynamically registered clients where scheme hijacking
        // is a concern.
        RedirectURIPolicyStrict RedirectURIPolicy = iota

        // RedirectURIPolicyAllowPrivateSchemes also allows private-use URI schemes
        // (e.g., cursor://, vscode://) per RFC 8252 Section 7.1.
        // This is appropriate for pre-registered/static clients where the administrator
        // explicitly configures trusted redirect URIs for native applications.
        RedirectURIPolicyAllowPrivateSchemes

        // RedirectURIPolicyAllowHTTP allows http:// for any host, bypassing the
        // loopback-only restriction. Intended for in-cluster Kubernetes deployments
        // where the OAuth callback runs over HTTP on a trusted internal network.
        RedirectURIPolicyAllowHTTP
)

// ValidateRedirectURI validates a redirect URI per RFC 6749 Section 3.1.2 and RFC 8252.
// The policy parameter controls whether private-use URI schemes are accepted.
//
// Validation rules applied:
//   - URI must not exceed MaxRedirectURILength (DoS protection)
//   - URI must be an absolute URI with a scheme (RFC 6749 Section 3.1.2)
//   - URI must not contain a fragment component (RFC 6749 Section 3.1.2)
//   - Scheme security per policy:
//   - Strict: only https or http-loopback (RFC 8252 Section 8.4)
//   - AllowPrivateSchemes: also allows private-use schemes (RFC 8252 Section 7.1)
func ValidateRedirectURI(uri string, policy RedirectURIPolicy) error {
        if len(uri) > MaxRedirectURILength {
                return fmt.Errorf("redirect_uri too long (maximum %d characters)", MaxRedirectURILength)
        }

        parsed, err := url.Parse(uri)
        if err != nil {
                return fmt.Errorf("invalid redirect_uri format: %w", err)
        }

        // RFC 6749 Section 3.1.2: must be absolute URI without fragment
        if !fosite.IsValidRedirectURI(parsed) {
                return fmt.Errorf("redirect_uri must be an absolute URI without a fragment")
        }

        // Apply scheme security policy
        switch policy {
        case RedirectURIPolicyStrict:
                // RFC 8252 Section 8.4: only https or http for loopback
                if !fosite.IsRedirectURISecureStrict(context.Background(), parsed) {
                        return fmt.Errorf("redirect_uri must use http (for loopback) or https scheme")
                }
        case RedirectURIPolicyAllowPrivateSchemes:
                // RFC 8252 Section 7.1: also allow private-use URI schemes
                if !fosite.IsRedirectURISecure(context.Background(), parsed) {
                        return fmt.Errorf("redirect_uri must use a secure scheme (https, http for loopback, or a private-use scheme)")
                }
        case RedirectURIPolicyAllowHTTP:
                // In-cluster mode: any http or https URI is accepted; only the basic
                // absolute-URI-without-fragment requirement from RFC 6749 Section 3.1.2 is enforced.
                if parsed.Scheme != "http" && parsed.Scheme != "https" {
                        return fmt.Errorf("redirect_uri must use http or https scheme")
                }
        default:
                return fmt.Errorf("unknown redirect URI policy: %d", policy)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package tokenexchange provides OAuth 2.0 Token Exchange (RFC 8693) support.
package tokenexchange

import (
        "context"
        "errors"
        "fmt"
        "net/http"
        "net/url"
        "strings"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/oauthproto"
)

// NormalizeTokenType converts a short token type name to its full URN.
// Accepts both short forms ("access_token", "id_token", "jwt") and full URNs.
// Returns the full URN or an error if the token type is invalid.
//
// This is primarily intended for CLI/user input processing. Internal APIs
// should use full URNs directly.
func NormalizeTokenType(tokenType string) (string, error) {
        // Empty string is valid (will use default)
        if tokenType == "" {
                return "", nil
        }

        // Check if already a full URN (backward compatibility)
        switch tokenType {
        case oauthproto.TokenTypeAccessToken, oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT:
                return tokenType, nil
        }

        // Convert short form to full URN
        switch tokenType {
        case "access_token":
                return oauthproto.TokenTypeAccessToken, nil
        case "id_token":
                return oauthproto.TokenTypeIDToken, nil
        case "jwt":
                return oauthproto.TokenTypeJWT, nil
        default:
                return "", fmt.Errorf("invalid token type %q: must be one of: access_token, id_token, jwt", tokenType)
        }
}

// actingParty represents the acting party in a token exchange delegation scenario.
// When present, it indicates that the actor token holder is acting on behalf of the subject token holder.
type actingParty struct {
        ActorToken     string
        ActorTokenType string
}

// exchangeRequest contains fields necessary to make an OAuth 2.0 token exchange.
// Based on RFC 8693: https://datatracker.ietf.org/doc/html/rfc8693
type exchangeRequest struct {
        // Required fields
        GrantType          string
        SubjectToken       string
        SubjectTokenType   string
        RequestedTokenType string

        // Optional fields
        Resource    string
        Audience    string
        Scope       []string
        ActingParty *actingParty
}

// clientAuthentication represents OAuth client credentials for token exchange.
type clientAuthentication struct {
        ClientID     string
        ClientSecret string
}

// ExchangeConfig holds the configuration for token exchange.
type ExchangeConfig struct {
        // TokenURL is the OAuth 2.0 token endpoint URL
        TokenURL string

        // ClientID is the OAuth 2.0 client identifier
        ClientID string

        // ClientSecret is the OAuth 2.0 client secret
        ClientSecret string

        // Audience is the target audience for the exchanged token (optional per RFC 8693)
        Audience string

        // Scopes is the list of scopes to request (optional per RFC 8693)
        Scopes []string

        // SubjectTokenType specifies the type of the subject token being exchanged.
        // Common values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT.
        // If empty, defaults to oauthproto.TokenTypeAccessToken.
        SubjectTokenType string

        // RequestedTokenType specifies the desired token type in the exchange response.
        // Defaults to oauthproto.TokenTypeAccessToken per RFC 8693. Set to any RFC 8693
        // token-type URN to request a different issued token type.
        RequestedTokenType string

        // Resource identifies the target resource per RFC 8693 Section 2.1 and RFC 8707.
        // Must be an absolute URI without a fragment. This implementation accepts a
        // single resource indicator; RFC 8707 permits multiple but multi-value support
        // is not yet exposed at this layer.
        Resource string

        // SubjectTokenProvider is a function that returns the subject token to exchange
        // we use a function to allow dynamic retrieval of the token (e.g. from request context)
        // and also to lazy-load the token only when needed, load from dynamic sources, etc.
        SubjectTokenProvider func() (string, error)

        // HTTPClient is the HTTP client to use for token exchange requests.
        // If nil, oauthproto.DefaultHTTPClient() will be used.
        HTTPClient *http.Client
}

// Validate checks if the ExchangeConfig contains all required fields.
//
// Side effect: when SubjectTokenType is provided as a short form
// ("access_token", "id_token", "jwt"), Validate normalizes it to the full
// RFC 8693 URN and writes the result back onto the receiver. Callers in
// pkg/vmcp/auth/strategies/tokenexchange.go and pkg/runner/config_builder.go
// read the normalized value after Validate returns, so this mutation is part
// of the documented contract.
func (c *ExchangeConfig) Validate() error {
        if c.TokenURL == "" {
                return fmt.Errorf("TokenURL is required")
        }

        if c.SubjectTokenProvider == nil {
                return fmt.Errorf("SubjectTokenProvider is required")
        }

        // ClientID is optional - some token exchange endpoints (like Google STS)
        // don't require client credentials and rely on the trust relationship
        // configured in the identity provider (e.g., Workload Identity Federation)

        // Validate and normalize SubjectTokenType if provided
        if c.SubjectTokenType != "" {
                normalized, err := NormalizeTokenType(c.SubjectTokenType)
                if err != nil {
                        return fmt.Errorf("invalid SubjectTokenType: %w", err)
                }
                // Update the config with the normalized value
                c.SubjectTokenType = normalized
        }

        // Validate Resource per RFC 8707 Section 2: absolute URI, no fragment.
        if c.Resource != "" {
                u, err := url.Parse(c.Resource)
                if err != nil {
                        return fmt.Errorf("invalid Resource URI: %w", err)
                }
                if !u.IsAbs() {
                        return fmt.Errorf("invalid Resource: must be an absolute URI (got %q)", c.Resource)
                }
                if u.Fragment != "" {
                        return fmt.Errorf("invalid Resource: must not include a fragment (RFC 8707 Section 2)")
                }
        }

        // Validate URL format
        _, err := url.Parse(c.TokenURL)
        if err != nil {
                return fmt.Errorf("TokenURL is not a valid URL: %w", err)
        }

        return nil
}

// tokenSource implements oauth2.TokenSource for token exchange.
type tokenSource struct {
        ctx  context.Context
        conf *ExchangeConfig
}

// Token implements oauth2.TokenSource interface.
// It performs the token exchange and returns an oauth2.Token.
func (ts *tokenSource) Token() (*oauth2.Token, error) {
        conf := ts.conf

        // Validate configuration
        if err := conf.Validate(); err != nil {
                return nil, fmt.Errorf("invalid config: %w", err)
        }

        // Get the subject token from the provider
        subjectToken, err := conf.SubjectTokenProvider()
        if err != nil {
                return nil, fmt.Errorf("failed to get subject token: %w", err)
        }

        // Determine subject token type (default to access_token if not specified)
        subjectTokenType := conf.SubjectTokenType
        if subjectTokenType == "" {
                subjectTokenType = oauthproto.TokenTypeAccessToken
        }

        // Build the token exchange request
        requestedTokenType := conf.RequestedTokenType
        if requestedTokenType == "" {
                requestedTokenType = oauthproto.TokenTypeAccessToken
        }

        request := &exchangeRequest{
                GrantType:          oauthproto.GrantTypeTokenExchange,
                Audience:           conf.Audience,
                Scope:              conf.Scopes,
                RequestedTokenType: requestedTokenType,
                Resource:           conf.Resource,
                SubjectToken:       subjectToken,
                SubjectTokenType:   subjectTokenType,
        }

        clientAuth := clientAuthentication{
                ClientID:     conf.ClientID,
                ClientSecret: conf.ClientSecret,
        }

        // Perform the exchange
        resp, err := exchangeToken(ts.ctx, conf.TokenURL, request, clientAuth, conf.HTTPClient)
        if err != nil {
                return nil, err
        }

        // RFC 8693 Section 2.2.1 requires token_type in the response. The shared
        // oauthproto.ParseTokenResponse is intentionally permissive on this field
        // (matching x/oauth2); token exchange tightens it back.
        if resp.Token.TokenType == "" {
                return nil, fmt.Errorf("token exchange: server returned empty token_type (required by RFC 8693)")
        }
        // RFC 8693 Section 2.2.1 requires issued_token_type in the response.
        if resp.IssuedTokenType == "" {
                return nil, fmt.Errorf("token exchange: server returned empty issued_token_type (required by RFC 8693)")
        }

        // Surface the RFC 8693 issued_token_type so callers can verify the server
        // returned the type they asked for (e.g., XAA's IdP exchange requires an ID-JAG).
        token := resp.Token.WithExtra(map[string]any{
                "issued_token_type": resp.IssuedTokenType,
        })

        return token, nil
}

// TokenSource returns an oauth2.TokenSource that performs token exchange.
func (c *ExchangeConfig) TokenSource(ctx context.Context) oauth2.TokenSource {
        return &tokenSource{
                ctx:  ctx,
                conf: c,
        }
}

// exchangeToken performs the actual HTTP request for token exchange.
// This is the internal implementation used by tokenSource.Token().
func exchangeToken(
        ctx context.Context,
        endpoint string,
        request *exchangeRequest,
        auth clientAuthentication,
        client *http.Client,
) (*oauthproto.TokenResponse, error) {
        data, err := buildFormData(request)
        if err != nil {
                return nil, err
        }

        req, err := oauthproto.NewFormRequest(ctx, endpoint, data, auth.ClientID, auth.ClientSecret)
        if err != nil {
                return nil, fmt.Errorf("tokenexchange: build request: %w", err)
        }

        resp, err := oauthproto.DoTokenRequest(client, req)
        if err != nil {
                // Preserve the pre-refactor behavior: scrub RetrieveError.Body so raw
                // upstream content cannot leak into error strings via err.Error().
                // pkg/oauthproto deliberately preserves Body for general-purpose
                // callers; tokenexchange opts back into the stricter behavior because
                // its errors propagate through vmcp / runner paths that may log them.
                var retrieveErr *oauth2.RetrieveError
                if errors.As(err, &retrieveErr) {
                        retrieveErr.Body = nil
                }
                return nil, err
        }
        return resp, nil
}

// buildFormData constructs the form data for a token exchange request according to RFC 8693.
func buildFormData(request *exchangeRequest) (url.Values, error) {
        data := url.Values{}

        // Grant type is always token exchange
        if request.GrantType == "" {
                request.GrantType = oauthproto.GrantTypeTokenExchange
        }
        data.Set("grant_type", request.GrantType)

        // Subject token is required
        if request.SubjectToken == "" {
                return nil, fmt.Errorf("subject_token is required")
        }
        data.Set("subject_token", request.SubjectToken)

        // Subject token type defaults to access_token if not specified
        if request.SubjectTokenType == "" {
                request.SubjectTokenType = oauthproto.TokenTypeAccessToken
        }
        data.Set("subject_token_type", request.SubjectTokenType)

        // Requested token type defaults to access_token if not specified
        if request.RequestedTokenType == "" {
                request.RequestedTokenType = oauthproto.TokenTypeAccessToken
        }
        data.Set("requested_token_type", request.RequestedTokenType)

        addOptionalFields(data, request)

        return data, nil
}

// addOptionalFields adds optional RFC 8693 fields to the form data.
func addOptionalFields(data url.Values, request *exchangeRequest) {
        if request.Audience != "" {
                data.Set("audience", request.Audience)
        }
        if len(request.Scope) > 0 {
                data.Set("scope", strings.Join(request.Scope, " "))
        }
        if request.Resource != "" {
                data.Set("resource", request.Resource)
        }

        // Actor token (for delegation scenarios)
        if request.ActingParty != nil && request.ActingParty.ActorToken != "" {
                data.Set("actor_token", request.ActingParty.ActorToken)
                if request.ActingParty.ActorTokenType != "" {
                        data.Set("actor_token_type", request.ActingParty.ActorTokenType)
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package tokenexchange

import (
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "os"

        "github.com/golang-jwt/jwt/v5"
        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/oauthproto"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// Middleware type constant
const (
        MiddlewareType = "tokenexchange"
)

// Header injection strategy constants
const (
        // HeaderStrategyReplace replaces the Authorization header with the exchanged token
        HeaderStrategyReplace = "replace"
        // HeaderStrategyCustom adds the exchanged token to a custom header
        HeaderStrategyCustom = "custom"
)

// Environment variable names
const (
        // EnvClientSecret is the environment variable name for the OAuth client secret
        // This corresponds to the "client_secret" field in the token exchange configuration
        //nolint:gosec // G101: this is an environment variable name, not a credential value
        EnvClientSecret = "TOOLHIVE_TOKEN_EXCHANGE_CLIENT_SECRET"
)

var errUnknownStrategy = errors.New("unknown token injection strategy")

// envGetter is a function that retrieves environment variables
// This can be overridden for testing
type envGetter func(string) string

// defaultEnvGetter is the default environment variable getter using os.Getenv
var defaultEnvGetter envGetter = os.Getenv

// MiddlewareParams represents the parameters for token exchange middleware
type MiddlewareParams struct {
        TokenExchangeConfig *Config `json:"token_exchange_config,omitempty"`
}

// Config holds configuration for token exchange middleware
type Config struct {
        // TokenURL is the OAuth 2.0 token endpoint URL
        TokenURL string `json:"token_url"`

        // ClientID is the OAuth 2.0 client identifier
        ClientID string `json:"client_id"`

        // ClientSecret is the OAuth 2.0 client secret
        ClientSecret string `json:"client_secret"`

        // Audience is the target audience for the exchanged token
        Audience string `json:"audience"`

        // Scopes is the list of scopes to request for the exchanged token
        Scopes []string `json:"scopes,omitempty"`

        // SubjectTokenType specifies the type of the subject token being exchanged.
        // Common values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT.
        // If empty, defaults to oauthproto.TokenTypeAccessToken.
        SubjectTokenType string `json:"subject_token_type,omitempty"`

        // HeaderStrategy determines how to inject the token
        // Valid values: HeaderStrategyReplace (default), HeaderStrategyCustom
        HeaderStrategy string `json:"header_strategy,omitempty"`

        // ExternalTokenHeaderName is the name of the custom header to use when HeaderStrategy is "custom"
        ExternalTokenHeaderName string `json:"external_token_header_name,omitempty"`
}

// Middleware wraps token exchange middleware functionality
type Middleware struct {
        middleware types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *Middleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (*Middleware) Close() error {
        // Token exchange middleware doesn't need cleanup
        return nil
}

// CreateMiddleware factory function for token exchange middleware
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal token exchange middleware parameters: %w", err)
        }

        // Token exchange config is required when this middleware type is specified
        if params.TokenExchangeConfig == nil {
                return fmt.Errorf("token exchange configuration is required but not provided")
        }

        // Validate configuration
        if err := validateTokenExchangeConfig(params.TokenExchangeConfig); err != nil {
                return fmt.Errorf("invalid token exchange configuration: %w", err)
        }

        middleware, err := createTokenExchangeMiddleware(*params.TokenExchangeConfig, nil, defaultEnvGetter)
        if err != nil {
                return fmt.Errorf("invalid token exchange middleware config: %w", err)
        }

        tokenExchangeMw := &Middleware{
                middleware: middleware,
        }

        // Add middleware to runner
        runner.AddMiddleware(config.Type, tokenExchangeMw)

        return nil
}

// validateTokenExchangeConfig validates the token exchange configuration
func validateTokenExchangeConfig(config *Config) error {
        if config.HeaderStrategy == HeaderStrategyCustom && config.ExternalTokenHeaderName == "" {
                return fmt.Errorf("external_token_header_name must be specified when header_strategy is '%s'", HeaderStrategyCustom)
        }

        if config.HeaderStrategy != "" &&
                config.HeaderStrategy != HeaderStrategyReplace &&
                config.HeaderStrategy != HeaderStrategyCustom {
                return fmt.Errorf("invalid header_strategy: %s (valid values: '%s', '%s')",
                        config.HeaderStrategy, HeaderStrategyReplace, HeaderStrategyCustom)
        }

        return nil
}

// injectionFunc is a function that injects a token into an HTTP request
type injectionFunc func(*http.Request, string) error

// createReplaceInjector creates an injection function that replaces the Authorization header
func createReplaceInjector() injectionFunc {
        return func(r *http.Request, token string) error {
                slog.Debug("Token exchange successful, replacing Authorization header")
                r.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
                return nil
        }
}

// createCustomInjector creates an injection function that adds the token to a custom header
func createCustomInjector(headerName string) injectionFunc {
        // Validate header name at creation time
        if headerName == "" {
                return func(_ *http.Request, _ string) error {
                        return fmt.Errorf("external_token_header_name must be specified when header_strategy is '%s'", HeaderStrategyCustom)
                }
        }

        return func(r *http.Request, token string) error {
                slog.Debug("Token exchange successful, adding token to custom header", "header", headerName)
                r.Header.Set(headerName, fmt.Sprintf("Bearer %s", token))
                return nil
        }
}

// SubjectTokenProvider is a function that provides the subject token for exchange.
// This is used when the token comes from an external source (e.g., OAuth flow)
// rather than from incoming request headers.
type SubjectTokenProvider func() (string, error)

// CreateMiddlewareFromHeader creates token exchange middleware that extracts
// the subject token from the incoming request's Authorization header.
// This is the recommended approach when the proxy receives authenticated requests
// and needs to exchange those tokens for backend access.
//
// For external authentication flows (OAuth/OIDC), use CreateMiddlewareFromTokenSource instead.
func CreateMiddlewareFromHeader(config Config) (types.MiddlewareFunction, error) {
        return createTokenExchangeMiddleware(config, nil, defaultEnvGetter)
}

// CreateMiddlewareFromTokenSource creates token exchange middleware using an oauth2.TokenSource.
// This is the recommended approach for external authentication flows (OAuth/OIDC).
//
// The middleware will automatically select the appropriate token based on config.SubjectTokenType:
//   - oauthproto.TokenTypeAccessToken: Uses token.AccessToken
//   - oauthproto.TokenTypeIDToken or oauthproto.TokenTypeJWT: Uses token.Extra("id_token")
//
// This moves the token selection logic into the middleware where it belongs,
// keeping the command layer focused on configuration.
func CreateMiddlewareFromTokenSource(
        config Config,
        tokenSource oauth2.TokenSource,
) (types.MiddlewareFunction, error) {
        if tokenSource == nil {
                return nil, fmt.Errorf("tokenSource cannot be nil")
        }

        // Validate SubjectTokenType early to catch configuration errors at startup
        if config.SubjectTokenType != "" &&
                config.SubjectTokenType != oauthproto.TokenTypeAccessToken &&
                config.SubjectTokenType != oauthproto.TokenTypeIDToken &&
                config.SubjectTokenType != oauthproto.TokenTypeJWT {
                return nil, fmt.Errorf("invalid SubjectTokenType: %s (must be one of: %s, %s, %s)",
                        config.SubjectTokenType, oauthproto.TokenTypeAccessToken, oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT)
        }

        // Create a SubjectTokenProvider that handles token selection based on config
        subjectTokenProvider := func() (string, error) {
                token, err := tokenSource.Token()
                if err != nil {
                        return "", fmt.Errorf("failed to get token: %w", err)
                }

                // Select appropriate token based on configured type
                switch config.SubjectTokenType {
                case oauthproto.TokenTypeIDToken:
                        // Extract ID token from Extra field (standard OIDC approach)
                        idToken, ok := token.Extra("id_token").(string)
                        if !ok || idToken == "" {
                                slog.Error("ID token not available in token response")
                                return "", errors.New("required token not available")
                        }
                        return idToken, nil

                case "", oauthproto.TokenTypeAccessToken:
                        // Use access token (default)
                        if token.AccessToken == "" {
                                slog.Error("Access token not available")
                                return "", errors.New("required token not available")
                        }
                        return token.AccessToken, nil

                default:
                        // This should never happen due to early validation, but handle defensively
                        slog.Error("Invalid subject token type", "type", config.SubjectTokenType)
                        return "", errors.New("invalid token configuration")
                }
        }

        return createTokenExchangeMiddleware(config, subjectTokenProvider, defaultEnvGetter)
}

// createTokenExchangeMiddleware is the internal implementation that accepts an envGetter
// This allows for dependency injection in tests
func createTokenExchangeMiddleware(
        config Config,
        subjectTokenProvider SubjectTokenProvider,
        getEnv envGetter,
) (types.MiddlewareFunction, error) {
        // Determine injection strategy at startup time
        strategy := config.HeaderStrategy
        if strategy == "" {
                strategy = HeaderStrategyReplace // Default to replace for backwards compatibility
        }

        var injectToken injectionFunc
        switch strategy {
        case HeaderStrategyReplace:
                injectToken = createReplaceInjector()
        case HeaderStrategyCustom:
                injectToken = createCustomInjector(config.ExternalTokenHeaderName)
        default:
                return nil, fmt.Errorf("%w: invalid header injection strategy %s", errUnknownStrategy, strategy)
        }

        // Resolve client secret from config or environment variable
        clientSecret := config.ClientSecret
        if clientSecret == "" {
                // If not provided in config, try to read from environment variable
                if envSecret := getEnv(EnvClientSecret); envSecret != "" {
                        clientSecret = envSecret
                        slog.Debug("Using client secret from environment variable")
                }
        }

        // Create base exchange config at startup time with all static fields
        baseExchangeConfig := ExchangeConfig{
                TokenURL:         config.TokenURL,
                ClientID:         config.ClientID,
                ClientSecret:     clientSecret,
                Audience:         config.Audience,
                Scopes:           config.Scopes,
                SubjectTokenType: config.SubjectTokenType,
                // SubjectTokenProvider will be set per request
        }

        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // Get identity from the auth middleware
                        identity, ok := auth.IdentityFromContext(r.Context())
                        if !ok {
                                slog.Debug("No identity found in context, proceeding without token exchange")
                                next.ServeHTTP(w, r)
                                return
                        }

                        // Extract claims from identity for token exchange
                        claims := jwt.MapClaims(identity.Claims)

                        var tokenProvider SubjectTokenProvider

                        // Determine token source based on whether external provider was given
                        if subjectTokenProvider != nil {
                                // if the subjectTokenProvider is provided, use it, e.g. for passing in id_tokens
                                slog.Debug("Using provided token source for token exchange")
                                tokenProvider = subjectTokenProvider
                        } else {
                                // otherwise, extract token from incoming request's Authorization header
                                subjectToken, err := auth.ExtractBearerToken(r)
                                if err != nil {
                                        slog.Debug("No valid Bearer token found, proceeding without token exchange", "error", err)
                                        next.ServeHTTP(w, r)
                                        return
                                }

                                tokenProvider = func() (string, error) {
                                        return subjectToken, nil
                                }
                        }

                        // Log some claim information for debugging
                        if sub, exists := claims["sub"]; exists {
                                slog.Debug("Performing token exchange for subject", "subject", sub)
                        }

                        // Create a copy of the base config with the request-specific subject token
                        exchangeConfig := baseExchangeConfig
                        exchangeConfig.SubjectTokenProvider = tokenProvider

                        // Get token from token source
                        tokenSource := exchangeConfig.TokenSource(r.Context())
                        exchangedToken, err := tokenSource.Token()
                        if err != nil {
                                slog.Warn("Token exchange failed", "error", err)
                                http.Error(w, "Token exchange failed", http.StatusUnauthorized)
                                return
                        }

                        // Token exchange succeeded
                        slog.Debug("Token exchange successful")

                        // Inject the exchanged token into the request using the preselected strategy
                        if err := injectToken(r, exchangedToken.AccessToken); err != nil {
                                slog.Warn("Failed to inject token", "error", err)
                                http.Error(w, "Token injection failed", http.StatusInternalServerError)
                                return
                        }

                        next.ServeHTTP(w, r)
                })
        }, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oauthproto

import (
        "net/http"
        "time"
)

// httpClientTimeout matches the timeout already used by NonCachingRefresher's
// http.Client and is appropriate for OAuth token endpoints, which respond
// quickly under healthy conditions.
const httpClientTimeout = 30 * time.Second

// UserAgentTransport is an http.RoundTripper that adds the ToolHive User-Agent
// header to outbound requests when no User-Agent is already set.
//
// It is intended for HTTP clients passed to golang.org/x/oauth2 via the
// oauth2.HTTPClient context value. Without it, the library falls back to
// http.DefaultClient and Go's default Go-http-client/2.0 User-Agent, which
// upstream servers and WAFs cannot attribute to ToolHive.
type UserAgentTransport struct {
        // Base is the underlying RoundTripper. If nil, http.DefaultTransport is used.
        Base http.RoundTripper
}

// Compile-time assertion that CloseIdleConnections stays present. Asserted
// against the local shape rather than networking.IdleConnectionCloser because
// pkg/networking imports this package (see CloseIdleConnections below).
var _ interface{ CloseIdleConnections() } = (*UserAgentTransport)(nil)

// RoundTrip implements http.RoundTripper. It clones the request before
// mutating headers, per the RoundTripper contract, and sets User-Agent only
// when the request has no User-Agent set so that callers layering another
// transport can override the value.
func (t *UserAgentTransport) RoundTrip(req *http.Request) (*http.Response, error) {
        base := t.Base
        if base == nil {
                base = http.DefaultTransport
        }
        clonedReq := req.Clone(req.Context())
        if clonedReq.Header == nil {
                clonedReq.Header = make(http.Header)
        }
        if clonedReq.Header.Get("User-Agent") == "" {
                clonedReq.Header.Set("User-Agent", UserAgent)
        }
        return base.RoundTrip(clonedReq)
}

// CloseIdleConnections forwards to the wrapped RoundTripper so that
// http.Client.CloseIdleConnections reaches the underlying connection pool
// instead of stopping at this wrapper. pkg/oauthproto is imported by
// pkg/networking, so it cannot use networking.ForwardCloseIdle here without an
// import cycle; the anonymous assertion below is the equivalent forward.
func (t *UserAgentTransport) CloseIdleConnections() {
        base := t.Base
        if base == nil {
                base = http.DefaultTransport
        }
        if c, ok := base.(interface{ CloseIdleConnections() }); ok {
                c.CloseIdleConnections()
        }
}

// NewHTTPClient returns an *http.Client whose transport sets the ToolHive
// User-Agent on every outbound request and that has a sensible timeout for
// OAuth token endpoint calls. Use it for HTTP clients passed to
// golang.org/x/oauth2 via the oauth2.HTTPClient context value.
func NewHTTPClient() *http.Client {
        return &http.Client{
                Transport: &UserAgentTransport{},
                Timeout:   httpClientTimeout,
        }
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package oidc

import (
        "strings"
        "time"
)

const (
        // DefaultScopes are the default OAuth scopes requested during login.
        DefaultScopes = "openid offline_access"
)

// ClientConfig holds the OIDC provider settings and cached token state shared
// by both registry OAuth and LLM gateway authentication flows. Token values
// are never stored here — only references and expiry metadata.
//
// Both pkg/config.RegistryOAuthConfig and pkg/llm.OIDCConfig are type aliases
// for this type, so validation logic and new fields stay in sync across both
// authentication flows.
type ClientConfig struct {
        Issuer       string   `yaml:"issuer,omitempty"        json:"issuer,omitempty"`
        ClientID     string   `yaml:"client_id,omitempty"     json:"client_id,omitempty"`
        Scopes       []string `yaml:"scopes,omitempty"        json:"scopes,omitempty"`
        Audience     string   `yaml:"audience,omitempty"      json:"audience,omitempty"`
        CallbackPort int      `yaml:"callback_port,omitempty" json:"callback_port,omitempty"`

        // CachedRefreshTokenRef is the secrets-provider key under which the refresh
        // token is stored (never the token value itself).
        CachedRefreshTokenRef string `yaml:"cached_refresh_token_ref,omitempty" json:"cached_refresh_token_ref,omitempty"`
        // CachedTokenExpiry is the expiry of the most recently cached access token,
        // used to surface helpful messages when the token is about to expire.
        CachedTokenExpiry time.Time `yaml:"cached_token_expiry,omitempty" json:"cached_token_expiry,omitempty"`
}

// EffectiveScopes returns the configured OIDC scopes, or the default scopes
// (openid, offline_access) if none are set.
func (c *ClientConfig) EffectiveScopes() []string {
        if len(c.Scopes) > 0 {
                return c.Scopes
        }
        return strings.Fields(DefaultScopes)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package plugins

import (
        "context"
        "errors"
        "net/http"

        "github.com/stacklok/toolhive-core/httperr"
        regtypes "github.com/stacklok/toolhive-core/registry/types"
        "github.com/stacklok/toolhive/pkg/skills"
        "github.com/stacklok/toolhive/pkg/skills/lockfile"
)

// ListOptions configures the behavior of the List operation. Alias for
// skills.ListOptions (identical shape: Scope, ClientApp, ProjectRoot, Group).
type ListOptions = skills.ListOptions

// InstallOptions configures the behavior of the Install operation. Mirrors
// skills.InstallOptions with the plugin-specific addition of Reference/Digest
// passthrough fields used by install-from-OCI flows.
type InstallOptions struct {
        // Name is the plugin name or OCI reference to install.
        Name string `json:"name"`
        // Version is the specific version to install. Empty means latest.
        Version string `json:"version,omitempty"`
        // Scope is the installation scope.
        Scope Scope `json:"scope,omitempty"`
        // Clients lists target clients (e.g., "claude-code").
        Clients []string `json:"clients,omitempty"`
        // Force allows overwriting unmanaged plugin directories.
        Force bool `json:"force,omitempty"`
        // ProjectRoot is the project root path for project-scoped installs.
        ProjectRoot string `json:"project_root,omitempty"`
        // Group is the group name to add the plugin to after installation.
        Group string `json:"group,omitempty"`
        // AllowUnsigned permits installing a project-scoped plugin whose
        // artifact carries no Sigstore signature. Without it, unsigned
        // artifacts are rejected; with it, the lock entry records the exception
        // as "unsigned: true". A plugin contributes hooks, agents, and MCP
        // servers to the client that loads it, so this is an explicit
        // per-install trust decision, never a default.
        AllowUnsigned bool `json:"allow_unsigned,omitempty"`
        // PublicKey is the base64-encoded DER SPKI cosign public key a
        // project-scoped install must verify the artifact against, for artifacts
        // signed with a cosign key pair rather than keylessly. Required on true
        // first use of such an artifact — the signing key is recoverable from
        // neither the artifact nor its bundle, so nothing else can supply the
        // trust anchor — and pinned into the lock entry, which supplies it on
        // every install thereafter. A value that conflicts with what the lock
        // already pins is rejected, never ignored.
        PublicKey string `json:"public_key,omitempty"`
        // LayerData is the tar.gz content from an OCI layer. Internal use only — NOT exposed via HTTP API.
        LayerData []byte `json:"-"`
        // Reference is the full OCI reference (e.g. ghcr.io/org/plugin:v1).
        Reference string `json:"-"`
        // Digest is the OCI digest for upgrade detection.
        Digest string `json:"-"`
        // Components is the plugin's component inventory, hydrated from the OCI
        // artifact config by install-from-OCI/git flows. Internal use only.
        Components ComponentInventory `json:"-"`
        // Dependencies is the plugin's external dependency list, hydrated from the
        // OCI artifact config. Internal use only.
        Dependencies []Dependency `json:"-"`
        // Tag is the OCI tag, hydrated from the resolved reference. Internal use only.
        Tag string `json:"-"`
        // Description is the plugin description, hydrated from the OCI artifact
        // config or git manifest. Internal use only.
        Description string `json:"-"`
        // LockSource overrides the value recorded as the lock entry's Source. When
        // empty, the entry's Source is Name as given by the caller before any
        // internal resolution. Set by Sync/Upgrade, which pass an already-resolved
        // Name that must not overwrite the entry's original Source. Internal use
        // only — NOT exposed via HTTP API.
        LockSource string `json:"-"`
        // LockResolvedReference overrides the value recorded as the lock entry's
        // ResolvedReference. When empty, the entry's ResolvedReference is
        // whatever this install actually resolved to. Set by Sync when
        // reinstalling at a pinned reference. Internal use only — NOT exposed
        // via HTTP API.
        LockResolvedReference string `json:"-"`
        // SyncRestore forces re-extraction to every existing client even when
        // Digest matches the currently-installed digest. Set by Sync when
        // reinstalling at a pinned reference: the whole point is repairing
        // on-disk drift that happened without the pinned digest changing, so the
        // normal "same digest means content is already correct" fast path must
        // not apply. Internal use only — NOT exposed via HTTP API.
        SyncRestore bool `json:"-"`
        // RefreshMetadata updates the installed record when content is already at
        // the requested digest and no client needs materialization. Upgrade sets it
        // for an explicitly allowed resolved-reference change. Internal use only.
        RefreshMetadata bool `json:"-"`
        // AllowSignerChange lets install-time verification re-record the
        // observed identity instead of enforcing the lock file's recorded one.
        // Internal use only — set by upgrade when its signer-change guard was
        // explicitly overridden. NOT exposed via HTTP API.
        AllowSignerChange bool `json:"-"`
        // ExpectedCanonicalName, when set, requires the resolved plugin/manifest name
        // to equal this value before any install mutation. Used by Sync/Upgrade so a
        // lock entry cannot be repaired under a different canonical identity.
        ExpectedCanonicalName string `json:"-"`
        // Unsigned records the trust decision that this install proceeded
        // without a verified signature (via AllowUnsigned). Set internally by
        // install-time verification; recorded as `unsigned: true` in the lock
        // entry. Internal use only — NOT exposed via HTTP API.
        Unsigned bool `json:"-"`
        // CatalogProvenance is the independently optional provenance constraints
        // declared by the plugin registry/catalog entry this install resolved
        // from. On true first use, install-time verification checks each non-empty
        // field against the observed signature. A lock entry, including a legacy
        // entry with no trust state, always takes precedence over this. Internal
        // use only — NOT exposed via HTTP API.
        CatalogProvenance *regtypes.Provenance `json:"-"`
        // Provenance carries the verified signer identity established during
        // install-time verification, for recording into the lock entry. Set by
        // the verification step, nil when the artifact is unsigned or
        // verification did not run. Unlike skills.InstallOptions, this is the
        // lock file's own shape: plugins have no API-facing provenance type to
        // convert through yet, and a conversion pair that exists only to be
        // round-tripped is a place for recorded trust data to get dropped.
        // Internal use only — NOT exposed via HTTP API.
        Provenance *lockfile.Provenance `json:"-"`
        // SigstoreBundle is the signature material backing Provenance, in core's
        // durable form (see verifier.Result.Bundle), persisted alongside the
        // install record so sync can re-verify offline against the artifact digest.
        // Internal use only — NOT exposed via HTTP API.
        SigstoreBundle []byte `json:"-"`
}

// ProvenanceInfo is the verified signer identity of an installed artifact,
// the in-memory mirror of the lock file's provenance block. Alias for
// skills.ProvenanceInfo: the RFC THV-0080 provenance contract is shared, and
// a parallel struct would only be a place for the two to drift.
type ProvenanceInfo = skills.ProvenanceInfo

// InstallResult contains the outcome of an Install operation.
type InstallResult struct {
        // Plugin is the installed plugin.
        Plugin InstalledPlugin `json:"plugin"`
        // PreExisting is the store record as it was before this install, or nil
        // when this install created the record. Rollback uses it to restore the
        // previous state instead of destructively deleting a record this call
        // did not create. Internal use only — NOT exposed via HTTP API.
        PreExisting *InstalledPlugin `json:"-"`
        // ContentDigest is the dirhash of the canonical plugin tree (ExtractPlugin
        // output), computed at install time for recording in the lock file.
        // Internal use only — NOT exposed via HTTP API.
        ContentDigest string `json:"-"`
        // RestoreFiles undoes this install's on-disk writes: dematerialize a
        // fresh install, or restore the previous tree after a failed upgrade.
        // Callers must join a non-nil error into the failure they are
        // compensating; discarding it can hide a partial restore. Internal use
        // only — NOT exposed via HTTP API.
        RestoreFiles func(context.Context) error `json:"-"`
        // Provenance is the verified signer identity this install recorded —
        // surfaced so callers can display what trust-on-first-use pinned.
        Provenance *ProvenanceInfo `json:"provenance,omitempty"`
        // Unsigned reports that the install was recorded as an explicit
        // unsigned exception.
        Unsigned bool `json:"unsigned,omitempty"`
}

// UninstallOptions configures the behavior of the Uninstall operation. Alias
// for skills.UninstallOptions (identical shape).
type UninstallOptions = skills.UninstallOptions

// InfoOptions configures the behavior of the Info operation. Alias for
// skills.InfoOptions.
type InfoOptions = skills.InfoOptions

// PluginInfo contains detailed information about an installed plugin.
type PluginInfo struct {
        // Metadata contains the plugin's metadata.
        Metadata PluginMetadata `json:"metadata"`
        // InstalledPlugin contains the full installation record.
        InstalledPlugin *InstalledPlugin `json:"installed_plugin,omitempty"`
        // UnmaterializedComponents lists, per client type, the component types the
        // plugin declares that the installed client adapter does NOT load. Populated
        // by Info by diffing InstalledPlugin.Components against each installed
        // client adapter's SupportedComponents.
        UnmaterializedComponents map[string][]ComponentType `json:"unmaterialized_components,omitempty"`
        // ProjectScopeDegradedClients lists the client types for which a
        // project-scoped install degraded (the adapter could only materialize at
        // user scope — e.g. Codex always writes to the user-scoped config.toml).
        // Populated by Info; empty for user-scoped installs. Recomputed at read
        // time from the stored scope + each adapter's capability, mirroring the
        // UnmaterializedComponents pattern (no persistence needed — the degradation
        // is deterministic from scope + client type).
        ProjectScopeDegradedClients []string `json:"project_scope_degraded_clients,omitempty"`
        // Provenance is the signer identity the project's lock file records for
        // this plugin, when project-scoped and lock-managed.
        Provenance *ProvenanceInfo `json:"provenance,omitempty"`
        // Unsigned reports that the lock file records an explicit unsigned
        // exception for this plugin.
        Unsigned bool `json:"unsigned,omitempty"`
        // TrustUnrecorded reports that the project's lock file has an entry for
        // this plugin which records neither a signer identity nor an unsigned
        // exception — an entry written before verification existed, or
        // hand-edited. It exists to keep that state distinguishable from having no
        // lock entry at all, which leaves Provenance and Unsigned equally empty:
        // sync reports this one as drift and can repair it, so Info must not
        // render it as if nothing were pinning the plugin.
        TrustUnrecorded bool `json:"trust_unrecorded,omitempty"`
}

// ContentOptions configures the behavior of the GetContent operation. Alias
// for skills.ContentOptions.
type ContentOptions = skills.ContentOptions

// BuildOptions configures the behavior of the Build operation. Alias for
// skills.BuildOptions (Path, Tag).
type BuildOptions = skills.BuildOptions

// BuildResult contains the outcome of a Build operation. Alias for
// skills.BuildResult (Reference).
type BuildResult = skills.BuildResult

// PushOptions configures the behavior of the Push operation. Alias for
// skills.PushOptions (identical shape: Reference, Key, IdentityToken,
// NoSign).
type PushOptions = skills.PushOptions

// ValidatePushSigning enforces the push endpoint's signing contract: exactly
// one of a cosign key, an OIDC identity token for keyless signing, or an
// explicit opt-out. Ambiguous or absent input is rejected with HTTP 400 before
// the artifact is pushed, rather than surfacing as a signing failure afterward.
//
// The HTTP handler runs this before dispatch and the service runs it again on
// the options it receives. Both call it so the API contract holds regardless
// of which PluginService implementation is wired in: a request naming only a
// reference must be a 400 from the endpoint itself, not from whichever
// service happens to answer. Mirrors skillsvc.validateSigningInputs; the
// error text names both the JSON fields and the plugin command's flags.
func ValidatePushSigning(opts PushOptions) error {
        methods := 0
        if opts.Key != "" {
                methods++
        }
        if opts.IdentityToken != "" {
                methods++
        }
        switch {
        case opts.NoSign && methods > 0:
                return httperr.WithCode(
                        errors.New("no_sign (--no-sign) cannot be combined with key (--key) or identity_token (--identity-token)"),
                        http.StatusBadRequest,
                )
        case !opts.NoSign && methods == 0:
                return httperr.WithCode(
                        errors.New("signing credential required: set key (--key), identity_token (--identity-token) "+
                                "for CI/OIDC keyless signing, or no_sign (--no-sign) to push unsigned"),
                        http.StatusBadRequest,
                )
        case !opts.NoSign && methods > 1:
                return httperr.WithCode(
                        errors.New("specify only one of key (--key) or identity_token (--identity-token)"),
                        http.StatusBadRequest,
                )
        }
        return nil
}

// SyncOptions configures a lock-file sync. Alias for skills.SyncOptions
// (identical shape including adoption's AllowUnsigned and PublicKey inputs).
type SyncOptions = skills.SyncOptions

// SyncResult is the outcome of a lock-file sync. Alias for skills.SyncResult.
type SyncResult = skills.SyncResult

// SyncFailure describes one plugin that failed to sync. Alias for
// skills.SyncFailure.
type SyncFailure = skills.SyncFailure

// FailureReason is a typed sync/upgrade failure reason. Alias for
// skills.FailureReason.
type FailureReason = skills.FailureReason

// Typed failure reasons for sync/upgrade operations. Aliased from
// skills.FailureReason because the RFC THV-0080 contract is shared.
const (
        FailureReasonRegistryUnreachable = skills.FailureReasonRegistryUnreachable
        FailureReasonDigestMissing       = skills.FailureReasonDigestMissing
        FailureReasonValidationRejected  = skills.FailureReasonValidationRejected
        FailureReasonLockWriteFailed     = skills.FailureReasonLockWriteFailed
        FailureReasonSignatureInvalid    = skills.FailureReasonSignatureInvalid
        FailureReasonSignerMismatch      = skills.FailureReasonSignerMismatch
        FailureReasonKeySigned           = skills.FailureReasonKeySigned
        FailureReasonUnsignedRejected    = skills.FailureReasonUnsignedRejected
        FailureReasonUnknown             = skills.FailureReasonUnknown

        // FailureReasonProvenanceFieldMismatch means the artifact verifies
        // against the recorded signer, but a pinned certificate field (the
        // repository ref or runner environment) no longer matches.
        FailureReasonProvenanceFieldMismatch = skills.FailureReasonProvenanceFieldMismatch
)

// UpgradeOptions configures a lock-file upgrade. Alias for
// skills.UpgradeOptions (identical shape including AllowRefChange,
// AllowSignerChange, and PublicKey).
type UpgradeOptions = skills.UpgradeOptions

// UpgradeResult is the outcome of a lock-file upgrade. Alias for
// skills.UpgradeResult.
type UpgradeResult = skills.UpgradeResult

// UpgradeOutcome describes one plugin considered for upgrade. Alias for
// skills.UpgradeOutcome.
type UpgradeOutcome = skills.UpgradeOutcome

// UpgradeStatus is the per-entry upgrade outcome. Alias for
// skills.UpgradeStatus.
type UpgradeStatus = skills.UpgradeStatus

// Per-entry upgrade outcomes. Aliased from skills.UpgradeStatus because the
// RFC THV-0080 contract is shared.
const (
        UpgradeStatusUpgraded            = skills.UpgradeStatusUpgraded
        UpgradeStatusUpToDate            = skills.UpgradeStatusUpToDate
        UpgradeStatusTrustUpdated        = skills.UpgradeStatusTrustUpdated
        UpgradeStatusNotUpgradable       = skills.UpgradeStatusNotUpgradable
        UpgradeStatusRefChangeBlocked    = skills.UpgradeStatusRefChangeBlocked
        UpgradeStatusSignerChangeBlocked = skills.UpgradeStatusSignerChangeBlocked
        UpgradeStatusFailed              = skills.UpgradeStatusFailed
)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package plugins

import (
        "bytes"
        "encoding/json"
        "errors"
        "fmt"
        "io/fs"
        "os"
        "path/filepath"
        "strings"
)

// ManifestPath is the required manifest file path for a plugin directory,
// matching the Claude Plugin manifest format.
const ManifestPath = ".claude-plugin/plugin.json"

// MaxManifestSize limits the plugin.json size to prevent JSON parsing attacks
// (e.g. billion laughs). Mirrors the packager's maxManifestSize.
const MaxManifestSize = 64 * 1024

// MaxComponentsPerGroup caps the number of entries allowed in a single
// component array (commands, agents, skills, hooks). Mirrors the skills
// parser's MaxDependencies bound and bounds the cost of ValidatePluginDir,
// which walks each bundled skill path. A 64KB manifest can otherwise pack
// thousands of "./x" entries cheaply.
const MaxComponentsPerGroup = 100

// ErrInvalidManifest indicates that the plugin manifest is malformed, missing,
// or fails toolhive-side strictness checks.
var ErrInvalidManifest = errors.New("invalid plugin manifest")

// Author represents the author field of a plugin manifest.
type Author struct {
        Name  string `json:"name,omitempty"`
        Email string `json:"email,omitempty"`
        URL   string `json:"url,omitempty"`
}

// PluginManifest represents the toolhive-readable fields of
// .claude-plugin/plugin.json. Unknown fields are preserved in Raw so the
// manifest can be round-tripped without loss.
type PluginManifest struct {
        Name        string `json:"name"`
        Version     string `json:"version,omitempty"`
        Description string `json:"description,omitempty"`
        Author      Author `json:"author,omitempty"`
        Homepage    string `json:"homepage,omitempty"`
        Repository  string `json:"repository,omitempty"`
        License     string `json:"license,omitempty"`
        // Keywords MUST be a JSON array (a string is a hard error — see strictStringSlice).
        Keywords strictStringSlice `json:"keywords"`
        // Component directories: each entry must be a relative path starting with "./".
        Commands   []string        `json:"commands,omitempty"`
        Agents     []string        `json:"agents,omitempty"`
        Skills     []string        `json:"skills,omitempty"`
        Hooks      []string        `json:"hooks,omitempty"`
        McpServers json.RawMessage `json:"mcpServers,omitempty"`
        LspServers json.RawMessage `json:"lspServers,omitempty"`
        // Raw is the full original document, preserved for round-tripping unknown fields.
        Raw json.RawMessage `json:"-"`
}

// strictStringSlice is a []string that rejects a JSON string value. Plugins
// require keywords to be a JSON array (unlike skills, which accept a
// space-delimited string fallback); a string here is a hard error.
type strictStringSlice []string

// UnmarshalJSON implements json.Unmarshaler. A JSON string is rejected;
// a JSON array is decoded into the slice; any other type is an error.
func (s *strictStringSlice) UnmarshalJSON(data []byte) error {
        trimmed := bytes.TrimSpace(data)
        if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) {
                *s = nil
                return nil
        }
        if trimmed[0] == '"' {
                return fmt.Errorf("%w: keywords must be an array, got string", ErrInvalidManifest)
        }
        var arr []string
        if err := json.Unmarshal(data, &arr); err != nil {
                return fmt.Errorf("%w: keywords must be an array: %w", ErrInvalidManifest, err)
        }
        *s = arr
        return nil
}

// ParsePluginManifest reads and parses .claude-plugin/plugin.json from
// pluginDir. It performs toolhive-specific pre-build strictness checks only;
// the packager re-walks the directory and re-validates at build time.
//
// Strictness (exit gate):
//   - keywords MUST be a JSON array; a string is a hard error.
//   - component paths (commands/agents/skills/hooks) must be relative,
//     start with "./", and contain no ".." traversal segments.
//
// The full document is preserved in .Raw (unknown fields preserved).
func ParsePluginManifest(pluginDir string) (*PluginManifest, error) {
        manifestPath := filepath.Join(pluginDir, ManifestPath)

        // Lstat before reading: reject a symlinked manifest (TOCTOU guard) and
        // avoid allocating an oversized file before the size check rejects it.
        fi, err := os.Lstat(manifestPath)
        if err != nil {
                if errors.Is(err, fs.ErrNotExist) {
                        return nil, fmt.Errorf("%s not found in plugin directory: %w", ManifestPath, ErrInvalidManifest)
                }
                return nil, fmt.Errorf("checking %s: %w", ManifestPath, err)
        }
        if fi.Mode()&os.ModeSymlink != 0 || !fi.Mode().IsRegular() {
                return nil, fmt.Errorf("%s must be a regular file: %w", ManifestPath, ErrInvalidManifest)
        }
        if fi.Size() > MaxManifestSize {
                return nil, fmt.Errorf("manifest size %d exceeds maximum of %d bytes: %w",
                        fi.Size(), MaxManifestSize, ErrInvalidManifest)
        }

        content, err := os.ReadFile(manifestPath) //#nosec G304 -- path joined from validated pluginDir; symlink rejected above
        if err != nil {
                return nil, fmt.Errorf("reading %s: %w", ManifestPath, err)
        }

        return parseManifestBytes(content)
}

// ParsePluginManifestFromBytes parses manifest bytes into a PluginManifest,
// applying the same strictness checks as ParsePluginManifest. Split out for
// callers (e.g. the git install flow) that already hold the manifest bytes
// rather than a directory on disk.
func ParsePluginManifestFromBytes(content []byte) (*PluginManifest, error) {
        return parseManifestBytes(content)
}

// parseManifestBytes parses manifest bytes into a PluginManifest, applying the
// strictness checks. Split from ParsePluginManifest for testability.
func parseManifestBytes(content []byte) (*PluginManifest, error) {
        content = bytes.TrimSpace(content)
        if len(content) == 0 {
                return nil, fmt.Errorf("%w: manifest is empty", ErrInvalidManifest)
        }
        if len(content) > MaxManifestSize {
                return nil, fmt.Errorf("manifest exceeds maximum size of %d bytes: %w", MaxManifestSize, ErrInvalidManifest)
        }

        // Decode into a json.Decoder withDisallowUnknownFields? No — we preserve
        // unknown fields in .Raw, so we use standard unmarshal which ignores them.
        var m PluginManifest
        if err := json.Unmarshal(content, &m); err != nil {
                return nil, fmt.Errorf("%w: parsing manifest JSON: %w", ErrInvalidManifest, err)
        }

        // Preserve the full document for round-tripping.
        m.Raw = make([]byte, len(content))
        copy(m.Raw, content)

        // Validate component paths after unmarshaling. Cap each group at
        // MaxComponentsPerGroup to bound the cost of downstream validation
        // (ValidatePluginDir walks each bundled skill path).
        for _, group := range [][]string{m.Commands, m.Agents, m.Skills, m.Hooks} {
                if len(group) > MaxComponentsPerGroup {
                        return nil, fmt.Errorf("%w: component group exceeds maximum of %d entries", ErrInvalidManifest, MaxComponentsPerGroup)
                }
                for _, p := range group {
                        if err := validateComponentPath(p); err != nil {
                                return nil, err
                        }
                }
        }

        return &m, nil
}

// validateComponentPath rejects empty, non-"./"-prefixed, absolute, or
// ".."-containing component paths. Component directories declared in the
// manifest must be relative paths scoped under the plugin directory to
// prevent path traversal during packaging/extraction.
func validateComponentPath(p string) error {
        if p == "" {
                return fmt.Errorf("%w: component path is empty", ErrInvalidManifest)
        }
        if filepath.IsAbs(p) {
                return fmt.Errorf("%w: component path %q must be relative (got absolute path)", ErrInvalidManifest, p)
        }
        if !strings.HasPrefix(p, "./") {
                return fmt.Errorf("%w: component path %q must be relative and start with ./", ErrInvalidManifest, p)
        }
        for _, segment := range strings.Split(filepath.ToSlash(p), "/") {
                if segment == ".." {
                        return fmt.Errorf("%w: component path %q must not contain '..' traversal", ErrInvalidManifest, p)
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package plugins

import (
        "fmt"
        "path/filepath"
        "strings"

        "github.com/stacklok/toolhive/pkg/skills"
)

// ValidatePluginName checks that a plugin name conforms to the kebab-case rule
// shared with skills (2-64 lowercase alphanumeric/hyphens, no consecutive
// hyphens). Re-exported from skills because the rule is identical.
func ValidatePluginName(name string) error { return skills.ValidateSkillName(name) }

// CheckFilesystem walks a plugin directory checking for symlinks and path
// traversal. Re-exported from skills (the filesystem safety check is generic).
var CheckFilesystem = skills.CheckFilesystem

// ValidatePluginDir validates a plugin directory at the given path. It is the
// plugin analogue of skills.ValidateSkillDir. Steps:
//  1. Clean + absolute path check.
//  2. CheckFilesystem (symlink + traversal walk).
//  3. ParsePluginManifest (strict keywords + component-path checks).
//  4. Validate name (ValidatePluginName) and assert it matches the dir basename.
//  5. For each manifest.Skills entry, reuse skills.ValidateSkillDir on the
//     bundled skill at <pluginDir>/<path-without-leading-./> — this reuses the
//     pkg/skills validator for bundled skills rather than duplicating it.
//
// I/O errors are returned as error; validation issues are returned in
// ValidationResult.
func ValidatePluginDir(path string) (*ValidationResult, error) {
        // Defense-in-depth: sanitize and validate the path before any filesystem
        // access. The caller (pluginsvc.Validate) also validates via
        // validateLocalPath, but we re-check here because ValidatePluginDir is
        // exported and may be called directly.
        path = filepath.Clean(path)
        if !filepath.IsAbs(path) {
                return nil, fmt.Errorf("path must be absolute, got %q", path)
        }

        var errs []string

        // Check for symlinks and path traversal in a single walk.
        if err := CheckFilesystem(path); err != nil {
                errs = append(errs, err.Error())
        }

        // Parse the manifest.
        manifest, err := ParsePluginManifest(path)
        if err != nil {
                errs = append(errs, fmt.Sprintf("invalid plugin manifest: %v", err))
                return &ValidationResult{
                        Valid:  false,
                        Errors: errs,
                }, nil
        }

        // Validate name + match directory basename.
        dirName := filepath.Base(path)
        if manifest.Name == "" {
                errs = append(errs, "plugin name is required in manifest")
        } else {
                if err := ValidatePluginName(manifest.Name); err != nil {
                        errs = append(errs, err.Error())
                }
                if manifest.Name != dirName {
                        errs = append(errs,
                                fmt.Sprintf("plugin name %q must match directory name %q", manifest.Name, dirName))
                }
        }

        // Reuse pkg/skills.Validator for bundled skills. Each manifest.Skills entry
        // is a relative "./..." path; strip the leading "./" to resolve it under
        // the plugin directory.
        for _, skillPath := range manifest.Skills {
                bundledPath := filepath.Join(path, strings.TrimPrefix(skillPath, "./"))
                skillResult, skillErr := skills.ValidateSkillDir(bundledPath)
                if skillErr != nil {
                        // I/O error (e.g. bundled path missing) — surface as a validation error,
                        // not a hard error, so the caller still gets a ValidationResult.
                        errs = append(errs, fmt.Sprintf("bundled skill %q: %v", skillPath, skillErr))
                        continue
                }
                if skillResult != nil && !skillResult.Valid {
                        for _, e := range skillResult.Errors {
                                errs = append(errs, fmt.Sprintf("bundled skill %q: %s", skillPath, e))
                        }
                }
        }

        return &ValidationResult{
                Valid:  len(errs) == 0,
                Errors: errs,
        }, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package process provides low level operations on OS processes
package process

import (
        "os"
)

// ToolHiveDetachedEnv is the environment variable used to indicate that the process is running in detached mode.
const ToolHiveDetachedEnv = "TOOLHIVE_DETACHED"

// ToolHiveDetachedValue is the expected value of ToolHiveDetachedEnv when set.
const ToolHiveDetachedValue = "1"

// IsDetached checks if the process is running in detached mode.
func IsDetached() bool {
        return os.Getenv(ToolHiveDetachedEnv) == ToolHiveDetachedValue
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

//go:build !windows

package process

import (
        "fmt"
        "os"
        "syscall"
)

// FindProcess finds a process by its ID and checks if it's running.
// This function works on Unix systems (Linux and macOS).
func FindProcess(pid int) (bool, error) {
        if pid <= 0 {
                return false, fmt.Errorf("invalid PID: %d", pid)
        }

        // On Unix systems, os.FindProcess always succeeds regardless of whether
        // the process exists or not. We need to send a signal to check if it's running.
        proc, err := os.FindProcess(pid)
        if err != nil {
                return false, fmt.Errorf("failed to find process: %w", err)
        }

        // Send signal 0 to check if the process exists
        // Signal 0 doesn't actually send a signal, but it checks if the process exists
        err = proc.Signal(syscall.Signal(0))

        // If there's no error, the process exists
        if err == nil {
                return true, nil
        }

        // If the error is "no such process", the process doesn't exist
        if err.Error() == "no such process" || err.Error() == "os: process already finished" {
                return false, nil
        }

        // For other errors (e.g., permission denied), return the error
        return false, fmt.Errorf("error checking process: %w", err)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

//go:build !windows

package process

import (
        "fmt"
        "os"
        "syscall"
)

// KillProcess kills a process by its ID
func KillProcess(pid int) error {
        if pid <= 0 {
                return fmt.Errorf("invalid PID: %d", pid)
        }

        // Check if the process exists
        process, err := os.FindProcess(pid)
        if err != nil {
                return fmt.Errorf("failed to find process: %w", err)
        }

        // Send a SIGTERM signal to the process
        if err := process.Signal(syscall.SIGTERM); err != nil {
                return fmt.Errorf("failed to send SIGTERM to process: %w", err)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package process

import (
        "strings"

        gopsutilprocess "github.com/shirou/gopsutil/v4/process"
)

const (
        // toolHiveBinaryName is the binary name used to identify ToolHive processes.
        // Used when TOOLHIVE_DETACHED cannot be read (e.g. platform restrictions).
        toolHiveBinaryName = "thv"
)

// containsThvBinary returns true if s appears to reference the thv binary
// (e.g. /usr/bin/thv, thv start, thv.exe), avoiding false positives like "toolhive".
func containsThvBinary(s string) bool {
        lower := strings.ToLower(s)
        return strings.Contains(lower, toolHiveBinaryName+".exe") ||
                strings.HasSuffix(lower, toolHiveBinaryName) ||
                strings.Contains(lower, "/"+toolHiveBinaryName) ||
                strings.Contains(lower, `\`+toolHiveBinaryName) ||
                strings.HasPrefix(lower, toolHiveBinaryName+" ") ||
                strings.Contains(lower, " "+toolHiveBinaryName+" ") ||
                strings.HasSuffix(lower, " "+toolHiveBinaryName)
}

// IsToolHiveProxyForWorkload returns true if the given PID belongs to the ToolHive
// proxy for the specified workload, so it is safe to kill when freeing a port.
// Returns false if the process is not that workload's proxy or if identity cannot
// be verified (fail-safe: do not kill).
//
// When workloadName is empty, only verifies it is a ToolHive process.
// When workloadName is non-empty, also verifies the process cmdline contains
// " start <workloadName> " (the detached proxy runs "thv start <name> --foreground").
//
// Verification checks, in order:
//  1. TOOLHIVE_DETACHED=1 in process environment (most reliable)
//  2. "thv" in executable path or command line (fallback when env unavailable)
//  3. workloadName in cmdline (when provided, avoids killing another workload's proxy)
func IsToolHiveProxyForWorkload(pid int, workloadName string) (bool, error) {
        if pid <= 0 {
                return false, nil
        }

        // PID fits in int32 on all supported platforms
        p, err := gopsutilprocess.NewProcess(int32(pid)) //nolint:gosec // G115
        if err != nil {
                return false, err
        }

        if !isToolHiveProcess(p) {
                return false, nil
        }

        if workloadName != "" {
                cmdline, err := p.Cmdline()
                if err != nil || !cmdlineContainsWorkload(cmdline, workloadName) {
                        return false, nil
                }
        }

        return true, nil
}

// isToolHiveProcess returns true if p is a ToolHive process (TOOLHIVE_DETACHED
// or thv binary in exe/cmdline).
func isToolHiveProcess(p *gopsutilprocess.Process) bool {
        if hasToolHiveDetachedEnv(p) {
                return true
        }
        if exe, err := p.Exe(); err == nil && containsThvBinary(exe) {
                return true
        }
        if cmdline, err := p.Cmdline(); err == nil && containsThvBinary(cmdline) {
                return true
        }
        return false
}

func hasToolHiveDetachedEnv(p *gopsutilprocess.Process) bool {
        env, err := p.Environ()
        if err != nil {
                return false
        }
        target := ToolHiveDetachedEnv + "=" + ToolHiveDetachedValue
        for _, e := range env {
                if e == target {
                        return true
                }
        }
        return false
}

// cmdlineContainsWorkload returns true if the cmdline indicates a "thv start <name> ..." process.
// Uses word boundaries to avoid partial matches (e.g. "g" matching "github").
func cmdlineContainsWorkload(cmdline, workloadName string) bool {
        if workloadName == "" {
                return false
        }
        pattern := " start " + workloadName
        if !strings.Contains(cmdline, pattern) {
                return false
        }
        // Ensure workloadName is not a prefix of a longer word: next char must be space, -, or end
        idx := strings.Index(cmdline, pattern)
        end := idx + len(pattern)
        if end >= len(cmdline) {
                return true
        }
        next := cmdline[end]
        return next == ' ' || next == '-'
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package process

import (
        "context"
        "fmt"
        "time"
)

// WaitForExit waits for the process with the given PID to exit.
// It polls FindProcess every 50ms until the process is no longer running
// or the context is cancelled. Callers should use context.WithTimeout to
// impose a deadline.
// Returns nil when the process has exited, or an error on context cancellation.
func WaitForExit(ctx context.Context, pid int) error {
        if pid <= 0 {
                return fmt.Errorf("invalid PID: %d", pid)
        }

        ticker := time.NewTicker(50 * time.Millisecond)
        defer ticker.Stop()

        for {
                alive, err := FindProcess(pid)
                if err != nil {
                        return err
                }
                if !alive {
                        return nil
                }

                select {
                case <-ctx.Done():
                        return ctx.Err()
                case <-ticker.C:
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package ratelimit

import (
        "math"
        "time"

        thvmcp "github.com/stacklok/toolhive/pkg/mcp"
)

const (
        // CodeRateLimited is the JSON-RPC error code for rate-limited requests,
        // mirroring HTTP 429 (Too Many Requests) the same way mcp.JSONRPCCodeDenied
        // mirrors 403. It sits outside the JSON-RPC reserved range
        // (-32768..-32000) so a single value is conformant on every MCP revision:
        // 2026-07-28 reserves -32020..-32099 exclusively for spec-defined codes
        // and directs application codes outside the reserved range, while
        // 2025-11-25 and earlier inherit plain JSON-RPC 2.0, which leaves the
        // space outside that range application-defined. Replaces -32029, chosen
        // by RFC THV-0057 when -32000..-32099 was uniformly implementation-defined
        // — a premise the 2026-07-28 partition invalidated (#6101).
        CodeRateLimited int64 = 429

        // MessageRateLimited is the error message for rate-limited requests.
        MessageRateLimited = "Rate limit exceeded"
)

// RateLimitedError reports that a request exceeded its configured rate limit.
type RateLimitedError struct {
        RetryAfter time.Duration
}

var _ thvmcp.CodedError = (*RateLimitedError)(nil)

func (*RateLimitedError) Error() string {
        return MessageRateLimited
}

// Code returns the ToolHive JSON-RPC-compatible code for rate-limited requests.
func (*RateLimitedError) Code() int64 {
        return CodeRateLimited
}

// Data returns structured retry metadata for transport adapters that cannot
// emit a custom JSON-RPC error object from the tool-handler seam.
func (e *RateLimitedError) Data() map[string]any {
        return map[string]any{
                "retryAfterSeconds": int(math.Ceil(e.RetryAfter.Seconds())),
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package bucket implements the token bucket algorithm backed by Redis.
package bucket

import (
        "context"
        "fmt"
        "math"
        "time"

        "github.com/redis/go-redis/v9"
)

// consumeAllScript atomically refills and attempts to consume one token from
// each of N buckets. Each bucket is a Redis hash with "tokens" and
// "last_refill" fields (microsecond precision).
//
// The script checks ALL buckets first, then only consumes from ALL if every
// bucket has sufficient tokens. This prevents draining a server-level bucket
// when a more restrictive per-tool bucket would reject.
//
// KEYS[1..N]   = bucket keys
// ARGV[1]      = number of buckets (N)
// ARGV[2..N+1] = maxTokens for each bucket
// ARGV[N+2..2N+1] = refillRate (tokens/sec, float) for each bucket
// ARGV[2N+2..3N+1] = expireSeconds for each bucket
//
// Returns: 0 if all consumed (allowed), or the 1-based index of the first
// bucket that rejected.
var consumeAllScript = redis.NewScript(`
local numBuckets = tonumber(ARGV[1])

-- Use Redis server clock for consistency across replicas
local timeResp = redis.call('TIME')
local now = tonumber(timeResp[1]) * 1000000 + tonumber(timeResp[2])

-- Phase 1: refill all buckets, check if each has >= 1 token
local refilled = {}
local rejectIdx = 0
for i = 1, numBuckets do
    local key = KEYS[i]
    local maxTokens = tonumber(ARGV[1 + i])
    local refillRate = tonumber(ARGV[1 + numBuckets + i])

    local data = redis.call('HMGET', key, 'tokens', 'last_refill')
    local tokens = tonumber(data[1])
    local lastRefill = tonumber(data[2])

    if tokens == nil then
        refilled[i] = maxTokens
    else
        local elapsedSec = math.max(0, (now - lastRefill) / 1000000)
        refilled[i] = math.min(maxTokens, tokens + elapsedSec * refillRate)
    end

    if refilled[i] < 1 and rejectIdx == 0 then
        rejectIdx = i
    end
end

-- Phase 2: if any bucket rejected, update state without consuming
if rejectIdx ~= 0 then
    for i = 1, numBuckets do
        local key = KEYS[i]
        local expireSec = tonumber(ARGV[1 + 2 * numBuckets + i])
        redis.call('HSET', key, 'tokens', refilled[i], 'last_refill', now)
        redis.call('EXPIRE', key, expireSec)
    end
    return rejectIdx
end

-- Phase 3: all buckets have tokens — consume one from each
for i = 1, numBuckets do
    local key = KEYS[i]
    local expireSec = tonumber(ARGV[1 + 2 * numBuckets + i])
    redis.call('HSET', key, 'tokens', refilled[i] - 1, 'last_refill', now)
    redis.call('EXPIRE', key, expireSec)
end
return 0
`)

// TokenBucket represents a single rate limit bucket backed by Redis.
type TokenBucket struct {
        key           string
        maxTokens     int32
        refillRate    float64 // tokens per second
        expireSeconds int64   // ceil(2 * refillPeriod) in seconds
}

// New creates a TokenBucket. The Redis key is derived from namespace, server
// name, and suffix (e.g., "shared" or "shared:tool:search").
func New(namespace, serverName, suffix string, maxTokens int32, refillPeriod time.Duration) *TokenBucket {
        refillSec := refillPeriod.Seconds()
        return &TokenBucket{
                key:           deriveKeyPrefix(namespace, serverName) + suffix,
                maxTokens:     maxTokens,
                refillRate:    float64(maxTokens) / refillSec,
                expireSeconds: int64(math.Ceil(2 * refillSec)),
        }
}

// deriveKeyPrefix creates the rate limit key prefix for a given namespace and server.
// Format: "thv:rl:{ns:name}:" where {ns:name} is a Redis hash tag ensuring all keys
// for a server land on the same Redis Cluster slot.
func deriveKeyPrefix(namespace, name string) string {
        return fmt.Sprintf("thv:rl:{%s:%s}:", namespace, name)
}

// ConsumeAll atomically checks and consumes one token from each bucket.
// All buckets are checked first; tokens are only consumed if ALL buckets
// have sufficient tokens.
//
// Returns the index of the first bucket that rejected (0-based), or -1 if
// all allowed.
func ConsumeAll(ctx context.Context, client redis.Cmdable, buckets []*TokenBucket) (int, error) {
        if len(buckets) == 0 {
                return -1, nil
        }

        keys := make([]string, len(buckets))
        // ARGV layout: [numBuckets, maxTokens..., refillRates..., expireSeconds...]
        args := make([]any, 1+3*len(buckets))
        args[0] = len(buckets)
        for i, b := range buckets {
                keys[i] = b.key
                args[1+i] = b.maxTokens
                args[1+len(buckets)+i] = fmt.Sprintf("%.6f", b.refillRate)
                args[1+2*len(buckets)+i] = b.expireSeconds
        }

        result, err := consumeAllScript.Run(ctx, client, keys, args...).Int64()
        if err != nil {
                return -1, fmt.Errorf("rate limit script error: %w", err)
        }
        if result == 0 {
                return -1, nil // all allowed
        }
        return int(result) - 1, nil // convert 1-based to 0-based index
}

// RetryAfter returns the minimum wait time for one token to become available.
func (b *TokenBucket) RetryAfter() time.Duration {
        d := time.Duration(float64(time.Second) / b.refillRate)
        // Round up to nearest second (minimum 1s) for Retry-After header
        if d < time.Second {
                return time.Second
        }
        return time.Duration(math.Ceil(d.Seconds())) * time.Second
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package ratelimit provides token-bucket rate limiting for MCP servers.
//
// The public API consists of the Limiter interface, Decision result type,
// and NewLimiter constructor. The token bucket implementation is internal.
package ratelimit

import (
        "context"
        "fmt"
        "time"

        "github.com/redis/go-redis/v9"
        "go.opentelemetry.io/otel"
        "go.opentelemetry.io/otel/metric"

        v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/ratelimit/internal/bucket"
)

// Limiter checks rate limits for an MCP server.
type Limiter interface {
        // Allow checks whether a request is permitted.
        // toolName is the MCP tool being called (empty for non-tool requests).
        // userID is the authenticated user (empty for unauthenticated requests).
        Allow(ctx context.Context, toolName, userID string) (*Decision, error)
}

// failOpenObserver records the policy decision made by production enforcement
// adapters after the built-in limiter returns an infrastructure error.
type failOpenObserver interface {
        recordFailOpen(context.Context)
}

// Decision holds the result of a rate limit check.
type Decision struct {
        // Allowed is true when the request may proceed.
        Allowed bool

        // RetryAfter is populated when Allowed is false.
        // It indicates the minimum wait before the next request may succeed.
        RetryAfter time.Duration
}

// Allow checks whether identity may call toolName through limiter. When the
// built-in limiter returns an infrastructure error, Allow records fail-open
// observability before preserving the error for the HTTP or vMCP adapter to log
// and apply its existing fail-open behavior.
func Allow(ctx context.Context, limiter Limiter, identity *auth.Identity, toolName string) error {
        if limiter == nil {
                return nil
        }

        // When no identity is present (unauthenticated), userID stays empty and
        // per-user buckets are skipped — only shared limits apply. CEL validation
        // ensures perUser rate limits require auth to be enabled.
        var userID string
        if identity != nil {
                userID = identity.Subject
        }

        decision, err := limiter.Allow(ctx, toolName, userID)
        if err != nil {
                if observer, ok := limiter.(failOpenObserver); ok {
                        observer.recordFailOpen(ctx)
                }
                return err
        }
        if !decision.Allowed {
                return &RateLimitedError{RetryAfter: decision.RetryAfter}
        }
        return nil
}

// NewLimiter constructs a Limiter from CRD configuration.
// Returns a no-op limiter (always allows) when crd is nil.
// namespace and name identify the MCP server for Redis key derivation.
func NewLimiter(client redis.Cmdable, namespace, name string, crd *v1beta1.RateLimitConfig) (Limiter, error) {
        return newLimiter(client, namespace, name, crd, otel.GetMeterProvider())
}

func newLimiter(
        client redis.Cmdable,
        namespace string,
        name string,
        crd *v1beta1.RateLimitConfig,
        meterProvider metric.MeterProvider,
) (Limiter, error) {
        if crd == nil {
                return noopLimiter{}, nil
        }

        telemetry, err := newRateLimitTelemetry(meterProvider, namespace, name)
        if err != nil {
                return nil, fmt.Errorf("rate limit telemetry: %w", err)
        }

        l := &limiter{
                client:    client,
                telemetry: telemetry,
        }

        if crd.Shared != nil {
                b, err := newBucket(namespace, name, "shared", crd.Shared)
                if err != nil {
                        return nil, fmt.Errorf("shared bucket: %w", err)
                }
                l.serverBucket = b
        }

        if crd.PerUser != nil {
                spec, err := newBucketSpec(namespace, name, crd.PerUser)
                if err != nil {
                        return nil, fmt.Errorf("perUser bucket: %w", err)
                }
                l.perUserSpec = &spec
        }

        for _, t := range crd.Tools {
                if t.Shared != nil {
                        b, err := newBucket(namespace, name, "shared:tool:"+t.Name, t.Shared)
                        if err != nil {
                                return nil, fmt.Errorf("tool %q shared bucket: %w", t.Name, err)
                        }
                        if l.toolBuckets == nil {
                                l.toolBuckets = make(map[string]*bucket.TokenBucket)
                        }
                        l.toolBuckets[t.Name] = b
                }
                if t.PerUser != nil {
                        spec, err := newBucketSpec(namespace, name, t.PerUser)
                        if err != nil {
                                return nil, fmt.Errorf("tool %q perUser bucket: %w", t.Name, err)
                        }
                        if l.perUserTools == nil {
                                l.perUserTools = make(map[string]bucketSpec)
                        }
                        l.perUserTools[t.Name] = spec
                }
        }

        return l, nil
}

// bucketSpec holds deferred bucket parameters for per-user buckets that are
// created on the fly in Allow() because the userID is not known at construction time.
type bucketSpec struct {
        namespace    string
        serverName   string
        maxTokens    int32
        refillPeriod time.Duration
}

// limitCheck keeps a bucket paired with its observability dimensions.
type limitCheck struct {
        bucket        *bucket.TokenBucket
        scope         string
        operationType string
}

func (c limitCheck) rejectionIdentifier() string {
        return c.scope + "_" + c.operationType
}

// limiter is the concrete implementation of Limiter.
type limiter struct {
        client       redis.Cmdable
        telemetry    *rateLimitTelemetry
        serverBucket *bucket.TokenBucket            // nil when no shared server limit
        toolBuckets  map[string]*bucket.TokenBucket // tool name -> shared bucket
        perUserSpec  *bucketSpec                    // nil when no server-level per-user limit
        perUserTools map[string]bucketSpec          // tool name -> per-user bucket spec; nil when none
}

var _ failOpenObserver = (*limiter)(nil)

func (l *limiter) recordFailOpen(ctx context.Context) {
        l.telemetry.recordFailOpen(ctx)
        recordRateLimitSpanOutcome(ctx, rateLimitDecisionAllowed, rateLimitRejectedByNone, true)
}

// Allow atomically checks all applicable rate limit buckets for the request.
// Tokens are only consumed if ALL buckets have sufficient capacity, preventing
// a rejected per-tool or per-user call from draining other budgets.
func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision, error) {
        // Collect applicable buckets in priority order.
        var checks []limitCheck
        if l.serverBucket != nil {
                checks = append(checks, limitCheck{
                        bucket:        l.serverBucket,
                        scope:         rateLimitScopeShared,
                        operationType: rateLimitOperationServer,
                })
        }
        if toolName != "" && l.toolBuckets != nil {
                if tb, ok := l.toolBuckets[toolName]; ok {
                        checks = append(checks, limitCheck{
                                bucket:        tb,
                                scope:         rateLimitScopeShared,
                                operationType: rateLimitOperationTool,
                        })
                }
        }

        // Per-user buckets are created on the fly because userID is request-scoped.
        // bucket.New only allocates a struct — all state lives in Redis, so creating
        // a new TokenBucket per request is safe (no local state to lose).
        //
        // Key prefixes deviate from RFC THV-0057 to prevent cross-type collisions:
        // RFC uses "user:{userId}:tool:{toolName}" for both scopes, but a userID
        // containing ":tool:" would collide with the per-tool key. Instead we use
        // distinct prefixes: "user:" for server-level, "user-tool:" for tool-level.
        if userID != "" {
                if l.perUserSpec != nil {
                        s := l.perUserSpec
                        checks = append(checks, limitCheck{
                                bucket: bucket.New(
                                        s.namespace, s.serverName,
                                        "user:"+userID,
                                        s.maxTokens, s.refillPeriod,
                                ),
                                scope:         rateLimitScopePerUser,
                                operationType: rateLimitOperationServer,
                        })
                }
                if toolName != "" && l.perUserTools != nil {
                        if s, ok := l.perUserTools[toolName]; ok {
                                // Key prefix "user-tool:" is distinct from "user:" to prevent
                                // collisions when a userID contains delimiter characters.
                                checks = append(checks, limitCheck{
                                        bucket: bucket.New(
                                                s.namespace, s.serverName,
                                                "user-tool:"+toolName+":"+userID,
                                                s.maxTokens, s.refillPeriod,
                                        ),
                                        scope:         rateLimitScopePerUser,
                                        operationType: rateLimitOperationTool,
                                })
                        }
                }
        }

        if len(checks) == 0 {
                recordRateLimitSpanOutcome(ctx, rateLimitDecisionAllowed, rateLimitRejectedByNone, false)
                return &Decision{Allowed: true}, nil
        }

        buckets := make([]*bucket.TokenBucket, len(checks))
        for i, check := range checks {
                buckets[i] = check.bucket
        }

        start := time.Now()
        rejectedIdx, err := bucket.ConsumeAll(ctx, l.client, buckets)
        l.telemetry.recordCheckLatency(ctx, time.Since(start))
        if err != nil {
                l.telemetry.recordRedisError(ctx, err)
                return nil, fmt.Errorf("rate limit check: %w", err)
        }
        if rejectedIdx >= 0 {
                l.telemetry.recordRejected(ctx, checks[rejectedIdx])
                recordRateLimitSpanOutcome(ctx, rateLimitDecisionRejected, checks[rejectedIdx].rejectionIdentifier(), false)
                return &Decision{
                        Allowed:    false,
                        RetryAfter: buckets[rejectedIdx].RetryAfter(),
                }, nil
        }

        l.telemetry.recordAllowed(ctx, checks)
        recordRateLimitSpanOutcome(ctx, rateLimitDecisionAllowed, rateLimitRejectedByNone, false)
        return &Decision{Allowed: true}, nil
}

// noopLimiter always allows requests.
type noopLimiter struct{}

func (noopLimiter) Allow(context.Context, string, string) (*Decision, error) {
        return &Decision{Allowed: true}, nil
}

// validateBucketCRD checks that a CRD bucket spec has valid parameters.
func validateBucketCRD(b *v1beta1.RateLimitBucket) (int32, time.Duration, error) {
        if b.MaxTokens < 1 {
                return 0, 0, fmt.Errorf("maxTokens must be >= 1, got %d", b.MaxTokens)
        }
        d := b.RefillPeriod.Duration
        if d <= 0 {
                return 0, 0, fmt.Errorf("refillPeriod must be positive, got %s", d)
        }
        return b.MaxTokens, d, nil
}

// newBucket validates a CRD bucket spec and creates a TokenBucket.
func newBucket(namespace, serverName, suffix string, b *v1beta1.RateLimitBucket) (*bucket.TokenBucket, error) {
        maxTokens, refillPeriod, err := validateBucketCRD(b)
        if err != nil {
                return nil, err
        }
        return bucket.New(namespace, serverName, suffix, maxTokens, refillPeriod), nil
}

// newBucketSpec validates a CRD bucket spec and creates a deferred bucketSpec
// for per-user buckets that are materialized at Allow() time.
func newBucketSpec(namespace, serverName string, b *v1beta1.RateLimitBucket) (bucketSpec, error) {
        maxTokens, refillPeriod, err := validateBucketCRD(b)
        if err != nil {
                return bucketSpec{}, err
        }
        return bucketSpec{
                namespace:    namespace,
                serverName:   serverName,
                maxTokens:    maxTokens,
                refillPeriod: refillPeriod,
        }, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package ratelimit

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "math"
        "net/http"
        "os"
        "time"

        "github.com/stacklok/toolhive-core/redisconn"
        v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/redisconfig"
        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // MiddlewareType is the type constant for the rate limit middleware.
        MiddlewareType = "ratelimit"

        // redisPasswordEnvVar is the environment variable containing the Redis password.
        // Shared with session storage — the operator injects it from the same Secret.
        redisPasswordEnvVar = "THV_SESSION_REDIS_PASSWORD" //nolint:gosec // G101: env var name, not a credential
)

// MiddlewareParams holds the parameters for the rate limit middleware factory.
type MiddlewareParams struct {
        Namespace  string                   `json:"namespace"`
        ServerName string                   `json:"server_name"`
        Config     *v1beta1.RateLimitConfig `json:"config"`
        RedisAddr  string                   `json:"redis_addr,omitempty"`
        RedisDB    int32                    `json:"redis_db,omitempty"`
        RedisTLS   *redisconfig.TLSConfig   `json:"redis_tls,omitempty"`
}

// rateLimitMiddleware wraps rate limiting functionality for the factory pattern.
type rateLimitMiddleware struct {
        handler types.MiddlewareFunction
        closer  io.Closer
}

// Handler returns the middleware function used by the proxy.
func (m *rateLimitMiddleware) Handler() types.MiddlewareFunction {
        return m.handler
}

// Close cleans up the Redis client.
func (m *rateLimitMiddleware) Close() error {
        if m.closer != nil {
                return m.closer.Close()
        }
        return nil
}

// NewRedisLimiter creates a Redis-backed rate limiter from typed params.
func NewRedisLimiter(params MiddlewareParams) (Limiter, io.Closer, error) {
        if params.RedisAddr == "" {
                return nil, nil, fmt.Errorf("rate limit middleware requires a Redis address")
        }

        tlsCfg, err := params.RedisTLS.Load()
        if err != nil {
                return nil, nil, fmt.Errorf("rate limit middleware TLS configuration: %w", err)
        }
        pingCtx, pingCancel := context.WithTimeout(context.Background(), 5*time.Second)
        defer pingCancel()
        client, err := redisconn.NewClient(pingCtx, &redisconn.Config{
                Addr:     params.RedisAddr,
                DB:       int(params.RedisDB),
                Password: os.Getenv(redisPasswordEnvVar),
                TLS:      tlsCfg,
        })
        if err != nil {
                return nil, nil, fmt.Errorf("rate limit middleware: failed to connect to Redis at %s: %w", params.RedisAddr, err)
        }

        limiter, err := NewLimiter(client, params.Namespace, params.ServerName, params.Config)
        if err != nil {
                _ = client.Close()
                return nil, nil, fmt.Errorf("failed to create rate limiter: %w", err)
        }

        return limiter, client, nil
}

// NewMiddleware creates a Redis-backed rate limit middleware from typed params.
func NewMiddleware(params MiddlewareParams) (types.Middleware, error) {
        limiter, closer, err := NewRedisLimiter(params)
        if err != nil {
                return nil, err
        }

        return &rateLimitMiddleware{
                handler: rateLimitHandler(limiter),
                closer:  closer,
        }, nil
}

// CreateMiddleware is the factory function for rate limit middleware.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal rate limit middleware parameters: %w", err)
        }

        mw, err := NewMiddleware(params)
        if err != nil {
                return err
        }
        runner.AddMiddleware(MiddlewareType, mw)
        return nil
}

// rateLimitHandler returns a middleware function that enforces rate limits
// on tools/call requests.
func rateLimitHandler(limiter Limiter) types.MiddlewareFunction {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // The configured chain MUST place MCP parsing before this middleware.
                        // A nil POST context may be an admitted response or a missing-parser
                        // misconfiguration. These controls do not apply to non-POST traffic;
                        // rate limits only apply to parsed tools/call requests.
                        parsed := mcp.GetParsedMCPRequest(r.Context())
                        if parsed == nil || parsed.Method != "tools/call" {
                                next.ServeHTTP(w, r)
                                return
                        }

                        identity, _ := auth.IdentityFromContext(r.Context())
                        err := Allow(r.Context(), limiter, identity, parsed.ResourceID)
                        var limited *RateLimitedError
                        if errors.As(err, &limited) {
                                writeRateLimited(w, parsed.ID, limited.RetryAfter)
                                return
                        }
                        if err != nil {
                                slog.Warn("rate limit check failed, allowing request", "error", err)
                                next.ServeHTTP(w, r)
                                return
                        }
                        next.ServeHTTP(w, r)
                })
        }
}

// writeRateLimited writes an HTTP 429 response with a JSON-RPC error body.
func writeRateLimited(w http.ResponseWriter, requestID any, retryAfter time.Duration) {
        retrySeconds := int(math.Ceil(retryAfter.Seconds()))
        w.Header().Set("Content-Type", "application/json")
        w.Header().Set("Retry-After", fmt.Sprintf("%d", retrySeconds))
        w.WriteHeader(http.StatusTooManyRequests)
        //nolint:gosec // G104: writing a static JSON error response to an HTTP client
        _, _ = w.Write(rateLimitedBody(requestID, retryAfter))
}

// rateLimitedBody returns the JSON-encoded body for a rate-limited JSON-RPC
// error.
//
// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error
// response as `id?: RequestId` where RequestId = string | number, so an
// absent id is encoded by omitting the "id" key, never as null. See
// session.HasJSONRPCID.
func rateLimitedBody(requestID any, retryAfter time.Duration) []byte {
        retrySeconds := math.Ceil(retryAfter.Seconds())
        resp := map[string]any{
                "jsonrpc": "2.0",
                "error": map[string]any{
                        "code":    CodeRateLimited,
                        "message": MessageRateLimited,
                        "data": map[string]any{
                                "retryAfterSeconds": retrySeconds,
                        },
                },
        }
        if session.HasJSONRPCID(requestID) {
                resp["id"] = requestID
        }
        data, err := json.Marshal(resp)
        if err != nil {
                return []byte(fmt.Sprintf(
                        `{"jsonrpc":"2.0","error":{"code":429,"message":"Rate limit exceeded","data":{"retryAfterSeconds":%.0f}}}`,
                        retrySeconds,
                ))
        }
        return data
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package ratelimit

import (
        "context"
        "errors"
        "fmt"
        "io"
        "net"
        "strings"
        "time"

        "github.com/redis/go-redis/v9"
        "go.opentelemetry.io/otel/attribute"
        "go.opentelemetry.io/otel/metric"
        "go.opentelemetry.io/otel/trace"

        "github.com/stacklok/toolhive/pkg/telemetry"
)

const (
        rateLimitInstrumentationName = "github.com/stacklok/toolhive/pkg/ratelimit"

        rateLimitDecisionAllowed  = "allowed"
        rateLimitDecisionRejected = "rejected"
        rateLimitRejectedByNone   = "none"

        rateLimitScopeShared  = "shared"
        rateLimitScopePerUser = "per_user"

        rateLimitOperationServer = "server"
        rateLimitOperationTool   = "tool"

        redisErrorTypeTimeout    = "timeout"
        redisErrorTypeConnection = "connection"
        redisErrorTypeAuth       = "auth"
        redisErrorTypeOther      = "other"
)

type rateLimitTelemetry struct {
        namespace  string
        serverName string

        decisions    metric.Int64Counter
        redisErrors  metric.Int64Counter
        failOpen     metric.Int64Counter
        checkLatency metric.Float64Histogram
}

func newRateLimitTelemetry(
        meterProvider metric.MeterProvider,
        namespace string,
        serverName string,
) (*rateLimitTelemetry, error) {
        if meterProvider == nil {
                return nil, nil
        }

        meter := meterProvider.Meter(rateLimitInstrumentationName)

        decisions, err := meter.Int64Counter(
                "toolhive_rate_limit_decisions",
                metric.WithDescription("Total number of rate limit bucket decisions"),
        )
        if err != nil {
                return nil, fmt.Errorf("create decision counter: %w", err)
        }

        redisErrors, err := meter.Int64Counter(
                "toolhive_rate_limit_redis_errors",
                metric.WithDescription("Total number of Redis errors during rate limit checks"),
        )
        if err != nil {
                return nil, fmt.Errorf("create Redis error counter: %w", err)
        }

        failOpen, err := meter.Int64Counter(
                "toolhive_rate_limit_fail_open",
                metric.WithDescription("Total number of rate limit checks allowed after an enforcement error"),
        )
        if err != nil {
                return nil, fmt.Errorf("create fail-open counter: %w", err)
        }

        checkLatency, err := meter.Float64Histogram(
                "toolhive_rate_limit_check_latency",
                metric.WithDescription("Duration of Redis Lua rate limit checks in seconds"),
                metric.WithUnit("s"),
                metric.WithExplicitBucketBoundaries(telemetry.MCPHistogramBuckets...),
        )
        if err != nil {
                return nil, fmt.Errorf("create check latency histogram: %w", err)
        }

        return &rateLimitTelemetry{
                namespace:    namespace,
                serverName:   serverName,
                decisions:    decisions,
                redisErrors:  redisErrors,
                failOpen:     failOpen,
                checkLatency: checkLatency,
        }, nil
}

func (t *rateLimitTelemetry) recordAllowed(ctx context.Context, checks []limitCheck) {
        if t == nil {
                return
        }
        for _, check := range checks {
                t.recordDecision(ctx, rateLimitDecisionAllowed, check)
        }
}

func (t *rateLimitTelemetry) recordRejected(ctx context.Context, check limitCheck) {
        if t == nil {
                return
        }
        t.recordDecision(ctx, rateLimitDecisionRejected, check)
}

func (t *rateLimitTelemetry) recordDecision(ctx context.Context, decision string, check limitCheck) {
        t.decisions.Add(ctx, 1, metric.WithAttributes(
                attribute.String("namespace", t.namespace),
                attribute.String("server", t.serverName),
                attribute.String("decision", decision),
                attribute.String("scope", check.scope),
                attribute.String("operation_type", check.operationType),
        ))
}

func (t *rateLimitTelemetry) recordRedisError(ctx context.Context, err error) {
        if t == nil {
                return
        }
        t.redisErrors.Add(ctx, 1, metric.WithAttributes(
                attribute.String("namespace", t.namespace),
                attribute.String("server", t.serverName),
                attribute.String("error_type", classifyRedisError(err)),
        ))
}

func (t *rateLimitTelemetry) recordFailOpen(ctx context.Context) {
        if t == nil {
                return
        }
        t.failOpen.Add(ctx, 1, metric.WithAttributes(
                attribute.String("namespace", t.namespace),
                attribute.String("server", t.serverName),
        ))
}

func (t *rateLimitTelemetry) recordCheckLatency(ctx context.Context, duration time.Duration) {
        if t == nil {
                return
        }
        t.checkLatency.Record(ctx, duration.Seconds(), metric.WithAttributes(
                attribute.String("namespace", t.namespace),
                attribute.String("server", t.serverName),
        ))
}

func recordRateLimitSpanOutcome(ctx context.Context, decision, rejectedBy string, failOpen bool) {
        trace.SpanFromContext(ctx).SetAttributes(
                attribute.String("rate_limit.decision", decision),
                attribute.String("rate_limit.rejected_by", rejectedBy),
                attribute.Bool("rate_limit.fail_open", failOpen),
        )
}

func classifyRedisError(err error) string {
        if redis.IsAuthError(err) {
                return redisErrorTypeAuth
        }

        if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, redis.ErrPoolTimeout) {
                return redisErrorTypeTimeout
        }
        var netErr net.Error
        if errors.As(err, &netErr) && netErr.Timeout() {
                return redisErrorTypeTimeout
        }

        if errors.Is(err, redis.ErrClosed) ||
                errors.Is(err, net.ErrClosed) ||
                errors.Is(err, io.EOF) ||
                errors.Is(err, io.ErrUnexpectedEOF) {
                return redisErrorTypeConnection
        }
        var opErr *net.OpError
        if errors.As(err, &opErr) {
                return redisErrorTypeConnection
        }

        message := strings.ToLower(err.Error())
        if strings.Contains(message, "connection refused") ||
                strings.Contains(message, "connection reset") ||
                strings.Contains(message, "broken pipe") ||
                strings.Contains(message, "no such host") {
                return redisErrorTypeConnection
        }

        return redisErrorTypeOther
}

//go:build !ignore_autogenerated

/*
Copyright 2025 Stacklok

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Code generated by controller-gen. DO NOT EDIT.

package types

import ()

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RateLimitBucket) DeepCopyInto(out *RateLimitBucket) {
        *out = *in
        out.RefillPeriod = in.RefillPeriod
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RateLimitBucket.
func (in *RateLimitBucket) DeepCopy() *RateLimitBucket {
        if in == nil {
                return nil
        }
        out := new(RateLimitBucket)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RateLimitConfig) DeepCopyInto(out *RateLimitConfig) {
        *out = *in
        if in.Shared != nil {
                in, out := &in.Shared, &out.Shared
                *out = new(RateLimitBucket)
                **out = **in
        }
        if in.PerUser != nil {
                in, out := &in.PerUser, &out.PerUser
                *out = new(RateLimitBucket)
                **out = **in
        }
        if in.Tools != nil {
                in, out := &in.Tools, &out.Tools
                *out = make([]ToolRateLimitConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RateLimitConfig.
func (in *RateLimitConfig) DeepCopy() *RateLimitConfig {
        if in == nil {
                return nil
        }
        out := new(RateLimitConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ToolRateLimitConfig) DeepCopyInto(out *ToolRateLimitConfig) {
        *out = *in
        if in.Shared != nil {
                in, out := &in.Shared, &out.Shared
                *out = new(RateLimitBucket)
                **out = **in
        }
        if in.PerUser != nil {
                in, out := &in.PerUser, &out.PerUser
                *out = new(RateLimitBucket)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ToolRateLimitConfig.
func (in *ToolRateLimitConfig) DeepCopy() *ToolRateLimitConfig {
        if in == nil {
                return nil
        }
        out := new(ToolRateLimitConfig)
        in.DeepCopyInto(out)
        return out
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package recovery adapts toolhive-core's panic recovery middleware to
// ToolHive's middleware factory, wiring in ToolHive's observability:
// OTel span error recording and Sentry issue reporting.
package recovery

import (
        "net/http"

        "go.opentelemetry.io/otel/codes"
        "go.opentelemetry.io/otel/trace"

        corerecovery "github.com/stacklok/toolhive-core/recovery"
        sentrypkg "github.com/stacklok/toolhive/pkg/sentry"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// MiddlewareType is the type constant for recovery middleware
const MiddlewareType = "recovery"

// panicHandler wires core's recovery hooks to ToolHive's OTel spans and
// Sentry reporting.
type panicHandler struct{}

// RecordError records a sanitized error on the request's span. The generic
// message is deliberate: panic values may embed credentials or internal
// state that must not reach external telemetry backends. Full details are
// in the log and in Sentry.
func (panicHandler) RecordError(r *http.Request, err error) {
        span := trace.SpanFromContext(r.Context())
        span.RecordError(err)
        span.SetStatus(codes.Error, "panic recovered")
}

// ReportPanic reports the raw panic value to Sentry. The Sentry span
// processor only creates transactions; reporting explicitly makes panics
// also appear as Issues in the Sentry Issues tab.
func (panicHandler) ReportPanic(r *http.Request, v any) {
        sentrypkg.RecoverPanic(r, v)
}

// Middleware is an HTTP middleware that recovers from panics.
// When a panic occurs, it logs the error and returns
// a 500 Internal Server Error response to the client.
func Middleware(next http.Handler) http.Handler {
        return corerecovery.Middleware(next, corerecovery.WithPanicHandler(panicHandler{}))
}

// FactoryMiddleware wraps recovery middleware functionality for the factory pattern.
type FactoryMiddleware struct{}

// Handler returns the middleware function used by the proxy.
func (FactoryMiddleware) Handler() types.MiddlewareFunction {
        return Middleware
}

// Close cleans up any resources used by the middleware.
func (FactoryMiddleware) Close() error {
        // Recovery middleware doesn't need cleanup
        return nil
}

// CreateMiddleware is the factory function for recovery middleware.
// It creates and registers the recovery middleware with the runner.
func CreateMiddleware(_ *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        recoveryMw := &FactoryMiddleware{}
        runner.AddMiddleware(MiddlewareType, recoveryMw)
        return nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package redisconfig provides serializable Redis connection settings.
//
// +groupName=toolhive.stacklok.dev
// +versionName=redisconfig
package redisconfig

import (
        "fmt"
        "os"

        "github.com/stacklok/toolhive-core/redisconn"
)

// TLSConfig configures TLS for a Redis connection. Presence enables TLS;
// an empty configuration verifies the server certificate using system roots.
// +kubebuilder:object:generate=true
// +gendoc
type TLSConfig struct {
        // InsecureSkipVerify skips server certificate and hostname verification.
        // +optional
        InsecureSkipVerify bool `json:"insecureSkipVerify,omitempty" yaml:"insecureSkipVerify,omitempty"`

        // CACertFile is the path to a PEM-encoded CA bundle. When configured, this
        // bundle replaces the system roots used to verify the Redis server.
        // +optional
        CACertFile string `json:"caCertFile,omitempty" yaml:"caCertFile,omitempty"`
}

// Load resolves and validates the CA bundle for a Redis connection. A nil
// receiver disables TLS. A configured CA file must contain valid certificates;
// errors are returned rather than falling back to system roots or plaintext.
func (c *TLSConfig) Load() (*redisconn.TLSConfig, error) {
        if c == nil {
                return nil, nil
        }
        cfg := &redisconn.TLSConfig{InsecureSkipVerify: c.InsecureSkipVerify}
        if c.CACertFile != "" {
                // #nosec G304 -- The CA file path comes from trusted runtime configuration.
                data, err := os.ReadFile(c.CACertFile)
                if err != nil {
                        return nil, fmt.Errorf("failed to read Redis CA cert file %q: %w", c.CACertFile, err)
                }
                if len(data) == 0 {
                        return nil, fmt.Errorf("redis CA cert file %q is empty", c.CACertFile)
                }
                cfg.CACert = data
        }
        if _, err := redisconn.BuildTLSConfig(cfg); err != nil {
                return nil, fmt.Errorf("invalid Redis TLS configuration: %w", err)
        }
        return cfg, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package legacyhint detects the deprecated ToolHive registry format and
// supplies a single migration message used by every parser/validator entry
// point. Lives in its own leaf package so pkg/registry and pkg/config can both
// import it without creating an import cycle (pkg/registry imports pkg/config).
package legacyhint

import "encoding/json"

// MigrationMessage is the user-facing text returned when a legacy ToolHive
// registry file is detected. Kept identical across the runtime parser,
// set-registry-file, set-registry-url, and remote provider validation paths.
const MigrationMessage = "registry file appears to be in the legacy ToolHive format; " +
        "run `thv registry convert --in <path> --in-place` to migrate to the upstream MCP format"

// Looks reports whether the JSON document has top-level "servers",
// "remote_servers", or "groups" — the markers of the legacy ToolHive registry
// layout. The upstream format wraps these under a top-level "data" object, so a
// match here means the input is legacy (or close enough that emitting the
// migration hint is more useful than a generic decode error).
//
// Used to short-circuit with a migration hint instead of the misleading
// "no servers" error that Go's JSON decoder produces when legacy fields are
// silently dropped during unmarshal into UpstreamRegistry.
func Looks(data []byte) bool {
        var probe struct {
                Servers       json.RawMessage `json:"servers"`
                RemoteServers json.RawMessage `json:"remote_servers"`
                Groups        json.RawMessage `json:"groups"`
        }
        if err := json.Unmarshal(data, &probe); err != nil {
                return false
        }
        return len(probe.Servers) > 0 || len(probe.RemoteServers) > 0 || len(probe.Groups) > 0
}

// IsUpstream reports whether the JSON document appears to use the upstream
// registry format. The discriminator is a top-level "data" object — only the
// upstream format wraps servers inside it. The "$schema" key alone is not
// sufficient because the legacy format also includes one.
func IsUpstream(data []byte) bool {
        var probe struct {
                Data json.RawMessage `json:"data"`
        }
        if err := json.Unmarshal(data, &probe); err != nil {
                return false
        }
        return len(probe.Data) > 0 && probe.Data[0] == '{'
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package runner provides functionality for running MCP servers
package runner

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"

        "github.com/stacklok/toolhive-core/permissions"
        v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
        "github.com/stacklok/toolhive/pkg/audit"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/awssts"
        "github.com/stacklok/toolhive/pkg/auth/remote"
        authsecrets "github.com/stacklok/toolhive/pkg/auth/secrets"
        "github.com/stacklok/toolhive/pkg/auth/upstreamswap"
        "github.com/stacklok/toolhive/pkg/authserver"
        "github.com/stacklok/toolhive/pkg/authz"
        "github.com/stacklok/toolhive/pkg/container"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/container/templates"
        "github.com/stacklok/toolhive/pkg/environment"
        "github.com/stacklok/toolhive/pkg/ignore"
        "github.com/stacklok/toolhive/pkg/labels"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto/tokenexchange"
        "github.com/stacklok/toolhive/pkg/redisconfig"
        "github.com/stacklok/toolhive/pkg/secrets"
        "github.com/stacklok/toolhive/pkg/state"
        "github.com/stacklok/toolhive/pkg/telemetry"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/webhook"
        workloadtypes "github.com/stacklok/toolhive/pkg/workloads/types"
)

// CurrentSchemaVersion is the current version of the RunConfig schema
// TODO: Set to "v1.0.0" when we clean up the middleware configuration.
const CurrentSchemaVersion = "v0.1.0"

// RunConfig contains all the configuration needed to run an MCP server
// It is serializable to JSON and YAML
// NOTE: This format is importable and exportable, and as a result should be
// considered part of ToolHive's API contract.
type RunConfig struct {
        // SchemaVersion is the version of the RunConfig schema
        SchemaVersion string `json:"schema_version" yaml:"schema_version"`

        // MCPServerGeneration is the K8s .metadata.generation of the MCPServer CR that rendered
        // this RunConfig. The Kubernetes runtime uses it as a monotonic version to prevent stale
        // rolling-update pods from overwriting a newer RunConfig's StatefulSet apply. Zero value
        // means unversioned (backward-compat with older operators, or non-operator callers).
        MCPServerGeneration int64 `json:"mcpserver_generation,omitempty" yaml:"mcpserver_generation,omitempty"`

        // Image is the Docker image to run
        Image string `json:"image" yaml:"image"`

        // RemoteURL is the URL of the remote MCP server (if running remotely)
        RemoteURL string `json:"remote_url,omitempty" yaml:"remote_url,omitempty"`

        // RegistryAPIURL is the registry API URL that served this server's metadata.
        // Empty when the server was not discovered via registry lookup.
        RegistryAPIURL string `json:"registry_api_url,omitempty" yaml:"registry_api_url,omitempty"`

        // RegistryURL is the registry URL that served this server's metadata.
        // Empty when the server was not discovered via registry lookup.
        RegistryURL string `json:"registry_url,omitempty" yaml:"registry_url,omitempty"`

        // RegistryServerName is the registry entry name used to look up this server's metadata.
        // Empty when the server was not discovered via registry lookup.
        RegistryServerName string `json:"registry_server_name,omitempty" yaml:"registry_server_name,omitempty"`

        // RemoteAuthConfig contains OAuth configuration for remote MCP servers
        RemoteAuthConfig *remote.Config `json:"remote_auth_config,omitempty" yaml:"remote_auth_config,omitempty"`

        // CmdArgs are the arguments to pass to the container
        CmdArgs []string `json:"cmd_args,omitempty" yaml:"cmd_args,omitempty"`

        // Name is the name of the MCP server
        Name string `json:"name" yaml:"name"`

        // ContainerName is the name of the container
        ContainerName string `json:"container_name,omitempty" yaml:"container_name,omitempty"`

        // BaseName is the base name used for the container (without prefixes)
        BaseName string `json:"base_name,omitempty" yaml:"base_name,omitempty"`

        // Transport is the transport mode (stdio, sse, or streamable-http)
        Transport types.TransportType `json:"transport" yaml:"transport" enums:"stdio,sse,streamable-http,inspector"`

        // Host is the host for the HTTP proxy
        Host string `json:"host" yaml:"host"`

        // Port is the port for the HTTP proxy to listen on (host port)
        Port int `json:"port" yaml:"port"`

        // TargetPort is the port for the container to expose (only applicable to SSE transport)
        TargetPort int `json:"target_port,omitempty" yaml:"target_port,omitempty"`

        // TargetHost is the host to forward traffic to (only applicable to SSE transport)
        TargetHost string `json:"target_host,omitempty" yaml:"target_host,omitempty"`

        // AllowedOrigins is the allowlist of values accepted on the HTTP Origin header,
        // used for DNS-rebinding protection per MCP 2025-11-25 §"Security Warning".
        // When empty and Host is loopback (127.0.0.1 / localhost / [::1]), a default
        // loopback-only allowlist is derived at middleware-wiring time.
        // When empty and Host is non-loopback, the middleware is disabled — operators
        // exposing the proxy publicly must configure an explicit allowlist.
        AllowedOrigins []string `json:"allowed_origins,omitempty" yaml:"allowed_origins,omitempty"`

        // Publish lists ports to publish to the host in format "hostPort:containerPort"
        Publish []string `json:"publish,omitempty" yaml:"publish,omitempty"`

        // PermissionProfileNameOrPath is the name or path of the permission profile
        PermissionProfileNameOrPath string `json:"permission_profile_name_or_path,omitempty" yaml:"permission_profile_name_or_path,omitempty"` //nolint:lll

        // PermissionProfile is the permission profile to use
        PermissionProfile *permissions.Profile `json:"permission_profile" yaml:"permission_profile" swaggerignore:"true"`

        // EnvVars are the parsed environment variables as key-value pairs
        EnvVars map[string]string `json:"env_vars,omitempty" yaml:"env_vars,omitempty"`

        // DEPRECATED: No longer appears to be used.
        // EnvFileDir is the directory path to load environment files from
        EnvFileDir string `json:"env_file_dir,omitempty" yaml:"env_file_dir,omitempty"`

        // Debug indicates whether debug mode is enabled
        Debug bool `json:"debug,omitempty" yaml:"debug,omitempty"`

        // Volumes are the directory mounts to pass to the container
        // Format: "host-path:container-path[:ro]"
        Volumes []string `json:"volumes,omitempty" yaml:"volumes,omitempty"`

        // ContainerLabels are the labels to apply to the container
        ContainerLabels map[string]string `json:"container_labels,omitempty" yaml:"container_labels,omitempty"`

        // DEPRECATED: Middleware configuration.
        // OIDCConfig contains OIDC configuration
        OIDCConfig *auth.TokenValidatorConfig `json:"oidc_config,omitempty" yaml:"oidc_config,omitempty"`

        // TokenExchangeConfig contains token exchange configuration for external authentication
        TokenExchangeConfig *tokenexchange.Config `json:"token_exchange_config,omitempty" yaml:"token_exchange_config,omitempty"`

        // UpstreamSwapConfig contains configuration for upstream token swap middleware.
        // When set along with EmbeddedAuthServerConfig, this middleware exchanges ToolHive JWTs
        // for upstream IdP tokens before forwarding requests to the MCP server.
        UpstreamSwapConfig *upstreamswap.Config `json:"upstream_swap_config,omitempty" yaml:"upstream_swap_config,omitempty"`

        // AWSStsConfig contains AWS STS token exchange configuration for accessing AWS services
        AWSStsConfig *awssts.Config `json:"aws_sts_config,omitempty" yaml:"aws_sts_config,omitempty"`

        // DEPRECATED: Middleware configuration.
        // AuthzConfig contains the authorization configuration
        AuthzConfig *authz.Config `json:"authz_config,omitempty" yaml:"authz_config,omitempty"`

        // DEPRECATED: Middleware configuration.
        // AuthzConfigPath is the path to the authorization configuration file
        AuthzConfigPath string `json:"authz_config_path,omitempty" yaml:"authz_config_path,omitempty"`

        // DEPRECATED: Middleware configuration.
        // AuditConfig contains the audit logging configuration
        AuditConfig *audit.Config `json:"audit_config,omitempty" yaml:"audit_config,omitempty"`

        // DEPRECATED: Middleware configuration.
        // AuditConfigPath is the path to the audit configuration file
        AuditConfigPath string `json:"audit_config_path,omitempty" yaml:"audit_config_path,omitempty"`

        // DEPRECATED: Middleware configuration.
        // TelemetryConfig contains the OpenTelemetry configuration
        TelemetryConfig *telemetry.Config `json:"telemetry_config,omitempty" yaml:"telemetry_config,omitempty"`

        // RateLimitConfig contains the CRD rate limiting configuration.
        // When set, rate limiting middleware is added to the proxy middleware chain.
        RateLimitConfig *v1beta1.RateLimitConfig `json:"rate_limit_config,omitempty" yaml:"rate_limit_config,omitempty"`

        // RateLimitNamespace is the Kubernetes namespace for Redis key derivation.
        RateLimitNamespace string `json:"rate_limit_namespace,omitempty" yaml:"rate_limit_namespace,omitempty"`

        // Secrets are the secret parameters to pass to the container
        // Format: "<secret name>,target=<target environment variable>"
        Secrets []string `json:"secrets,omitempty" yaml:"secrets,omitempty"`

        // K8sPodTemplatePatch is a JSON string to patch the Kubernetes pod template
        // Only applicable when using Kubernetes runtime
        K8sPodTemplatePatch string `json:"k8s_pod_template_patch,omitempty" yaml:"k8s_pod_template_patch,omitempty"`

        // Deployer is the container runtime to use (not serialized)
        Deployer rt.Deployer `json:"-" yaml:"-"`

        // buildContext indicates whether this config is being built for CLI or operator use (not serialized)
        buildContext BuildContext

        // IsolateNetwork indicates whether to isolate the network for the container
        IsolateNetwork bool `json:"isolate_network,omitempty" yaml:"isolate_network,omitempty"`

        // AllowDockerGateway permits outbound connections to Docker gateway addresses
        // (all daemon bridge gateways, host.docker.internal, and gateway.docker.internal). These are
        // blocked by default in the egress proxy even when InsecureAllowAll is set.
        // Only applicable to Docker deployments with network isolation enabled.
        // Gateways are snapshotted at deployment; redeploy to include bridges created later.
        // Gateway access is port-independent: it ignores the permission profile's
        // allowed ports, so once enabled the gateway is reachable on any port.
        AllowDockerGateway bool `json:"allow_docker_gateway,omitempty" yaml:"allow_docker_gateway,omitempty"`

        // TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies
        TrustProxyHeaders bool `json:"trust_proxy_headers,omitempty" yaml:"trust_proxy_headers,omitempty"`

        // StrictProtocolValidation enables strict MCP-Protocol-Version validation
        // on the streamable HTTP proxy: a request whose header names an unknown
        // MCP revision is rejected with HTTP 400. Default false accepts any
        // version string (an absent header is always accepted in either mode).
        StrictProtocolValidation bool `json:"strict_protocol_validation,omitempty" yaml:"strict_protocol_validation,omitempty"`

        // Stateless indicates the server only supports POST (no SSE/GET).
        // When true, the proxy returns 405 for incoming GET requests and uses a
        // POST-based health check instead of the default GET probe.
        // Applies to both remote URLs and local container workloads.
        Stateless bool `json:"stateless,omitempty" yaml:"stateless,omitempty"`

        // SessionTTL is the inactivity timeout for proxy sessions, expressed as a Go
        // duration string (e.g. "30m", "2h", "168h"). Empty uses the transport
        // default (2h). Negative durations and values that fail time.ParseDuration
        // are rejected at runtime.
        // String (not time.Duration) keeps the wire format unit-explicit: a
        // time.Duration field serializes as nanoseconds in JSON.
        SessionTTL string `json:"session_ttl,omitempty" yaml:"session_ttl,omitempty" example:"2h"`

        // MaxRequestBodySize is the maximum inbound MCP proxy request body size in bytes.
        // Zero uses the default limit of 8 MiB. Negative values are rejected
        // when the RunConfig is built or used at runtime.
        MaxRequestBodySize int64 `json:"max_request_body_size,omitempty" yaml:"max_request_body_size,omitempty"`

        // ProxyReadTimeout bounds reading the entire request (headers + body) on the
        // proxy HTTP server, expressed as a Go duration string (e.g. "30s", "1m").
        // Empty uses the proxy default (30s). Negative durations and values that fail
        // time.ParseDuration are rejected at runtime. Applies to all HTTP transports.
        // String (not time.Duration) keeps the wire format unit-explicit.
        ProxyReadTimeout string `json:"proxy_read_timeout,omitempty" yaml:"proxy_read_timeout,omitempty" example:"30s"`

        // ProxyMode is the effective HTTP protocol the proxy uses.
        // For stdio transports, this is the configured mode (sse or streamable-http).
        // For direct transports (sse/streamable-http), this matches the transport type.
        // Note: "sse" is deprecated; use "streamable-http" instead.
        ProxyMode types.ProxyMode `json:"proxy_mode,omitempty" yaml:"proxy_mode,omitempty" enums:"sse,streamable-http"`

        // DEPRECATED: No longer appears to be used.
        // ThvCABundle is the path to the CA certificate bundle for ToolHive HTTP operations
        ThvCABundle string `json:"thv_ca_bundle,omitempty" yaml:"thv_ca_bundle,omitempty"`

        // DEPRECATED: No longer appears to be used.
        // JWKSAuthTokenFile is the path to file containing auth token for JWKS/OIDC requests
        JWKSAuthTokenFile string `json:"jwks_auth_token_file,omitempty" yaml:"jwks_auth_token_file,omitempty"`

        // Group is the name of the group this workload belongs to, if any
        Group string `json:"group,omitempty" yaml:"group,omitempty"`

        // DEPRECATED: Middleware configuration.
        // ToolsFilter is the list of tools to filter
        ToolsFilter []string `json:"tools_filter,omitempty" yaml:"tools_filter,omitempty"`

        // DEPRECATED: Middleware configuration.
        // ToolsOverride is a map from an actual tool to its overridden name and/or description
        ToolsOverride map[string]ToolOverride `json:"tools_override,omitempty" yaml:"tools_override,omitempty"`

        // IgnoreConfig contains configuration for ignore processing
        IgnoreConfig *ignore.Config `json:"ignore_config,omitempty" yaml:"ignore_config,omitempty"`

        // MiddlewareConfigs contains the list of middleware to apply to the transport
        // and the configuration for each middleware.
        MiddlewareConfigs []types.MiddlewareConfig `json:"middleware_configs,omitempty" yaml:"middleware_configs,omitempty"`

        // AdditionalMiddlewareConfigs carries pre-built middleware configs injected by
        // external-auth handlers (reached via *[]RunConfigBuilderOption) rather than
        // derived from typed RunConfig fields. PopulateMiddlewareConfigs splices these
        // into the chain in the backend-egress group — after auth and before recovery —
        // instead of discarding them. Upstream carries these configs verbatim and never
        // inspects their parameters; the middleware type identity (e.g. an enterprise
        // auth type) is supplied by the caller via types.MiddlewareConfig.Type.
        //
        // Each entry's Type is expected to be a NEW egress middleware type (e.g. OBO),
        // not one already produced from a typed RunConfig field (auth, authz, audit,
        // tokenExchange, awssts, …). Dispatch in the proxyrunner is purely by Type
        // string, so an injected Type that shadows a typed-field type would add a
        // second instance of that middleware to the chain; the seam does not validate
        // against this.
        AdditionalMiddlewareConfigs []types.MiddlewareConfig `json:"additional_middleware_configs,omitempty" yaml:"additional_middleware_configs,omitempty"` //nolint:lll

        // ValidatingWebhooks contains the configuration for validating webhook middleware.
        ValidatingWebhooks []webhook.Config `json:"validating_webhooks,omitempty" yaml:"validating_webhooks,omitempty"`

        // MutatingWebhooks contains the configuration for mutating webhook middleware.
        // Mutating webhooks run before validating webhooks, per RFC THV-0017 ordering.
        MutatingWebhooks []webhook.Config `json:"mutating_webhooks,omitempty" yaml:"mutating_webhooks,omitempty"`

        // existingPort is the port from an existing workload being updated (not serialized)
        // Used during port validation to allow reusing the same port
        existingPort int

        // EndpointPrefix is an explicit prefix to prepend to SSE endpoint URLs.
        // This is used to handle path-based ingress routing scenarios.
        EndpointPrefix string `json:"endpoint_prefix,omitempty" yaml:"endpoint_prefix,omitempty"`

        // RuntimeConfig allows overriding the default runtime configuration
        // for this specific workload (base images and packages)
        RuntimeConfig *templates.RuntimeConfig `json:"runtime_config,omitempty" yaml:"runtime_config,omitempty"`

        // HeaderForward contains configuration for injecting headers into requests to remote servers.
        HeaderForward *HeaderForwardConfig `json:"header_forward,omitempty" yaml:"header_forward,omitempty"`

        // EmbeddedAuthServerConfig contains configuration for the embedded OAuth2/OIDC authorization server.
        // When set, the proxy runner will start an embedded auth server that delegates to upstream IDPs.
        // This is the serializable RunConfig; secrets are referenced by file paths or env var names.
        EmbeddedAuthServerConfig *authserver.RunConfig `json:"embedded_auth_server_config,omitempty" yaml:"embedded_auth_server_config,omitempty"` //nolint:lll

        // ScalingConfig contains configuration for horizontal scaling of the proxy runner.
        // Only applicable when running in Kubernetes with the ToolHive operator.
        // When nil, no scaling configuration is applied (single-replica default behavior).
        ScalingConfig *ScalingConfig `json:"scaling_config,omitempty" yaml:"scaling_config,omitempty"`
}

// ScalingConfig contains configuration for horizontal scaling of the proxy runner backend.
// It is intentionally kept as a separate struct so future scaling knobs (MinReplicas,
// MaxReplicas, ScaleDownWindow, etc.) can be added here without growing the top-level
// RunConfig shape.
type ScalingConfig struct {
        // BackendReplicas is the desired StatefulSet replica count for the proxy runner backend.
        // When nil, replicas are unmanaged (preserving HPA or manual kubectl control).
        // When set (including 0), the value is an explicit replica count.
        BackendReplicas *int32 `json:"backend_replicas,omitempty" yaml:"backend_replicas,omitempty"`

        // SessionRedis holds non-sensitive Redis connection parameters for distributed session storage.
        // Populated only when MCPServer.spec.sessionStorage.provider == "redis".
        // The Redis password is not included — it is injected as env var THV_SESSION_REDIS_PASSWORD.
        // +optional
        SessionRedis *SessionRedisConfig `json:"session_redis,omitempty" yaml:"session_redis,omitempty"`
}

// SessionRedisConfig contains non-sensitive Redis connection parameters used for distributed
// session storage when the operator is configured with sessionStorage.provider == "redis".
// The Redis password is excluded and injected separately as env var THV_SESSION_REDIS_PASSWORD.
type SessionRedisConfig struct {
        // Address is the Redis server address (host:port).
        Address string `json:"address,omitempty" yaml:"address,omitempty"`

        // DB is the Redis database number.
        DB int32 `json:"db,omitempty" yaml:"db,omitempty"`

        // KeyPrefix is an optional prefix applied to all Redis keys used by ToolHive.
        KeyPrefix string `json:"key_prefix,omitempty" yaml:"key_prefix,omitempty"`

        // TLS enables TLS for session storage and rate limiting. An empty
        // configuration verifies the Redis server using system roots.
        TLS *redisconfig.TLSConfig `json:"tls,omitempty" yaml:"tls,omitempty"`
}

// NormalizeProxyMode sets ProxyMode to the effective value based on the
// transport type, so downstream readers always see the actual HTTP protocol.
func (c *RunConfig) NormalizeProxyMode() {
        c.ProxyMode = types.EffectiveProxyMode(c.Transport, c.ProxyMode)
}

// WriteJSON serializes the RunConfig to JSON and writes it to the provided writer
func (c *RunConfig) WriteJSON(w io.Writer) error {
        // Ensure the schema version is set
        if c.SchemaVersion == "" {
                c.SchemaVersion = CurrentSchemaVersion
        }
        encoder := json.NewEncoder(w)
        encoder.SetIndent("", "  ")
        return encoder.Encode(c)
}

// ReadJSON deserializes the RunConfig from JSON read from the provided reader
func ReadJSON(r io.Reader) (*RunConfig, error) {
        var config RunConfig
        if err := state.ReadJSON(r, &config); err != nil {
                return nil, err
        }

        // Initialize maps if they're nil after deserialization
        if config.EnvVars == nil {
                config.EnvVars = make(map[string]string)
        }
        if config.ContainerLabels == nil {
                config.ContainerLabels = make(map[string]string)
        }

        // Initialize slices if they're nil after deserialization
        if config.CmdArgs == nil {
                config.CmdArgs = []string{}
        }
        if config.Volumes == nil {
                config.Volumes = []string{}
        }
        if config.Secrets == nil {
                config.Secrets = []string{}
        }

        // Set the default schema version if not set
        if config.SchemaVersion == "" {
                config.SchemaVersion = CurrentSchemaVersion
        }

        // Migrate plain text OAuth client secrets to CLI format
        if err := migrateOAuthClientSecret(&config); err != nil {
                return nil, fmt.Errorf("failed to migrate OAuth client secret: %w", err)
        }

        // Migrate plain text bearer tokens to CLI format
        if err := migrateBearerToken(&config); err != nil {
                return nil, fmt.Errorf("failed to migrate bearer token: %w", err)
        }

        // Repair legacy serialized auth middleware in memory without rewriting the source config.
        if err := canonicalizeOIDCMiddlewareConfig(&config); err != nil {
                return nil, fmt.Errorf("invalid OIDC middleware configuration: %w", err)
        }

        // Normalize proxyMode so pre-existing configs always reflect the effective protocol
        config.NormalizeProxyMode()

        // Self-heal legacy on-disk configs that persisted isolate_network=true with a
        // non-bridge network mode. Isolation is only enforceable in bridge mode, so
        // reflect reality in memory for status/list. See issue #5775.
        degradeNetworkIsolation(&config)

        return &config, nil
}

// degradeNetworkIsolation drops network isolation from a loaded config when its
// resolved network mode cannot enforce it, so status/list reflect reality.
//
// ReadJSON is reached from read-only paths (thv list, status, export, API GETs)
// as well as deploy. The "none" case logs at DEBUG since it's merely redundant
// (already maximally confined) and would otherwise fire on every read of a
// static, already-known condition. The host/custom case stays at WARN even on
// read paths: it means the user's --isolate-network request is being silently
// ignored, which they should keep seeing until they either drop the flag or
// switch network modes. See #5775.
func degradeNetworkIsolation(config *RunConfig) {
        if !config.IsolateNetwork {
                return
        }
        mode := ""
        if config.PermissionProfile != nil && config.PermissionProfile.Network != nil {
                mode = config.PermissionProfile.Network.Mode
        }
        if networking.IsBridgeMode(mode) {
                return
        }
        config.IsolateNetwork = false
        if mode == "none" {
                slog.Debug(networking.NetworkIsolationNoneRedundantMsg, "network_mode", mode)
                return
        }
        slog.Warn(networking.NetworkIsolationHostDroppedMsg, "network_mode", mode)
}

// migrateOAuthClientSecret migrates plain text OAuth client secrets to CLI format
// This handles the transition from storing plain text secrets to CLI format references
func migrateOAuthClientSecret(config *RunConfig) error {
        if config.RemoteAuthConfig == nil {
                return nil // No OAuth config to migrate
        }

        if config.RemoteAuthConfig.ClientSecret == "" {
                return nil
        }

        // Check if the client secret is already in CLI format
        if _, err := secrets.ParseSecretParameter(config.RemoteAuthConfig.ClientSecret); err == nil {
                return nil // Already in CLI format, no migration needed
        }

        // The client secret is in plain text format - migrate it
        cliFormatSecret, err := authsecrets.ProcessSecret(
                config.Name,
                config.RemoteAuthConfig.ClientSecret,
                authsecrets.TokenTypeOAuthClientSecret,
        )
        if err != nil {
                return fmt.Errorf("failed to process OAuth client secret: %w", err)
        }

        // Update the RunConfig to use the CLI format reference
        config.RemoteAuthConfig.ClientSecret = cliFormatSecret

        // Save the migrated RunConfig back to disk so migration only happens once
        if err := config.SaveState(context.Background()); err != nil {
                // Log error without potentially sensitive details - only log error type and message
                slog.Warn("failed to save migrated RunConfig for workload", "name", config.Name, "error", err)
                // Don't fail the migration - the secret is already stored and the config is updated in memory
        }

        return nil
}

// migrateBearerToken migrates plain text bearer tokens to CLI format
// This handles the transition from storing plain text tokens to CLI format references
func migrateBearerToken(config *RunConfig) error {
        if config.RemoteAuthConfig == nil {
                return nil // No remote auth config to migrate
        }

        if config.RemoteAuthConfig.BearerToken == "" {
                return nil
        }

        // Check if the bearer token is already in CLI format
        if _, err := secrets.ParseSecretParameter(config.RemoteAuthConfig.BearerToken); err == nil {
                return nil // Already in CLI format, no migration needed
        }

        // The bearer token is in plain text format - migrate it
        cliFormatToken, err := authsecrets.ProcessSecret(
                config.Name,
                config.RemoteAuthConfig.BearerToken,
                authsecrets.TokenTypeBearerToken,
        )
        if err != nil {
                return fmt.Errorf("failed to process bearer token: %w", err)
        }

        // Update the RunConfig to use the CLI format reference
        config.RemoteAuthConfig.BearerToken = cliFormatToken

        // Save the migrated RunConfig back to disk so migration only happens once
        if err := config.SaveState(context.Background()); err != nil {
                // Log error without potentially sensitive details - only log error type and message
                slog.Warn("failed to save migrated RunConfig for workload", "name", config.Name, "error", err)
                // Don't fail the migration - the secret is already stored and the config is updated in memory
        }

        return nil
}

// NewRunConfig creates a new RunConfig with default values
func NewRunConfig() *RunConfig {
        return &RunConfig{
                ContainerLabels: make(map[string]string),
                EnvVars:         make(map[string]string),
        }
}

// WithAuthz adds authorization configuration to the RunConfig
func (c *RunConfig) WithAuthz(config *authz.Config) *RunConfig {
        c.AuthzConfig = config
        return c
}

// WithAudit adds audit configuration to the RunConfig
func (c *RunConfig) WithAudit(config *audit.Config) *RunConfig {
        c.AuditConfig = config
        return c
}

// WithMiddlewareConfig adds middleware configuration to the RunConfig
func (c *RunConfig) WithMiddlewareConfig(middlewareConfig []types.MiddlewareConfig) *RunConfig {
        c.MiddlewareConfigs = middlewareConfig
        return c
}

// WithTransport parses and sets the transport type
func (c *RunConfig) WithTransport(t string) (*RunConfig, error) {
        transportType, err := types.ParseTransportType(t)
        if err != nil {
                return c, fmt.Errorf("invalid transport mode: %s. Valid modes are: sse, streamable-http, stdio", t)
        }
        c.Transport = transportType
        return c, nil
}

// WithPorts configures the host and target ports
func (c *RunConfig) WithPorts(proxyPort, targetPort int) (*RunConfig, error) {
        // Skip port validation for operator context - ports will be used in containers, not on operator host
        if c.buildContext == BuildContextOperator {
                c.Port = proxyPort
                c.TargetPort = targetPort
                return c, nil
        }

        // CLI context: perform port validation as before
        var selectedPort int
        var err error

        // If the user requested an explicit proxy port, check if it's available.
        // If not available - treat as an error, since picking a random port here
        // is going to lead to confusion.
        if proxyPort != 0 {
                // Skip validation if reusing the same port from existing workload (during update)
                if proxyPort == c.existingPort && c.existingPort > 0 {
                        slog.Debug("reusing existing port", "port", proxyPort)
                        selectedPort = proxyPort
                } else if !networking.IsAvailable(proxyPort) {
                        return c, fmt.Errorf("requested proxy port %d is not available", proxyPort)
                } else {
                        slog.Debug("using requested port", "port", proxyPort)
                        selectedPort = proxyPort
                }
        } else {
                // Otherwise - pick a random available port.
                selectedPort, err = networking.FindOrUsePort(proxyPort)
                if err != nil {
                        return c, err
                }
        }
        c.Port = selectedPort

        // Select a target port for the container if using SSE or Streamable HTTP transport.
        // Target ports are container-internal; do not validate them against host availability.
        if c.Transport == types.TransportTypeSSE || c.Transport == types.TransportTypeStreamableHTTP {
                if targetPort != 0 {
                        slog.Debug("using target port", "port", targetPort)
                        c.TargetPort = targetPort
                } else {
                        selectedTargetPort, err := networking.FindOrUsePort(0)
                        if err != nil {
                                return c, fmt.Errorf("target port error: %w", err)
                        }
                        slog.Debug("using target port", "port", selectedTargetPort)
                        c.TargetPort = selectedTargetPort
                }
        }

        return c, nil
}

// WithEnvironmentVariables sets environment variables
func (c *RunConfig) WithEnvironmentVariables(envVars map[string]string) (*RunConfig, error) {
        // Initialize EnvVars if it's nil
        if c.EnvVars == nil {
                c.EnvVars = make(map[string]string)
        }

        // Merge the provided environment variables with existing ones
        for key, value := range envVars {
                c.EnvVars[key] = value
        }

        // Set transport-specific environment variables
        environment.SetTransportEnvironmentVariables(c.EnvVars, string(c.Transport), c.TargetPort)
        return c, nil
}

// ValidateSecrets checks if the secrets can be parsed and are valid
func (c *RunConfig) ValidateSecrets(ctx context.Context, userProvider secrets.Provider) error {
        if len(c.Secrets) > 0 {
                _, err := environment.ParseSecretParameters(ctx, c.Secrets, userProvider)
                if err != nil {
                        return fmt.Errorf("failed to get secrets: %w", err)
                }
        }
        if c.RemoteAuthConfig != nil && c.RemoteAuthConfig.ClientSecret != "" {
                _, err := secrets.ParseSecretParameter(c.RemoteAuthConfig.ClientSecret)
                if err != nil {
                        return fmt.Errorf("failed to get secrets: %w", err)
                }
        }

        return nil
}

// WithSecrets processes secrets and adds them to environment variables.
// systemProvider is used for system-managed secrets (auth tokens, registry credentials).
// userProvider is used for user-managed secrets (--secret flags, header secrets).
func (c *RunConfig) WithSecrets(
        ctx context.Context,
        systemProvider secrets.Provider,
        userProvider secrets.Provider,
) (*RunConfig, error) {
        // Process regular secrets if provided — these are user-managed (from --secret flags)
        if len(c.Secrets) > 0 {
                secretVariables, err := environment.ParseSecretParameters(ctx, c.Secrets, userProvider)
                if err != nil {
                        return c, fmt.Errorf("failed to get secrets: %w", err)
                }

                // Initialize EnvVars if it's nil
                if c.EnvVars == nil {
                        c.EnvVars = make(map[string]string)
                }

                // Add secret variables to environment variables
                for key, value := range secretVariables {
                        c.EnvVars[key] = value
                }
        }

        if c.RemoteAuthConfig != nil {
                c.RemoteAuthConfig.ClearResolvedSecrets()
        }

        // Resolve RemoteAuthConfig.ClientSecret while retaining its persistable reference.
        if c.RemoteAuthConfig != nil && c.RemoteAuthConfig.ClientSecret != "" {
                if secretParam, err := secrets.ParseSecretParameter(c.RemoteAuthConfig.ClientSecret); err == nil {
                        actualSecret, err := systemProvider.GetSecret(ctx, secretParam.Name)
                        if err != nil {
                                return c, fmt.Errorf("failed to resolve OAuth client secret '%s': %w", secretParam.Name, err)
                        }
                        c.RemoteAuthConfig.SetResolvedClientSecret(actualSecret)
                }
        }

        // Resolve RemoteAuthConfig.BearerToken while retaining its persistable reference.
        if c.RemoteAuthConfig != nil && c.RemoteAuthConfig.BearerToken != "" {
                if secretParam, err := secrets.ParseSecretParameter(c.RemoteAuthConfig.BearerToken); err == nil {
                        actualToken, err := systemProvider.GetSecret(ctx, secretParam.Name)
                        if err != nil {
                                return c, fmt.Errorf("failed to resolve bearer token '%s': %w", secretParam.Name, err)
                        }
                        c.RemoteAuthConfig.SetResolvedBearerToken(actualToken)
                }
        }

        // Process HeaderForward.AddHeadersFromSecret — user-managed secrets
        if err := c.resolveHeaderForwardSecrets(ctx, userProvider); err != nil {
                return c, err
        }

        return c, nil
}

// resolveHeaderForwardSecrets resolves secret references in HeaderForward.AddHeadersFromSecret
// and builds the merged resolvedHeaders map for middleware consumption.
// Only the secret references are persisted to disk; actual values exist only in memory
// via the non-serialized resolvedHeaders field.
func (c *RunConfig) resolveHeaderForwardSecrets(ctx context.Context, userProvider secrets.Provider) error {
        if c.HeaderForward == nil || len(c.HeaderForward.AddHeadersFromSecret) == 0 {
                return nil
        }
        // Build merged map: start with plaintext headers, then overlay resolved secrets.
        merged := make(map[string]string, len(c.HeaderForward.AddPlaintextHeaders)+len(c.HeaderForward.AddHeadersFromSecret))
        for k, v := range c.HeaderForward.AddPlaintextHeaders {
                merged[k] = v
        }
        for headerName, secretName := range c.HeaderForward.AddHeadersFromSecret {
                actualValue, err := userProvider.GetSecret(ctx, secretName)
                if err != nil {
                        return fmt.Errorf("failed to resolve header secret %q: %w", secretName, err)
                }
                merged[headerName] = actualValue
        }
        c.HeaderForward.resolvedHeaders = merged
        return nil
}

// mergeEnvVars is a helper method to merge environment variables into RunConfig
func (c *RunConfig) mergeEnvVars(envVars map[string]string) *RunConfig {
        // Initialize EnvVars if it's nil
        if c.EnvVars == nil {
                c.EnvVars = make(map[string]string)
        }

        // Add env vars to environment variables
        for key, value := range envVars {
                c.EnvVars[key] = value
        }

        return c
}

// WithEnvFilesFromDirectory processes environment files from a directory and adds them to environment variables
func (c *RunConfig) WithEnvFilesFromDirectory(dirPath string) (*RunConfig, error) {
        envVars, err := processEnvFilesDirectory(dirPath)
        if err != nil {
                return c, fmt.Errorf("failed to process env files from %s: %w", dirPath, err)
        }

        return c.mergeEnvVars(envVars), nil
}

// WithEnvFile processes a single environment file and adds it to environment variables
func (c *RunConfig) WithEnvFile(filePath string) (*RunConfig, error) {
        envVars, err := processEnvFile(filePath)
        if err != nil {
                return c, fmt.Errorf("failed to process env file %s: %w", filePath, err)
        }

        return c.mergeEnvVars(envVars), nil
}

// WithContainerName generates container name if not already set
// Returns the config and a boolean indicating if the name was sanitized
func (c *RunConfig) WithContainerName() (*RunConfig, bool) {
        var wasModified bool

        if c.ContainerName == "" {
                if c.Image != "" {
                        // For container-based servers
                        // Sanitize the name if provided to ensure it's safe for file paths
                        safeName := ""
                        if c.Name != "" {
                                safeName, wasModified = workloadtypes.SanitizeWorkloadName(c.Name)
                        }
                        containerName, baseName := container.GetOrGenerateContainerName(safeName, c.Image)
                        c.ContainerName = containerName
                        c.BaseName = baseName
                } else if c.RemoteURL != "" && c.Name != "" {
                        // For remote servers, sanitize the provided name to ensure it's safe for file paths
                        c.BaseName, wasModified = workloadtypes.SanitizeWorkloadName(c.Name)
                        c.ContainerName = c.Name
                }
        }
        return c, wasModified
}

// WithStandardLabels adds standard labels to the container
func (c *RunConfig) WithStandardLabels() *RunConfig {
        if c.ContainerLabels == nil {
                c.ContainerLabels = make(map[string]string)
        }
        // Use Name if ContainerName is not set
        containerName := c.ContainerName
        if containerName == "" {
                containerName = c.Name
        }

        transportLabel := c.Transport.String()
        // Use the Group field from the RunConfig
        labels.AddStandardLabels(c.ContainerLabels, containerName, c.BaseName, transportLabel, c.Port)
        return c
}

// GetBaseName returns the base name for the run configuration
func (c *RunConfig) GetBaseName() string {
        return c.BaseName
}

// SaveState saves the run configuration to the state store
func (c *RunConfig) SaveState(ctx context.Context) error {
        return state.SaveRunConfig(ctx, c)
}

// LoadState loads a run configuration from the state store
func LoadState(ctx context.Context, name string) (*RunConfig, error) {
        return state.LoadRunConfig(ctx, name, ReadJSON)
}

// ToolOverride represents a tool override.
// Both Name and Description can be overridden independently, but
// they can't be both empty.
type ToolOverride struct {
        // Name is the redefined name of the tool
        Name string `json:"name,omitempty"`
        // Description is the redefined description of the tool
        Description string `json:"description,omitempty"`
}

// HeaderForwardConfig defines configuration for injecting headers into requests to remote servers.
// Headers are added server-side, so clients don't need to configure them individually.
type HeaderForwardConfig struct {
        // AddPlaintextHeaders is a map of header names to literal values to inject into requests.
        // WARNING: These values are stored in plaintext in the configuration.
        // For sensitive values (API keys, tokens), use AddHeadersFromSecret instead.
        AddPlaintextHeaders map[string]string `json:"add_plaintext_headers,omitempty" yaml:"add_plaintext_headers,omitempty"`

        // AddHeadersFromSecret is a map of header names to secret names.
        // The key is the header name, the value is the secret name in ToolHive's secrets manager.
        // Resolved at runtime via WithSecrets() into resolvedHeaders.
        // The actual secret value is only held in memory, never persisted.
        AddHeadersFromSecret map[string]string `json:"add_headers_from_secret,omitempty" yaml:"add_headers_from_secret,omitempty"`

        // resolvedHeaders holds the merged set of headers (plaintext + resolved secrets)
        // for middleware consumption. Never serialized to disk.
        resolvedHeaders map[string]string
}

// ResolvedHeaders returns the merged set of headers for middleware use.
// After WithSecrets() has run, this includes both plaintext and secret-backed headers.
// If secrets have not been resolved yet, returns only AddPlaintextHeaders.
// Safe to call on a nil receiver.
func (h *HeaderForwardConfig) ResolvedHeaders() map[string]string {
        if h == nil {
                return nil
        }
        if h.resolvedHeaders != nil {
                return h.resolvedHeaders
        }
        return h.AddPlaintextHeaders
}

// HasHeaders returns true if any headers are configured (plaintext or secret-backed).
// Safe to call on a nil receiver.
func (h *HeaderForwardConfig) HasHeaders() bool {
        if h == nil {
                return false
        }
        return len(h.AddPlaintextHeaders) > 0 || len(h.AddHeadersFromSecret) > 0
}

// DefaultCallbackPort is the default port for the OAuth callback server
const DefaultCallbackPort = 8666

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "maps"
        "net/url"
        "os"
        "path/filepath"
        "slices"
        "strings"
        "time"

        "github.com/stacklok/toolhive-core/permissions"
        regtypes "github.com/stacklok/toolhive-core/registry/types"
        v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
        "github.com/stacklok/toolhive/pkg/audit"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/awssts"
        "github.com/stacklok/toolhive/pkg/auth/remote"
        "github.com/stacklok/toolhive/pkg/authserver"
        "github.com/stacklok/toolhive/pkg/authserver/server/registration"
        "github.com/stacklok/toolhive/pkg/authz"
        appconfig "github.com/stacklok/toolhive/pkg/config"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/container/templates"
        "github.com/stacklok/toolhive/pkg/ignore"
        "github.com/stacklok/toolhive/pkg/labels"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/oauthproto/tokenexchange"
        "github.com/stacklok/toolhive/pkg/recovery"
        "github.com/stacklok/toolhive/pkg/telemetry"
        "github.com/stacklok/toolhive/pkg/transport"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/usagemetrics"
        "github.com/stacklok/toolhive/pkg/webhook"
)

// BuildContext defines the context in which the RunConfigBuilder is being used
type BuildContext int

const (
        // BuildContextCLI indicates the builder is being used in CLI context with full validation
        BuildContextCLI BuildContext = iota
        // BuildContextOperator indicates the builder is being used in Kubernetes operator context
        BuildContextOperator
)

// runConfigBuilder provides a fluent interface for building RunConfig instances
type runConfigBuilder struct {
        config *RunConfig
        // Store transport string separately to avoid type confusion
        transportString string
        // Store ports separately for proper validation
        port       int
        targetPort int
        // registryProxyPort is the proxy port from the registry metadata (remote servers).
        // Used as a fallback when port is 0 (not set by CLI).
        registryProxyPort int
        // Store network mode to apply to permission profile after it's loaded
        networkMode string
        // isolateNetworkExplicit records whether the user explicitly requested network
        // isolation (e.g. --isolate-network was passed) rather than it being the default.
        // It controls whether an incompatible network mode fails fast or silently degrades.
        isolateNetworkExplicit bool
        // Build context determines which validation and features are enabled
        buildContext BuildContext
}

// RunConfigBuilderOption is a function that modifies the RunConfigBuilder
type RunConfigBuilderOption func(*runConfigBuilder) error

// WithRuntime sets the container runtime
func WithRuntime(deployer rt.Deployer) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if b.buildContext == BuildContextCLI {
                        b.config.Deployer = deployer
                }
                return nil
        }
}

// WithImage sets the Docker image
func WithImage(image string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Image = image
                return nil
        }
}

// WithMCPServerGeneration sets the MCPServer generation as the monotonic version stamp.
func WithMCPServerGeneration(gen int64) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.MCPServerGeneration = gen
                return nil
        }
}

// WithRuntimeConfig sets the runtime configuration (base images and packages)
func WithRuntimeConfig(runtimeConfig *templates.RuntimeConfig) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.RuntimeConfig = runtimeConfig
                return nil
        }
}

// WithRemoteURL sets the remote URL for the MCP server
func WithRemoteURL(remoteURL string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.RemoteURL = remoteURL
                return nil
        }
}

// WithRegistrySourceURLs records the registry URLs that served this server's metadata.
func WithRegistrySourceURLs(apiURL, registryURL string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.RegistryAPIURL = apiURL
                b.config.RegistryURL = registryURL
                return nil
        }
}

// ResolveRegistrySourceURLs returns the registry API URL and registry URL from
// config when the server was discovered via registry lookup (non-nil metadata).
// Both values are empty when metadata is nil (direct image reference or protocol scheme)
// or when appConfig is nil.
func ResolveRegistrySourceURLs(serverMetadata regtypes.ServerMetadata, appConfig *appconfig.Config) (apiURL, registryURL string) {
        if serverMetadata == nil || appConfig == nil {
                return "", ""
        }
        return appConfig.RegistryApiUrl, appConfig.RegistryUrl
}

// WithRegistryServerName records the registry entry name used to look up this server's metadata.
func WithRegistryServerName(name string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.RegistryServerName = name
                return nil
        }
}

// ResolveRegistryServerName returns the registry entry name from server metadata
// when the server was discovered via registry lookup (non-nil metadata).
// Returns empty string when metadata is nil (direct image reference or protocol scheme).
func ResolveRegistryServerName(serverMetadata regtypes.ServerMetadata) string {
        if serverMetadata == nil {
                return ""
        }
        return serverMetadata.GetName()
}

// WithRegistryProxyPort sets the proxy port from registry metadata.
// This is used as a fallback when the CLI --proxy-port flag is not set.
func WithRegistryProxyPort(port int) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.registryProxyPort = port
                return nil
        }
}

// WithRemoteAuth sets the remote authentication configuration
func WithRemoteAuth(config *remote.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if config == nil {
                        config = &remote.Config{
                                CallbackPort: remote.DefaultCallbackPort,
                        }
                }
                b.config.RemoteAuthConfig = config
                return nil
        }
}

// WithName sets the MCP server name
func WithName(name string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Name = name
                return nil
        }
}

// WithCmdArgs sets the command arguments
func WithCmdArgs(args []string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.CmdArgs = args
                return nil
        }
}

// WithHost sets the host (applies default if empty)
func WithHost(host string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if host == "" {
                        host = transport.LocalhostIPv4
                }
                b.config.Host = host
                return nil
        }
}

// WithTargetHost sets the target host (applies default if empty)
func WithTargetHost(targetHost string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if b.config.RemoteURL != "" {
                        remoteURL, err := url.Parse(b.config.RemoteURL)
                        if err == nil {
                                targetHost = remoteURL.Host
                        } else {
                                slog.Warn("Failed to parse remote URL", "error", err)
                                targetHost = transport.LocalhostIPv4
                        }
                } else if targetHost == "" {
                        targetHost = transport.LocalhostIPv4
                }
                b.config.TargetHost = targetHost
                return nil
        }
}

// WithPublish sets the published ports
func WithPublish(publish []string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Publish = publish
                return nil
        }
}

// WithDebug sets debug mode
func WithDebug(debug bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Debug = debug
                return nil
        }
}

// WithVolumes sets the volume mounts
func WithVolumes(volumes []string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Volumes = volumes
                return nil
        }
}

// WithSecrets sets the secrets list
func WithSecrets(secrets []string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Secrets = secrets
                return nil
        }
}

// WithAuthzConfigPath sets the authorization config path
func WithAuthzConfigPath(path string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.AuthzConfigPath = path
                return nil
        }
}

// WithAuthzConfig sets the authorization config data
func WithAuthzConfig(config *authz.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.AuthzConfig = config
                return nil
        }
}

// WithValidatingWebhooks sets the validating webhook configurations.
// These webhooks run after mutating webhooks and can accept or deny requests.
func WithValidatingWebhooks(webhooks []webhook.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.ValidatingWebhooks = webhooks
                return nil
        }
}

// WithMutatingWebhooks sets the mutating webhook configurations.
// These webhooks run before validating webhooks and can transform requests.
func WithMutatingWebhooks(webhooks []webhook.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.MutatingWebhooks = webhooks
                return nil
        }
}

// WithAuditConfigPath sets the audit config path
func WithAuditConfigPath(path string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.AuditConfigPath = path
                return nil
        }
}

// WithPermissionProfileNameOrPath sets the permission profile name or path.
// If called multiple times or mixed with WithPermissionProfile,
// the last call takes precedence.
func WithPermissionProfileNameOrPath(profile string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.PermissionProfileNameOrPath = profile
                b.config.PermissionProfile = nil // Clear any existing profile
                return nil
        }
}

// WithPermissionProfile sets the permission profile directly.
// If called multiple times or mixed with WithPermissionProfile,
// the last call takes precedence.
func WithPermissionProfile(profile *permissions.Profile) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.PermissionProfile = profile
                b.config.PermissionProfileNameOrPath = "" // Clear any existing name or path
                return nil
        }
}

// WithNetworkIsolation sets network isolation.
//
// Network isolation is only enforceable in bridge mode. When it is requested
// together with a non-bridge network mode (host/none/custom), the build
// reconciles the conflict: without the WithNetworkIsolationExplicit companion
// (i.e. the user did not actively pass --isolate-network), a host/custom
// conflict degrades to isolation=false with a warning rather than failing fast.
// See WithNetworkIsolationExplicit and issue #5775.
func WithNetworkIsolation(isolate bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.IsolateNetwork = isolate
                return nil
        }
}

// WithNetworkIsolationExplicit records whether network isolation was explicitly
// requested by the user (i.e. --isolate-network was actually passed) rather than
// defaulted. It exists to drive the fail-fast-vs-degrade decision when isolation
// conflicts with a non-bridge network mode: when true, a host/custom conflict
// fails the build so the user learns their request cannot be honored; when false
// (defaulted), the conflict degrades to isolation=false with a warning. See #5775.
func WithNetworkIsolationExplicit(explicit bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.isolateNetworkExplicit = explicit
                return nil
        }
}

// WithAllowDockerGateway sets whether to allow outbound connections to Docker gateway addresses
func WithAllowDockerGateway(allow bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.AllowDockerGateway = allow
                return nil
        }
}

// WithAllowedOrigins sets the HTTP Origin-header allowlist used for
// DNS-rebinding protection (MCP 2025-11-25 §"Security Warning").
// An empty slice defers the choice to middleware wiring, which derives a
// loopback-only default when the bind host is loopback and otherwise leaves
// the middleware disabled.
func WithAllowedOrigins(origins []string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.AllowedOrigins = origins
                return nil
        }
}

// WithTrustProxyHeaders sets whether to trust X-Forwarded-* headers from reverse proxies
func WithTrustProxyHeaders(trust bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.TrustProxyHeaders = trust
                return nil
        }
}

// WithStrictProtocolValidation sets whether the streamable HTTP proxy rejects
// requests carrying an unknown/unsupported MCP-Protocol-Version header with
// HTTP 400. Default false accepts any version string.
func WithStrictProtocolValidation(strict bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.StrictProtocolValidation = strict
                return nil
        }
}

// WithStateless declares the server is stateless (POST-only, no SSE).
func WithStateless(stateless bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Stateless = stateless
                return nil
        }
}

// WithEndpointPrefix sets the path prefix for SSE endpoint URLs
func WithEndpointPrefix(prefix string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.EndpointPrefix = prefix
                return nil
        }
}

// WithSessionTTL sets the inactivity timeout for proxy sessions.
// Zero is valid and means "use the transport default" (2h).
// Negative values return an error.
//
// The value is stored as a Go duration string on RunConfig so it survives a
// JSON/YAML round-trip in the runconfig API contract; a time.Duration field
// would serialize as nanoseconds.
func WithSessionTTL(ttl time.Duration) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if ttl < 0 {
                        return fmt.Errorf("session-ttl must be non-negative, got %s", ttl)
                }
                if ttl == 0 {
                        b.config.SessionTTL = ""
                        return nil
                }
                b.config.SessionTTL = ttl.String()
                return nil
        }
}

// WithMaxRequestBodySize sets the maximum inbound MCP proxy request body size in bytes.
// Zero uses the default limit of 8 MiB. Negative values return an error.
func WithMaxRequestBodySize(maxBytes int64) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if maxBytes < 0 {
                        return fmt.Errorf("max-request-body-size must be non-negative, got %d", maxBytes)
                }
                b.config.MaxRequestBodySize = maxBytes
                return nil
        }
}

// WithProxyReadTimeout sets http.Server.ReadTimeout on the proxy, bounding how
// long the server will spend reading a request (headers + body). Zero is valid
// and means "use the proxy default" (30s). Negative values return an error.
//
// The value is stored as a Go duration string on RunConfig so it survives a
// JSON/YAML round-trip; a time.Duration field would serialize as nanoseconds.
func WithProxyReadTimeout(d time.Duration) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if d < 0 {
                        return fmt.Errorf("proxy-read-timeout must be non-negative, got %s", d)
                }
                if d == 0 {
                        b.config.ProxyReadTimeout = ""
                        return nil
                }
                b.config.ProxyReadTimeout = d.String()
                return nil
        }
}

// WithNetworkMode sets the network mode for the container.
// The network mode will be applied to the permission profile after it is loaded.
func WithNetworkMode(networkMode string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.networkMode = networkMode
                return nil
        }
}

// WithK8sPodPatch sets the Kubernetes pod template patch
func WithK8sPodPatch(patch string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.K8sPodTemplatePatch = patch
                return nil
        }
}

// WithProxyMode sets the proxy mode
func WithProxyMode(mode types.ProxyMode) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.ProxyMode = mode
                return nil
        }
}

// WithGroup sets the group name for the workload
func WithGroup(groupName string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.Group = groupName
                return nil
        }
}

// WithLabels sets custom labels from command-line flags
func WithLabels(labelStrings []string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if len(labelStrings) == 0 {
                        return nil
                }

                // Initialize ContainerLabels if it's nil
                if b.config.ContainerLabels == nil {
                        b.config.ContainerLabels = make(map[string]string)
                }

                // Parse and add each label
                for _, labelString := range labelStrings {
                        key, value, err := labels.ParseLabel(labelString)
                        if err != nil {
                                slog.Warn("Skipping invalid label", "label", labelString, "error", err)
                                continue
                        }
                        b.config.ContainerLabels[key] = value
                }

                return nil
        }
}

// WithTransportAndPorts sets transport and port configuration
func WithTransportAndPorts(mcpTransport string, port, targetPort int) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.transportString = mcpTransport
                b.port = port
                b.targetPort = targetPort
                return nil
        }
}

// WithExistingPort sets the existing port for update operations
// This allows port reuse during workload updates by skipping validation for the same port
func WithExistingPort(port int) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.existingPort = port
                return nil
        }
}

// WithAuditEnabled configures audit settings
func WithAuditEnabled(enableAudit bool, auditConfigPath string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if enableAudit && auditConfigPath == "" {
                        b.config.AuditConfig = audit.DefaultConfig()
                }
                return nil
        }
}

// WithTokenValidatorConfig sets a copy of the canonical OIDC token validator configuration.
func WithTokenValidatorConfig(config *auth.TokenValidatorConfig) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if config == nil {
                        b.config.OIDCConfig = nil
                        return nil
                }

                configCopy := *config
                configCopy.Scopes = slices.Clone(config.Scopes)
                b.config.OIDCConfig = &configCopy
                return nil
        }
}

// WithOIDCConfig configures OIDC settings
func WithOIDCConfig(
        oidcIssuer string,
        oidcAudience string,
        oidcJwksURL string,
        oidcIntrospectionURL string,
        oidcClientID string,
        oidcClientSecret string,
        thvCABundle string,
        jwksAuthTokenFile string,
        resourceURL string,
        jwksAllowPrivateIP bool,
        insecureAllowHTTP bool,
        scopes []string,
) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if oidcIssuer != "" || oidcAudience != "" || oidcJwksURL != "" || oidcIntrospectionURL != "" ||
                        oidcClientID != "" || oidcClientSecret != "" {
                        b.config.OIDCConfig = &auth.TokenValidatorConfig{
                                Issuer:            oidcIssuer,
                                Audience:          oidcAudience,
                                JWKSURL:           oidcJwksURL,
                                IntrospectionURL:  oidcIntrospectionURL,
                                ClientID:          oidcClientID,
                                ClientSecret:      oidcClientSecret,
                                CACertPath:        thvCABundle,
                                AuthTokenFile:     jwksAuthTokenFile,
                                AllowPrivateIP:    jwksAllowPrivateIP,
                                InsecureAllowHTTP: insecureAllowHTTP,
                                Scopes:            scopes,
                        }
                }

                // Set JWKS-related configuration
                b.config.ThvCABundle = thvCABundle
                b.config.JWKSAuthTokenFile = jwksAuthTokenFile

                // Set ResourceURL if OIDCConfig exists or if resourceURL is not empty
                if b.config.OIDCConfig != nil {
                        b.config.OIDCConfig.ResourceURL = resourceURL
                } else if resourceURL != "" {
                        // Create OIDCConfig just for ResourceURL if it doesn't exist but resourceURL is provided
                        b.config.OIDCConfig = &auth.TokenValidatorConfig{
                                ResourceURL: resourceURL,
                        }
                }

                return nil
        }
}

// WithTokenExchangeConfig sets the token exchange configuration
func WithTokenExchangeConfig(config *tokenexchange.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.TokenExchangeConfig = config
                return nil
        }
}

// WithAWSStsConfig sets the AWS STS token exchange configuration
func WithAWSStsConfig(config *awssts.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.AWSStsConfig = config
                return nil
        }
}

// WithTelemetryConfigFromFlags configures telemetry settings (legacy - custom attributes handled via middleware)
func WithTelemetryConfigFromFlags(
        otelEndpoint string,
        otelEnablePrometheusMetricsPath bool,
        otelTracingEnabled bool,
        otelMetricsEnabled bool,
        otelServiceName string,
        otelSamplingRate float64,
        otelHeaders []string,
        otelInsecure bool,
        otelEnvironmentVariables []string,
        otelUseLegacyAttributes bool,
) RunConfigBuilderOption {
        config := telemetry.MaybeMakeConfig(
                otelEndpoint,
                otelEnablePrometheusMetricsPath,
                otelTracingEnabled,
                otelMetricsEnabled,
                otelServiceName,
                otelSamplingRate,
                otelHeaders,
                otelInsecure,
                otelEnvironmentVariables,
                otelUseLegacyAttributes,
        )
        return WithTelemetryConfig(config)
}

// WithMetricsOnTransportPort records whether /metrics should also be served on the
// main transport port. Nil means the caller expressed no preference and the current
// default applies at startup; see telemetry.DefaultMetricsOnTransportPort.
//
// Must be applied after the telemetry configuration; it is a no-op when telemetry is
// disabled, since there is no config to record it on.
func WithMetricsOnTransportPort(v *bool) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if v == nil || b.config.TelemetryConfig == nil {
                        return nil
                }
                b.config.TelemetryConfig.MetricsOnTransportPort = v
                return nil
        }
}

// WithTelemetryConfig sets the telemetry configuration
func WithTelemetryConfig(config *telemetry.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.TelemetryConfig = config
                return nil
        }
}

// WithRateLimitConfig sets the rate limiting configuration.
func WithRateLimitConfig(namespace string, config *v1beta1.RateLimitConfig) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.RateLimitConfig = config
                b.config.RateLimitNamespace = namespace
                return nil
        }
}

// WithToolsFilter sets the tools filter
func WithToolsFilter(toolsFilter []string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.ToolsFilter = toolsFilter
                return nil
        }
}

// WithToolsOverride sets the tool override map for the RunConfig
// This method is mutually exclusive with WithToolOverrideFile
func WithToolsOverride(toolOverride map[string]ToolOverride) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.ToolsOverride = toolOverride
                return nil
        }
}

// WithIgnoreConfig sets the ignore configuration
func WithIgnoreConfig(ignoreConfig *ignore.Config) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.IgnoreConfig = ignoreConfig
                return nil
        }
}

// WithMiddlewareFromFlags creates middleware configurations directly from flag values
func WithMiddlewareFromFlags(
        oidcConfig *auth.TokenValidatorConfig,
        tokenExchangeConfig *tokenexchange.Config,
        toolsFilter []string,
        toolsOverride map[string]ToolOverride,
        telemetryConfig *telemetry.Config,
        authzConfigPath string,
        enableAudit bool,
        auditConfigPath string,
        serverName string,
        transportType string,
        disableUsageMetrics bool,
) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                var middlewareConfigs []types.MiddlewareConfig

                // NOTE: order matters here. Specifically, these routines use append
                // to add new middleware configs. The proxy wraps the handler in
                // reverse slice order (see applyMiddlewares), so the first entry
                // added here is the OUTERMOST wrapper and runs first at HTTP request
                // time; the last entry runs closest to the handler.
                //
                // We should avoid doing this and a better pattern would be to let the
                // actual proxy determine the order of application of middlewares, since
                // the types of middleware are known at compile time.

                // Audit middleware (if enabled) goes first so it is the outermost
                // wrapper after the body-limit prepended in runner.Run: every request
                // that passes the size cap produces an audit event no matter which
                // middleware rejects it — authentication (401), webhook denials, and
                // authorization (403, outcome "denied") included. Identity and parsed
                // MCP data are read back from the inner auth/parser middlewares via
                // the holder carriers (auth.IdentityHolder, mcp.ParsedRequestHolder).
                middlewareConfigs = addAuditMiddleware(middlewareConfigs, enableAudit, auditConfigPath, serverName, transportType)

                // Add tool filter middlewares
                middlewareConfigs = addToolFilterMiddlewares(middlewareConfigs, toolsFilter, toolsOverride)

                // Add core middlewares (always present)
                embeddedAuthServerIssuer := func() string {
                        if b.config.EmbeddedAuthServerConfig != nil {
                                return b.config.EmbeddedAuthServerConfig.Issuer
                        }
                        return ""
                }()
                middlewareConfigs = addCoreMiddlewares(
                        middlewareConfigs,
                        oidcConfig,
                        tokenExchangeConfig,
                        embeddedAuthServerIssuer,
                        disableUsageMetrics,
                )

                // NOTE: Header forward middleware is NOT added here because secret-backed
                // headers are not yet resolved at builder time. It is added in Runner.Run()
                // after WithSecrets() resolves all secret references.

                // NOTE: AWS STS middleware is NOT added here because it is only configured
                // through the operator path via PopulateMiddlewareConfigs(), not via CLI flags.
                //
                // NOTE: addCoreMiddlewares also injects usage metrics before webhook insertion here,
                // which differs slightly from PopulateMiddlewareConfigs where usage metrics is added
                // after webhooks. Since mutating webhooks republish the parsed request
                // (mcp.RepublishParsedMCPRequest), the two paths now disagree when a webhook rewrites
                // the JSON-RPC method: this path counts the method as received, the operator path counts
                // the method as mutated. Only tool-call counts are affected, and only for a webhook that
                // patches "method" itself. The broader ordering TODO remains to unify these paths.

                // Add Mutating webhooks before Validating webhooks
                var err error
                middlewareConfigs, err = addMutatingWebhookMiddleware(middlewareConfigs, b.config)
                if err != nil {
                        return err
                }

                // Add Validating webhooks
                middlewareConfigs, err = addValidatingWebhookMiddleware(middlewareConfigs, b.config)
                if err != nil {
                        return err
                }

                // Add optional middlewares. Audit was added at the top of the chain
                // so authorization denials (403) still produce an audit event with
                // outcome "denied".
                middlewareConfigs = addTelemetryMiddleware(middlewareConfigs, telemetryConfig, serverName, transportType)
                var authzErr error
                middlewareConfigs, authzErr = addAuthzMiddleware(middlewareConfigs, authzConfigPath, b.config.EmbeddedAuthServerConfig)
                if authzErr != nil {
                        return authzErr
                }

                // Add recovery middleware (always present, added last so it is the
                // innermost wrapper, executing closest to the handler — matching
                // PopulateMiddlewareConfigs)
                middlewareConfigs = addRecoveryMiddleware(middlewareConfigs)

                // Set the populated middleware configs
                b.config.MiddlewareConfigs = middlewareConfigs
                return nil
        }
}

// WithAdditionalMiddlewareConfigs appends pre-built middleware configs to the
// RunConfig's AdditionalMiddlewareConfigs slice.
//
// Unlike the typed-field middleware (auth, token exchange, AWS STS, ...) that
// PopulateMiddlewareConfigs derives from RunConfig fields, these configs are
// supplied already-built by external-auth handlers reached via
// *[]RunConfigBuilderOption. PopulateMiddlewareConfigs splices them into the
// backend-egress group — after auth and before recovery — so they survive
// population and reach the serialized RunConfig the proxyrunner reads.
//
// The carrier is generic: upstream moves these configs verbatim and never
// interprets their parameters; the middleware type identity is supplied by the
// caller via the config's Type. Multiple calls and multiple arguments are
// additive; nil entries are skipped.
//
// NOTE: the injected configs are consumed only by PopulateMiddlewareConfigs (the
// operator build path). The CLI flag path (WithMiddlewareFromFlags) builds
// MiddlewareConfigs directly and never reads AdditionalMiddlewareConfigs, so
// combining this option with that path would silently drop the injected config.
//
// SECURITY: a config's Parameters are serialized verbatim into the RunConfig the
// operator writes to a ConfigMap — plaintext, not a Secret — readable by anyone
// with ConfigMap read access in the workload namespace. Injected parameters must
// therefore NEVER embed raw credentials (client secrets, refresh tokens,
// service-account keys, …). Reference such secrets by environment-variable name
// and resolve them at runtime instead, mirroring the typed token-exchange
// middleware, which leaves its serialized client_secret empty and reads
// TOOLHIVE_TOKEN_EXCHANGE_CLIENT_SECRET at startup. Because the seam carries
// parameters opaquely it cannot enforce this — it is a contract the injecting
// handler must uphold.
func WithAdditionalMiddlewareConfigs(configs ...*types.MiddlewareConfig) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                for _, c := range configs {
                        if c == nil {
                                continue
                        }
                        b.config.AdditionalMiddlewareConfigs = append(b.config.AdditionalMiddlewareConfigs, *c)
                }
                return nil
        }
}

// WithEnvVars sets environment variables from a map
func WithEnvVars(envVars map[string]string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if b.config.EnvVars == nil {
                        b.config.EnvVars = make(map[string]string)
                }
                for key, value := range envVars {
                        b.config.EnvVars[key] = value
                }
                return nil
        }
}

// addToolFilterMiddlewares adds tool filter middlewares if tools filter is provided
func addToolFilterMiddlewares(
        middlewareConfigs []types.MiddlewareConfig,
        toolsFilter []string,
        toolsOverride map[string]ToolOverride,
) []types.MiddlewareConfig {
        if len(toolsFilter) == 0 && len(toolsOverride) == 0 {
                return middlewareConfigs
        }

        overrides := make(map[string]mcp.ToolOverride)
        for actualName, tool := range toolsOverride {
                overrides[actualName] = mcp.ToolOverride{
                        Name:        tool.Name,
                        Description: tool.Description,
                }
        }

        toolFilterParams := mcp.ToolFilterMiddlewareParams{
                FilterTools:   toolsFilter,
                ToolsOverride: overrides,
        }

        // Add tool filter middleware
        if toolFilterConfig, err := types.NewMiddlewareConfig(mcp.ToolFilterMiddlewareType, toolFilterParams); err == nil {
                middlewareConfigs = append(middlewareConfigs, *toolFilterConfig)
        }

        // Add tool call filter middleware
        if toolCallFilterConfig, err := types.NewMiddlewareConfig(mcp.ToolCallFilterMiddlewareType, toolFilterParams); err == nil {
                middlewareConfigs = append(middlewareConfigs, *toolCallFilterConfig)
        }

        return middlewareConfigs
}

// addCoreMiddlewares adds core middlewares that are always present
func addCoreMiddlewares(
        middlewareConfigs []types.MiddlewareConfig,
        oidcConfig *auth.TokenValidatorConfig,
        tokenExchangeConfig *tokenexchange.Config,
        embeddedAuthServerIssuer string,
        disableUsageMetrics bool,
) []types.MiddlewareConfig {
        // Authentication middleware (always present)
        authParams := auth.MiddlewareParams{
                OIDCConfig:               oidcConfig,
                EmbeddedAuthServerIssuer: embeddedAuthServerIssuer,
        }
        if authConfig, err := types.NewMiddlewareConfig(auth.MiddlewareType, authParams); err == nil {
                middlewareConfigs = append(middlewareConfigs, *authConfig)
        }

        // Token Exchange middleware (conditionally present)
        if tokenExchangeConfig != nil {
                tokenExchangeParams := tokenexchange.MiddlewareParams{
                        TokenExchangeConfig: tokenExchangeConfig,
                }
                if tokenExchangeMwConfig, err := types.NewMiddlewareConfig(tokenexchange.MiddlewareType, tokenExchangeParams); err == nil {
                        middlewareConfigs = append(middlewareConfigs, *tokenExchangeMwConfig)
                } else {
                        slog.Warn("Failed to create token exchange middleware config", "error", err)
                }
        }

        // MCP Parser middleware (always present)
        mcpParserParams := mcp.ParserMiddlewareParams{}
        if mcpParserConfig, err := types.NewMiddlewareConfig(mcp.ParserMiddlewareType, mcpParserParams); err == nil {
                middlewareConfigs = append(middlewareConfigs, *mcpParserConfig)
        }

        // Usage metrics middleware (if enabled)
        if usagemetrics.ShouldEnableMetrics(disableUsageMetrics) {
                usageMetricsParams := usagemetrics.MiddlewareParams{}
                if usageMetricsConfig, err := types.NewMiddlewareConfig(usagemetrics.MiddlewareType, usageMetricsParams); err == nil {
                        middlewareConfigs = append(middlewareConfigs, *usageMetricsConfig)
                }
        }

        return middlewareConfigs
}

// addTelemetryMiddleware adds telemetry middleware if enabled
func addTelemetryMiddleware(
        middlewareConfigs []types.MiddlewareConfig,
        telemetryConfig *telemetry.Config,
        serverName, transportType string,
) []types.MiddlewareConfig {
        if telemetryConfig == nil {
                return middlewareConfigs
        }

        telemetryParams := telemetry.FactoryMiddlewareParams{
                Config:     telemetryConfig,
                ServerName: serverName,
                Transport:  transportType,
        }
        if telemetryMwConfig, err := types.NewMiddlewareConfig(telemetry.MiddlewareType, telemetryParams); err == nil {
                middlewareConfigs = append(middlewareConfigs, *telemetryMwConfig)
        }

        return middlewareConfigs
}

// addAuthzMiddleware adds authorization middleware if a config path is provided.
// When embeddedAuthServerCfg is non-nil, the loaded Cedar config is enriched
// with the upstream provider name so Cedar policies evaluate claims from the
// upstream IDP token rather than the ToolHive-issued JWT.
func addAuthzMiddleware(
        middlewareConfigs []types.MiddlewareConfig,
        authzConfigPath string,
        embeddedAuthServerCfg *authserver.RunConfig,
) ([]types.MiddlewareConfig, error) {
        if authzConfigPath == "" {
                return middlewareConfigs, nil
        }

        // Load the authz config eagerly so that a malformed file produces a clear
        // error now rather than a confusing failure at request time.
        authzConfigData, err := authz.LoadConfig(authzConfigPath)
        if err != nil {
                return nil, fmt.Errorf("failed to load authorization config from %q: %w", authzConfigPath, err)
        }

        enriched, err := injectUpstreamProviderIfNeeded(authzConfigData, embeddedAuthServerCfg)
        if err != nil {
                return nil, fmt.Errorf("failed to inject upstream provider into authz config: %w", err)
        }

        authzParams := authz.FactoryMiddlewareParams{
                ConfigPath: authzConfigPath, // Keep for backwards compatibility.
                ConfigData: enriched,
        }

        authzConfig, err := types.NewMiddlewareConfig(authz.MiddlewareType, authzParams)
        if err != nil {
                return nil, fmt.Errorf("failed to create authorization middleware config: %w", err)
        }
        return append(middlewareConfigs, *authzConfig), nil
}

// addAuditMiddleware adds audit middleware if enabled or config path is provided
func addAuditMiddleware(
        middlewareConfigs []types.MiddlewareConfig,
        enableAudit bool,
        auditConfigPath, serverName, transportType string,
) []types.MiddlewareConfig {
        if !enableAudit && auditConfigPath == "" {
                return middlewareConfigs
        }

        auditParams := audit.MiddlewareParams{
                ConfigPath:    auditConfigPath, // Keep for backwards compatibility
                Component:     serverName,      // Use server name as component
                TransportType: transportType,   // Pass the actual transport type
        }

        // Read audit config contents if path is provided
        if auditConfigPath != "" {
                if auditConfigData, err := audit.LoadFromFile(auditConfigPath); err == nil {
                        auditParams.ConfigData = auditConfigData
                }
                // Note: We keep ConfigPath set for backwards compatibility
        }

        if auditConfig, err := types.NewMiddlewareConfig(audit.MiddlewareType, auditParams); err == nil {
                middlewareConfigs = append(middlewareConfigs, *auditConfig)
        }

        return middlewareConfigs
}

// addRecoveryMiddleware adds recovery middleware (always present, added last).
// The proxy wraps the handler in reverse slice order, so the last entry is the
// INNERMOST wrapper: it catches panics from the handler itself, but not from
// middleware added earlier in the slice (which wrap it).
func addRecoveryMiddleware(middlewareConfigs []types.MiddlewareConfig) []types.MiddlewareConfig {
        recoveryConfig, err := types.NewMiddlewareConfig(recovery.MiddlewareType, nil)
        if err != nil {
                slog.Warn("failed to create recovery middleware", "error", err)
                return middlewareConfigs
        }
        middlewareConfigs = append(middlewareConfigs, *recoveryConfig)
        return middlewareConfigs
}

// NewOperatorRunConfigBuilder creates a new RunConfigBuilder configured for operator use
func NewOperatorRunConfigBuilder(
        ctx context.Context,
        imageMetadata *regtypes.ImageMetadata,
        envVars map[string]string,
        envVarValidator EnvVarValidator,
        runConfigOptions ...RunConfigBuilderOption,
) (*RunConfig, error) {
        return internalRunConfigBuilder(ctx,
                &runConfigBuilder{
                        config: &RunConfig{
                                ContainerLabels: make(map[string]string),
                                EnvVars:         make(map[string]string),
                        },
                        buildContext: BuildContextOperator,
                }, imageMetadata, envVars, envVarValidator, runConfigOptions...)
}

// NewRunConfigBuilder creates the final RunConfig instance with validation and processing
func NewRunConfigBuilder(
        ctx context.Context,
        imageMetadata *regtypes.ImageMetadata,
        envVars map[string]string,
        envVarValidator EnvVarValidator,
        runConfigOptions ...RunConfigBuilderOption,
) (*RunConfig, error) {
        return internalRunConfigBuilder(ctx,
                &runConfigBuilder{
                        config: &RunConfig{
                                ContainerLabels: make(map[string]string),
                                EnvVars:         make(map[string]string),
                        },
                        buildContext: BuildContextCLI,
                }, imageMetadata, envVars, envVarValidator, runConfigOptions...)
}

func internalRunConfigBuilder(
        ctx context.Context,
        b *runConfigBuilder,
        imageMetadata *regtypes.ImageMetadata,
        envVars map[string]string,
        envVarValidator EnvVarValidator,
        runConfigOptions ...RunConfigBuilderOption,
) (*RunConfig, error) {
        // Set the build context on the config to control validation behavior
        b.config.buildContext = b.buildContext

        // Apply all the options
        for _, option := range runConfigOptions {
                if err := option(b); err != nil {
                        return nil, fmt.Errorf("failed to apply run config option: %w", err)
                }
        }

        // Resolve the OTel service name from the workload name when not explicitly set.
        // This ensures the service name is always populated before persisting the config.
        telemetry.ResolveServiceName(b.config.TelemetryConfig, b.config.Name)

        // When the embedded auth server is configured and the proxy has no
        // explicit PRM scopes, derive them from the AS's ScopesSupported.
        // This ensures the PRM advertises the same scopes the AS supports
        // (including offline_access for refresh tokens).
        // If the AS also has no explicit scopes, use the auth server's
        // default scopes (registration.DefaultScopes).
        if b.config.EmbeddedAuthServerConfig != nil &&
                b.config.OIDCConfig != nil &&
                len(b.config.OIDCConfig.Scopes) == 0 {
                scopes := b.config.EmbeddedAuthServerConfig.ScopesSupported
                if len(scopes) == 0 {
                        scopes = registration.DefaultScopes
                }
                b.config.OIDCConfig.Scopes = scopes
        }

        if err := canonicalizeOIDCMiddlewareConfig(b.config); err != nil {
                return nil, fmt.Errorf("invalid OIDC middleware configuration: %w", err)
        }

        // When using the CLI validation strategy, this is where the prompting for
        // missing environment variables will happen.
        processedEnvVars := envVars
        if envVarValidator != nil {
                validatedEnvVars, err := envVarValidator.Validate(ctx, imageMetadata, b.config, envVars)
                if err != nil {
                        return nil, fmt.Errorf("failed to validate environment variables: %w", err)
                }
                processedEnvVars = validatedEnvVars
        }

        // Do some final validation which can only be done after everything else is set.
        // Apply image metadata overrides if needed.
        if err := b.validateConfig(imageMetadata); err != nil {
                return nil, fmt.Errorf("failed to validate run config: %w", err)
        }

        // Now set environment variables with the correct transport and ports resolved
        if _, err := b.config.WithEnvironmentVariables(processedEnvVars); err != nil {
                return nil, fmt.Errorf("failed to set environment variables: %w", err)
        }

        // Set schema version.
        b.config.SchemaVersion = CurrentSchemaVersion

        // Normalize proxyMode to the effective value before returning.
        b.config.NormalizeProxyMode()

        return b.config, nil
}

// validateConfig ensures the RunConfig is valid and sets up some of the final
// configuration details which can only be applied after all other flags are added.
// This function also handles setting missing values based on the image metadata (if present).
//
//nolint:gocyclo // This function needs to be refactored to reduce cyclomatic complexity.
func (b *runConfigBuilder) validateConfig(imageMetadata *regtypes.ImageMetadata) error {
        c := b.config
        var err error

        // The old logic claimed to override the name with the name from the registry
        // but didn't. Instead, it used the name passed in from the CLI.
        // See: https://github.com/stacklok/toolhive/blob/2873152b62bf61698cbcdd0aba1707a046151e67/cmd/thv/app/run.go#L425
        // The following code implements what I believe was the intended behavior:
        if c.Name == "" && imageMetadata != nil {
                c.Name = imageMetadata.Name
        }

        // Check to see if the mcpTransport is defined in the metadata.
        // Use this value if it was not set by the user.
        // Else, default to stdio.
        mcpTransport := b.transportString
        if mcpTransport == "" {
                if imageMetadata != nil && imageMetadata.Transport != "" {
                        slog.Debug("Using registry transport", "transport", imageMetadata.Transport)
                        mcpTransport = imageMetadata.Transport
                } else {
                        slog.Debug("Defaulting mcpTransport to stdio")
                        mcpTransport = types.TransportTypeStdio.String()
                }
        }
        // Set mcpTransport
        if _, err = c.WithTransport(mcpTransport); err != nil {
                return err
        }

        // Use registry ports if not overridden and if the mcpTransport is HTTP-based.
        proxyPort := b.port
        targetPort := b.targetPort
        if imageMetadata != nil {
                isHTTPServer := mcpTransport == types.TransportTypeSSE.String() ||
                        mcpTransport == types.TransportTypeStreamableHTTP.String()
                if isHTTPServer {
                        // Use registry proxy port if not set by CLI
                        if proxyPort == 0 && imageMetadata.ProxyPort > 0 {
                                slog.Debug("Using registry proxy port", "port", imageMetadata.ProxyPort)
                                proxyPort = imageMetadata.ProxyPort
                        }
                        // Use registry target port if not set by CLI
                        if targetPort == 0 && imageMetadata.TargetPort > 0 {
                                slog.Debug("Using registry target port", "port", imageMetadata.TargetPort)
                                targetPort = imageMetadata.TargetPort
                        }
                }
        }
        // Use registry proxy port from remote server metadata if not set by CLI
        if proxyPort == 0 && b.registryProxyPort > 0 {
                slog.Debug("Using remote server registry proxy port", "port", b.registryProxyPort)
                proxyPort = b.registryProxyPort
        }
        // Configure ports and target host
        if _, err = c.WithPorts(proxyPort, targetPort); err != nil {
                return err
        }

        // Load or default the permission profile
        // NOTE: This must be done before processing volume mounts
        permissionProfile, err := b.loadPermissionProfile(imageMetadata)
        if err != nil {
                return err
        }
        c.PermissionProfile = clonePermissionProfile(permissionProfile)

        // Apply network mode to permission profile if specified
        if b.networkMode != "" {
                // Ensure Network permissions struct exists
                if c.PermissionProfile.Network == nil {
                        c.PermissionProfile.Network = &permissions.NetworkPermissions{}
                }
                c.PermissionProfile.Network.Mode = b.networkMode
                slog.Debug("Setting network mode on permission profile", "network_mode", b.networkMode)
        }

        // Reconcile network isolation against the resolved network mode. Isolation is
        // only enforceable in bridge mode; for host/none/custom modes the isolation
        // sidecars have no route out and the workload bypasses them (issue #5775).
        // Key off the resolved profile mode so this also covers modes coming from a
        // permission-profile file rather than the --network flag.
        if err = b.reconcileNetworkIsolation(); err != nil {
                return err
        }

        // Process volume mounts
        if err = b.processVolumeMounts(); err != nil {
                return err
        }

        // Generate container name if not already set
        _, wasModified := c.WithContainerName()
        if wasModified && c.Name != "" {
                slog.Warn("The provided name contained invalid characters and was sanitized", "name", c.Name)
        }

        // Add standard labels
        c.WithStandardLabels()

        // Add authorization configuration if provided
        if c.AuthzConfigPath != "" {
                authzConfig, err := authz.LoadConfig(c.AuthzConfigPath)
                if err != nil {
                        return fmt.Errorf("failed to load authorization configuration: %w", err)
                }
                c.WithAuthz(authzConfig)
        }

        // Add audit configuration if provided
        if c.AuditConfigPath != "" {
                auditConfig, err := audit.LoadFromFile(c.AuditConfigPath)
                if err != nil {
                        return fmt.Errorf("failed to load audit configuration: %w", err)
                }
                c.WithAudit(auditConfig)
        }
        // Note: AuditConfig is already set from --enable-audit flag if provided

        if imageMetadata != nil && len(imageMetadata.Args) > 0 {
                if len(c.CmdArgs) == 0 {
                        // No user args provided, use registry defaults
                        slog.Debug("Using registry default args", "args", imageMetadata.Args)
                        c.CmdArgs = append(c.CmdArgs, imageMetadata.Args...)
                }
        }

        for toolName, tool := range c.ToolsOverride {
                if tool.Name == "" && tool.Description == "" {
                        return fmt.Errorf("tool override for %s must have either Name or Description set", toolName)
                }
        }

        if c.ToolsOverride != nil && imageMetadata != nil && imageMetadata.Tools != nil {
                slog.Debug("Using tools override", "tools", c.ToolsOverride)
                for toolName := range c.ToolsOverride {
                        if !slices.Contains(imageMetadata.Tools, toolName) {
                                return fmt.Errorf("tool %s not found in registry", toolName)
                        }
                }
        }

        if c.ToolsFilter != nil && imageMetadata != nil && imageMetadata.Tools != nil {
                slog.Debug("Using tools filter", "filter", c.ToolsFilter)
                for _, tool := range c.ToolsFilter {
                        name := tool

                        if c.ToolsOverride != nil {
                                for actualName, toolOverride := range c.ToolsOverride {
                                        if toolOverride.Name == tool {
                                                name = actualName
                                                break
                                        }
                                }
                        }

                        if !slices.Contains(imageMetadata.Tools, name) {
                                return fmt.Errorf("tool %s not found in registry", name)
                        }
                }
        }

        return nil
}

// reconcileNetworkIsolation drops network isolation when the resolved network
// mode cannot enforce it. In bridge mode isolation is left untouched. For "none"
// (already maximally confined) isolation is silently dropped. For host/other
// non-bridge modes isolation is dropped with a warning, unless the user
// explicitly requested isolation, in which case the build fails fast.
func (b *runConfigBuilder) reconcileNetworkIsolation() error {
        c := b.config
        if !c.IsolateNetwork {
                return nil
        }

        // Resolve the effective network mode from the permission profile.
        mode := ""
        if c.PermissionProfile != nil && c.PermissionProfile.Network != nil {
                mode = c.PermissionProfile.Network.Mode
        }
        if networking.IsBridgeMode(mode) {
                return nil
        }

        if mode == "none" {
                // "none" is already maximally confined; isolation is merely redundant.
                c.IsolateNetwork = false
                slog.Debug(networking.NetworkIsolationNoneRedundantMsg, "network_mode", mode)
                return nil
        }

        // host (or other non-bridge) is less restrictive than isolation, so dropping
        // it reduces confinement. Fail fast when it was explicitly requested.
        //
        // The mode may come from --permission-profile rather than --network, so the
        // message names the resolved mode instead of assuming a flag the user may
        // never have passed. See #5794 review discussion.
        if b.isolateNetworkExplicit {
                return fmt.Errorf(
                        "network isolation cannot be enforced with the resolved network mode %q: "+
                                "use a bridge network mode to keep enforced isolation, or pass --isolate-network=false to keep %q networking",
                        mode, mode)
        }

        c.IsolateNetwork = false
        slog.Warn(networking.NetworkIsolationHostDroppedMsg, "network_mode", mode)
        return nil
}

func (b *runConfigBuilder) loadPermissionProfile(imageMetadata *regtypes.ImageMetadata) (*permissions.Profile, error) {
        // The permission profile object takes precedence over the name or path.
        if b.config.PermissionProfile != nil {
                return b.config.PermissionProfile, nil
        }

        // Try to load the permission profile by name or path.
        if b.config.PermissionProfileNameOrPath != "" {
                switch b.config.PermissionProfileNameOrPath {
                case permissions.ProfileNone, "stdio":
                        return permissions.BuiltinNoneProfile(), nil
                case permissions.ProfileNetwork:
                        return permissions.BuiltinNetworkProfile(), nil
                default:
                        // Try to load from file
                        return permissions.FromFile(b.config.PermissionProfileNameOrPath)
                }
        }

        // If a profile was not set by name or path, check the image metadata.
        if imageMetadata != nil && imageMetadata.Permissions != nil {

                slog.Debug("Using registry permission profile", "permissions", imageMetadata.Permissions)
                return imageMetadata.Permissions, nil
        }

        // If no metadata is available, use the network permission profile as default.
        slog.Debug("Using default permission profile", "profile", permissions.ProfileNetwork)
        return permissions.BuiltinNetworkProfile(), nil
}

// clonePermissionProfile returns a deep copy that the builder can safely mutate.
func clonePermissionProfile(profile *permissions.Profile) *permissions.Profile {
        if profile == nil {
                return nil
        }

        cloned := *profile
        cloned.Read = slices.Clone(profile.Read)
        cloned.Write = slices.Clone(profile.Write)

        if profile.Network != nil {
                network := *profile.Network
                cloned.Network = &network

                if profile.Network.Outbound != nil {
                        outbound := *profile.Network.Outbound
                        outbound.AllowHost = slices.Clone(profile.Network.Outbound.AllowHost)
                        outbound.AllowPort = slices.Clone(profile.Network.Outbound.AllowPort)
                        network.Outbound = &outbound
                }

                if profile.Network.Inbound != nil {
                        inbound := *profile.Network.Inbound
                        inbound.AllowHost = slices.Clone(profile.Network.Inbound.AllowHost)
                        network.Inbound = &inbound
                }
        }

        return &cloned
}

// processVolumeMounts processes volume mounts and adds them to the permission profile
func (b *runConfigBuilder) processVolumeMounts() error {

        // Skip if no volumes to process
        if len(b.config.Volumes) == 0 {
                return nil
        }

        // Ensure permission profile is loaded
        if b.config.PermissionProfile == nil {
                return fmt.Errorf("permission profile is required when using volume mounts")
        }

        // Create a map of existing mount targets for quick lookup
        existingMounts := make(map[string]string)

        // Add existing read mounts to the map
        for _, m := range b.config.PermissionProfile.Read {
                source, target, _ := m.Parse()
                existingMounts[target] = source
        }

        // Add existing write mounts to the map
        for _, m := range b.config.PermissionProfile.Write {
                source, target, _ := m.Parse()
                existingMounts[target] = source
        }

        // Process each volume mount
        for _, volume := range b.config.Volumes {
                // Parse read-only flag
                readOnly := strings.HasSuffix(volume, ":ro")
                volumeSpec := volume
                if readOnly {
                        volumeSpec = strings.TrimSuffix(volume, ":ro")
                }

                // Create and parse mount declaration
                mount := permissions.MountDeclaration(volumeSpec)
                source, target, err := mount.Parse()
                if err != nil {
                        return fmt.Errorf("invalid volume format: %s (%w)", volume, err)
                }

                // Validate source path exists on the host (CLI context only)
                if b.buildContext == BuildContextCLI {
                        absSource := source
                        if !filepath.IsAbs(source) {
                                cwd, err := os.Getwd()
                                if err != nil {
                                        return fmt.Errorf("failed to resolve relative volume path %s: %w", source, err)
                                }
                                absSource = filepath.Join(cwd, source)
                        }
                        if _, err := os.Stat(absSource); err != nil {
                                if errors.Is(err, os.ErrNotExist) {
                                        return fmt.Errorf("volume source path does not exist: %s", absSource)
                                }
                                return fmt.Errorf("failed to access volume source path %s: %w", absSource, err)
                        }
                }

                // Check for duplicate mount target
                if existingSource, isDuplicate := existingMounts[target]; isDuplicate {
                        slog.Warn("Skipping duplicate mount target", "target", target, "existing_source", existingSource)
                        continue
                }

                // Add the mount to the appropriate permission list
                if readOnly {
                        b.config.PermissionProfile.Read = append(b.config.PermissionProfile.Read, mount)
                } else {
                        b.config.PermissionProfile.Write = append(b.config.PermissionProfile.Write, mount)
                }

                // Add to the map of existing mounts
                existingMounts[target] = source

                slog.Debug("Adding volume mount", "source", source, "target", target,
                        "mode", map[bool]string{true: "read-only", false: "read-write"}[readOnly])
        }

        return nil
}

// BuildForOperator creates a RunConfig for operator use, using the same validation as CLI
func (b *runConfigBuilder) BuildForOperator() (*RunConfig, error) {
        if b.buildContext != BuildContextOperator {
                return nil, fmt.Errorf("BuildForOperator can only be used with BuildContextOperator")
        }

        // Set build context on the config to control validation behavior
        b.config.buildContext = BuildContextOperator

        // Use the same validation logic as CLI, but without image metadata (pass nil)
        if err := b.validateConfig(nil); err != nil {
                return nil, fmt.Errorf("failed to validate run config: %w", err)
        }

        // Set schema version
        b.config.SchemaVersion = CurrentSchemaVersion

        return b.config, nil
}

// WithEnvVars sets environment variables from a map
func (b *runConfigBuilder) WithEnvVars(envVars map[string]string) *runConfigBuilder {
        if b.config.EnvVars == nil {
                b.config.EnvVars = make(map[string]string)
        }
        for key, value := range envVars {
                b.config.EnvVars[key] = value
        }
        return b
}

// WithEnvFile adds environment variables from a single file
func WithEnvFile(filePath string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if _, err := b.config.WithEnvFile(filePath); err != nil {
                        return err
                }
                return nil
        }
}

// WithEnvFilesFromDirectory adds environment variables from all files in a directory
func WithEnvFilesFromDirectory(dirPath string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if _, err := b.config.WithEnvFilesFromDirectory(dirPath); err != nil {
                        return err
                }
                return nil
        }
}

// WithHeaderForward adds plaintext header forward entries for remote MCP servers.
// The headers parameter contains literal header values (non-sensitive, stored as-is in RunConfig).
// Multiple calls are additive; later values for the same header name overwrite earlier ones.
func WithHeaderForward(headers map[string]string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if len(headers) == 0 {
                        return nil
                }
                hf := b.ensureHeaderForward()
                if hf.AddPlaintextHeaders == nil {
                        hf.AddPlaintextHeaders = make(map[string]string, len(headers))
                }
                maps.Copy(hf.AddPlaintextHeaders, headers)
                return nil
        }
}

// WithHeaderForwardSecrets adds secret-backed header forward entries for remote MCP servers.
// The headers parameter maps header names to secret names in the secrets manager.
// Secret values are resolved at runtime via WithSecrets() and never persisted to disk.
// Multiple calls are additive; later values for the same header name overwrite earlier ones.
func WithHeaderForwardSecrets(headers map[string]string) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                if len(headers) == 0 {
                        return nil
                }
                hf := b.ensureHeaderForward()
                if hf.AddHeadersFromSecret == nil {
                        hf.AddHeadersFromSecret = make(map[string]string, len(headers))
                }
                maps.Copy(hf.AddHeadersFromSecret, headers)
                return nil
        }
}

func (b *runConfigBuilder) ensureHeaderForward() *HeaderForwardConfig {
        if b.config.HeaderForward == nil {
                b.config.HeaderForward = &HeaderForwardConfig{}
        }
        return b.config.HeaderForward
}

// WithEmbeddedAuthServerConfig sets the embedded auth server configuration.
// The config is a RunConfig with file paths and env var names for secrets.
func WithEmbeddedAuthServerConfig(config *authserver.RunConfig) RunConfigBuilderOption {
        return func(b *runConfigBuilder) error {
                b.config.EmbeddedAuthServerConfig = config
                return nil
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "context"
        "fmt"
        "net/http"
        "time"

        "github.com/stacklok/toolhive/pkg/diagnostics"
        "github.com/stacklok/toolhive/pkg/telemetry"
        "github.com/stacklok/toolhive/pkg/transport"
)

// diagnosticsStopTimeout bounds the graceful shutdown of the diagnostics
// listener on Run's error path, where no caller-supplied context is available.
const diagnosticsStopTimeout = 10 * time.Second

// startDiagnosticsServer starts the diagnostics listener when the telemetry
// middleware enabled the Prometheus metrics path, and is a no-op otherwise.
//
// The metrics endpoint is deliberately kept off the application listener so it
// can be governed by port: NetworkPolicy cannot filter on HTTP path, so a shared
// port makes "allow MCP, deny scraping" unexpressible. This does not make the
// endpoint authenticated — the diagnostics listener carries no middleware. See
// pkg/diagnostics for the full rationale and its limits.
func (r *Runner) startDiagnosticsServer() error {
        if r.prometheusHandler == nil {
                return nil
        }

        server, err := diagnostics.New(r.bindHost(), diagnosticsPort(r.Config.TelemetryConfig), r.prometheusHandler)
        if err != nil {
                return fmt.Errorf("failed to create diagnostics server: %w", err)
        }
        if err := server.Start(); err != nil {
                return fmt.Errorf("failed to start diagnostics server: %w", err)
        }

        r.diagnosticsServer = server
        return nil
}

// bindHost returns the host the runner's side listeners (diagnostics, auth
// server TLS) bind to.
//
// They bind to the same host as the proxy so a deployment that reaches the
// workload can reach them, but on separate ports that deployments are not
// expected to route publicly. Note this means the operator's 0.0.0.0 default
// applies here too: the endpoints stay reachable from other pods, so
// restricting them is a NetworkPolicy job, which is what the separate ports
// make possible. Mirror the builder's host default so a config assembled
// without WithHost does not bind to every interface.
func (r *Runner) bindHost() string {
        if r.Config.Host == "" {
                return transport.LocalhostIPv4
        }
        return r.Config.Host
}

// diagnosticsPort resolves the port the diagnostics listener should request.
// An unset port falls back to diagnostics.DefaultPort so a scraper has a
// predictable target rather than an arbitrary one.
func diagnosticsPort(cfg *telemetry.Config) int {
        if cfg == nil {
                return diagnostics.DefaultPort
        }
        return diagnostics.ResolvePort(cfg.PrometheusPort)
}

// mountPrometheusHandlerOnTransportPort reports whether Run should hand the
// Prometheus handler to the transport, which is what makes the proxies also serve
// /metrics on the transport port during the migration window (see
// telemetry.DefaultMetricsOnTransportPort).
func mountPrometheusHandlerOnTransportPort(handler http.Handler, cfg *telemetry.Config) bool {
        return handler != nil && cfg.ServeMetricsOnTransportPort()
}

// stopDiagnosticsServer shuts the diagnostics listener down. It is safe to call
// when no diagnostics server was started.
func (r *Runner) stopDiagnosticsServer(ctx context.Context) error {
        if r.diagnosticsServer == nil {
                return nil
        }

        server := r.diagnosticsServer
        r.diagnosticsServer = nil
        if err := server.Stop(ctx); err != nil {
                return fmt.Errorf("failed to stop diagnostics server: %w", err)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "bufio"
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "os"
        "strings"

        "golang.org/x/term"

        "github.com/stacklok/toolhive-core/registry/types"
        "github.com/stacklok/toolhive/pkg/config"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// EnvVarValidator defines the interface for checking that the expected
// environment variables and secrets have been supplied when creating a
// workload. This is implemented as a strategy pattern since the handling
// is different for the CLI vs the API and k8s.
type EnvVarValidator interface {
        // Validate checks that all required environment variables and secrets are provided
        // and returns the processed environment variables to be set.
        Validate(
                ctx context.Context,
                metadata *registry.ImageMetadata,
                runConfig *RunConfig,
                suppliedEnvVars map[string]string,
        ) (map[string]string, error)
}

// DetachedEnvVarValidator implements the EnvVarValidator interface for
// scenarios where the user cannot be prompted for input. Any missing,
// mandatory variables will result in an error being returned.
type DetachedEnvVarValidator struct{}

// Validate checks that all required environment variables and secrets are provided
// and returns the processed environment variables to be set.
func (*DetachedEnvVarValidator) Validate(
        _ context.Context,
        metadata *registry.ImageMetadata,
        runConfig *RunConfig,
        suppliedEnvVars map[string]string,
) (map[string]string, error) {
        // Check variables in metadata if we are processing an image from our registry.
        if metadata != nil {
                secretsList := runConfig.Secrets
                registryEnvVars := metadata.EnvVars
                for _, envVar := range registryEnvVars {
                        if isEnvVarProvided(envVar.Name, suppliedEnvVars, secretsList) {
                                continue
                        } else if envVar.Required {
                                return nil, missingRequiredEnvVarError(envVar, nil)
                        } else if envVar.Default != "" {
                                addAsEnvironmentVariable(envVar, envVar.Default, &suppliedEnvVars)
                        }
                }
        }

        return suppliedEnvVars, nil
}

// CLIEnvVarValidator implements the EnvVarValidator interface for
// CLI usage. If any missing, mandatory variables are found, this code will
// prompt the user to supply them through stdin.
type CLIEnvVarValidator struct {
        configProvider config.Provider
}

// NewCLIEnvVarValidator creates a new CLI environment variable validator with the given config provider.
func NewCLIEnvVarValidator(configProvider config.Provider) *CLIEnvVarValidator {
        return &CLIEnvVarValidator{
                configProvider: configProvider,
        }
}

// Validate checks that all required environment variables and secrets are provided
// and returns the processed environment variables to be set.
func (v *CLIEnvVarValidator) Validate(
        ctx context.Context,
        metadata *registry.ImageMetadata,
        runConfig *RunConfig,
        suppliedEnvVars map[string]string,
) (map[string]string, error) {
        envVars := make(map[string]string)

        // Copy the supplied environment variables
        for k, v := range suppliedEnvVars {
                envVars[k] = v
        }

        // If we are processing an image from our registry, we need to check the
        // variables defined in the metadata.
        if metadata != nil {
                secretsConfig := runConfig.Secrets
                // Create new slice for extra secrets
                secretsList := make([]string, 0, len(secretsConfig))

                // Copy existing secrets
                secretsList = append(secretsList, secretsConfig...)
                registryEnvVars := metadata.EnvVars

                // Initialize secrets manager if needed
                secretsManager := v.initializeSecretsManagerIfNeeded(registryEnvVars)

                // One reader for every prompt in this call: a fresh bufio.Reader per prompt
                // would throw away any line it had already buffered from os.Stdin, which is
                // how a piped answer for the next variable would be lost.
                stdinReader := bufio.NewReader(os.Stdin)

                // Process each environment variable from the registry
                for _, envVar := range registryEnvVars {
                        if isEnvVarProvided(envVar.Name, envVars, secretsList) {
                                continue
                        }

                        if envVar.Required {

                                if envVar.Secret {
                                        // Check if secrets manager is available before attempting to retrieve secret.
                                        // Falls back to prompt if unavailable or secret not found.
                                        if secretsManager != nil {
                                                value, err := secretsManager.GetSecret(ctx, envVar.Name)
                                                if err != nil {
                                                        slog.Warn("unable to find secret in the secrets manager", "name", envVar.Name, "error", err)
                                                } else {
                                                        addNewVariable(ctx, envVar, value, secretsManager, &envVars, &secretsList)
                                                        continue
                                                }
                                        } else {
                                                slog.Warn("secrets manager not configured (setup incomplete or missing provider) - " +
                                                        "falling back to prompt")
                                        }

                                        // If secrets manager unavailable or secret not found, fall through to prompt
                                }

                                value, err := promptForEnvironmentVariable(envVar, stdinReader)
                                if err != nil {
                                        // The variable is required: without a value the workload starts
                                        // misconfigured, so fail instead of warning and skipping it.
                                        return nil, missingRequiredEnvVarError(envVar, err)
                                }
                                if value == "" {
                                        return nil, missingRequiredEnvVarError(envVar, nil)
                                }
                                addNewVariable(ctx, envVar, value, secretsManager, &envVars, &secretsList)
                        } else if envVar.Default != "" {
                                addNewVariable(ctx, envVar, envVar.Default, secretsManager, &envVars, &secretsList)
                        }
                }

                runConfig.Secrets = secretsList
        }

        return envVars, nil
}

// promptForEnvironmentVariable prompts the user for an environment variable value
// and reads it as a whole line. Required registry variables routinely take values
// containing spaces (a path with a space, a space-separated OAuth scope list),
// which fmt.Scanln cannot read: it stops after the first whitespace-delimited
// token and reports "expected newline" for the rest.
func promptForEnvironmentVariable(envVar *registry.EnvVar, stdin *bufio.Reader) (string, error) {
        var byteValue []byte
        var err error
        if envVar.Secret {
                fmt.Printf("Required secret environment variable: %s (%s)", envVar.Name, envVar.Description)
                fmt.Printf("Enter value for %s (input will be hidden): ", envVar.Name)
                byteValue, err = term.ReadPassword(int(os.Stdin.Fd())) //nolint:gosec // G115: stdin fd is always small
                fmt.Println()                                          // Move to the next line after hidden input
                if err != nil {
                        return "", fmt.Errorf("failed to read hidden input: %w", err)
                }
        } else {
                fmt.Printf("Required environment variable: %s (%s)", envVar.Name, envVar.Description)
                fmt.Printf("Enter value for %s: ", envVar.Name)
                // Read the full line, matching the other interactive prompts in this repo.
                // A value that arrives without a trailing newline still counts: ReadString
                // returns the data alongside io.EOF in that case.
                line, readErr := stdin.ReadString('\n')
                // EOF is the normal end of a line that arrived without a trailing
                // newline. Any other error, or EOF before a byte arrived, means the
                // value is not complete: reject it even when a partial line came back.
                if readErr != nil && (!errors.Is(readErr, io.EOF) || line == "") {
                        return "", fmt.Errorf("failed to read input: %w", readErr)
                }
                byteValue = []byte(line)
        }

        return strings.TrimSpace(string(byteValue)), nil
}

// missingRequiredEnvVarError explains how to supply a required variable without
// a prompt, so a non-interactive run does not fail with only io.EOF.
func missingRequiredEnvVarError(envVar *registry.EnvVar, cause error) error {
        message := fmt.Sprintf(
                "missing required environment variable %s (pass it with --env %s=<value> or --secret <name>,target=%s)",
                envVar.Name, envVar.Name, envVar.Name,
        )
        if cause != nil {
                return fmt.Errorf("%s: %w", message, cause)
        }
        return errors.New(message)
}

// addNewVariable adds an environment variable or secret to the appropriate list
func addNewVariable(
        ctx context.Context,
        envVar *registry.EnvVar,
        value string,
        secretsManager secrets.Provider,
        envVars *map[string]string,
        secretsList *[]string,
) {
        if envVar.Secret && secretsManager != nil {
                addAsSecret(ctx, envVar, value, secretsManager, secretsList, envVars)
        } else {
                addAsEnvironmentVariable(envVar, value, envVars)
        }
}

// addAsSecret stores the value as a secret and adds a secret reference
func addAsSecret(
        ctx context.Context,
        envVar *registry.EnvVar,
        value string,
        secretsManager secrets.Provider,
        secretsList *[]string,
        envVars *map[string]string,
) {
        var secretName string
        if envVar.Required {
                secretName = fmt.Sprintf("registry-user-%s", strings.ToLower(envVar.Name))
        } else {
                secretName = fmt.Sprintf("registry-default-%s", strings.ToLower(envVar.Name))
        }

        if err := secretsManager.SetSecret(ctx, secretName, value); err != nil {
                slog.Warn("failed to store secret", "secret_name", secretName, "error", err)
                slog.Warn("falling back to environment variable", "name", envVar.Name)
                (*envVars)[envVar.Name] = value
                slog.Debug("added environment variable (secret fallback)", "name", envVar.Name)
        } else {
                // Create secret reference for RunConfig
                secretEntry := fmt.Sprintf("%s,target=%s", secretName, envVar.Name)
                *secretsList = append(*secretsList, secretEntry)
                if envVar.Required {
                        slog.Debug("created secret", "name", envVar.Name, "secret_name", secretName)
                } else {
                        slog.Debug("created secret with default value", "name", envVar.Name, "secret_name", secretName)
                }
        }
}

// initializeSecretsManagerIfNeeded initializes the secrets manager if there are secret environment variables
func (v *CLIEnvVarValidator) initializeSecretsManagerIfNeeded(registryEnvVars []*registry.EnvVar) secrets.Provider {
        // Check if we have any secret environment variables
        hasSecrets := false
        for _, envVar := range registryEnvVars {
                if envVar.Secret {
                        hasSecrets = true
                        break
                }
        }

        if !hasSecrets {
                return nil
        }

        secretsManager, err := v.getSecretsManager()
        if err != nil {
                slog.Warn("failed to initialize secrets manager", "error", err)
                slog.Warn("secret environment variables will be stored as regular environment variables")
                return nil
        }

        return secretsManager
}

// Duplicated from cmd/thv/app/app.go
// It may be possible to de-duplicate this in future.
func (v *CLIEnvVarValidator) getSecretsManager() (secrets.Provider, error) {
        cfg := v.configProvider.GetConfig()

        // Check if secrets setup has been completed
        if !cfg.Secrets.SetupCompleted {
                return nil, secrets.ErrSecretsNotSetup
        }

        providerType, err := cfg.Secrets.GetProviderType()
        if err != nil {
                return nil, fmt.Errorf("failed to get secrets provider type: %w", err)
        }

        manager, err := secrets.CreateProvider(providerType, secrets.WithScope(secrets.ScopeWorkloads))
        if err != nil {
                return nil, fmt.Errorf("failed to create secrets manager: %w", err)
        }

        return manager, nil
}

// Shared Logic follows

// isEnvVarProvided checks if an environment variable is already provided
func isEnvVarProvided(name string, envVars map[string]string, secretsConfig []string) bool {
        // Check if the environment variable is already provided in the command line
        if _, exists := envVars[name]; exists {
                return true
        }

        // Check if the environment variable is provided as a secret
        return findEnvironmentVariableFromSecrets(secretsConfig, name)
}

func findEnvironmentVariableFromSecrets(secs []string, envVarName string) bool {
        for _, secret := range secs {
                if isSecretReferenceEnvVar(secret, envVarName) {
                        return true
                }
        }

        return false
}

func isSecretReferenceEnvVar(secret, envVarName string) bool {
        parts := strings.Split(secret, ",")
        if len(parts) != 2 {
                return false
        }

        targetSplit := strings.Split(parts[1], "=")
        if len(targetSplit) != 2 {
                return false
        }

        if targetSplit[1] == envVarName {
                return true
        }

        return false
}

// addAsEnvironmentVariable adds the value as a regular environment variable
func addAsEnvironmentVariable(
        envVar *registry.EnvVar,
        value string,
        envVars *map[string]string,
) {
        (*envVars)[envVar.Name] = value

        if envVar.Secret {
                if envVar.Required {
                        slog.Debug("added secret as environment variable (no secrets manager)", "name", envVar.Name)
                } else {
                        slog.Debug("added default secret as environment variable (no secrets manager)", "name", envVar.Name)
                }
        } else {
                if envVar.Required {
                        slog.Debug("added environment variable", "name", envVar.Name)
                } else {
                        slog.Debug("using default value", "name", envVar.Name, "default_value", value)
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "strings"

        "github.com/stacklok/toolhive/pkg/environment"
)

// processEnvFilesDirectory detects and processes environment files from a directory
// Returns a map of environment variables to be merged with RunConfig.EnvVars
func processEnvFilesDirectory(dirPath string) (map[string]string, error) {
        // Check if directory exists
        entries, err := os.ReadDir(dirPath)
        if err != nil {
                if os.IsNotExist(err) {
                        slog.Debug("Env files directory does not exist", "dir", dirPath)
                        return make(map[string]string), nil // Return empty map, not an error
                }
                return nil, fmt.Errorf("failed to read env files directory %s: %w", dirPath, err)
        }

        slog.Debug("Env files directory detected, processing environment files", "dir", dirPath)

        allEnvVars := make(map[string]string)
        processedCount := 0

        for _, entry := range entries {
                // Skip directories
                if entry.IsDir() {
                        continue
                }

                // Skip hidden files
                if strings.HasPrefix(entry.Name(), ".") {
                        continue
                }

                filePath := filepath.Join(dirPath, entry.Name())
                fileEnvVars, err := processEnvFile(filePath)
                if err != nil {
                        slog.Warn("Failed to process env file", "name", entry.Name(), "error", err)
                        continue
                }

                // Merge env vars, with later files potentially overriding earlier ones
                for key, value := range fileEnvVars {
                        allEnvVars[key] = value
                }
                processedCount++
        }

        slog.Debug("Processed env files", "files_processed", processedCount, "vars_extracted", len(allEnvVars))
        return allEnvVars, nil
}

// processEnvFile reads and processes a single environment file
// Uses existing ToolHive environment parsing utilities
func processEnvFile(path string) (map[string]string, error) {
        content, err := os.ReadFile(path) // #nosec G304 - path is controlled internally, validated by caller
        if err != nil {
                return nil, fmt.Errorf("failed to read file: %w", err)
        }

        // Convert content to slice of KEY=VALUE lines for existing parser
        lines := strings.Split(string(content), "\n")
        var envLines []string

        for _, line := range lines {
                line = strings.TrimSpace(line)

                // Skip empty lines and comments
                if line == "" || strings.HasPrefix(line, "#") {
                        continue
                }

                // Handle export statements (common in shell env files)
                line = strings.TrimPrefix(line, "export ")

                // Only process lines that contain '=' (KEY=VALUE format)
                if strings.Contains(line, "=") {
                        envLines = append(envLines, line)
                }
        }

        if len(envLines) == 0 {
                slog.Debug("No environment variables found", "file", filepath.Base(path))
                return make(map[string]string), nil
        }

        // Use existing ToolHive utility to parse KEY=VALUE format
        envVars, err := environment.ParseEnvironmentVariables(envLines)
        if err != nil {
                return nil, fmt.Errorf("failed to parse environment variables in %s: %w", filepath.Base(path), err)
        }

        slog.Debug("Extracted environment variables", "count", len(envVars), "file", filepath.Base(path))
        return envVars, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "bytes"
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "slices"

        "github.com/stacklok/toolhive/pkg/audit"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/awssts"
        "github.com/stacklok/toolhive/pkg/auth/obo"
        "github.com/stacklok/toolhive/pkg/auth/upstreamswap"
        "github.com/stacklok/toolhive/pkg/authserver"
        "github.com/stacklok/toolhive/pkg/authz"
        "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar"
        "github.com/stacklok/toolhive/pkg/bodylimit"
        cfg "github.com/stacklok/toolhive/pkg/config"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/oauthproto/tokenexchange"
        "github.com/stacklok/toolhive/pkg/ratelimit"
        "github.com/stacklok/toolhive/pkg/recovery"
        "github.com/stacklok/toolhive/pkg/telemetry"
        headerfwd "github.com/stacklok/toolhive/pkg/transport/middleware"
        "github.com/stacklok/toolhive/pkg/transport/middleware/origin"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/usagemetrics"
        "github.com/stacklok/toolhive/pkg/webhook/mutating"
        "github.com/stacklok/toolhive/pkg/webhook/validating"
)

// GetSupportedMiddlewareFactories returns a map of supported middleware types to their factory functions
func GetSupportedMiddlewareFactories() map[string]types.MiddlewareFactory {
        return map[string]types.MiddlewareFactory{
                auth.MiddlewareType:                   auth.CreateMiddleware,
                tokenexchange.MiddlewareType:          tokenexchange.CreateMiddleware,
                upstreamswap.MiddlewareType:           upstreamswap.CreateMiddleware,
                awssts.MiddlewareType:                 awssts.CreateMiddleware,
                obo.MiddlewareType:                    obo.CreateMiddleware,
                bodylimit.MiddlewareType:              bodylimit.CreateMiddleware,
                mcp.ParserMiddlewareType:              mcp.CreateParserMiddleware,
                mcp.ToolFilterMiddlewareType:          mcp.CreateToolFilterMiddleware,
                mcp.ToolCallFilterMiddlewareType:      mcp.CreateToolCallFilterMiddleware,
                ratelimit.MiddlewareType:              ratelimit.CreateMiddleware,
                usagemetrics.MiddlewareType:           usagemetrics.CreateMiddleware,
                telemetry.MiddlewareType:              telemetry.CreateMiddleware,
                authz.MiddlewareType:                  authz.CreateMiddleware,
                audit.MiddlewareType:                  audit.CreateMiddleware,
                recovery.MiddlewareType:               recovery.CreateMiddleware,
                headerfwd.HeaderForwardMiddlewareName: headerfwd.CreateMiddleware,
                headerfwd.StripAuthMiddlewareName:     headerfwd.CreateStripAuthMiddleware,
                origin.MiddlewareType:                 origin.CreateMiddleware,
                validating.MiddlewareType:             validating.CreateMiddleware,
                mutating.MiddlewareType:               mutating.CreateMiddleware,
        }
}

// PopulateMiddlewareConfigs populates the MiddlewareConfigs slice based on the RunConfig settings
// This function serves as a bridge between the old configuration style and the new generic middleware system
//
//nolint:gocyclo // Function complexity is acceptable for middleware configuration
func PopulateMiddlewareConfigs(config *RunConfig) error {
        var middlewareConfigs []types.MiddlewareConfig
        // TODO: Consider extracting other middleware setup into helper functions like addUsageMetricsMiddleware
        //
        // NOTE: Origin-validation middleware is intentionally NOT added here. It is
        // wired centrally in runner.Run (via prependOriginMiddleware) for both the
        // operator/proxyrunner path (this function) and the CLI path
        // (WithMiddlewareFromFlags), because that is the only place where the
        // effective Host/Port/AllowedOrigins are fully resolved.

        // Body size limit middleware (always present, outermost). See addBodyLimitMiddleware.
        middlewareConfigs, err := addBodyLimitMiddleware(middlewareConfigs, config.MaxRequestBodySize)
        if err != nil {
                return err
        }

        // Audit middleware (if enabled). Added directly inside body-limit so it
        // wraps the REST of the chain: every request that passes the size cap
        // produces an audit event no matter which middleware rejects it.
        // Authentication (401) and authorization/webhook denials (403) map to
        // outcome "denied"; other rejections such as rate limiting (429) map to
        // "failure" (see audit.determineOutcome). Identity and parsed MCP data are
        // read back from the inner auth/parser middlewares via the holder
        // carriers (see auth.IdentityHolder, mcp.ParsedRequestHolder).
        if config.AuditConfig != nil {
                auditParams := audit.MiddlewareParams{
                        ConfigPath:    config.AuditConfigPath, // Keep for backwards compatibility
                        ConfigData:    config.AuditConfig,     // Use the loaded config data
                        Component:     config.AuditConfig.Component,
                        TransportType: config.Transport.String(), // Pass the actual transport type
                }
                auditConfig, err := types.NewMiddlewareConfig(audit.MiddlewareType, auditParams)
                if err != nil {
                        return fmt.Errorf("failed to create audit middleware config: %w", err)
                }
                middlewareConfigs = append(middlewareConfigs, *auditConfig)
        }

        // Authentication middleware (always present)
        embeddedAuthServerIssuer := func() string {
                if config.EmbeddedAuthServerConfig != nil {
                        return config.EmbeddedAuthServerConfig.Issuer
                }
                return ""
        }()
        authParams := auth.MiddlewareParams{
                OIDCConfig:               config.OIDCConfig,
                EmbeddedAuthServerIssuer: embeddedAuthServerIssuer,
        }
        authConfig, authErr := types.NewMiddlewareConfig(auth.MiddlewareType, authParams)
        if authErr != nil {
                return fmt.Errorf("failed to create auth middleware config: %w", authErr)
        }
        middlewareConfigs = append(middlewareConfigs, *authConfig)

        // Upstream swap middleware (if embedded auth server is configured)
        // This exchanges ToolHive JWTs for upstream IdP tokens when embedded auth server is used.
        // IMPORTANT: Must run BEFORE token exchange middleware so it can read the `tsid` claim
        // from the original ToolHive JWT before any token modification occurs.
        middlewareConfigs, err = addUpstreamSwapMiddleware(middlewareConfigs, config)
        if err != nil {
                return err
        }

        // Token exchange middleware (if configured)
        // Runs after upstream swap so that if both are configured, upstream swap can first
        // inject the upstream IdP token, then token exchange can further transform it if needed.
        middlewareConfigs, err = addTokenExchangeMiddleware(middlewareConfigs, config.TokenExchangeConfig)
        if err != nil {
                return err
        }

        // Tools filter and override middleware (if enabled)
        if len(config.ToolsFilter) > 0 || len(config.ToolsOverride) > 0 {
                // Prepare overrides map (convert runner.ToolOverride -> mcp.ToolOverride)
                overrides := make(map[string]mcp.ToolOverride)
                for actualName, tool := range config.ToolsOverride {
                        overrides[actualName] = mcp.ToolOverride{
                                Name:        tool.Name,
                                Description: tool.Description,
                        }
                }

                // Add tool filter middleware with both filter and overrides
                toolFilterParams := mcp.ToolFilterMiddlewareParams{
                        FilterTools:   config.ToolsFilter,
                        ToolsOverride: overrides,
                }
                toolFilterConfig, err := types.NewMiddlewareConfig(mcp.ToolFilterMiddlewareType, toolFilterParams)
                if err != nil {
                        return fmt.Errorf("failed to create tool filter middleware config: %w", err)
                }
                middlewareConfigs = append(middlewareConfigs, *toolFilterConfig)

                // Add tool call filter middleware with same params
                toolCallFilterConfig, err := types.NewMiddlewareConfig(mcp.ToolCallFilterMiddlewareType, toolFilterParams)
                if err != nil {
                        return fmt.Errorf("failed to create tool call filter middleware config: %w", err)
                }
                middlewareConfigs = append(middlewareConfigs, *toolCallFilterConfig)
        }

        // MCP Parser middleware (always present)
        mcpParserParams := mcp.ParserMiddlewareParams{}
        mcpParserConfig, err := types.NewMiddlewareConfig(mcp.ParserMiddlewareType, mcpParserParams)
        if err != nil {
                return fmt.Errorf("failed to create MCP parser middleware config: %w", err)
        }
        middlewareConfigs = append(middlewareConfigs, *mcpParserConfig)

        // Telemetry middleware (if enabled).
        // Positioned after MCP parsing and before rate limiting so the limiter can
        // annotate an active request span. MCP identity is finalized after the inner
        // chain returns, using the parse republished by any mutating webhook.
        if config.TelemetryConfig != nil {
                telemetryParams := telemetry.FactoryMiddlewareParams{
                        Config:     config.TelemetryConfig,
                        ServerName: config.Name,
                        Transport:  config.Transport.String(),
                }
                telemetryConfig, err := types.NewMiddlewareConfig(telemetry.MiddlewareType, telemetryParams)
                if err != nil {
                        return fmt.Errorf("failed to create telemetry middleware config: %w", err)
                }
                middlewareConfigs = append(middlewareConfigs, *telemetryConfig)
        }

        // Rate limit middleware (if configured)
        // Positioned after MCP parser (needs tool name from context) and telemetry
        // (needs the active request span for rate-limit attributes).
        middlewareConfigs, err = addRateLimitMiddleware(middlewareConfigs, config)
        if err != nil {
                return err
        }

        // Mutating Webhooks middleware (if configured).
        // Must run before validating webhooks and authorization. Parsing, telemetry,
        // and rate limiting remain outside both webhook layers; republishing through
        // the shared parsed-request holder lets telemetry observe the final request
        // after the inner chain returns.
        middlewareConfigs, err = addMutatingWebhookMiddleware(middlewareConfigs, config)
        if err != nil {
                return err
        }

        // Validating Webhooks middleware (if configured)
        middlewareConfigs, err = addValidatingWebhookMiddleware(middlewareConfigs, config)
        if err != nil {
                return err
        }

        // Load application config for global settings
        configProvider := cfg.NewDefaultProvider()
        appConfig := configProvider.GetConfig()

        // Usage metrics middleware (if enabled)
        middlewareConfigs, err = addUsageMetricsMiddleware(middlewareConfigs, appConfig.DisableUsageMetrics)
        if err != nil {
                return err
        }

        // Authorization middleware (if enabled)
        if config.AuthzConfig != nil {
                authzCfgData, err := injectUpstreamProviderIfNeeded(config.AuthzConfig, config.EmbeddedAuthServerConfig)
                if err != nil {
                        return fmt.Errorf("failed to inject upstream provider into authorization config: %w", err)
                }
                authzParams := authz.FactoryMiddlewareParams{
                        ConfigPath: config.AuthzConfigPath, // Keep for backwards compatibility
                        ConfigData: authzCfgData,           // Use the (possibly-enriched) config data
                }
                authzConfig, err := types.NewMiddlewareConfig(authz.MiddlewareType, authzParams)
                if err != nil {
                        return fmt.Errorf("failed to create authorization middleware config: %w", err)
                }
                middlewareConfigs = append(middlewareConfigs, *authzConfig)
        }

        // AWS STS middleware (if configured)
        // Placed after audit/authz so that authorization is checked before exchanging
        // credentials, and close to the backend so SigV4 signing happens as late as
        // possible — minimizing the chance of subsequent middleware invalidating the signature.
        middlewareConfigs, err = addAWSStsMiddleware(middlewareConfigs, config)
        if err != nil {
                return err
        }

        // Header forward middleware (if configured for remote servers).
        // Added near the end so it executes closest to the backend handler (innermost).
        // By this point, WithSecrets() has resolved any secret-backed headers
        // into resolvedHeaders, so we pass the merged map to the factory.
        middlewareConfigs, err = addHeaderForwardMiddleware(middlewareConfigs, config)
        if err != nil {
                return err
        }

        // Additional middleware configs injected by external-auth handlers via
        // WithAdditionalMiddlewareConfigs (e.g. the enterprise OBO handler). Upstream
        // carries these pre-built configs verbatim and does not interpret their
        // parameters. They are spliced into the backend-egress group here: after auth
        // (so the authenticated identity is available) and after the awssts /
        // header-forward middleware, and before recovery — giving an injected egress
        // middleware the final say on the outbound request to the backend. Appending an
        // empty slice is a no-op, so this is harmless when nothing was injected.
        //
        // config.AdditionalMiddlewareConfigs is intentionally NOT cleared after the
        // splice. Two reasons: (1) this overwrite-based build is idempotent — re-running
        // PopulateMiddlewareConfigs rebuilds the local slice from scratch and re-reads
        // the carrier, so the injected entry appears exactly once each time rather than
        // accumulating; and (2) the carrier stays serialized as a fallback, so a config
        // persisted before population (empty MiddlewareConfigs) still gets the entry when
        // the proxyrunner re-populates via the len(MiddlewareConfigs)==0 guard. The cost
        // is that a fully-populated ConfigMap carries the entry in both slices; the
        // duplicate is inert because the proxyrunner reads MiddlewareConfigs.
        middlewareConfigs = append(middlewareConfigs, config.AdditionalMiddlewareConfigs...)

        // Recovery middleware (always present, added last). The proxy transports
        // apply this slice in reverse order (see applyMiddlewares), so the
        // last-appended entry becomes the INNERMOST wrapper and executes closest to
        // the handler. Recovery therefore catches panics from the handler and the
        // inner middleware; panics raised in middleware that wrap it (earlier
        // entries, such as body-limit and auth) are not caught here.
        recoveryConfig, err := types.NewMiddlewareConfig(recovery.MiddlewareType, nil)
        if err != nil {
                return fmt.Errorf("failed to create recovery middleware config: %w", err)
        }
        middlewareConfigs = append(middlewareConfigs, *recoveryConfig)

        // Set the populated middleware configs
        config.MiddlewareConfigs = middlewareConfigs
        return nil
}

var errOIDCAuthenticationMiddlewareCount = errors.New("OIDC configuration requires exactly one authentication middleware")

// canonicalizeOIDCMiddlewareConfig synchronizes the serialized authentication
// middleware with the canonical OIDC configuration when the chain is populated.
func canonicalizeOIDCMiddlewareConfig(config *RunConfig) error {
        if config.OIDCConfig == nil || len(config.MiddlewareConfigs) == 0 {
                return nil
        }

        authMiddlewareIndex := -1
        for i, middlewareConfig := range config.MiddlewareConfigs {
                if middlewareConfig.Type != auth.MiddlewareType {
                        continue
                }
                if authMiddlewareIndex >= 0 {
                        return errOIDCAuthenticationMiddlewareCount
                }
                authMiddlewareIndex = i
        }
        if authMiddlewareIndex < 0 {
                return errOIDCAuthenticationMiddlewareCount
        }

        var params auth.MiddlewareParams
        parameters := config.MiddlewareConfigs[authMiddlewareIndex].Parameters
        if string(bytes.TrimSpace(parameters)) == "null" {
                return errors.New("authentication middleware parameters cannot be null")
        }
        if err := json.Unmarshal(parameters, &params); err != nil {
                return fmt.Errorf("failed to decode authentication middleware parameters: %w", err)
        }

        if embeddedAuthServerConfig := config.EmbeddedAuthServerConfig; embeddedAuthServerConfig != nil &&
                embeddedAuthServerConfig.Issuer != "" {
                params.EmbeddedAuthServerIssuer = embeddedAuthServerConfig.Issuer
        }

        canonicalConfig := *config.OIDCConfig
        canonicalConfig.Scopes = slices.Clone(config.OIDCConfig.Scopes)
        params.OIDCConfig = &canonicalConfig
        middlewareConfig, err := types.NewMiddlewareConfig(auth.MiddlewareType, params)
        if err != nil {
                return fmt.Errorf("failed to encode authentication middleware parameters: %w", err)
        }
        config.MiddlewareConfigs[authMiddlewareIndex] = *middlewareConfig
        return nil
}

// addMutatingWebhookMiddleware configures the mutating webhook middleware if any webhooks are defined.
// It must be called before addValidatingWebhookMiddleware to preserve the RFC-specified ordering.
func addMutatingWebhookMiddleware(configs []types.MiddlewareConfig, runConfig *RunConfig) ([]types.MiddlewareConfig, error) {
        if len(runConfig.MutatingWebhooks) == 0 {
                return configs, nil
        }

        params := mutating.FactoryMiddlewareParams{
                MiddlewareParams: mutating.MiddlewareParams{
                        Webhooks: runConfig.MutatingWebhooks,
                },
                ServerName: runConfig.Name,
                Transport:  runConfig.Transport.String(),
        }

        config, err := types.NewMiddlewareConfig(mutating.MiddlewareType, params)
        if err != nil {
                return nil, fmt.Errorf("failed to create mutating webhook middleware config: %w", err)
        }

        return append(configs, *config), nil
}

// addValidatingWebhookMiddleware configures the validating webhook middleware if any webhooks are defined
func addValidatingWebhookMiddleware(configs []types.MiddlewareConfig, runConfig *RunConfig) ([]types.MiddlewareConfig, error) {
        if len(runConfig.ValidatingWebhooks) == 0 {
                return configs, nil
        }

        params := validating.FactoryMiddlewareParams{
                MiddlewareParams: validating.MiddlewareParams{
                        Webhooks: runConfig.ValidatingWebhooks,
                },
                ServerName: runConfig.Name,
                Transport:  runConfig.Transport.String(),
        }

        config, err := types.NewMiddlewareConfig(validating.MiddlewareType, params)
        if err != nil {
                return nil, fmt.Errorf("failed to create validating webhook middleware config: %w", err)
        }

        return append(configs, *config), nil
}

// addTokenExchangeMiddleware adds token exchange middleware if configured
func addTokenExchangeMiddleware(
        middlewares []types.MiddlewareConfig,
        tokenExchangeConfig *tokenexchange.Config,
) ([]types.MiddlewareConfig, error) {
        if tokenExchangeConfig == nil {
                return middlewares, nil
        }

        tokenExchangeParams := tokenexchange.MiddlewareParams{
                TokenExchangeConfig: tokenExchangeConfig,
        }
        tokenExchangeMwConfig, err := types.NewMiddlewareConfig(
                tokenexchange.MiddlewareType,
                tokenExchangeParams,
        )
        if err != nil {
                return nil, fmt.Errorf("failed to create token exchange middleware config: %w", err)
        }
        return append(middlewares, *tokenExchangeMwConfig), nil
}

// addBodyLimitMiddleware ensures the body-size-limit middleware is present as the
// outermost entry (index 0) of the chain. It is the symmetric counterpart to recovery
// ("always present, innermost"): body-limit is "always present, outermost", so an
// oversized request body is rejected with 413 before auth, the MCP parser, or any
// handler buffers it via io.ReadAll — regardless of which builder assembled the chain.
//
// Existing body-limit entries are replaced so the typed RunConfig field remains
// authoritative. The relative order of every other middleware is preserved.
// Zero defaults to bodylimit.DefaultMaxRequestBodySize; negative values are rejected.
func addBodyLimitMiddleware(
        middlewares []types.MiddlewareConfig,
        maxRequestBodySize int64,
) ([]types.MiddlewareConfig, error) {
        if maxRequestBodySize < 0 {
                return nil, fmt.Errorf("max_request_body_size must be non-negative, got %d", maxRequestBodySize)
        }
        if maxRequestBodySize == 0 {
                maxRequestBodySize = bodylimit.DefaultMaxRequestBodySize
        }
        bodyLimitConfig, err := types.NewMiddlewareConfig(bodylimit.MiddlewareType, bodylimit.MiddlewareParams{
                MaxBytes: maxRequestBodySize,
        })
        if err != nil {
                return nil, fmt.Errorf("failed to create body limit middleware config: %w", err)
        }

        result := make([]types.MiddlewareConfig, 0, len(middlewares)+1)
        result = append(result, *bodyLimitConfig)
        for _, middleware := range middlewares {
                if middleware.Type != bodylimit.MiddlewareType {
                        result = append(result, middleware)
                }
        }
        return result, nil
}

// addHeaderForwardMiddleware adds header forward middleware if configured for remote servers
func addHeaderForwardMiddleware(middlewares []types.MiddlewareConfig, config *RunConfig) ([]types.MiddlewareConfig, error) {
        if config.RemoteURL == "" || !config.HeaderForward.HasHeaders() {
                return middlewares, nil
        }

        headerForwardParams := headerfwd.HeaderForwardMiddlewareParams{
                AddHeaders: config.HeaderForward.ResolvedHeaders(),
        }
        headerForwardConfig, err := types.NewMiddlewareConfig(headerfwd.HeaderForwardMiddlewareName, headerForwardParams)
        if err != nil {
                return nil, fmt.Errorf("failed to create header forward middleware config: %w", err)
        }
        return append(middlewares, *headerForwardConfig), nil
}

// addUsageMetricsMiddleware adds usage metrics middleware if enabled
func addUsageMetricsMiddleware(middlewares []types.MiddlewareConfig, configDisabled bool) ([]types.MiddlewareConfig, error) {
        if !usagemetrics.ShouldEnableMetrics(configDisabled) {
                return middlewares, nil
        }

        usageMetricsParams := usagemetrics.MiddlewareParams{}
        usageMetricsConfig, err := types.NewMiddlewareConfig(usagemetrics.MiddlewareType, usageMetricsParams)
        if err != nil {
                return nil, fmt.Errorf("failed to create usage metrics middleware config: %w", err)
        }
        return append(middlewares, *usageMetricsConfig), nil
}

// addUpstreamSwapMiddleware adds upstream swap middleware if the embedded auth server is configured.
// This middleware exchanges ToolHive JWTs for upstream IdP tokens.
// The middleware is only added when EmbeddedAuthServerConfig is set and
// DisableUpstreamTokenInjection is false. If UpstreamSwapConfig is nil,
// default configuration values are used.
func addUpstreamSwapMiddleware(
        middlewares []types.MiddlewareConfig,
        config *RunConfig,
) ([]types.MiddlewareConfig, error) {
        // Only add middleware if embedded auth server is configured
        if config.EmbeddedAuthServerConfig == nil {
                return middlewares, nil
        }

        // A token-only embedded auth server has no upstream credential to swap. Like
        // DisableUpstreamTokenInjection, it strips client credentials only after auth.
        // An explicit swap configuration is invalid rather than silently ignored.
        if len(config.EmbeddedAuthServerConfig.Upstreams) == 0 {
                if config.UpstreamSwapConfig != nil {
                        return nil, fmt.Errorf("upstream swap cannot be configured without upstreams")
                }
                if err := validateCredentialInjectionConflicts(config, "token-only embedded auth server"); err != nil {
                        return nil, err
                }
                return addAuthHeaderStripMiddleware(middlewares)
        }

        // When upstream token injection is disabled, strip the client's credential
        // headers (Authorization, Cookie, Proxy-Authorization) so they never reach
        // the upstream server. Two ordering invariants apply, pinned by
        // TestPopulateMiddlewareConfigs_StripAuthOrdering:
        //   - strip-auth is appended after the auth middleware, so the client JWT
        //     is fully validated (and the identity stored in the request context
        //     for authz/audit) before the header is removed;
        //   - token-injecting middlewares (token exchange, AWS STS) run closer to
        //     the backend and would re-add an Authorization header after the
        //     strip, silently defeating the flag — that contradiction is rejected
        //     here instead.
        if config.EmbeddedAuthServerConfig.DisableUpstreamTokenInjection {
                if config.UpstreamSwapConfig != nil {
                        return nil, fmt.Errorf("disableUpstreamTokenInjection cannot be combined with upstream swap: " +
                                "upstream swap would re-add an Authorization header after strip-auth removes it")
                }
                if err := validateCredentialInjectionConflicts(config, "disableUpstreamTokenInjection"); err != nil {
                        return nil, err
                }
                return addAuthHeaderStripMiddleware(middlewares)
        }

        // Use provided config or defaults
        upstreamSwapConfig := config.UpstreamSwapConfig
        if upstreamSwapConfig == nil {
                upstreamSwapConfig = &upstreamswap.Config{}
        }

        // Derive ProviderName from the upstream config if not explicitly set
        if upstreamSwapConfig.ProviderName == "" {
                var names []string
                if embeddedCfg := config.EmbeddedAuthServerConfig; embeddedCfg != nil {
                        names = make([]string, len(embeddedCfg.Upstreams))
                        for i, u := range embeddedCfg.Upstreams {
                                names[i] = u.Name
                        }
                }
                upstreamSwapConfig.ProviderName = authserver.ResolveFirstUpstreamName(names)
        }

        upstreamSwapParams := upstreamswap.MiddlewareParams{
                Config: upstreamSwapConfig,
        }
        upstreamSwapMwConfig, err := types.NewMiddlewareConfig(
                upstreamswap.MiddlewareType,
                upstreamSwapParams,
        )
        if err != nil {
                return nil, fmt.Errorf("failed to create upstream swap middleware config: %w", err)
        }
        return append(middlewares, *upstreamSwapMwConfig), nil
}

// injectUpstreamProviderIfNeeded enriches an authz.Config with the
// PrimaryUpstreamProvider derived from the embedded auth server config.
// When the embedded auth server is active, Cedar policies should evaluate
// claims from the upstream IDP token rather than the ToolHive-issued JWT.
// If embeddedCfg is nil the original authzCfg is returned unchanged.
func injectUpstreamProviderIfNeeded(
        authzCfg *authz.Config,
        embeddedCfg *authserver.RunConfig,
) (*authz.Config, error) {
        if embeddedCfg == nil {
                return authzCfg, nil
        }

        if len(embeddedCfg.Upstreams) == 0 {
                return authzCfg, nil
        }

        names := make([]string, len(embeddedCfg.Upstreams))
        for i, u := range embeddedCfg.Upstreams {
                names[i] = u.Name
        }
        providerName := authserver.ResolveFirstUpstreamName(names)

        return cedar.InjectUpstreamProvider(authzCfg, providerName)
}

// validateCredentialStrippingMiddleware verifies that a pre-built middleware chain
// cannot bypass the credential-stripping guarantees required when no upstream
// credential is available or its injection is disabled. The chain must include
// exactly one strip-auth middleware after authentication.
func validateCredentialStrippingMiddleware(
        middlewares []types.MiddlewareConfig,
        config *RunConfig,
) error {
        if !requiresCredentialStripping(config) {
                return nil
        }
        if err := validateCredentialStrippingConfig(config); err != nil {
                return err
        }

        authIndex, stripAuthIndex, err := credentialStrippingMiddlewareIndices(middlewares)
        if err != nil {
                return err
        }
        if stripAuthIndex == -1 {
                return fmt.Errorf("credential stripping requires strip-auth middleware")
        }
        if authIndex == -1 {
                return fmt.Errorf("credential stripping requires auth middleware")
        }
        if stripAuthIndex < authIndex {
                return fmt.Errorf("credential stripping requires strip-auth middleware after auth middleware")
        }
        return nil
}

func requiresCredentialStripping(config *RunConfig) bool {
        return config.EmbeddedAuthServerConfig != nil &&
                (len(config.EmbeddedAuthServerConfig.Upstreams) == 0 ||
                        config.EmbeddedAuthServerConfig.DisableUpstreamTokenInjection)
}

func validateCredentialStrippingConfig(config *RunConfig) error {
        if config.UpstreamSwapConfig != nil {
                return fmt.Errorf("credential stripping cannot be combined with upstream swap")
        }
        if config.TokenExchangeConfig != nil {
                return fmt.Errorf("credential stripping cannot be combined with token exchange")
        }
        if config.AWSStsConfig != nil {
                return fmt.Errorf("credential stripping cannot be combined with AWS STS")
        }
        if hasCredentialInjectingAdditionalMiddleware(config.AdditionalMiddlewareConfigs) {
                return fmt.Errorf("credential stripping cannot be combined with OBO middleware")
        }
        return nil
}

func credentialStrippingMiddlewareIndices(middlewares []types.MiddlewareConfig) (int, int, error) {
        authIndex, stripAuthIndex := -1, -1
        for i, middleware := range middlewares {
                if err := credentialInjectingMiddlewareError(middleware.Type); err != nil {
                        return 0, 0, err
                }
                if middleware.Type == auth.MiddlewareType && authIndex == -1 {
                        authIndex = i
                }
                if middleware.Type != headerfwd.StripAuthMiddlewareName {
                        continue
                }
                if stripAuthIndex != -1 {
                        return 0, 0, fmt.Errorf("credential stripping requires exactly one strip-auth middleware")
                }
                stripAuthIndex = i
        }
        return authIndex, stripAuthIndex, nil
}

func credentialInjectingMiddlewareError(middlewareType string) error {
        switch middlewareType {
        case upstreamswap.MiddlewareType:
                return fmt.Errorf("credential stripping cannot be combined with upstream swap middleware")
        case tokenexchange.MiddlewareType:
                return fmt.Errorf("credential stripping cannot be combined with token exchange middleware")
        case awssts.MiddlewareType:
                return fmt.Errorf("credential stripping cannot be combined with AWS STS middleware")
        case obo.MiddlewareType:
                return fmt.Errorf("credential stripping cannot be combined with OBO middleware")
        default:
                return nil
        }
}

func validateCredentialInjectionConflicts(config *RunConfig, reason string) error {
        if hasCredentialInjectingAdditionalMiddleware(config.AdditionalMiddlewareConfigs) {
                return fmt.Errorf(
                        "%s cannot be combined with OBO middleware: "+
                                "OBO would re-add credentials after strip-auth removes them", reason)
        }
        if config.TokenExchangeConfig != nil {
                return fmt.Errorf(
                        "%s cannot be combined with token exchange: "+
                                "token exchange would re-add an Authorization header after strip-auth removes it", reason)
        }
        if config.AWSStsConfig != nil {
                return fmt.Errorf(
                        "%s cannot be combined with AWS STS: "+
                                "SigV4 signing would re-add credentials after strip-auth removes them", reason)
        }
        return nil
}

func hasCredentialInjectingAdditionalMiddleware(middlewares []types.MiddlewareConfig) bool {
        return slices.ContainsFunc(middlewares, func(middleware types.MiddlewareConfig) bool {
                return middleware.Type == obo.MiddlewareType
        })
}

// addAuthHeaderStripMiddleware adds the strip-auth middleware
// (pkg/transport/middleware), which removes the client's credential headers
// before forwarding to the upstream. This prevents the client's ToolHive JWT,
// cookies, and proxy credentials from leaking to upstream servers that don't
// expect them.
func addAuthHeaderStripMiddleware(
        middlewares []types.MiddlewareConfig,
) ([]types.MiddlewareConfig, error) {
        mwConfig, err := types.NewMiddlewareConfig(headerfwd.StripAuthMiddlewareName, struct{}{})
        if err != nil {
                return nil, fmt.Errorf("failed to create strip-auth middleware config: %w", err)
        }
        return append(middlewares, *mwConfig), nil
}

// addAWSStsMiddleware adds AWS STS middleware if configured.
// Returns an error if AWSStsConfig is set but RemoteURL is empty, because
// SigV4 signing is only meaningful for remote MCP servers.
func addAWSStsMiddleware(middlewares []types.MiddlewareConfig, config *RunConfig) ([]types.MiddlewareConfig, error) {
        if config.AWSStsConfig == nil {
                return middlewares, nil
        }

        if config.RemoteURL == "" {
                return nil, fmt.Errorf("AWS STS middleware requires a remote URL: SigV4 signing is only meaningful for remote MCP servers")
        }

        awsStsParams := awssts.MiddlewareParams{
                AWSStsConfig: config.AWSStsConfig,
                TargetURL:    config.RemoteURL, // Use remote URL as the target for SigV4 signing
        }
        awsStsMwConfig, err := types.NewMiddlewareConfig(awssts.MiddlewareType, awsStsParams)
        if err != nil {
                return nil, fmt.Errorf("failed to create AWS STS middleware config: %w", err)
        }
        return append(middlewares, *awsStsMwConfig), nil
}

// prependOriginMiddleware prepends Origin-header validation middleware for
// DNS-rebind protection per MCP 2025-11-25 §"Security Warning". It is placed at
// the front of the chain so disallowed Origin values are rejected before
// authentication or any business logic runs. Default-derivation logic lives in
// origin.ResolveAllowedOrigins so the standalone `thv proxy` command and the
// runner path agree on behavior.
//
// This is called from runner.Run after both middleware-population paths
// (PopulateMiddlewareConfigs and WithMiddlewareFromFlags) have run, because
// that is the only point where the effective Host/Port/AllowedOrigins are
// fully resolved — the CLI builder defers port resolution to validateConfig.
//
// When the effective allowlist is empty — which happens when the operator
// binds to a non-loopback host without supplying --allowed-origins — the
// middleware is skipped entirely and a WARN is logged so the security-disabled
// state is visible in operator logs. A follow-up PR hardens the non-loopback
// path by requiring an explicit opt-in flag (see audit row 22).
func prependOriginMiddleware(middlewares []types.MiddlewareConfig, config *RunConfig) ([]types.MiddlewareConfig, error) {
        allowed := origin.ResolveAllowedOrigins(config.Host, config.Port, config.AllowedOrigins)
        if len(allowed) == 0 {
                slog.Warn("Origin validation disabled — no allowlist configured for non-loopback bind",
                        "host", config.Host,
                        "port", config.Port,
                        "hint", "pass --allowed-origins=https://your-client.example to enable DNS-rebind protection",
                )
                return middlewares, nil
        }

        params := origin.MiddlewareParams{AllowedOrigins: allowed}
        mwCfg, err := types.NewMiddlewareConfig(origin.MiddlewareType, params)
        if err != nil {
                return nil, fmt.Errorf("failed to create origin middleware config: %w", err)
        }
        // Prepend so Origin validation is the outermost wrapper (runs first at
        // request time). Build a new slice to avoid mutating the caller's backing
        // array.
        return append([]types.MiddlewareConfig{*mwCfg}, middlewares...), nil
}

// addRateLimitMiddleware adds rate limit middleware if configured.
func addRateLimitMiddleware(middlewares []types.MiddlewareConfig, config *RunConfig) ([]types.MiddlewareConfig, error) {
        if config.RateLimitConfig == nil {
                return middlewares, nil
        }

        if config.ScalingConfig == nil || config.ScalingConfig.SessionRedis == nil {
                return nil, fmt.Errorf("rate limiting requires sessionStorage with provider redis")
        }
        redisAddr := config.ScalingConfig.SessionRedis.Address
        redisDB := config.ScalingConfig.SessionRedis.DB

        params := ratelimit.MiddlewareParams{
                Namespace:  config.RateLimitNamespace,
                ServerName: config.Name,
                Config:     config.RateLimitConfig,
                RedisAddr:  redisAddr,
                RedisDB:    redisDB,
                RedisTLS:   config.ScalingConfig.SessionRedis.TLS,
        }
        mwConfig, err := types.NewMiddlewareConfig(ratelimit.MiddlewareType, params)
        if err != nil {
                return nil, fmt.Errorf("failed to create rate limit middleware config: %w", err)
        }
        return append(middlewares, *mwConfig), nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "strings"
)

// This was moved from the CLI to allow it to be shared with the lifecycle manager.
// It will likely be moved elsewhere in a future PR.

// CleanupTempPermissionProfile removes a temporary permission profile file if it was created by toolhive
func CleanupTempPermissionProfile(permissionProfilePath string) error {
        if permissionProfilePath == "" {
                return nil
        }

        // Check if this is a temporary file created by toolhive
        if !isTempPermissionProfile(permissionProfilePath) {
                //nolint:gosec // G706: path is user-provided file, not secret
                slog.Debug("Permission profile is not a temporary file, skipping cleanup", "path", permissionProfilePath)
                return nil
        }

        // Check if the file exists
        // #nosec G703 -- permissionProfilePath is validated by isTempPermissionProfile above
        if _, err := os.Stat(permissionProfilePath); os.IsNotExist(err) {
                //nolint:gosec // G706: path is validated by isTempPermissionProfile
                slog.Debug("Temporary permission profile file does not exist, skipping cleanup", "path", permissionProfilePath)
                return nil
        }

        // Remove the temporary file
        // #nosec G703 -- permissionProfilePath is validated by isTempPermissionProfile above
        if err := os.Remove(permissionProfilePath); err != nil {
                return fmt.Errorf("failed to remove temporary permission profile file %s: %w", permissionProfilePath, err)
        }

        //nolint:gosec // G706: path is validated by isTempPermissionProfile
        slog.Debug("Removed temporary permission profile file", "path", permissionProfilePath)
        return nil
}

// isTempPermissionProfile checks if a file path is a temporary permission profile created by toolhive
func isTempPermissionProfile(filePath string) bool {
        if filePath == "" {
                return false
        }

        // Get the base name of the file
        fileName := filepath.Base(filePath)

        // Check if it matches the pattern: toolhive-*-permissions-*.json
        if !strings.HasPrefix(fileName, "toolhive-") ||
                !strings.Contains(fileName, "-permissions-") ||
                !strings.HasSuffix(fileName, ".json") {
                return false
        }

        // Check if it's in a temporary directory (os.TempDir() or similar)
        tempDir := os.TempDir()
        fileDir := filepath.Dir(filePath)

        // Check if the file is in the system temp directory or a subdirectory of it
        relPath, err := filepath.Rel(tempDir, fileDir)
        if err != nil {
                return false
        }

        // If the relative path doesn't start with "..", then it's within the temp directory
        return !strings.HasPrefix(relPath, "..")
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "context"
        "sync"
)

// PolicyGate is called before server creation operations to allow external
// policy enforcement. Additional methods (e.g., CheckStopServer) may be added
// in future issues; downstream implementations should embed a NoopPolicyGate
// to remain forward-compatible.
type PolicyGate interface {
        // CheckCreateServer is called before a local workload container is set up.
        // Return a non-nil error to block server creation.
        CheckCreateServer(ctx context.Context, cfg *RunConfig) error
}

// NoopPolicyGate is a policy gate that allows all operations. Downstream
// implementations should embed this struct to remain forward-compatible when
// new methods are added to the PolicyGate interface.
type NoopPolicyGate struct{}

// CheckCreateServer implements PolicyGate by allowing all create operations.
func (NoopPolicyGate) CheckCreateServer(_ context.Context, _ *RunConfig) error {
        return nil
}

// allowAllGate is the default policy gate used when no gate has been registered.
type allowAllGate struct {
        NoopPolicyGate
}

var (
        policyGateMu sync.RWMutex
        policyGate   PolicyGate = allowAllGate{}
)

// RegisterPolicyGate replaces the active policy gate with g. It is safe to
// call from multiple goroutines, though it is intended to be called once at
// program startup before any runners are created.
func RegisterPolicyGate(g PolicyGate) {
        policyGateMu.Lock()
        defer policyGateMu.Unlock()
        policyGate = g
}

// ActivePolicyGate returns the currently registered policy gate under the
// package-level mutex. It is exported for use by other toolhive packages
// (e.g. retriever) that enforce policy outside Runner.Run; it is not
// intended for external consumers.
func ActivePolicyGate() PolicyGate {
        policyGateMu.RLock()
        defer policyGateMu.RUnlock()
        return policyGate
}

// EagerCheckCreateServer calls CheckCreateServer on the currently registered
// policy gate. Callers in the CLI or API layer should call this before saving
// state or spawning a detached worker so that policy violations surface
// synchronously in the main process, not silently in the background.
func EagerCheckCreateServer(ctx context.Context, cfg *RunConfig) error {
        return ActivePolicyGate().CheckCreateServer(ctx, cfg)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "context"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "regexp"
        "strings"
        "time"

        nameref "github.com/google/go-containerregistry/pkg/name"

        "github.com/stacklok/toolhive/pkg/certs"
        "github.com/stacklok/toolhive/pkg/config"
        "github.com/stacklok/toolhive/pkg/container/images"
        "github.com/stacklok/toolhive/pkg/container/templates"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// Protocol schemes
const (
        UVXScheme = "uvx://"
        NPXScheme = "npx://"
        GOScheme  = "go://"
)

// HandleProtocolScheme checks if the serverOrImage string contains a protocol scheme (uvx://, npx://, or go://)
// and builds a Docker image for it if needed.
// Returns the Docker image name to use and any error encountered.
func HandleProtocolScheme(
        ctx context.Context,
        imageManager images.ImageManager,
        serverOrImage string,
        caCertPath string,
        runtimeOverride *templates.RuntimeConfig,
) (string, error) {
        return BuildFromProtocolSchemeWithName(ctx, imageManager, serverOrImage, caCertPath, "", nil, runtimeOverride, false)
}

// BuildFromProtocolSchemeWithName checks if the serverOrImage string contains a protocol scheme (uvx://, npx://, or go://)
// and builds a Docker image for it if needed with a custom image name.
// If imageName is empty, a default name will be generated.
// buildArgs are baked into the container's ENTRYPOINT at build time (e.g., required subcommands).
// If dryRun is true, returns the Dockerfile content instead of building the image.
// Returns the Docker image name (or Dockerfile content if dryRun) and any error encountered.
func BuildFromProtocolSchemeWithName(
        ctx context.Context,
        imageManager images.ImageManager,
        serverOrImage string,
        caCertPath string,
        imageName string,
        buildArgs []string,
        runtimeOverride *templates.RuntimeConfig,
        dryRun bool,
) (string, error) {
        transportType, packageName, err := ParseProtocolScheme(serverOrImage)
        if err != nil {
                return "", err
        }

        templateData, err := createTemplateData(transportType, packageName, caCertPath, buildArgs, runtimeOverride)
        if err != nil {
                return "", err
        }

        // If dry-run, just return the Dockerfile content
        if dryRun {
                dockerfileContent, err := templates.GetDockerfileTemplate(transportType, templateData)
                if err != nil {
                        return "", fmt.Errorf("failed to get Dockerfile template: %w", err)
                }
                return dockerfileContent, nil
        }

        return buildImageFromTemplateWithName(ctx, imageManager, transportType, packageName, templateData, imageName)
}

// ParseProtocolScheme extracts the transport type and package name from the protocol scheme.
func ParseProtocolScheme(serverOrImage string) (templates.TransportType, string, error) {
        if strings.HasPrefix(serverOrImage, UVXScheme) {
                return templates.TransportTypeUVX, strings.TrimPrefix(serverOrImage, UVXScheme), nil
        }
        if strings.HasPrefix(serverOrImage, NPXScheme) {
                return templates.TransportTypeNPX, strings.TrimPrefix(serverOrImage, NPXScheme), nil
        }
        if strings.HasPrefix(serverOrImage, GOScheme) {
                return templates.TransportTypeGO, strings.TrimPrefix(serverOrImage, GOScheme), nil
        }
        return "", "", fmt.Errorf("unsupported protocol scheme: %s", serverOrImage)
}

// validateBuildArgs ensures buildArgs don't contain single quotes which would break
// shell quoting in the UVX template. Single quotes cannot be escaped within single-quoted
// strings in shell, making them the only character that can enable command injection.
// NPX and GO use JSON array ENTRYPOINTs without shell interpretation, so they're safe.
func validateBuildArgs(buildArgs []string) error {
        for _, arg := range buildArgs {
                if strings.Contains(arg, "'") {
                        return fmt.Errorf("buildArg cannot contain single quotes: %s", arg)
                }
        }
        return nil
}

// packageNameChars allows the characters that npm, PyPI and Go package
// references legitimately need: scopes and versions (@, /), Go module paths
// (/, .), local Go paths (./, ../, /), PyPI pins and extras (=, [, ]) and the
// usual separators. Every character a shell treats specially is excluded, so a
// validated name cannot break out of the RUN instructions it is rendered into.
var packageNameChars = regexp.MustCompile(`^[A-Za-z0-9@/:._+=~\[\]-]+$`)

// validatePackageName rejects package names that could inject shell syntax into
// the generated Dockerfile. The name is interpolated into RUN instructions in
// npx.tmpl, uvx.tmpl and go.tmpl, so it has to be constrained before it reaches
// the template data.
func validatePackageName(packageName string) error {
        if packageName == "" {
                return fmt.Errorf("package name cannot be empty")
        }
        if !packageNameChars.MatchString(packageName) {
                return fmt.Errorf(
                        "invalid package name %q: only letters, digits and the characters @/:._+=~[]- are allowed",
                        packageName,
                )
        }
        return nil
}

// createTemplateData creates the template data with optional CA certificate and build arguments.
func createTemplateData(
        transportType templates.TransportType, packageName, caCertPath string, buildArgs []string,
        runtimeOverride *templates.RuntimeConfig,
) (templates.TemplateData, error) {
        // Validate the package name to prevent shell injection in the generated Dockerfile
        if err := validatePackageName(packageName); err != nil {
                return templates.TemplateData{}, err
        }

        // Validate buildArgs to prevent shell injection in templates that use sh -c
        if err := validateBuildArgs(buildArgs); err != nil {
                return templates.TemplateData{}, err
        }

        // Check if this is a local path (for Go packages only)
        isLocalPath := transportType == templates.TransportTypeGO && isLocalGoPath(packageName)

        templateData := templates.TemplateData{
                MCPPackage:  packageName,
                IsLocalPath: isLocalPath,
                BuildArgs:   buildArgs,
        }

        if caCertPath != "" {
                if err := addCACertToTemplate(caCertPath, &templateData); err != nil {
                        return templateData, err
                }
        }

        // Load build environment variables from configuration
        if err := addBuildEnvToTemplate(&templateData); err != nil {
                return templateData, err
        }

        // Load build auth files from configuration and secrets
        if err := addBuildAuthFilesToTemplate(&templateData); err != nil {
                return templateData, err
        }

        // Load runtime configuration (base images and packages)
        runtimeConfig, err := loadRuntimeConfig(transportType, runtimeOverride)
        if err != nil {
                return templateData, err
        }
        templateData.RuntimeConfig = runtimeConfig

        return templateData, nil
}

// loadRuntimeConfig loads the runtime configuration for a given transport type.
// Priority order:
// 1. Override provided as parameter (merged with defaults, then validated)
// 2. User configuration from config file
// 3. Default configuration for the transport type
//
// When an override is provided, empty fields fall back to the defaults and
// AdditionalPackages are merged (defaults first, then unique override entries).
func loadRuntimeConfig(
        transportType templates.TransportType,
        override *templates.RuntimeConfig,
) (*templates.RuntimeConfig, error) {
        defaults := templates.GetDefaultRuntimeConfig(transportType)

        // If override is provided, merge with defaults before validating
        if override != nil {
                merged := defaults.WithOverrides(override)
                if err := merged.ValidateFor(transportType); err != nil {
                        return nil, fmt.Errorf("invalid runtime config override: %w", err)
                }
                return merged, nil
        }

        // Try loading from user config (merge with defaults, then validate)
        provider := config.NewProvider()
        if userConfig, err := provider.GetRuntimeConfig(string(transportType)); err == nil && userConfig != nil {
                merged := defaults.WithOverrides(userConfig)
                if err := merged.ValidateFor(transportType); err != nil {
                        return nil, fmt.Errorf("invalid runtime config in config file for %s: %w", transportType, err)
                }
                return merged, nil
        }

        // Fall back to defaults. GetDefaultRuntimeConfig already returns a value
        // detached from the package-global RuntimeDefaults, so no further clone
        // is needed here.
        if err := defaults.ValidateFor(transportType); err != nil {
                return nil, fmt.Errorf("invalid default runtime config for %s: %w", transportType, err)
        }
        return &defaults, nil
}

// addBuildEnvToTemplate loads build environment variables from config and adds them to template data.
// It resolves values from three sources:
// 1. Literal values stored in BuildEnv
// 2. Values from ToolHive secrets (BuildEnvFromSecrets)
// 3. Values from the current shell environment (BuildEnvFromShell)
func addBuildEnvToTemplate(templateData *templates.TemplateData) error {
        provider := config.NewProvider()
        resolvedEnv := make(map[string]string)

        // 1. Add literal values
        literalEnv := provider.GetAllBuildEnv()
        for k, v := range literalEnv {
                resolvedEnv[k] = v
        }

        // 2. Resolve values from secrets
        secretRefs := provider.GetAllBuildEnvFromSecrets()
        if len(secretRefs) > 0 {
                secretValues, err := resolveSecretsForBuildEnv(secretRefs)
                if err != nil {
                        return fmt.Errorf("failed to resolve secrets for build env: %w", err)
                }
                for k, v := range secretValues {
                        resolvedEnv[k] = v
                }
        }

        // 3. Resolve values from shell environment
        shellRefs := provider.GetAllBuildEnvFromShell()
        for _, key := range shellRefs {
                value := os.Getenv(key)
                if value == "" {
                        slog.Warn("Build env variable configured to read from shell, but not set in environment", "key", key)
                        continue
                }
                resolvedEnv[key] = value
        }

        if len(resolvedEnv) > 0 {
                templateData.BuildEnv = resolvedEnv
                slog.Debug("Loaded build environment variable(s) (redacted for security)", "count", len(resolvedEnv))
        }

        return nil
}

// addBuildAuthFilesToTemplate loads build auth files from config and secrets, adding them to template data.
func addBuildAuthFilesToTemplate(templateData *templates.TemplateData) error {
        provider := config.NewProvider()
        configuredFiles := provider.GetConfiguredBuildAuthFiles()

        if len(configuredFiles) == 0 {
                return nil
        }

        // Resolve auth file content from secrets
        authFiles, err := resolveBuildAuthFilesFromSecrets(configuredFiles)
        if err != nil {
                return err
        }

        if len(authFiles) > 0 {
                templateData.BuildAuthFiles = authFiles
                slog.Debug("Loaded build auth file(s)", "count", len(authFiles))
        }

        return nil
}

// resolveBuildAuthFilesFromSecrets resolves auth file content from the secrets provider.
func resolveBuildAuthFilesFromSecrets(configuredFiles []string) (map[string]string, error) {
        ctx := context.Background()
        configProvider := config.NewProvider()
        cfg := configProvider.GetConfig()

        // Check if secrets are set up
        if !cfg.Secrets.SetupCompleted {
                return nil, secrets.ErrSecretsNotSetup
        }

        providerType, err := cfg.Secrets.GetProviderType()
        if err != nil {
                return nil, fmt.Errorf("failed to get secrets provider type: %w", err)
        }

        manager, err := secrets.CreateProvider(providerType, secrets.WithScope(secrets.ScopeWorkloads))
        if err != nil {
                return nil, fmt.Errorf("failed to create secrets provider: %w", err)
        }

        resolved := make(map[string]string, len(configuredFiles))
        for _, fileType := range configuredFiles {
                secretName := config.BuildAuthFileSecretName(fileType)
                content, err := manager.GetSecret(ctx, secretName)
                if err != nil {
                        return nil, fmt.Errorf("failed to get secret '%s' for auth file %s: %w", secretName, fileType, err)
                }
                resolved[fileType] = content
        }

        return resolved, nil
}

// resolveSecretsForBuildEnv resolves secret references to their actual values.
func resolveSecretsForBuildEnv(secretRefs map[string]string) (map[string]string, error) {
        ctx := context.Background()
        configProvider := config.NewProvider()
        cfg := configProvider.GetConfig()

        // Check if secrets are set up
        if !cfg.Secrets.SetupCompleted {
                return nil, secrets.ErrSecretsNotSetup
        }

        providerType, err := cfg.Secrets.GetProviderType()
        if err != nil {
                return nil, fmt.Errorf("failed to get secrets provider type: %w", err)
        }

        manager, err := secrets.CreateProvider(providerType, secrets.WithUserFacing())
        if err != nil {
                return nil, fmt.Errorf("failed to create secrets provider: %w", err)
        }

        resolved := make(map[string]string, len(secretRefs))
        for key, secretName := range secretRefs {
                value, err := manager.GetSecret(ctx, secretName)
                if err != nil {
                        return nil, fmt.Errorf("failed to get secret '%s' for build env variable %s: %w", secretName, key, err)
                }

                // Validate the secret value doesn't contain dangerous characters
                if err := config.ValidateBuildEnvValue(value); err != nil {
                        return nil, fmt.Errorf("secret '%s' contains invalid value for build env variable %s: %w", secretName, key, err)
                }

                resolved[key] = value
        }

        return resolved, nil
}

// addCACertToTemplate reads and validates a CA certificate, adding it to the template data.
func addCACertToTemplate(caCertPath string, templateData *templates.TemplateData) error {
        slog.Debug("Using custom CA certificate", "path", caCertPath)

        // Read the CA certificate file
        // #nosec G304 -- This is a user-provided file path that we need to read
        caCertContent, err := os.ReadFile(caCertPath)
        if err != nil {
                return fmt.Errorf("failed to read CA certificate file: %w", err)
        }

        // Validate that the file contains a valid PEM certificate
        if err := certs.ValidateCACertificate(caCertContent); err != nil {
                return fmt.Errorf("invalid CA certificate: %w", err)
        }

        // Add the CA certificate content to the template data
        templateData.CACertContent = string(caCertContent)
        slog.Debug("Successfully validated and loaded CA certificate")
        return nil
}

// buildContext represents a Docker build context with cleanup functionality.
type buildContext struct {
        Dir            string
        DockerfilePath string
        CleanupFunc    func()
}

// setupBuildContext sets up the appropriate build context directory based on whether
// we're dealing with a local path or remote package.
func setupBuildContext(packageName string, isLocalPath bool) (*buildContext, error) {
        if isLocalPath {
                return setupLocalBuildContext(packageName)
        }
        return setupTempBuildContext()
}

// setupLocalBuildContext sets up a build context using the local directory directly.
func setupLocalBuildContext(packageName string) (*buildContext, error) {
        absPath, err := filepath.Abs(packageName)
        if err != nil {
                return nil, fmt.Errorf("failed to get absolute path for %s: %w", packageName, err)
        }

        // Check if the source path exists
        if _, err := os.Stat(absPath); err != nil {
                return nil, fmt.Errorf("source path does not exist: %s: %w", absPath, err)
        }

        // For Go projects, use the current working directory as the build context
        // to ensure go.mod and the entire project structure is available
        currentDir, err := os.Getwd()
        if err != nil {
                return nil, fmt.Errorf("failed to get current working directory: %w", err)
        }

        dockerfilePath := filepath.Join(currentDir, "Dockerfile")

        slog.Debug("Using current working directory as build context", "dir", currentDir)

        return &buildContext{
                Dir:            currentDir,
                DockerfilePath: dockerfilePath,
                CleanupFunc: func() {
                        // Clean up the temporary Dockerfile only if we created it
                        if _, err := os.Stat(dockerfilePath); err == nil {
                                // Check if this is our generated Dockerfile by reading the first few lines
                                // #nosec G304 -- This is a controlled file read operation for cleanup verification
                                if content, readErr := os.ReadFile(dockerfilePath); readErr == nil {
                                        if strings.Contains(string(content), "# Generated by ToolHive") {
                                                if err := os.Remove(dockerfilePath); err != nil {
                                                        slog.Debug("Failed to remove temporary Dockerfile", "error", err)
                                                }
                                        }
                                }
                        }
                },
        }, nil
}

// setupTempBuildContext sets up a temporary build context directory.
func setupTempBuildContext() (*buildContext, error) {
        tempDir, err := os.MkdirTemp("", "toolhive-docker-build-")
        if err != nil {
                return nil, fmt.Errorf("failed to create temporary directory: %w", err)
        }

        dockerfilePath := filepath.Join(tempDir, "Dockerfile")

        slog.Debug("Using temporary directory as build context", "dir", tempDir)

        return &buildContext{
                Dir:            tempDir,
                DockerfilePath: dockerfilePath,
                CleanupFunc: func() {
                        if err := os.RemoveAll(tempDir); err != nil {
                                slog.Debug("Failed to remove temporary directory", "error", err)
                        }
                },
        }, nil
}

// writeDockerfile writes the Dockerfile content to the build context.
// For local paths, it checks if a Dockerfile already exists and avoids overwriting it.
func writeDockerfile(dockerfilePath, dockerfileContent string, isLocalPath bool) error {
        if isLocalPath {
                // Check if a Dockerfile already exists
                if _, err := os.Stat(dockerfilePath); err == nil {
                        slog.Debug("Dockerfile already exists, using existing Dockerfile", "path", dockerfilePath)
                        return nil // Use the existing Dockerfile
                }
        }

        // Add a comment marker to identify our generated Dockerfile
        markedContent := "# Generated by ToolHive - temporary file\n" + dockerfileContent

        if err := os.WriteFile(dockerfilePath, []byte(markedContent), 0600); err != nil {
                return fmt.Errorf("failed to write Dockerfile: %w", err)
        }

        if isLocalPath {
                slog.Debug("Created temporary Dockerfile", "path", dockerfilePath)
        }

        return nil
}

// writeCACertificate writes the CA certificate to the build context if provided.
func writeCACertificate(buildContextDir, caCertContent string, isLocalPath bool) (func(), error) {
        if caCertContent == "" {
                return func() {}, nil
        }

        caCertFilePath := filepath.Join(buildContextDir, "ca-cert.crt")
        if err := os.WriteFile(caCertFilePath, []byte(caCertContent), 0600); err != nil {
                return nil, fmt.Errorf("failed to write CA certificate file: %w", err)
        }

        slog.Debug("Added CA certificate to build context", "path", caCertFilePath)

        var cleanupFunc func()
        if isLocalPath {
                // For local paths, clean up the CA certificate file after build
                cleanupFunc = func() {
                        if err := os.Remove(caCertFilePath); err != nil {
                                slog.Debug("Failed to remove temporary CA certificate", "error", err)
                        }
                }
        } else {
                // For temp directories, no specific cleanup needed (handled by build context cleanup)
                cleanupFunc = func() {}
        }

        return cleanupFunc, nil
}

// writeAuthFiles writes auth files to the build context.
// Returns a cleanup function to remove the files after build.
func writeAuthFiles(buildContextDir string, authFiles map[string]string, isLocalPath bool) (func(), error) {
        if len(authFiles) == 0 {
                return func() {}, nil
        }

        // Map of auth file types to their filenames in the build context
        authFileNames := map[string]string{
                "npmrc":  ".npmrc",
                "netrc":  ".netrc",
                "yarnrc": ".yarnrc",
        }

        var writtenFiles []string
        for fileType, content := range authFiles {
                filename, ok := authFileNames[fileType]
                if !ok {
                        continue
                }

                filePath := filepath.Join(buildContextDir, filename)
                if err := os.WriteFile(filePath, []byte(content), 0600); err != nil {
                        // Clean up any files we've written so far
                        for _, f := range writtenFiles {
                                _ = os.Remove(f)
                        }
                        return nil, fmt.Errorf("failed to write auth file %s: %w", filename, err)
                }
                writtenFiles = append(writtenFiles, filePath)
                slog.Debug("Added auth file to build context", "path", filePath)
        }

        var cleanupFunc func()
        if isLocalPath {
                cleanupFunc = func() {
                        for _, f := range writtenFiles {
                                if err := os.Remove(f); err != nil {
                                        slog.Debug("Failed to remove temporary auth file", "path", f, "error", err)
                                }
                        }
                }
        } else {
                cleanupFunc = func() {}
        }

        return cleanupFunc, nil
}

// generateImageName generates a unique Docker image name based on the package and transport type.
func generateImageName(transportType templates.TransportType, packageName string) string {
        tag := time.Now().Format("20060102150405")
        return strings.ToLower(fmt.Sprintf("toolhivelocal/%s-%s:%s",
                string(transportType),
                PackageNameToImageName(packageName),
                tag))
}

// buildImageFromTemplateWithName builds a Docker image from the template data with a custom image name.
// If imageName is empty, a default name will be generated.
func buildImageFromTemplateWithName(
        ctx context.Context,
        imageManager images.ImageManager,
        transportType templates.TransportType,
        packageName string,
        templateData templates.TemplateData,
        imageName string,
) (string, error) {

        // Get the Dockerfile content
        dockerfileContent, err := templates.GetDockerfileTemplate(transportType, templateData)
        if err != nil {
                return "", fmt.Errorf("failed to get Dockerfile template: %w", err)
        }

        // Set up the build context
        buildCtx, err := setupBuildContext(packageName, templateData.IsLocalPath)
        if err != nil {
                return "", err
        }
        defer buildCtx.CleanupFunc()

        // Write the Dockerfile
        if err := writeDockerfile(buildCtx.DockerfilePath, dockerfileContent, templateData.IsLocalPath); err != nil {
                return "", err
        }

        // Write CA certificate if provided
        caCertCleanup, err := writeCACertificate(buildCtx.Dir, templateData.CACertContent, templateData.IsLocalPath)
        if err != nil {
                return "", err
        }
        defer caCertCleanup()

        // Write auth files if provided
        authFilesCleanup, err := writeAuthFiles(buildCtx.Dir, templateData.BuildAuthFiles, templateData.IsLocalPath)
        if err != nil {
                return "", err
        }
        defer authFilesCleanup()

        // Use provided image name or generate one
        finalImageName := imageName
        if finalImageName == "" {
                finalImageName = generateImageName(transportType, packageName)
        } else {
                // Validate the provided image name using go-containerregistry
                ref, err := nameref.ParseReference(finalImageName)
                if err != nil {
                        return "", fmt.Errorf("invalid image name format '%s': %w", finalImageName, err)
                }
                // Use the normalized reference string
                finalImageName = ref.String()
                slog.Debug("Using validated image name", "image", finalImageName)
        }

        // Log the build process
        slog.Debug("Building Docker image for package", "transport_type", transportType, "package", packageName)
        slog.Debug("Using Dockerfile", "dockerfile_content", dockerfileContent)

        if err := imageManager.BuildImage(ctx, buildCtx.Dir, finalImageName); err != nil {
                return "", fmt.Errorf("failed to build Docker image: %w", err)
        }
        slog.Debug("Successfully built Docker image", "image", finalImageName)

        return finalImageName, nil
}

// PackageNameToImageName replaces slashes with dashes to create a valid Docker image name. If there
// is a version in the package name, the @ is replaced with a dash.
// For local paths, we clean up the path to make it a valid image name.
func PackageNameToImageName(packageName string) string {
        imageName := packageName

        // Handle local paths by cleaning them up
        imageName = strings.TrimPrefix(imageName, "./")
        imageName = strings.TrimPrefix(imageName, "../")

        // Replace problematic characters
        imageName = strings.ReplaceAll(imageName, "/", "-")
        imageName = strings.ReplaceAll(imageName, "@", "-")
        imageName = strings.ReplaceAll(imageName, ".", "-")

        // Ensure the name doesn't start with a dash
        imageName = strings.TrimPrefix(imageName, "-")

        // If the name is empty after cleaning, use a default
        if imageName == "" || imageName == "-" {
                imageName = "toolhive-container"
        }

        return imageName
}

// isLocalGoPath checks if the given path is a local Go path that should be copied into the container.
// Local paths start with "." (relative) or "/" (absolute).
func isLocalGoPath(path string) bool {
        return strings.HasPrefix(path, "./") || strings.HasPrefix(path, "../") || strings.HasPrefix(path, "/") || path == "."
}

// IsImageProtocolScheme checks if the serverOrImage string contains a protocol scheme (uvx://, npx://, or go://)
func IsImageProtocolScheme(serverOrImage string) bool {
        return strings.HasPrefix(serverOrImage, UVXScheme) ||
                strings.HasPrefix(serverOrImage, NPXScheme) ||
                strings.HasPrefix(serverOrImage, GOScheme)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package runner provides functionality for running MCP servers
package runner

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

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive-core/redisconn"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/remote"
        authsecrets "github.com/stacklok/toolhive/pkg/auth/secrets"
        "github.com/stacklok/toolhive/pkg/auth/upstreamtoken"
        authserverrunner "github.com/stacklok/toolhive/pkg/authserver/runner"
        "github.com/stacklok/toolhive/pkg/authserver/server/keys"
        "github.com/stacklok/toolhive/pkg/client"
        "github.com/stacklok/toolhive/pkg/config"
        ct "github.com/stacklok/toolhive/pkg/container"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/diagnostics"
        "github.com/stacklok/toolhive/pkg/labels"
        "github.com/stacklok/toolhive/pkg/process"
        "github.com/stacklok/toolhive/pkg/runtime"
        "github.com/stacklok/toolhive/pkg/secrets"
        "github.com/stacklok/toolhive/pkg/state"
        "github.com/stacklok/toolhive/pkg/telemetry"
        "github.com/stacklok/toolhive/pkg/transport"
        headerfwd "github.com/stacklok/toolhive/pkg/transport/middleware"
        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/workloads/statuses"
)

// ErrContainerExitedRestartNeeded is returned when a container exits and needs to be restarted
var ErrContainerExitedRestartNeeded = errors.New("container exited, restart needed")

// probeProtocolVersion is the MCP protocol revision the readiness probe advertises
// in the JSON-RPC params.protocolVersion of its initialize handshake.
//
// It is deliberately pinned to 2025-11-25 — the current stable revision and the
// newest one that still defines the `initialize` method — rather than tracking
// mcp.LATEST_PROTOCOL_VERSION. The upcoming 2026-07-28 revision (draft; subject to
// change until it ships) removes `initialize` entirely (replaced by server/discover
// + per-request _meta, SEP-2575), so a probe that sent method:"initialize" with a
// 2026-07-28 version string would be internally inconsistent. When ToolHive adopts
// 2026-07-28 (issue #5754) the probe must switch to server/discover — which, like
// all 2026-07-28 Streamable HTTP POSTs, carries the required Mcp-Method/Mcp-Name
// headers — and fall back to this initialize path for older backends. Pinning here
// keeps the probe's method and version consistent regardless of what the SDK later
// declares "latest".
//
// The value is sent only in the request body, NOT as an MCP-Protocol-Version
// header: the spec scopes that header to requests made AFTER initialization
// (carrying the negotiated version), and requires a server to reject an unsupported
// header value with HTTP 400. Sending it on the initialize request itself would let
// a backend that only supports an older revision 400 the probe — a false
// "not-ready" — whereas body-based version negotiation degrades gracefully (the
// server answers HTTP 200 with its own supported version).
const probeProtocolVersion = "2025-11-25"

// Runner is responsible for running an MCP server with the provided configuration
type Runner struct {
        // Config is the configuration for the runner
        Config *RunConfig

        // telemetryProvider is the OpenTelemetry provider for cleanup
        telemetryProvider *telemetry.Provider

        // supportedMiddleware is a map of supported middleware types to their factory functions.
        supportedMiddleware map[string]types.MiddlewareFactory

        // middlewares is a slice of created middleware instances for cleanup
        middlewares []types.Middleware

        // namedMiddlewares is a slice of named middleware to apply to the transport
        namedMiddlewares []types.NamedMiddleware

        // authInfoHandler is the authentication info handler set by auth middleware
        authInfoHandler http.Handler

        // prometheusHandler is the Prometheus metrics handler set by telemetry middleware
        prometheusHandler http.Handler

        // diagnosticsServer serves the Prometheus /metrics endpoint on its own
        // listener, keeping it off the application listener where it would bypass
        // the middleware chain. Nil unless the metrics path is enabled.
        diagnosticsServer *diagnostics.Server

        statusManager statuses.StatusManager

        // authenticatedTokenSource is the wrapped token source for remote workloads with authentication monitoring
        authenticatedTokenSource *auth.MonitoredTokenSource

        // monitoringCtx is the context for background authentication monitoring
        // It is cancelled during Cleanup() to stop monitoring
        monitoringCtx    context.Context
        monitoringCancel context.CancelFunc

        // embeddedAuthServer is the embedded OAuth/OIDC authorization server.
        // Only initialized when Config.EmbeddedAuthServerConfig is set.
        embeddedAuthServer *authserverrunner.EmbeddedAuthServer

        // upstreamTokenReader provides read-only access to upstream tokens for
        // identity enrichment in auth middleware. Set when the embedded auth
        // server is initialized in Run().
        // Nil when no embedded auth server is configured.
        upstreamTokenReader upstreamtoken.TokenReader

        // keyProvider provides in-process JWKS key lookups from the embedded
        // auth server, eliminating self-referential HTTP calls.
        // Nil when no embedded auth server is configured.
        keyProvider keys.PublicKeyProvider
}

// statusManagerAdapter adapts statuses.StatusManager to auth.StatusUpdater interface
type statusManagerAdapter struct {
        sm statuses.StatusManager
}

func (a *statusManagerAdapter) SetWorkloadStatus(
        ctx context.Context,
        workloadName string,
        status rt.WorkloadStatus,
        reason string,
) error {
        slog.Debug("setting workload status", "workload", workloadName, "status", status, "reason", reason)
        return a.sm.SetWorkloadStatus(ctx, workloadName, status, reason)
}

// NewRunner creates a new Runner with the provided configuration
func NewRunner(runConfig *RunConfig, statusManager statuses.StatusManager) *Runner {
        return &Runner{
                Config:              runConfig,
                statusManager:       statusManager,
                supportedMiddleware: GetSupportedMiddlewareFactories(),
        }
}

// AddMiddleware adds a middleware instance and its function to the runner with a name
func (r *Runner) AddMiddleware(name string, middleware types.Middleware) {
        r.middlewares = append(r.middlewares, middleware)
        r.namedMiddlewares = append(r.namedMiddlewares, types.NamedMiddleware{
                Name:     name,
                Function: middleware.Handler(),
        })
}

// SetAuthInfoHandler sets the authentication info handler
func (r *Runner) SetAuthInfoHandler(handler http.Handler) {
        r.authInfoHandler = handler
}

// SetPrometheusHandler sets the Prometheus metrics handler
func (r *Runner) SetPrometheusHandler(handler http.Handler) {
        r.prometheusHandler = handler
}

// GetConfig returns a config interface for middleware to access runner configuration
func (r *Runner) GetConfig() types.RunnerConfig {
        return r.Config
}

// GetUpstreamTokenReader returns the UpstreamTokenReader for identity
// enrichment in the auth middleware. Returns nil if no embedded auth
// server is configured.
func (r *Runner) GetUpstreamTokenReader() upstreamtoken.TokenReader {
        return r.upstreamTokenReader
}

// GetKeyProvider returns the embedded auth server's public key provider
// for in-process JWKS key lookups. Returns nil if no embedded auth server
// is configured.
func (r *Runner) GetKeyProvider() keys.PublicKeyProvider {
        return r.keyProvider
}

// GetName returns the name of the mcp-service from the runner config (implements types.RunnerConfig)
func (c *RunConfig) GetName() string {
        return c.Name
}

// GetPort returns the port from the runner config (implements types.RunnerConfig)
func (c *RunConfig) GetPort() int {
        return c.Port
}

// parseProxyTimeout parses an optional Go duration string from RunConfig. An
// empty string yields 0, which the proxy treats as "use the package default".
// Negative durations and unparsable values are rejected.
func parseProxyTimeout(name, value string) (time.Duration, error) {
        if value == "" {
                return 0, nil
        }
        d, err := time.ParseDuration(value)
        if err != nil {
                return 0, fmt.Errorf("invalid %s %q: %w", name, value, err)
        }
        if d < 0 {
                return 0, fmt.Errorf("%s must be non-negative, got %s", name, d)
        }
        return d, nil
}

// Run runs the MCP server with the provided configuration
//
//nolint:gocyclo // This function is complex but manageable
func (r *Runner) Run(ctx context.Context) error {
        // Validate serialized RunConfig input before any provider or runtime work.
        // The builder and middleware helper repeat this check at their own boundaries.
        if r.Config.MaxRequestBodySize < 0 {
                return fmt.Errorf("max_request_body_size must be non-negative, got %d", r.Config.MaxRequestBodySize)
        }

        // Resolve session TTL once so both the transport proxy and Redis storage use
        // the same effective value, rather than each applying their own zero-fallback
        // independently. SessionTTL is stored as a Go duration string so the
        // runconfig wire format does not depend on nanosecond integers.
        effectiveSessionTTL := session.DefaultSessionTTL
        if r.Config.SessionTTL != "" {
                parsed, err := time.ParseDuration(r.Config.SessionTTL)
                if err != nil {
                        return fmt.Errorf("invalid session_ttl %q: %w", r.Config.SessionTTL, err)
                }
                if parsed < 0 {
                        return fmt.Errorf("session_ttl must be non-negative, got %s", parsed)
                }
                if parsed > 0 {
                        effectiveSessionTTL = parsed
                }
        }

        // Resolve the proxy HTTP server read timeout. Empty/zero means "use the proxy
        // default"; the proxy option ignores non-positive values.
        proxyReadTimeout, parseErr := parseProxyTimeout("proxy_read_timeout", r.Config.ProxyReadTimeout)
        if parseErr != nil {
                return parseErr
        }

        // Create transport with runtime
        transportConfig := types.Config{
                Type:                     r.Config.Transport,
                ProxyPort:                r.Config.Port,
                TargetPort:               r.Config.TargetPort,
                Host:                     r.Config.Host,
                TargetHost:               r.Config.TargetHost,
                Deployer:                 r.Config.Deployer,
                Debug:                    r.Config.Debug,
                TrustProxyHeaders:        r.Config.TrustProxyHeaders,
                StrictProtocolValidation: r.Config.StrictProtocolValidation,
                EndpointPrefix:           r.Config.EndpointPrefix,
                SessionTTL:               effectiveSessionTTL,
                ReadTimeout:              proxyReadTimeout,
                StripConsentCookie:       r.Config.EmbeddedAuthServerConfig != nil,
        }

        // Set proxy mode for stdio transport
        transportConfig.ProxyMode = r.Config.ProxyMode

        // Process secrets before middleware population so that resolved values
        // (e.g., header forward secrets) are available to middleware factories.
        hasRegularSecrets := len(r.Config.Secrets) > 0
        hasRemoteAuthSecret := r.Config.RemoteAuthConfig != nil &&
                (r.Config.RemoteAuthConfig.ClientSecret != "" || r.Config.RemoteAuthConfig.BearerToken != "")
        hasHeaderForwardSecrets := r.Config.HeaderForward != nil && len(r.Config.HeaderForward.AddHeadersFromSecret) > 0

        slog.Debug("secret processing check",
                "has_regular_secrets", hasRegularSecrets,
                "has_remote_auth_secret", hasRemoteAuthSecret,
                "has_header_forward_secrets", hasHeaderForwardSecrets)
        if hasRemoteAuthSecret {
                if r.Config.RemoteAuthConfig.ClientSecret != "" {
                        slog.Debug("remote auth config has client secret configured")
                }
                if r.Config.RemoteAuthConfig.BearerToken != "" {
                        slog.Debug("remote auth config has bearer token configured")
                }
        }

        if hasRegularSecrets || hasRemoteAuthSecret || hasHeaderForwardSecrets {
                slog.Debug("calling WithSecrets to process secrets")
                cfgprovider := config.NewDefaultProvider()
                cfg := cfgprovider.GetConfig()

                providerType, err := cfg.Secrets.GetProviderType()
                if err != nil {
                        return fmt.Errorf("error determining secrets provider type: %w", err)
                }

                systemProvider, err := secrets.CreateProvider(providerType, secrets.WithScope(secrets.ScopeWorkloads))
                if err != nil {
                        return fmt.Errorf("error instantiating system secret manager: %w", err)
                }
                userProvider, err := secrets.CreateProvider(providerType, secrets.WithUserFacing())
                if err != nil {
                        return fmt.Errorf("error instantiating user secret manager: %w", err)
                }

                // Process secrets (including RemoteAuthConfig and header forward secret resolution)
                if _, err = r.Config.WithSecrets(ctx, systemProvider, userProvider); err != nil {
                        return err
                }
        }

        // Populate default middlewares from config fields if not already populated.
        // This runs after WithSecrets so resolved values are available.
        if len(r.Config.MiddlewareConfigs) == 0 {
                if err := PopulateMiddlewareConfigs(r.Config); err != nil {
                        return fmt.Errorf("failed to populate middleware configs: %w", err)
                }
        } else {
                // MiddlewareConfigs was pre-populated (e.g., by WithMiddlewareFromFlags).
                // Header forward is appended here (consistent with PopulateMiddlewareConfigs
                // which also places it at the end) after secret resolution, because
                // secret-backed header values are not available at builder time.
                var err error
                r.Config.MiddlewareConfigs, err = addHeaderForwardMiddleware(r.Config.MiddlewareConfigs, r.Config)
                if err != nil {
                        return fmt.Errorf("failed to add header forward middleware: %w", err)
                }
        }

        if err := canonicalizeOIDCMiddlewareConfig(r.Config); err != nil {
                return fmt.Errorf("invalid OIDC middleware configuration: %w", err)
        }

        if err := validateCredentialStrippingMiddleware(r.Config.MiddlewareConfigs, r.Config); err != nil {
                return fmt.Errorf("invalid credential stripping middleware configuration: %w", err)
        }

        // Origin-header validation (DNS-rebinding protection per MCP 2025-11-25
        // §"Security Warning") is wired here, after both middleware-population
        // paths, because it is the single place where Host/Port/AllowedOrigins are
        // fully resolved: the CLI builder (WithMiddlewareFromFlags) defers port
        // resolution to validateConfig, so the effective port is not known at
        // builder time.
        var err error
        r.Config.MiddlewareConfigs, err = prependOriginMiddleware(r.Config.MiddlewareConfigs, r.Config)
        if err != nil {
                return fmt.Errorf("failed to add origin middleware: %w", err)
        }

        // Body-size limit is always the outermost middleware, regardless of how the
        // chain was assembled (PopulateMiddlewareConfigs above, or WithMiddlewareFromFlags
        // which pre-populates the slice and takes the else branch). Existing body-limit
        // entries are replaced so the typed RunConfig value is authoritative.
        r.Config.MiddlewareConfigs, err = addBodyLimitMiddleware(
                r.Config.MiddlewareConfigs,
                r.Config.MaxRequestBodySize,
        )
        if err != nil {
                return fmt.Errorf("failed to add body limit middleware: %w", err)
        }

        // Initialize embedded auth server if configured.
        // This must happen before middleware creation so that the upstream token
        // service is available to middleware factories (e.g., upstreamswap).
        if r.Config.EmbeddedAuthServerConfig != nil {
                // Proxy runner supports only single-upstream configs; multi-upstream
                // requires VirtualMCPServer.
                if len(r.Config.EmbeddedAuthServerConfig.Upstreams) > 1 {
                        return fmt.Errorf(
                                "proxy runner does not support multiple upstream providers (found %d); "+
                                        "use VirtualMCPServer for multi-upstream deployments",
                                len(r.Config.EmbeddedAuthServerConfig.Upstreams),
                        )
                }

                var err error
                r.embeddedAuthServer, err = authserverrunner.NewEmbeddedAuthServer(
                        ctx, r.Config.EmbeddedAuthServerConfig, authserverrunner.WithListenerHost(r.bindHost()),
                )
                if err != nil {
                        return fmt.Errorf("failed to create embedded auth server: %w", err)
                }
                defer func() {
                        if err := r.embeddedAuthServer.Close(); err != nil {
                                slog.Warn("failed to close embedded auth server", "error", err)
                        }
                }()
                slog.Debug("embedded authorization server initialized")

                // Create the upstream token service eagerly now that the auth server exists.
                // IDPTokenStorage is guaranteed non-nil after successful construction.
                // UpstreamTokenRefresher may be nil if no upstream IDP is configured;
                // InProcessService handles this gracefully (returns ErrNoRefreshToken).
                stor := r.embeddedAuthServer.IDPTokenStorage()
                refresher := r.embeddedAuthServer.UpstreamTokenRefresher()
                r.upstreamTokenReader = upstreamtoken.NewInProcessService(stor, refresher)

                // Expose key provider for in-process JWKS lookups (avoids self-referential HTTP)
                r.keyProvider = r.embeddedAuthServer.KeyProvider()

                // Mount auth server routes at specific prefixes to avoid conflicts with MCP endpoints
                // (e.g., /.well-known/oauth-protected-resource is an MCP endpoint, not auth server)
                transportConfig.PrefixHandlers = r.embeddedAuthServer.Routes()
        }

        // Create middleware from the MiddlewareConfigs instances in the RunConfig.
        for _, middlewareConfig := range r.Config.MiddlewareConfigs {
                // First, get the correct factory function for the middleware type.
                factory, ok := r.supportedMiddleware[middlewareConfig.Type]
                if !ok {
                        return fmt.Errorf("unsupported middleware type: %s", middlewareConfig.Type)
                }

                // Create the middleware instance using the factory function.
                // The factory will add the middleware to the runner and handle any special configuration.
                if err := factory(&middlewareConfig, r); err != nil {
                        return fmt.Errorf("failed to create middleware of type %s: %w", middlewareConfig.Type, err)
                }
        }

        // Set all named middleware and handlers on transport config
        transportConfig.Middlewares = r.namedMiddlewares
        transportConfig.AuthInfoHandler = r.authInfoHandler

        if err := r.startEmbeddedAuthServerListener(); err != nil {
                return err
        }

        // Metrics are served on a dedicated diagnostics listener so access can be
        // restricted by port; see pkg/diagnostics.
        if err := r.startDiagnosticsServer(); err != nil {
                return err
        }

        // During the deprecation window /metrics is ALSO served on the transport port,
        // so an existing scrape configuration keeps working while it is moved. Handing
        // the handler to the transport is what mounts it there; leaving it nil is what
        // keeps the proxies from mounting it at all. See
        // telemetry.DefaultMetricsOnTransportPort for the cutover.
        if mountPrometheusHandlerOnTransportPort(r.prometheusHandler, r.Config.TelemetryConfig) {
                transportConfig.PrometheusHandler = r.prometheusHandler
                slog.Warn("serving prometheus metrics on the transport port as well as the diagnostics port; "+
                        "this is deprecated and the transport-port copy will be removed; see #6384 for the timeline. "+
                        "Move scrapers to the "+
                        "diagnostics address logged above, then set metricsOnTransportPort: false to verify",
                        "transport_port", r.Config.Port)
        }

        // Release the listener on every exit from Run, not just the happy path.
        // Cleanup runs only via stopMCPServer, which several paths skip: the early
        // error returns below, and the graceful container-exit branch that returns
        // nil after the transport already stopped. Leaking here would strand a
        // goroutine and a bound port, and under workloads.Manager's restart loop each
        // attempt would strand another and push the next onto a different port —
        // silently breaking the scrape target this change exists to stabilise.
        //
        // A duplicate stop is harmless: stopDiagnosticsServer is nil-safe and
        // idempotent, so the paths that do reach Cleanup are unaffected.
        defer func() {
                stopCtx, cancel := context.WithTimeout(context.Background(), diagnosticsStopTimeout)
                defer cancel()
                if err := r.stopDiagnosticsServer(stopCtx); err != nil {
                        slog.Warn("failed to stop diagnostics server", "error", err)
                }
        }()

        // Set up the transport
        slog.Debug("setting up transport", "transport", r.Config.Transport)

        // Prepare transport options based on workload type
        var transportOpts []transport.Option
        var setupResult *runtime.SetupResult

        // Check policy gate before creating the server (applies to both local and remote)
        if err := ActivePolicyGate().CheckCreateServer(ctx, r.Config); err != nil {
                return fmt.Errorf("server creation blocked by policy: %w", err)
        }

        if r.Config.RemoteURL == "" {
                // For local workloads, deploy the container using runtime.Setup first
                var scalingConfig *rt.ScalingConfig
                if r.Config.ScalingConfig != nil {
                        scalingConfig = &rt.ScalingConfig{
                                BackendReplicas: r.Config.ScalingConfig.BackendReplicas,
                        }
                }
                result, err := runtime.Setup(
                        ctx,
                        r.Config.Transport,
                        r.Config.Deployer,
                        r.Config.ContainerName,
                        r.Config.Image,
                        r.Config.CmdArgs,
                        r.Config.EnvVars,
                        r.Config.ContainerLabels,
                        r.Config.PermissionProfile,
                        r.Config.K8sPodTemplatePatch,
                        r.Config.IsolateNetwork,
                        r.Config.AllowDockerGateway,
                        r.Config.IgnoreConfig,
                        r.Config.Host,
                        r.Config.TargetPort,
                        r.Config.TargetHost,
                        r.Config.Publish,
                        scalingConfig,
                        r.Config.MCPServerGeneration,
                )
                if err != nil {
                        return fmt.Errorf("failed to set up workload: %w", err)
                }
                setupResult = result

                // Configure the transport with the setup results using options
                transportOpts = append(transportOpts, transport.WithContainerName(setupResult.ContainerName))
                if setupResult.TargetURI != "" {
                        transportOpts = append(transportOpts, transport.WithTargetURI(setupResult.TargetURI))
                }
        }

        // When Redis session storage is configured, create a Redis-backed session store
        // so sessions are shared across proxy replicas instead of being pod-local.
        if r.Config.ScalingConfig != nil && r.Config.ScalingConfig.SessionRedis != nil {
                redisCfg := r.Config.ScalingConfig.SessionRedis
                keyPrefix := redisCfg.KeyPrefix
                if keyPrefix == "" {
                        keyPrefix = "thv:proxy:session:"
                }
                tlsCfg, err := redisCfg.TLS.Load()
                if err != nil {
                        return fmt.Errorf("redis session storage TLS configuration: %w", err)
                }
                storage, err := session.NewRedisStorage(ctx, redisconn.Config{
                        Addr:     redisCfg.Address,
                        Password: os.Getenv(session.RedisPasswordEnvVar),
                        DB:       int(redisCfg.DB),
                        TLS:      tlsCfg,
                }, keyPrefix, effectiveSessionTTL)
                if err != nil {
                        return fmt.Errorf("failed to create Redis session storage: %w", err)
                }
                slog.Info("using Redis session storage",
                        "store", redisCfg.Address,
                        "address", redisCfg.Address,
                        "db", redisCfg.DB,
                        "key_prefix", keyPrefix,
                )
                transportConfig.SessionStorage = storage
        }

        // Create transport with options
        transportHandler, err := transport.NewFactory().Create(transportConfig, transportOpts...)
        if err != nil {
                return fmt.Errorf("failed to create transport: %w", err)
        }

        // For remote MCP servers, set the remote URL on HTTP transports
        if r.Config.RemoteURL != "" {
                transportHandler.SetRemoteURL(r.Config.RemoteURL)

                // Handle remote authentication if configured
                tokenSource, err := r.handleRemoteAuthentication(ctx)
                if err != nil {
                        return fmt.Errorf("failed to authenticate to remote server: %w", err)
                }

                // Wrap the token source with authentication monitoring for remote workloads
                if tokenSource != nil {
                        // Create a child context for monitoring that can be cancelled during cleanup
                        r.monitoringCtx, r.monitoringCancel = context.WithCancel(ctx)
                        // Create adapter to bridge statuses.StatusManager to auth.StatusUpdater
                        adapter := &statusManagerAdapter{sm: r.statusManager}

                        // Capture the upstream issuer and resolved client_id so the DCR
                        // remediation warning emitted by isTransientNetworkError on a
                        // permanent 4xx (indicating a stale RFC 7591 registration)
                        // carries enough context for an operator to identify which
                        // upstream AS + client_id to re-register. Precedence (cached
                        // CIMD > cached DCR > static) lives next to
                        // resolveClientCredentials in pkg/auth/remote so both call
                        // sites stay in sync.
                        upstream, clientID := r.Config.RemoteAuthConfig.LogContext()

                        r.authenticatedTokenSource = auth.NewMonitoredTokenSource(
                                r.monitoringCtx,
                                tokenSource,
                                r.Config.BaseName,
                                upstream,
                                clientID,
                                adapter,
                        )
                        tokenSource = r.authenticatedTokenSource
                        r.authenticatedTokenSource.StartBackgroundMonitoring()
                }

                // Set the token source on the transport
                transportHandler.SetTokenSource(tokenSource)

                // Set the health check failure callback for remote servers
                transportHandler.SetOnHealthCheckFailed(func() {
                        slog.Warn("health check failed for remote server, marking as unhealthy", "server", r.Config.BaseName)
                        // Use Background context for status update callback - this is triggered by health check
                        // failure and is independent of any request context. The callback is fired asynchronously
                        // and needs its own lifecycle separate from the transport's parent context.
                        if err := r.statusManager.SetWorkloadStatus(
                                context.Background(),
                                r.Config.BaseName,
                                rt.WorkloadStatusUnhealthy,
                                "Health check failed",
                        ); err != nil {
                                slog.Error("failed to update workload status", "error", err)
                        }
                })

                // Set the unauthorized response callback for bearer token authentication
                errorMsg := "Bearer token authentication failed. Please restart the server with a new token"
                transportHandler.SetOnUnauthorizedResponse(func() {
                        slog.Warn("received 401 Unauthorized response for remote server, marking as unauthenticated", "server", r.Config.BaseName)
                        // Use Background context for status update callback - this is triggered by 401 response
                        // and is independent of any request context. The callback is fired asynchronously
                        // and needs its own lifecycle separate from the transport's parent context.
                        if err := r.statusManager.SetWorkloadStatus(
                                context.Background(),
                                r.Config.BaseName,
                                rt.WorkloadStatusUnauthenticated,
                                errorMsg,
                        ); err != nil {
                                slog.Error("failed to update workload status", "error", err)
                        }
                })
        }

        // Configure stateless mode if requested. Stateless mode applies to any
        // streamable-HTTP server (remote or local container) where the upstream
        // only accepts POST and does not support SSE-based sessions.
        if r.Config.Stateless {
                httpT, ok := transportHandler.(*transport.HTTPTransport)
                if !ok {
                        return fmt.Errorf("--stateless requires streamable-HTTP or SSE transport, got %T", transportHandler)
                }
                httpT.SetStateless(true)
        }

        // Start the transport (which also starts the container and monitoring)
        slog.Debug("starting transport", "transport", r.Config.Transport, "container", r.Config.ContainerName)
        if err := transportHandler.Start(ctx); err != nil {
                return fmt.Errorf("failed to start transport: %w", err)
        }

        slog.Debug("mcp server started successfully", "container", r.Config.ContainerName)

        // Wait for the MCP server to accept initialize requests before updating client configurations.
        // This prevents timing issues where clients try to connect before the server is fully ready.
        // We repeatedly call initialize until it succeeds (up to 5 minutes).
        // Note: We skip this check for pure STDIO transport because STDIO servers may reject
        // multiple initialize calls (see #1982).
        transportType := labels.GetTransportType(r.Config.ContainerLabels)
        serverURL := transport.GenerateMCPServerURL(
                transportType,
                string(r.Config.ProxyMode),
                "localhost",
                r.Config.Port,
                r.Config.ContainerName,
                r.Config.RemoteURL)

        // Only wait for initialization on non-STDIO transports
        // STDIO servers communicate directly via stdin/stdout and calling initialize multiple times
        // can cause issues as the behavior is not specified by the MCP spec
        if transportType != "stdio" {
                // Repeatedly try calling initialize until it succeeds (up to 5 minutes)
                // Some servers (like mcp-optimizer) can take significant time to start up.
                // When OIDC auth is configured, the local proxy rejects the unauthenticated
                // probe with 401/403, which still indicates the server is ready.
                authExpected := r.Config.OIDCConfig != nil
                if err := waitForInitializeSuccess(ctx, serverURL, transportType, authExpected, 5*time.Minute); err != nil {
                        slog.Warn("initialize not successful, but continuing", "error", err)
                        // Continue anyway to maintain backward compatibility, but log a warning
                }
        } else {
                slog.Debug("skipping initialize check for STDIO transport")
        }

        // Update client configurations with the MCP server URL.
        // Note that this function checks the configuration to determine which
        // clients should be updated, if any.
        clientManager, err := client.NewManager(ctx)
        if err != nil {
                slog.Warn("failed to create client manager", "error", err)
        } else {
                if err := clientManager.AddServerToClients(ctx, r.Config.ContainerName, serverURL, transportType, r.Config.Group); err != nil {
                        slog.Warn("failed to add server to client configurations", "error", err)
                }
        }

        // Define a function to stop the MCP server
        stopMCPServer := func(reason string) {
                // Use Background context for cleanup operations. The parent context may already be
                // cancelled when this cleanup function runs (e.g., on graceful shutdown or context
                // cancellation). We need a fresh context with its own timeout to ensure cleanup
                // operations complete successfully regardless of the parent context state.
                cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 1*time.Minute)
                defer cleanupCancel()
                slog.Debug("stopping MCP server", "reason", reason)

                // Stop the transport (which also stops the container, monitoring, and handles removal)
                slog.Debug("stopping transport", "transport", r.Config.Transport)
                if err := transportHandler.Stop(cleanupCtx); err != nil {
                        slog.Warn("failed to stop transport", "error", err)
                }

                // Cleanup telemetry provider
                if err := r.Cleanup(cleanupCtx); err != nil {
                        slog.Warn("failed to cleanup telemetry", "error", err)
                }

                // Remove the PID file if it exists. Use PID-guarded reset so that a
                // dying process does not clobber the PID of a replacement process that
                // started in the meantime (e.g. during thv rm + thv run).
                if err := r.statusManager.ResetWorkloadPIDIfMatch(cleanupCtx, r.Config.BaseName, os.Getpid()); err != nil {
                        slog.Warn("failed to reset workload PID", "container", r.Config.ContainerName, "error", err)
                }

                slog.Debug("mcp server stopped", "container", r.Config.ContainerName)
        }

        if err := r.statusManager.SetWorkloadPID(ctx, r.Config.BaseName, os.Getpid()); err != nil {
                slog.Warn("failed to set workload PID", "error", err)
        }

        if process.IsDetached() {
                // We're a detached process running in foreground mode
                // Write the PID to a file so the stop command can kill the process
                slog.Info("running as detached process", "pid", os.Getpid())
        } else {
                // Notify that user that the workload has started successfully when using --foreground
                slog.Info("workload started successfully, press Ctrl+C to stop")
        }

        // Create a done channel to signal when the server has been stopped
        doneCh := make(chan struct{})

        // Start a goroutine to monitor the transport's running state
        go func() {
                for {
                        // Safely check if transportHandler is nil
                        if transportHandler == nil {
                                slog.Debug("transport handler is nil, exiting monitoring routine")
                                close(doneCh)
                                return
                        }

                        // Check if the transport is still running
                        running, err := transportHandler.IsRunning()
                        if err != nil {
                                slog.Error("error checking transport status", "error", err)
                                // Don't exit immediately on error, try again after pause
                                time.Sleep(1 * time.Second)
                                continue
                        }
                        if !running {
                                // Transport is no longer running (container exited or was stopped)
                                slog.Warn("transport is no longer running, attempting automatic restart")
                                close(doneCh)
                                return
                        }

                        // Sleep for a short time before checking again
                        time.Sleep(1 * time.Second)
                }
        }()

        // At this point, we can consider the workload started successfully.
        // However, we should preserve unauthenticated status if it was already set
        // (e.g., if bearer token authentication failed during initialization)
        currentWorkload, err := r.statusManager.GetWorkload(ctx, r.Config.BaseName)
        if err != nil && !errors.Is(err, rt.ErrWorkloadNotFound) {
                slog.Warn("failed to get current workload status", "error", err)
        }

        // Only set status to running if it's not already unauthenticated
        // This preserves the unauthenticated state when bearer token authentication fails
        if err == nil && currentWorkload.Status == rt.WorkloadStatusUnauthenticated {
                slog.Debug("preserving unauthenticated status for workload", "workload", r.Config.BaseName)
        } else {
                if err := r.statusManager.SetWorkloadStatus(ctx, r.Config.BaseName, rt.WorkloadStatusRunning, ""); err != nil {
                        // If we can't set the status to `running` - treat it as a fatal error.
                        return fmt.Errorf("failed to set workload status: %w", err)
                }
        }

        // Wait for either a signal or the done channel to be closed
        select {
        case <-ctx.Done():
                stopMCPServer("Context cancelled")
        case <-doneCh:
                // The transport has already been stopped (likely by the container exit)
                // Remove the old PID from the state file. Use PID-guarded reset to
                // avoid clobbering a replacement process's PID.
                if err := r.statusManager.ResetWorkloadPIDIfMatch(ctx, r.Config.BaseName, os.Getpid()); err != nil {
                        slog.Warn("failed to reset workload PID", "workload", r.Config.BaseName, "error", err)
                }

                // Check if workload still exists (using status manager and runtime)
                // If it doesn't exist, it was removed - clean up client config
                // If it exists, it exited unexpectedly - signal restart needed
                exists, checkErr := r.doesWorkloadExist(ctx, r.Config.BaseName)
                if checkErr != nil {
                        slog.Warn("failed to check if workload exists", "error", checkErr)
                        // Assume restart needed if we can't check
                } else if !exists {
                        // Workload doesn't exist in `thv ls` - it was removed
                        slog.Debug("Workload no longer exists, removing from client configurations",
                                "workload", r.Config.BaseName)
                        clientManager, clientErr := client.NewManager(ctx)
                        if clientErr == nil {
                                removeErr := clientManager.RemoveServerFromClients(
                                        ctx,
                                        r.Config.ContainerName,
                                        r.Config.Group,
                                )
                                if removeErr != nil {
                                        slog.Warn("failed to remove from client config", "error", removeErr)
                                } else {
                                        slog.Debug("Successfully removed from client configurations",
                                                "container", r.Config.ContainerName)
                                }
                        }
                        slog.Debug("MCP server stopped and cleaned up", "container", r.Config.ContainerName)
                        return nil // Exit gracefully, no restart
                }

                // Workload still exists - signal restart needed
                slog.Debug("MCP server stopped, restart needed", "container", r.Config.ContainerName)
                return ErrContainerExitedRestartNeeded
        }

        return nil
}

// doesWorkloadExist checks if a workload exists in the status manager and runtime.
// For remote workloads, it trusts the status manager.
// For container workloads, it verifies the container exists in the runtime.
func (r *Runner) doesWorkloadExist(ctx context.Context, workloadName string) (bool, error) {
        // Check if workload exists by trying to get it from status manager
        workload, err := r.statusManager.GetWorkload(ctx, workloadName)
        if err != nil {
                if errors.Is(err, rt.ErrWorkloadNotFound) {
                        return false, nil
                }
                return false, fmt.Errorf("failed to check if workload exists: %w", err)
        }

        // If remote workload, check if it should exist
        if workload.Remote {
                // For remote workloads, trust the status manager
                return workload.Status != rt.WorkloadStatusError, nil
        }

        // For container workloads, verify the container actually exists in the runtime
        // Create a runtime instance to check if container exists
        backend, err := ct.NewFactory().Create(ctx)
        if err != nil {
                slog.Warn("Failed to create runtime to check container existence", "error", err)
                // Fall back to status manager only
                return workload.Status != rt.WorkloadStatusError, nil
        }

        // Check if container exists in the runtime (not just running)
        // GetWorkloadInfo will return an error if the container doesn't exist
        _, err = backend.GetWorkloadInfo(ctx, workloadName)
        if err != nil {
                // Container doesn't exist
                slog.Debug("Container not found in runtime", "workload", workloadName, "error", err)
                return false, nil
        }

        // Container exists (may be running or stopped)
        return true, nil
}

// handleRemoteAuthentication handles authentication for remote MCP servers
func (r *Runner) handleRemoteAuthentication(ctx context.Context) (oauth2.TokenSource, error) {
        if r.Config.RemoteAuthConfig == nil {
                return nil, nil
        }

        // Get the secret manager for token storage
        secretManager, err := authsecrets.GetSecretsManager()
        if err != nil {
                // Secret manager not available - log warning but continue
                // OAuth will work but tokens won't be persisted across restarts
                slog.Warn("Secret manager not available, OAuth tokens will not be persisted", "error", err)
        }

        // Create remote authentication handler
        authHandler := remote.NewHandler(r.Config.RemoteAuthConfig)

        // Set the secret provider for retrieving cached tokens
        if secretManager != nil {
                authHandler.SetSecretProvider(secretManager)
        }

        // Set up token persister to save tokens across restarts
        if secretManager != nil {
                authHandler.SetTokenPersister(func(refreshToken string, expiry time.Time) error {
                        return r.persistRefreshToken(ctx, secretManager, refreshToken, expiry)
                })

                // Set up client credentials persister for DCR (Dynamic Client Registration)
                authHandler.SetClientCredentialsPersister(func(
                        clientID, clientSecret string,
                        secretExpiry time.Time,
                        regAccessToken, regClientURI string,
                        tokenEndpointAuthMethod string,
                        registeredCallbackPort int,
                ) error {
                        return r.persistClientCredentials(
                                ctx, secretManager, clientID, clientSecret,
                                secretExpiry, regAccessToken, regClientURI, tokenEndpointAuthMethod, registeredCallbackPort)
                })
        }

        // Perform authentication
        tokenSource, err := authHandler.Authenticate(ctx, r.Config.RemoteURL)
        if err != nil {
                return nil, fmt.Errorf("remote authentication failed: %w", err)
        }

        return tokenSource, nil
}

// decodeStoredRunConfig decodes stored state without ReadJSON: its migrations
// save state themselves and its normalizations are meant to stay in memory.
func decodeStoredRunConfig(reader io.Reader) (*RunConfig, error) {
        var stored *RunConfig
        if err := json.NewDecoder(reader).Decode(&stored); err != nil {
                // Custom unmarshalers (remote.Config, time.Time) can echo input.
                return nil, fmt.Errorf("invalid stored run configuration: %T", err)
        }
        if stored == nil {
                return nil, errors.New("stored run configuration is empty")
        }
        return stored, nil
}

// loadAndScrubRunConfig returns the stored state with legacy resolved secret
// values removed. It changes neither the live config nor the state on disk.
func (r *Runner) loadAndScrubRunConfig(ctx context.Context) (*RunConfig, error) {
        stored, err := state.LoadRunConfig(ctx, r.Config.BaseName, decodeStoredRunConfig)
        if err != nil {
                return nil, fmt.Errorf("failed to load stored run configuration for %q: %w", r.Config.BaseName, err)
        }
        for i, reference := range stored.Secrets {
                parameter, err := secrets.ParseSecretParameter(reference)
                if err != nil {
                        return nil, fmt.Errorf("invalid stored secret reference at index %d", i)
                }
                delete(stored.EnvVars, parameter.Target)
        }
        if stored.HeaderForward != nil {
                if err := scrubStoredHeaderParameters(stored.MiddlewareConfigs, stored.HeaderForward.AddHeadersFromSecret); err != nil {
                        return nil, err
                }
        }
        return stored, nil
}

// scrubStoredHeaderParameters removes secret headers from header-forward
// entries, modifying the middlewares slice in place; pass a freshly decoded copy.
// Params are decoded into the struct the middleware factory uses, so duplicate
// or case-variant add_headers members resolve as they do at runtime.
func scrubStoredHeaderParameters(middlewares []types.MiddlewareConfig, secretHeaders map[string]string) error {
        secret := make(map[string]bool, len(secretHeaders))
        for name := range secretHeaders {
                secret[http.CanonicalHeaderKey(name)] = true
        }
        for i := range middlewares {
                if middlewares[i].Type != headerfwd.HeaderForwardMiddlewareName || len(middlewares[i].Parameters) == 0 {
                        continue
                }
                var params headerfwd.HeaderForwardMiddlewareParams
                if err := json.Unmarshal(middlewares[i].Parameters, &params); err != nil {
                        return fmt.Errorf("invalid stored header forwarding parameters: %T", err)
                }
                maps.DeleteFunc(params.AddHeaders, func(name, _ string) bool {
                        return secret[http.CanonicalHeaderKey(name)]
                })
                encoded, err := json.Marshal(params)
                if err != nil {
                        return fmt.Errorf("failed to encode stored header forwarding parameters: %T", err)
                }
                middlewares[i].Parameters = encoded
        }
        return nil
}

func (r *Runner) persistRefreshToken(
        ctx context.Context,
        secretManager secrets.Provider,
        refreshToken string,
        expiry time.Time,
) error {
        updatedConfig, err := r.loadAndScrubRunConfig(ctx)
        if err != nil {
                return err
        }
        secretName, err := authsecrets.GenerateUniqueSecretNameWithPrefix(
                r.Config.Name,
                "OAUTH_REFRESH_TOKEN_",
                secretManager,
        )
        if err != nil {
                return fmt.Errorf("failed to generate secret name: %w", err)
        }

        if err := authsecrets.StoreSecretInManagerWithProvider(ctx, secretName, refreshToken, secretManager); err != nil {
                return fmt.Errorf("failed to store refresh token: %w", err)
        }

        updatedRemoteAuthConfig := *r.Config.RemoteAuthConfig
        updatedConfig.RemoteAuthConfig = &updatedRemoteAuthConfig
        updatedRemoteAuthConfig.CachedRefreshTokenRef = secretName
        updatedRemoteAuthConfig.CachedTokenExpiry = expiry

        if err := updatedConfig.SaveState(ctx); err != nil {
                return fmt.Errorf("failed to save config with token reference: %w", err)
        }

        // Preserve pointer identity so the auth handler and runner observe the same
        // complete config only after the durable write succeeds.
        *r.Config.RemoteAuthConfig = updatedRemoteAuthConfig

        slog.Debug("Stored OAuth refresh token in secret manager", "secret_name", secretName)
        return nil
}

func (r *Runner) persistClientCredentials(
        ctx context.Context,
        secretManager secrets.Provider,
        clientID, clientSecret string,
        secretExpiry time.Time,
        regAccessToken, regClientURI string,
        tokenEndpointAuthMethod string,
        registeredCallbackPort int,
) error {
        updatedConfig, err := r.loadAndScrubRunConfig(ctx)
        if err != nil {
                return err
        }
        updatedRemoteAuthConfig := remote.Config{}
        if r.Config.RemoteAuthConfig != nil {
                updatedRemoteAuthConfig = *r.Config.RemoteAuthConfig
        }
        updatedConfig.RemoteAuthConfig = &updatedRemoteAuthConfig
        updatedRemoteAuthConfig.CachedClientID = clientID

        if clientSecret != "" {
                clientSecretSecretName := updatedRemoteAuthConfig.CachedClientSecretRef
                if clientSecretSecretName == "" {
                        clientSecretSecretName, err = authsecrets.GenerateUniqueSecretNameWithPrefix(
                                r.Config.Name,
                                "OAUTH_CLIENT_SECRET_",
                                secretManager,
                        )
                        if err != nil {
                                return fmt.Errorf("failed to generate client secret secret name: %w", err)
                        }
                }

                if err := authsecrets.StoreSecretInManagerWithProvider(ctx, clientSecretSecretName, clientSecret, secretManager); err != nil {
                        return fmt.Errorf("failed to store client secret: %w", err)
                }
                updatedRemoteAuthConfig.CachedClientSecretRef = clientSecretSecretName
        }

        updatedRemoteAuthConfig.CachedSecretExpiry = secretExpiry

        if regAccessToken != "" {
                regTokenSecretName := updatedRemoteAuthConfig.CachedRegTokenRef
                if regTokenSecretName == "" {
                        regTokenSecretName, err = authsecrets.GenerateUniqueSecretNameWithPrefix(r.Config.Name, "OAUTH_REG_TOKEN_", secretManager)
                        if err != nil {
                                return fmt.Errorf("failed to generate registration token secret name: %w", err)
                        }
                }

                if err := authsecrets.StoreSecretInManagerWithProvider(ctx, regTokenSecretName, regAccessToken, secretManager); err != nil {
                        return fmt.Errorf("failed to store registration access token: %w", err)
                }
                updatedRemoteAuthConfig.CachedRegTokenRef = regTokenSecretName
                slog.Debug("Stored DCR registration access token for RFC 7592 operations")
        }

        updatedRemoteAuthConfig.CachedRegClientURI = regClientURI
        updatedRemoteAuthConfig.CachedTokenEndpointAuthMethod = tokenEndpointAuthMethod
        updatedRemoteAuthConfig.CachedDCRCallbackPort = registeredCallbackPort

        if err := updatedConfig.SaveState(ctx); err != nil {
                return fmt.Errorf("failed to save config with client credentials: %w", err)
        }

        // Preserve pointer identity so the auth handler and runner continue to
        // observe the same complete config after the durable write succeeds.
        if r.Config.RemoteAuthConfig == nil {
                r.Config.RemoteAuthConfig = &updatedRemoteAuthConfig
        } else {
                *r.Config.RemoteAuthConfig = updatedRemoteAuthConfig
        }

        slog.Debug("Stored DCR client credentials", "client_id", clientID,
                "has_expiry", !secretExpiry.IsZero(),
                "has_reg_token", regAccessToken != "",
                "has_reg_uri", regClientURI != "")
        return nil
}

// Cleanup performs cleanup operations for the runner, including shutting down all middleware.
func (r *Runner) Cleanup(ctx context.Context) error {
        // For simplicity, return the last error we encounter during cleanup.
        var lastErr error

        // Clean up all middleware instances
        for i, middleware := range r.middlewares {
                if err := middleware.Close(); err != nil {
                        slog.Warn("Failed to close middleware", "index", i, "error", err)
                        lastErr = err
                }
        }

        // Close embedded auth server
        if r.embeddedAuthServer != nil {
                if err := r.embeddedAuthServer.Close(); err != nil {
                        slog.Warn("Failed to close embedded auth server", "error", err)
                        if lastErr == nil {
                                lastErr = err
                        }
                }
        }

        // Stop the diagnostics listener before the telemetry provider it scrapes
        // from is shut down.
        if err := r.stopDiagnosticsServer(ctx); err != nil {
                slog.Warn("Failed to stop diagnostics server", "error", err)
                if lastErr == nil {
                        lastErr = err
                }
        }

        // Legacy telemetry provider cleanup (will be removed when telemetry middleware handles it)
        if r.telemetryProvider != nil {
                slog.Debug("Shutting down telemetry provider")
                if err := r.telemetryProvider.Shutdown(ctx); err != nil {
                        slog.Warn("failed to shutdown telemetry provider", "error", err)
                        lastErr = err
                }
        }

        // Stop background authentication monitoring for remote workloads
        // Cancel the monitoring context to stop the background goroutine
        if r.monitoringCancel != nil {
                r.monitoringCancel()
                r.monitoringCancel = nil
        }

        return lastErr
}

// isReadyStatus reports whether an HTTP status code from the readiness probe indicates
// the MCP server is ready. A 200 always means ready. When authExpected is true, a 401 or
// 403 also means ready: it proves the local proxy listener is up and its auth middleware
// is enforcing credentials against the unauthenticated probe.
//
// Today ToolHive's OIDC validator rejects the tokenless probe with 401; 403 is accepted
// defensively to cover upstream IdP or edge behavior, not any current ToolHive code path.
func isReadyStatus(statusCode int, authExpected bool) bool {
        if statusCode == http.StatusOK {
                return true
        }
        return authExpected && (statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden)
}

// waitForInitializeSuccess repeatedly checks if the MCP server is ready to accept requests.
// This prevents timing issues where clients try to connect before the server is fully ready.
// It makes repeated attempts with exponential backoff up to a maximum timeout.
//
// The probe is unauthenticated and targets ToolHive's own local proxy. When authExpected
// is true (the workload has OIDC configured), an HTTP 401 or 403 is treated as "ready":
// it proves the proxy listener is up and the auth middleware is enforcing credentials. When
// authExpected is false, only HTTP 200 is accepted.
//
// Note that in the OIDC case readiness reflects only the proxy/auth layer being up — the
// auth middleware rejects the probe before the request reaches the backend, so this probe
// does not confirm backend MCP server initialization.
//
// Note: This function should not be called for STDIO transport.
func waitForInitializeSuccess(
        ctx context.Context,
        serverURL, transportType string,
        authExpected bool,
        maxWaitTime time.Duration,
) error {
        // Determine the endpoint and method to use based on transport type
        var endpoint string
        var method string
        var payload string

        switch transportType {
        case "streamable-http", "streamable":
                // For streamable-http, send initialize request to /mcp endpoint
                // Format: http://localhost:port/mcp
                endpoint = serverURL
                method = "POST"
                payload = fmt.Sprintf(
                        `{"jsonrpc":"2.0","method":"initialize","id":"toolhive-init-check",`+
                                `"params":{"protocolVersion":%q,"capabilities":{},`+
                                `"clientInfo":{"name":"toolhive","version":"1.0"}}}`,
                        probeProtocolVersion,
                )
        case "sse":
                // For SSE, just check if the SSE endpoint is available
                // We can't easily call initialize without establishing a full SSE connection,
                // so we just verify the endpoint responds.
                // Format: http://localhost:port/sse#container-name -> http://localhost:port/sse
                endpoint = serverURL
                // Remove fragment if present (everything after #)
                if idx := strings.Index(endpoint, "#"); idx != -1 {
                        endpoint = endpoint[:idx]
                }
                method = "GET"
                payload = ""
        default:
                // For other transports, no HTTP check is needed
                slog.Debug("Skipping readiness check for transport type", "transport", transportType)
                return nil
        }

        // Setup retry logic with exponential backoff
        startTime := time.Now()
        attempt := 0
        delay := 100 * time.Millisecond
        maxDelay := 2 * time.Second // Cap at 2 seconds between retries

        // Per-attempt outcomes below log at DEBUG, which is invisible in a default
        // deployment — a workload stuck in `starting` for minutes then shows no
        // trace of WHY the probe kept failing. Surface a periodic INFO progress
        // line (long-running operation) carrying the last observed outcome.
        // Every loop iteration assigns lastObserved before it is read.
        var lastObserved string
        const progressInterval = 30 * time.Second
        nextProgressLog := progressInterval

        slog.Info("Waiting for MCP server to be ready", "endpoint", endpoint, "timeout", maxWaitTime)

        // Create HTTP client with a reasonable timeout for requests
        httpClient := &http.Client{
                Timeout: 10 * time.Second,
        }

        for {
                attempt++

                // Make the readiness check request
                var req *http.Request
                var err error
                if payload != "" {
                        req, err = http.NewRequestWithContext(ctx, method, endpoint, bytes.NewBufferString(payload))
                } else {
                        req, err = http.NewRequestWithContext(ctx, method, endpoint, nil)
                }

                if err != nil {
                        slog.Debug("Failed to create request", "attempt", attempt, "error", err)
                        lastObserved = fmt.Sprintf("failed to create request: %v", err)
                } else {
                        if method == "POST" {
                                req.Header.Set("Content-Type", "application/json")
                                req.Header.Set("Accept", "application/json, text/event-stream")
                                // No MCP-Protocol-Version header: it is scoped to post-initialize
                                // requests (carrying the negotiated version) and a server MUST reject
                                // an unsupported value with HTTP 400. The initialize body's
                                // protocolVersion negotiates gracefully instead. See probeProtocolVersion.
                        }

                        resp, err := httpClient.Do(req) // #nosec G704 -- endpoint is the local MCP server readiness URL
                        if err == nil {
                                //nolint:errcheck // Ignoring close error on response body in error path
                                defer resp.Body.Close()

                                if isReadyStatus(resp.StatusCode, authExpected) {
                                        elapsed := time.Since(startTime)
                                        slog.Debug("MCP server is ready", //nolint:gosec // G706: status code and attempt are integers
                                                "elapsed", elapsed, "attempt", attempt, "status_code", resp.StatusCode)
                                        return nil
                                }

                                slog.Debug("Server returned status", //nolint:gosec // G706: status code and attempt are integers
                                        "status_code", resp.StatusCode, "attempt", attempt)
                                lastObserved = fmt.Sprintf("HTTP %d", resp.StatusCode)
                        } else {
                                slog.Debug("Failed to reach endpoint", "attempt", attempt, "error", err)
                                lastObserved = fmt.Sprintf("unreachable: %v", err)
                        }
                }

                // Check if we've exceeded the maximum wait time
                elapsed := time.Since(startTime)
                if elapsed >= maxWaitTime {
                        return fmt.Errorf("initialize not successful after %v (%d attempts, last observed: %s)",
                                elapsed, attempt, lastObserved)
                }

                if elapsed >= nextProgressLog {
                        slog.Info("Still waiting for MCP server to be ready", //nolint:gosec // G706: attempt is an integer
                                "endpoint", endpoint, "elapsed", elapsed.Round(time.Second), "attempt", attempt, "last_observed", lastObserved)
                        nextProgressLog += progressInterval
                }

                // Wait before retrying
                select {
                case <-ctx.Done():
                        return fmt.Errorf("context cancelled while waiting for initialize")
                case <-time.After(delay):
                        // Continue to next attempt
                }

                // Update delay for next iteration with exponential backoff
                delay *= 2
                if delay > maxDelay {
                        delay = maxDelay
                }
        }
}

// startEmbeddedAuthServerListener starts the embedded auth server's TLS
// listener when one is configured, and is a no-op otherwise. The listener port
// is fixed and part of the issuer URL, so a clash with the proxy or diagnostics
// port is reported instead of being resolved by picking another port.
func (r *Runner) startEmbeddedAuthServerListener() error {
        if r.embeddedAuthServer == nil || r.Config.EmbeddedAuthServerConfig == nil ||
                r.Config.EmbeddedAuthServerConfig.TLSListener == nil {
                return nil
        }
        if r.Config.Port == authserverrunner.TLSListenerPort {
                return fmt.Errorf("proxy port %d conflicts with the auth server TLS listener port", r.Config.Port)
        }
        if r.prometheusHandler != nil && diagnosticsPort(r.Config.TelemetryConfig) == authserverrunner.TLSListenerPort {
                return fmt.Errorf("diagnostics port %d conflicts with the auth server TLS listener port", authserverrunner.TLSListenerPort)
        }
        if err := r.embeddedAuthServer.Start(); err != nil {
                return fmt.Errorf("failed to start auth server TLS listener: %w", err)
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package runner

import (
        "log/slog"
        "strings"

        appconfig "github.com/stacklok/toolhive/pkg/config"
        "github.com/stacklok/toolhive/pkg/telemetry"
)

// otelHeadersEnvVar is the standard OpenTelemetry environment variable carrying
// OTLP export headers as a comma-separated list (e.g. "key1=val1,key2=val2").
// A config source may deliver headers through this variable in the app config's
// EnvVars instead of as header flags.
const otelHeadersEnvVar = "OTEL_EXPORTER_OTLP_HEADERS"

// BuildTelemetryConfigFromAppConfig builds a *telemetry.Config from an
// application OpenTelemetry config (read from ~/.config/toolhive/config.yaml
// or supplied by a config source).
//
// Both the "thv run" CLI and the "POST /api/v1/workloads" API path call this
// so workloads created via either surface inherit the same telemetry
// settings. Export headers come from two sources: the caller's header flags
// (the CLI surfaces "--otel-headers"; the API passes none) and an
// OTEL_EXPORTER_OTLP_HEADERS entry in the config's EnvVars. Flags take
// precedence. customAttributes has no equivalent in the application config
// and is passed through by the caller.
//
// Returns nil when no telemetry should be enabled: either no endpoint and no
// Prometheus metrics path, or both tracing and metrics disabled with no
// Prometheus path.
//
// CLI-style defaults are applied: TracingEnabled, MetricsEnabled, and
// UseLegacyAttributes default to true when unset in the config so that
// configuring just an endpoint produces a working setup.
func BuildTelemetryConfigFromAppConfig(
        otel appconfig.OpenTelemetryConfig,
        serviceName string,
        headers []string,
        customAttributes string,
) *telemetry.Config {
        if otel.Endpoint == "" && !otel.EnablePrometheusMetricsPath {
                return nil
        }

        tracingEnabled := true
        if otel.TracingEnabled != nil {
                tracingEnabled = *otel.TracingEnabled
        }
        metricsEnabled := true
        if otel.MetricsEnabled != nil {
                metricsEnabled = *otel.MetricsEnabled
        }
        useLegacyAttributes := true
        if otel.UseLegacyAttributes != nil {
                useLegacyAttributes = *otel.UseLegacyAttributes
        }

        if !tracingEnabled && !metricsEnabled && !otel.EnablePrometheusMetricsPath {
                return nil
        }

        parsedHeaders := mergeTelemetryHeaders(headers, otel.EnvVars)

        var processedEnvVars []string
        for _, entry := range otel.EnvVars {
                for ev := range strings.SplitSeq(entry, ",") {
                        trimmed := strings.TrimSpace(ev)
                        // Entries are host environment variable names. A "NAME=value"
                        // assignment (e.g. OTEL_EXPORTER_OTLP_HEADERS, parsed above into
                        // Headers) may itself contain commas, so drop it and the rest of
                        // the entry: later pieces may be fragments of its value.
                        if name, _, isAssignment := strings.Cut(trimmed, "="); isAssignment {
                                // Never log the entry itself: it may carry a credential.
                                if name != otelHeadersEnvVar {
                                        slog.Warn("Ignoring otel env-vars entry that is not an environment variable name")
                                }
                                break
                        }
                        if trimmed != "" {
                                processedEnvVars = append(processedEnvVars, trimmed)
                        }
                }
        }

        customAttrs, err := telemetry.ParseCustomAttributes(customAttributes)
        if err != nil {
                slog.Warn("Failed to parse custom attributes", "error", err)
                customAttrs = nil
        }

        cfg := &telemetry.Config{
                Endpoint:                    otel.Endpoint,
                ServiceName:                 serviceName,
                ServiceVersion:              "", // resolved at runtime in NewProvider()
                TracingEnabled:              tracingEnabled,
                MetricsEnabled:              metricsEnabled,
                Headers:                     parsedHeaders,
                Insecure:                    otel.Insecure,
                EnablePrometheusMetricsPath: otel.EnablePrometheusMetricsPath,
                EnvironmentVariables:        processedEnvVars,
                CustomAttributes:            customAttrs,
                UseLegacyAttributes:         useLegacyAttributes,
        }
        cfg.SetSamplingRateFromFloat(otel.SamplingRate)
        return cfg
}

// mergeTelemetryHeaders builds the OTLP export header map from two sources:
// an OTEL_EXPORTER_OTLP_HEADERS entry in the config's EnvVars (used when the
// caller passes no header flags, such as the API path) and caller-supplied
// header flags ("key=value", from the CLI's --otel-headers). Env headers are
// applied first so flags override them on conflict. Values are split on the
// first '=' only, so a value containing '=' (e.g. base64 padding) or spaces
// (e.g. "Basic <token>") is preserved intact.
func mergeTelemetryHeaders(flagHeaders, envVars []string) map[string]string {
        headers := make(map[string]string)
        for _, entry := range envVars {
                name, value, ok := strings.Cut(entry, "=")
                if !ok || name != otelHeadersEnvVar {
                        continue
                }
                for pair := range strings.SplitSeq(value, ",") {
                        if key, val, ok := strings.Cut(strings.TrimSpace(pair), "="); ok && key != "" {
                                headers[key] = val
                        }
                }
        }

        for _, h := range flagHeaders {
                if key, val, ok := strings.Cut(h, "="); ok {
                        headers[key] = val
                }
        }
        return headers
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package runtime provides workload deployment setup functionality
// that was previously part of the transport package.
package runtime

import (
        "context"
        "fmt"
        "log/slog"

        "github.com/stacklok/toolhive-core/permissions"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/ignore"
        "github.com/stacklok/toolhive/pkg/networking"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

var transportEnvMap = map[types.TransportType]string{
        types.TransportTypeSSE:            "sse",
        types.TransportTypeStreamableHTTP: "streamable-http",
        types.TransportTypeStdio:          "stdio",
}

// SetupResult contains the results of setting up a workload
type SetupResult struct {
        ContainerName string
        TargetURI     string
        TargetPort    int
        TargetHost    string
}

// Setup prepares and deploys a workload for use with a transport.
// The runtime parameter provides access to container operations.
// The permissionProfile is used to configure container permissions (including network mode).
// The k8sPodTemplatePatch is a JSON string to patch the Kubernetes pod template.
// Returns the container name and target URI for configuring the transport.
func Setup(
        ctx context.Context,
        transportType types.TransportType,
        runtime rt.Deployer,
        containerName string,
        image string,
        cmdArgs []string,
        envVars, labels map[string]string,
        permissionProfile *permissions.Profile,
        k8sPodTemplatePatch string,
        isolateNetwork bool,
        allowDockerGateway bool,
        ignoreConfig *ignore.Config,
        host string,
        targetPort int,
        targetHost string,
        publishedPorts []string,
        scalingConfig *rt.ScalingConfig,
        runConfigMCPServerGeneration int64,
) (*SetupResult, error) {
        // Add transport-specific environment variables
        env, ok := transportEnvMap[transportType]
        if !ok && transportType != types.TransportTypeStdio {
                return nil, fmt.Errorf("unsupported transport type: %s", transportType)
        }

        // For stdio transport, env is already set above
        if transportType == types.TransportTypeStdio {
                envVars["MCP_TRANSPORT"] = "stdio"
        } else {
                envVars["MCP_TRANSPORT"] = env

                // Use the target port for the container's environment variables
                envVars["MCP_PORT"] = fmt.Sprintf("%d", targetPort)
                envVars["FASTMCP_PORT"] = fmt.Sprintf("%d", targetPort)
                envVars["MCP_HOST"] = targetHost
        }

        // Create workload options
        containerOptions := rt.NewDeployWorkloadOptions()
        containerOptions.K8sPodTemplatePatch = k8sPodTemplatePatch
        containerOptions.IgnoreConfig = ignoreConfig

        // Process published ports
        for _, portSpec := range publishedPorts {
                hostPort, containerPort, err := networking.ParsePortSpec(portSpec)
                if err != nil {
                        return nil, fmt.Errorf("failed to parse published port '%s': %w", portSpec, err)
                }

                // Add to exposed ports
                containerPortStr := fmt.Sprintf("%d/tcp", containerPort)
                containerOptions.ExposedPorts[containerPortStr] = struct{}{}

                // Add to port bindings
                // Check if we already have bindings for this port
                bindings := containerOptions.PortBindings[containerPortStr]
                bindings = append(bindings, rt.PortBinding{
                        HostPort: hostPort,
                })
                containerOptions.PortBindings[containerPortStr] = bindings
        }
        containerOptions.ScalingConfig = scalingConfig
        containerOptions.AllowDockerGateway = allowDockerGateway
        containerOptions.RunConfigMCPServerGeneration = runConfigMCPServerGeneration

        if transportType == types.TransportTypeStdio {
                containerOptions.AttachStdio = true
        } else {
                // Expose the target port in the container
                containerPortStr := fmt.Sprintf("%d/tcp", targetPort)
                containerOptions.ExposedPorts[containerPortStr] = struct{}{}

                // Create host port bindings (configurable through the --host flag)
                portBindings := []rt.PortBinding{
                        {
                                HostIP:   host,
                                HostPort: fmt.Sprintf("%d", targetPort),
                        },
                }

                // Set the port bindings
                // Note: if the user explicitly publishes the target port using --publish,
                // we append the default transport binding to the list of bindings for that port.
                if _, ok := containerOptions.PortBindings[containerPortStr]; ok {
                        containerOptions.PortBindings[containerPortStr] = append(containerOptions.PortBindings[containerPortStr], portBindings...)
                } else {
                        containerOptions.PortBindings[containerPortStr] = portBindings
                }
        }

        // Create the container
        slog.Debug("deploying workload", "container", containerName, "image", image)
        exposedPort, err := runtime.DeployWorkload(
                ctx,
                image,
                containerName,
                cmdArgs,
                envVars,
                labels,
                permissionProfile,
                transportType.String(),
                containerOptions,
                isolateNetwork,
        )
        if err != nil {
                return nil, fmt.Errorf("failed to create container: %w", err)
        }
        slog.Debug("container created", "container", containerName)

        result := &SetupResult{
                ContainerName: containerName,
                TargetHost:    targetHost,
                TargetPort:    targetPort,
        }

        // For stdio transport, there's no target URI
        if transportType == types.TransportTypeStdio {
                result.TargetURI = ""
        } else {
                // Update target host and port if needed (for Kubernetes)
                if (transportType == types.TransportTypeSSE || transportType == types.TransportTypeStreamableHTTP) && rt.IsKubernetesRuntime() {
                        // If the MCPServiceName is set, use it as the target host
                        if containerOptions.MCPServiceName != "" {
                                result.TargetHost = containerOptions.MCPServiceName
                        }
                }

                // we don't want to override the targetPort in a Kubernetes deployment. Because
                // by default the Kubernetes container deployer returns `0` for the exposedPort
                // therefore causing the "target port not set" error when it is assigned to the targetPort.
                // Issues:
                // - https://github.com/stacklok/toolhive/issues/902
                // - https://github.com/stacklok/toolhive/issues/924
                if !rt.IsKubernetesRuntime() {
                        result.TargetPort = exposedPort
                }

                // Construct target URI
                result.TargetURI = fmt.Sprintf("http://%s:%d", result.TargetHost, result.TargetPort)
        }

        return result, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "context"
        "fmt"
        "os"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/secrets/clients"
)

//go:generate mockgen -destination=mocks/mock_onepassword.go -package=mocks -source=1password.go OPSecretsService

// Err1PasswordReadOnly indicates that the 1Password secrets manager is read-only.
// Is it returned by operations which attempt to change values in 1Password.
var Err1PasswordReadOnly = fmt.Errorf("1Password secrets manager is read-only, write operations are not supported")

// OnePasswordManager manages secrets in 1Password.
type OnePasswordManager struct {
        client clients.OnePasswordClient
}

var timeout = 5 * time.Second

// GetSecret retrieves a secret from 1Password.
func (o *OnePasswordManager) GetSecret(ctx context.Context, path string) (string, error) {
        if !strings.Contains(path, "op://") {
                return "", fmt.Errorf("invalid secret path: %s", path)
        }

        secret, err := o.client.Resolve(ctx, path)
        if err != nil {
                return "", fmt.Errorf("error resolving secret: %w", err)
        }

        return secret, nil
}

// SetSecret is not supported for 1Password unless there is
// demand for it.
func (*OnePasswordManager) SetSecret(_ context.Context, _, _ string) error {
        return Err1PasswordReadOnly
}

// DeleteSecret is not supported for 1Password unless there is
// demand for it.
func (*OnePasswordManager) DeleteSecret(_ context.Context, _ string) error {
        return Err1PasswordReadOnly
}

// ListSecrets lists the paths to the secrets in 1Password.
// 1Password has a hierarchy of vaults, items, and fields.
// Each secret is represented as a path in the format:
// op://<vault>/<item>/<field>
func (o *OnePasswordManager) ListSecrets(ctx context.Context) ([]SecretDescription, error) {
        // First, grab the list of vaults we have access to.
        vaults, err := o.client.ListVaults(ctx)
        if err != nil {
                return nil, fmt.Errorf("error retrieving vaults from 1password API: %w", err)
        }

        var secrets []SecretDescription
        // For each vault...
        for _, vault := range vaults {
                items, err := o.client.ListItems(ctx, vault.ID)
                if err != nil {
                        return nil, fmt.Errorf("error retrieving secrets from 1password API: %w", err)
                }

                // For each item in the vault...
                for _, item := range items {
                        details, err := o.client.GetItem(ctx, vault.ID, item.ID)
                        if err != nil {
                                return nil, fmt.Errorf("error retrieving item details from 1password API: %w", err)
                        }
                        // For each field in the item...
                        for _, field := range details.Fields {
                                // Create a path and human-readable name for each field.
                                description := SecretDescription{
                                        Key:         fmt.Sprintf("op://%s/%s/%s", item.VaultID, item.ID, field.ID),
                                        Description: fmt.Sprintf("%s :: %s :: %s", vault.Title, item.Title, field.Title),
                                }
                                secrets = append(secrets, description)
                        }
                }
        }

        return secrets, nil
}

// DeleteSecrets is a no-op for the 1Password provider (read-only).
func (*OnePasswordManager) DeleteSecrets(_ context.Context, _ []string) error {
        return nil
}

// Cleanup is not needed for 1Password.
func (*OnePasswordManager) Cleanup() error {
        return nil
}

// Capabilities returns the capabilities of the 1Password provider.
// Read-only provider with listing support.
func (*OnePasswordManager) Capabilities() ProviderCapabilities {
        return ProviderCapabilities{
                CanRead:    true,
                CanWrite:   false, // 1Password is read-only for now
                CanDelete:  false, // 1Password is read-only for now
                CanList:    true,  // Listing is now supported
                CanCleanup: false, // Not applicable for 1Password
        }
}

// NewOnePasswordManager creates an instance of OnePasswordManager.
func NewOnePasswordManager() (Provider, error) {
        token := os.Getenv("OP_SERVICE_ACCOUNT_TOKEN")
        if token == "" {
                return nil, fmt.Errorf("OP_SERVICE_ACCOUNT_TOKEN is not set")
        }

        ctx, cancel := context.WithTimeout(context.Background(), timeout)
        defer cancel()

        client, err := clients.NewOnePasswordClient(ctx, token)
        if err != nil {
                return nil, fmt.Errorf("error creating 1Password client: %w", err)
        }

        return &OnePasswordManager{
                client: client,
        }, nil
}

// NewOnePasswordManagerWithClient creates an instance of OnePasswordManager with a provided 1password client.
// This function is primarily intended for testing purposes.
func NewOnePasswordManagerWithClient(client clients.OnePasswordClient) *OnePasswordManager {
        return &OnePasswordManager{
                client: client,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package aes contains functions for encrypting and decrypting data using AES-GCM
package aes

import (
        "crypto/aes"
        "crypto/cipher"
        "crypto/rand"
        "errors"
        "io"
)

const maxPlaintextSize = 32 * 1024 * 1024

// ErrExceedsMaxSize is returned when the plaintext is too large to encrypt.
var ErrExceedsMaxSize = errors.New("plaintext is too large, limited to 32MiB")

// Encrypt encrypts data using 256-bit AES-GCM.  This both hides the content of
// the data and provides a check that it hasn't been altered. Output takes the
// form nonce|ciphertext|tag where '|' indicates concatenation.
func Encrypt(plaintext []byte, key []byte) ([]byte, error) {
        if len(plaintext) > maxPlaintextSize {
                return nil, ErrExceedsMaxSize
        }

        block, err := aes.NewCipher(key[:])
        if err != nil {
                return nil, err
        }

        gcm, err := cipher.NewGCM(block)
        if err != nil {
                return nil, err
        }

        nonce := make([]byte, gcm.NonceSize())
        _, err = io.ReadFull(rand.Reader, nonce)
        if err != nil {
                return nil, err
        }

        return gcm.Seal(nonce, nonce, plaintext, nil), nil
}

// Decrypt decrypts data using 256-bit AES-GCM.  This both hides the content of
// the data and provides a check that it hasn't been altered. Expects input
// form nonce|ciphertext|tag where '|' indicates concatenation.
func Decrypt(ciphertext []byte, key []byte) ([]byte, error) {
        block, err := aes.NewCipher(key[:])
        if err != nil {
                return nil, err
        }

        gcm, err := cipher.NewGCM(block)
        if err != nil {
                return nil, err
        }

        if len(ciphertext) < gcm.NonceSize() {
                return nil, errors.New("malformed ciphertext")
        }

        return gcm.Open(nil,
                ciphertext[:gcm.NonceSize()],
                ciphertext[gcm.NonceSize():],
                nil,
        )
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package clients contains code for connecting to secret provider APIs.
package clients

import (
        "context"
        "fmt"

        "github.com/1password/onepassword-sdk-go"
)

//go:generate mockgen -destination=mocks/mock_onepassword.go -package=mocks -source=1password.go OnePasswordClient

// OnePasswordClient defines the subset of the 1Password SDK that we use.
type OnePasswordClient interface {
        Resolve(ctx context.Context, secretReference string) (string, error)
        ListItems(ctx context.Context, vaultID string, filters ...onepassword.ItemListFilter) ([]onepassword.ItemOverview, error)
        ListVaults(ctx context.Context) ([]onepassword.VaultOverview, error)
        GetItem(ctx context.Context, vaultID, itemID string) (onepassword.Item, error)
}

// NewOnePasswordClient creates a OnePasswordClient from the 1Password SDK
func NewOnePasswordClient(ctx context.Context, token string) (OnePasswordClient, error) {
        client, err := onepassword.NewClient(
                ctx,
                onepassword.WithServiceAccountToken(token),
                onepassword.WithIntegrationInfo(onepassword.DefaultIntegrationName, onepassword.DefaultIntegrationVersion),
        )
        if err != nil {
                return nil, fmt.Errorf("error creating 1Password client: %w", err)
        }

        return &onePasswordClient{client: client}, nil
}

// defaultOnePasswordClient implements the OnePasswordClient interface.
// Note that the methods we need are from two different interfaces in the SDK.
// This implementation presents them in a single interface for ease of mocking.
type onePasswordClient struct {
        client *onepassword.Client
}

func (opc *onePasswordClient) Resolve(ctx context.Context, secretReference string) (string, error) {
        secret, err := opc.client.Secrets().Resolve(ctx, secretReference)
        if err != nil {
                return "", fmt.Errorf("error resolving secret: %w", err)
        }
        return secret, nil
}

func (opc *onePasswordClient) ListItems(
        ctx context.Context,
        vaultID string,
        filters ...onepassword.ItemListFilter,
) ([]onepassword.ItemOverview, error) {
        return opc.client.Items().List(ctx, vaultID, filters...)
}

func (opc *onePasswordClient) ListVaults(ctx context.Context) ([]onepassword.VaultOverview, error) {
        return opc.client.Vaults().List(ctx)
}

func (opc *onePasswordClient) GetItem(ctx context.Context, vaultID, itemID string) (onepassword.Item, error) {
        return opc.client.Items().Get(ctx, vaultID, itemID)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "bytes"
        "context"
        "crypto/sha256"
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "maps"
        "os"
        "path"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/secrets/aes"
)

// On-disk framing for the encrypted secrets file:
//
//        magic    "THVSEC"  6 bytes
//        version  0x01      1 byte
//        salt               16 bytes
//        body               nonce|ciphertext|tag, as produced by aes.Encrypt
//
// The version selects the Argon2id cost parameters, which are fixed in code
// (see kdf.go) rather than stored here.
//
// Files without the magic prefix predate this framing: they are the bare output
// of aes.Encrypt under an unsalted SHA-256 of the password. Ordinary operations
// preserve that format for compatibility; UpgradeProtection explicitly upgrades it.
const (
        secretsFileMagic   = "THVSEC"
        secretsFileVersion = 0x01

        headerLength = len(secretsFileMagic) + 1 + saltLength
)

// ErrMalformedSecretsFile is returned when the secrets file header cannot be parsed.
var ErrMalformedSecretsFile = errors.New("malformed secrets file")

// ErrUnsupportedSecretsFileVersion is returned when the secrets file uses a
// framed format version this binary does not understand.
var ErrUnsupportedSecretsFileVersion = errors.New("unsupported secrets file format version")

// ErrDecryptionFailed is returned when the secrets file cannot be decrypted,
// which almost always means the password is wrong.
var ErrDecryptionFailed = errors.New("unable to decrypt secrets file")

type fileFormatVersion uint8

const (
        fileFormatEmpty fileFormatVersion = iota
        fileFormatLegacy
        fileFormatV1
)

// fileFormat describes the on-disk format returned by a read. Its salt is set
// only for framed version 1 files.
type fileFormat struct {
        version fileFormatVersion
        salt    []byte
}

// EncryptedManager stores secrets in an encrypted file.
// AES-256-GCM is used for encryption, with the key derived from the password
// using Argon2id and a per-file salt stored in the file header.
type EncryptedManager struct {
        filePath string
        // Password the encryption key is derived from.
        password []byte
}

// fileStructure is the structure of the secrets file.
type fileStructure struct {
        Secrets map[string]string `json:"secrets"`
}

// GetSecret retrieves a secret from the secret store.
//
// The file is read and decrypted on every call so that changes written by
// other processes are immediately visible.
func (e *EncryptedManager) GetSecret(_ context.Context, name string) (string, error) {
        if name == "" {
                return "", errors.New("secret name cannot be empty")
        }

        secrets, _, err := e.readFileSecrets()
        if err != nil {
                return "", fmt.Errorf("reading secrets: %w", err)
        }
        value, ok := secrets[name]
        if !ok {
                return "", fmt.Errorf("%w: %s", ErrSecretNotFound, name)
        }
        return value, nil
}

// SetSecret stores a secret in the secret store.
func (e *EncryptedManager) SetSecret(_ context.Context, name, value string) error {
        if name == "" {
                return errors.New("secret name cannot be empty")
        }

        return fileutils.WithFileLock(e.filePath, func() error {
                // Re-read the file inside the lock to avoid overwriting changes
                // made by other processes since this manager was created.
                secrets, format, err := e.readFileSecrets()
                if err != nil {
                        return err
                }
                secrets[name] = value
                return e.writeFileSecrets(secrets, format)
        })
}

// DeleteSecret removes a secret from the secret store.
func (e *EncryptedManager) DeleteSecret(_ context.Context, name string) error {
        if name == "" {
                return errors.New("secret name cannot be empty")
        }

        return fileutils.WithFileLock(e.filePath, func() error {
                // Re-read the file inside the lock so the existence check
                // reflects the current on-disk state.
                secrets, format, err := e.readFileSecrets()
                if err != nil {
                        return err
                }
                if _, ok := secrets[name]; !ok {
                        return fmt.Errorf("%w: %s", ErrSecretNotFound, name)
                }
                delete(secrets, name)
                return e.writeFileSecrets(secrets, format)
        })
}

// ListSecrets returns a list of all secret names stored in the manager.
func (e *EncryptedManager) ListSecrets(_ context.Context) ([]SecretDescription, error) {
        secrets, _, err := e.readFileSecrets()
        if err != nil {
                return nil, fmt.Errorf("reading secrets: %w", err)
        }
        result := make([]SecretDescription, 0, len(secrets))
        for key := range secrets {
                result = append(result, SecretDescription{Key: key})
        }
        return result, nil
}

// DeleteSecrets removes all named keys from the store.
func (e *EncryptedManager) DeleteSecrets(_ context.Context, keys []string) error {
        return fileutils.WithFileLock(e.filePath, func() error {
                // Re-read the file inside the lock to avoid losing changes made
                // by other processes since this manager was created.
                current, format, err := e.readFileSecrets()
                if err != nil {
                        return err
                }
                for _, key := range keys {
                        delete(current, key)
                }
                return e.writeFileSecrets(current, format)
        })
}

// Cleanup removes all secrets managed by this manager.
func (e *EncryptedManager) Cleanup() error {
        return fileutils.WithFileLock(e.filePath, func() error {
                // Re-read under the lock so the replacement retains the current format.
                _, format, err := e.readFileSecrets()
                if err != nil {
                        return err
                }
                return e.writeFileSecrets(make(map[string]string), format)
        })
}

// Capabilities returns the capabilities of the encrypted provider.
func (*EncryptedManager) Capabilities() ProviderCapabilities {
        return ProviderCapabilities{
                CanRead:    true,
                CanWrite:   true,
                CanDelete:  true,
                CanList:    true,
                CanCleanup: true,
        }
}

// UpgradeProtection explicitly upgrades a legacy store to the current framed
// format. It authenticates the original bytes, installs the upgraded file
// atomically, and verifies the installed file before returning success.
func (e *EncryptedManager) UpgradeProtection() error {
        return fileutils.WithFileLock(e.filePath, func() error {
                original, err := os.ReadFile(e.filePath) // #nosec G304: File path is not configurable at this time.
                if err != nil {
                        return fmt.Errorf("reading secrets file for upgrade: %w", err)
                }
                secrets, format, err := e.decodeSecretsFile(original)
                if err != nil {
                        return fmt.Errorf("authenticating secrets file for upgrade: %w", err)
                }
                if format.version != fileFormatLegacy {
                        return nil
                }
                newSalt, err := generateSalt()
                if err != nil {
                        return err
                }
                if err := e.installUpgrade(secrets, fileFormat{version: fileFormatV1, salt: newSalt}); err != nil {
                        return err
                }
                return e.verifyUpgradeOrRollback(original, secrets)
        })
}

// installUpgrade writes secrets in the given (already-migrated) format. Split
// out from UpgradeProtection so tests can install a real upgrade and then
// corrupt the result on disk before exercising verifyUpgradeOrRollback,
// which is the only way to deterministically reach the rollback branch
// without a production-only test seam.
func (e *EncryptedManager) installUpgrade(secrets map[string]string, newFormat fileFormat) error {
        return e.writeFileSecrets(secrets, newFormat)
}

// verifyUpgradeOrRollback re-reads what installUpgrade just wrote and confirms
// it decrypts back to want under the new format. On any mismatch — a failed
// read, a version that isn't fileFormatV1, or decrypted secrets that differ —
// it restores original, the pre-upgrade encrypted bytes, so a verification
// failure never leaves the store in a state only the failed write can read.
func (e *EncryptedManager) verifyUpgradeOrRollback(original []byte, want map[string]string) error {
        verified, verifiedFormat, err := e.readFileSecrets()
        if err == nil && verifiedFormat.version == fileFormatV1 && maps.Equal(want, verified) {
                return nil
        }
        verifyErr := err
        if verifyErr == nil {
                verifyErr = errors.New("decrypted secrets do not match")
        }
        rollbackErr := fileutils.AtomicWriteFile(e.filePath, original, 0600)
        if rollbackErr != nil {
                return fmt.Errorf("verifying upgraded secrets file: %w; restoring original encrypted bytes: %v", verifyErr, rollbackErr)
        }
        return fmt.Errorf("verifying upgraded secrets file: %w; restored original encrypted bytes", verifyErr)
}

// NewEncryptedManager creates an instance of EncryptedManager.
//
// password is the user's password, not a derived key. Callers that previously
// passed sha256(password) still compile against this signature but will fail to
// decrypt, because that value is now hashed again as if it were a password.
//
// The manager takes the password rather than a key because the key derivation
// salt lives in the secrets file itself and is only known once the file is read.
//
// Opening an existing store authenticates that it is readable, but does not
// change its format. Use UpgradeProtection to explicitly upgrade legacy files.
func NewEncryptedManager(filePath string, password []byte) (Provider, error) {
        return newEncryptedManager(filePath, password)
}

func newEncryptedManager(filePath string, password []byte) (*EncryptedManager, error) {
        if len(password) == 0 {
                return nil, errors.New("password cannot be empty")
        }

        filePath = path.Clean(filePath)

        // Ensure the file exists (create if needed).
        // #nosec G304: File path is not configurable at this time.
        f, err := os.OpenFile(filePath, os.O_CREATE|os.O_RDWR, 0600)
        if err != nil {
                return nil, fmt.Errorf("failed to open secrets file: %w", err)
        }
        if err := f.Close(); err != nil {
                slog.Warn("Failed to close secrets file", "error", err)
        }

        manager := &EncryptedManager{
                filePath: filePath,
                password: bytes.Clone(password),
        }

        // Validate the file is readable and correctly encrypted at startup.
        stat, err := os.Stat(filePath)
        if err != nil {
                return nil, fmt.Errorf("failed to stat secrets file: %w", err)
        }
        if stat.Size() > 0 {
                if _, _, err := manager.readFileSecrets(); err != nil {
                        printSecretsFileHint(filePath, err)
                        return nil, err
                }
        }

        return manager, nil
}

// printSecretsFileHint explains why the secrets file could not be opened. The
// underlying error surfaces deep inside a CLI command, where it is easy to
// mistake for an unrelated failure.
func printSecretsFileHint(filePath string, err error) {
        switch {
        case errors.Is(err, ErrDecryptionFailed):
                fmt.Fprintf(os.Stderr, "\nSecrets file decryption failed: this usually means the password "+
                        "is incorrect or the secrets file has been corrupted.\n"+
                        "If your keyring was recently reset, try again with your original password.\n"+
                        "If the secrets file is corrupted, delete it at %s and run 'thv secret setup' to start fresh.\n\n",
                        filePath)
        case errors.Is(err, ErrUnsupportedSecretsFileVersion):
                fmt.Fprintf(os.Stderr, "\nThe secrets file at %s was written by a newer version of thv. "+
                        "Upgrade thv to access it.\n\n", filePath)
        case errors.Is(err, ErrMalformedSecretsFile):
                fmt.Fprintf(os.Stderr, "\nThe secrets file at %s is not in a format this version of thv understands.\n"+
                        "If it was written by a newer version of thv, upgrade. Otherwise the file has been corrupted; "+
                        "delete it and run 'thv secret setup' to start fresh.\n\n", filePath)
        }
}

// readFileSecrets reads and decrypts the on-disk secrets, returning the format
// required to preserve it on an ordinary write.
func (e *EncryptedManager) readFileSecrets() (map[string]string, fileFormat, error) {
        // #nosec G304: File path is not configurable at this time.
        data, err := os.ReadFile(e.filePath)
        if err != nil {
                if errors.Is(err, os.ErrNotExist) {
                        return make(map[string]string), fileFormat{version: fileFormatEmpty}, nil
                }
                return nil, fileFormat{}, fmt.Errorf("failed to read secrets file: %w", err)
        }
        return e.decodeSecretsFile(data)
}

// decodeSecretsFile authenticates and decodes encrypted file bytes.
func (e *EncryptedManager) decodeSecretsFile(data []byte) (map[string]string, fileFormat, error) {
        if len(data) == 0 {
                return make(map[string]string), fileFormat{version: fileFormatEmpty}, nil
        }
        key, body, format, err := e.decryptionKey(data)
        if err != nil {
                return nil, fileFormat{}, err
        }
        decrypted, err := aes.Decrypt(body, key)
        if err != nil {
                return nil, fileFormat{}, fmt.Errorf("%w: %w", ErrDecryptionFailed, err)
        }
        var contents fileStructure
        if err := json.Unmarshal(decrypted, &contents); err != nil {
                return nil, fileFormat{}, fmt.Errorf("failed to decode secrets file: %w", err)
        }
        if contents.Secrets == nil {
                return make(map[string]string), format, nil
        }
        return contents.Secrets, format, nil
}

func (e *EncryptedManager) decryptionKey(data []byte) ([]byte, []byte, fileFormat, error) {
        if !bytes.HasPrefix(data, []byte(secretsFileMagic)) {
                sum := sha256.Sum256(e.password)
                return sum[:], data, fileFormat{version: fileFormatLegacy}, nil
        }
        salt, body, err := parseSecretsHeader(data)
        if err != nil {
                return nil, nil, fileFormat{}, err
        }
        format := fileFormat{version: fileFormatV1, salt: salt}
        return deriveKeyCached(e.filePath, e.password, salt), body, format, nil
}

// writeFileSecrets encrypts and atomically writes the secrets map using format.
// It must be called while holding the file lock.
func (e *EncryptedManager) writeFileSecrets(secrets map[string]string, format fileFormat) error {
        contents, err := json.Marshal(fileStructure{Secrets: secrets})
        if err != nil {
                return fmt.Errorf("failed to marshal secrets: %w", err)
        }
        if format.version == fileFormatLegacy {
                sum := sha256.Sum256(e.password)
                body, err := aes.Encrypt(contents, sum[:])
                if err != nil {
                        return fmt.Errorf("failed to encrypt secrets: %w", err)
                }
                if err := fileutils.AtomicWriteFile(e.filePath, body, 0600); err != nil {
                        return fmt.Errorf("failed to write secrets to file: %w", err)
                }
                return nil
        }
        salt := format.salt
        if format.version == fileFormatEmpty {
                salt, err = generateSalt()
                if err != nil {
                        return err
                }
        }
        body, err := aes.Encrypt(contents, deriveKeyCached(e.filePath, e.password, salt))
        if err != nil {
                return fmt.Errorf("failed to encrypt secrets: %w", err)
        }
        if err := fileutils.AtomicWriteFile(e.filePath, append(encodeSecretsHeader(salt), body...), 0600); err != nil {
                return fmt.Errorf("failed to write secrets to file: %w", err)
        }
        return nil
}

// encodeSecretsHeader builds the fixed-size header that precedes the body.
func encodeSecretsHeader(salt []byte) []byte {
        header := make([]byte, 0, headerLength)
        header = append(header, secretsFileMagic...)
        header = append(header, secretsFileVersion)
        return append(header, salt...)
}

// parseSecretsHeader parses a framed secrets file, returning its salt and the
// encrypted body. The caller must have already checked the magic prefix.
func parseSecretsHeader(data []byte) ([]byte, []byte, error) {
        if len(data) < headerLength {
                return nil, nil, fmt.Errorf("%w: header is truncated (%d bytes)", ErrMalformedSecretsFile, len(data))
        }
        if version := data[len(secretsFileMagic)]; version != secretsFileVersion {
                return nil, nil, fmt.Errorf("%w: %d", ErrUnsupportedSecretsFileVersion, version)
        }
        // Clone the salt so that the retained copy does not pin the whole file buffer.
        return bytes.Clone(data[len(secretsFileMagic)+1 : headerLength]), data[headerLength:], nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "context"
        "errors"
        "fmt"
        "os"
)

// EnvironmentProvider reads secrets from environment variables
type EnvironmentProvider struct {
        prefix string
}

// NewEnvironmentProvider creates a new environment variable secrets provider
func NewEnvironmentProvider() Provider {
        return &EnvironmentProvider{
                prefix: EnvVarPrefix,
        }
}

// GetSecret retrieves a secret from environment variables
func (e *EnvironmentProvider) GetSecret(_ context.Context, name string) (string, error) {
        if name == "" {
                return "", errors.New("secret name cannot be empty")
        }

        envVar := e.prefix + name
        value := os.Getenv(envVar)
        if value == "" {
                return "", fmt.Errorf("%w: %s", ErrSecretNotFound, name)
        }

        return value, nil
}

// SetSecret is not supported for environment variables
func (*EnvironmentProvider) SetSecret(_ context.Context, name, _ string) error {
        if name == "" {
                return errors.New("secret name cannot be empty")
        }
        return errors.New("environment provider is read-only")
}

// DeleteSecret is not supported for environment variables
func (*EnvironmentProvider) DeleteSecret(_ context.Context, name string) error {
        if name == "" {
                return errors.New("secret name cannot be empty")
        }
        return errors.New("environment provider is read-only")
}

// ListSecrets is not supported for environment variables for security reasons
func (*EnvironmentProvider) ListSecrets(_ context.Context) ([]SecretDescription, error) {
        return nil, errors.New("environment provider does not support listing secrets for security reasons")
}

// DeleteSecrets is a no-op for the environment provider (read-only).
func (*EnvironmentProvider) DeleteSecrets(_ context.Context, _ []string) error {
        return nil
}

// Cleanup is a no-op for environment provider
func (*EnvironmentProvider) Cleanup() error {
        return nil
}

// Capabilities returns the capabilities of the environment provider
func (*EnvironmentProvider) Capabilities() ProviderCapabilities {
        return ProviderCapabilities{
                CanRead:    true,
                CanWrite:   false,
                CanDelete:  false,
                CanList:    false,
                CanCleanup: false,
        }
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "context"
        "crypto/rand"
        "encoding/base64"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "os"
        "sync"

        "github.com/adrg/xdg"
        "golang.org/x/term"

        "github.com/stacklok/toolhive-core/httperr"
        "github.com/stacklok/toolhive/pkg/process"
        "github.com/stacklok/toolhive/pkg/secrets/keyring"
)

const (
        // PasswordEnvVar is the environment variable used to specify the password for encrypting and decrypting secrets.
        PasswordEnvVar = "TOOLHIVE_SECRETS_PASSWORD"

        // ProviderEnvVar is the environment variable used to specify the secrets provider type.
        ProviderEnvVar = "TOOLHIVE_SECRETS_PROVIDER"

        keyringService = "toolhive"
)

var (
        keyringProvider keyring.Provider
        keyringOnce     sync.Once
)

func getKeyringProvider() keyring.Provider {
        keyringOnce.Do(func() {
                keyringProvider = keyring.NewCompositeProvider()
        })
        return keyringProvider
}

// ProviderType represents an enum of the types of available secrets providers.
type ProviderType string

const (
        // EncryptedType represents the encrypted secret provider.
        EncryptedType ProviderType = "encrypted"

        // OnePasswordType represents the 1Password secret provider.
        OnePasswordType ProviderType = "1password"

        // EnvironmentType represents the environment variable secret provider
        EnvironmentType ProviderType = "environment"
)

// ErrUnknownManagerType is returned when an invalid value for ProviderType is specified.
var ErrUnknownManagerType = httperr.WithCode(
        errors.New("unknown secret manager type"),
        http.StatusBadRequest,
)

// ErrSecretsNotSetup is returned when secrets functionality is used before running setup.
var ErrSecretsNotSetup = httperr.WithCode(
        errors.New("secrets provider not configured. "+
                "Please run 'thv secret setup' to configure a secrets provider first"),
        http.StatusNotFound,
)

// SetupResult contains the result of a provider setup operation
type SetupResult struct {
        ProviderType ProviderType
        Success      bool
        Message      string
        Error        error
}

// ValidateProvider validates that a provider can be created and performs basic functionality tests
func ValidateProvider(ctx context.Context, providerType ProviderType) *SetupResult {
        return ValidateProviderWithPassword(ctx, providerType, "")
}

// ValidateProviderWithPassword validates that a provider can be created and performs basic functionality tests.
// If password is provided for encrypted provider, it uses that password instead of reading from stdin.
func ValidateProviderWithPassword(ctx context.Context, providerType ProviderType, password string) *SetupResult {
        result := &SetupResult{
                ProviderType: providerType,
                Success:      false,
        }

        // Test that we can create the provider
        provider, err := CreateSecretProviderWithPassword(providerType, password)
        if err != nil {
                result.Error = fmt.Errorf("failed to create provider: %w", err)
                result.Message = fmt.Sprintf("Failed to initialize %s provider", providerType)
                return result
        }

        // Perform provider-specific validation
        switch providerType {
        case EncryptedType:
                return validateEncryptedProvider(ctx, provider, result)
        case OnePasswordType:
                return validateOnePasswordProvider(ctx, provider, result)
        case EnvironmentType:
                return ValidateEnvironmentProvider(ctx, provider, result)
        default:
                result.Error = fmt.Errorf("unknown provider type: %s", providerType)
                result.Message = "Unknown provider type"
                return result
        }
}

// ValidateEnvironmentProvider tests environment provider functionality
func ValidateEnvironmentProvider(ctx context.Context, provider Provider, result *SetupResult) *SetupResult {
        // Test basic functionality by attempting to get a test secret
        // We don't expect it to exist, but we should get a proper "not found" error
        _, err := provider.GetSecret(ctx, "nonexistent-test-secret")
        if err == nil {
                result.Error = fmt.Errorf("expected error for nonexistent secret, but got none")
                result.Message = "Environment provider validation failed"
                return result
        }

        // Check that we get the expected not-found error
        if !IsNotFoundError(err) {
                result.Error = fmt.Errorf("unexpected error format: %w", err)
                result.Message = "Environment provider validation failed"
                return result
        }

        result.Success = true
        result.Message = "Environment provider validation successful"
        return result
}

// validateEncryptedProvider tests basic encrypted provider functionality
func validateEncryptedProvider(ctx context.Context, provider Provider, result *SetupResult) *SetupResult {
        // Test basic functionality with a temporary secret
        testKey := "setup-validation-test"
        testValue := "test-value"

        // Test setting a secret
        if err := provider.SetSecret(ctx, testKey, testValue); err != nil {
                result.Error = fmt.Errorf("failed to test secret storage: %w", err)
                result.Message = "Failed to store test secret"
                return result
        }

        // Test retrieving the secret
        retrievedValue, err := provider.GetSecret(ctx, testKey)
        if err != nil {
                result.Error = fmt.Errorf("failed to test secret retrieval: %w", err)
                result.Message = "Failed to retrieve test secret"
                return result
        }

        // Verify the value matches
        if retrievedValue != testValue {
                result.Error = fmt.Errorf("secret test failed: expected %s, got %s", testValue, retrievedValue)
                result.Message = "Secret value mismatch during validation"
                return result
        }

        // Clean up test secret
        _ = provider.DeleteSecret(ctx, testKey)

        result.Success = true
        result.Message = "Encrypted provider validation successful"
        return result
}

// validateOnePasswordProvider tests 1Password provider connectivity
func validateOnePasswordProvider(ctx context.Context, provider Provider, result *SetupResult) *SetupResult {
        // Test basic functionality by attempting to list secrets
        _, err := provider.ListSecrets(ctx)
        if err != nil {
                result.Error = fmt.Errorf("failed to connect to 1Password: %w", err)
                result.Message = "Failed to connect to 1Password service"
                return result
        }

        result.Success = true
        result.Message = "1Password provider validation successful"
        return result
}

// ErrKeyringNotAvailable is returned when the OS keyring is not available for the encrypted provider.
var ErrKeyringNotAvailable = httperr.WithCode(
        errors.New("OS keyring is not available. "+
                "The encrypted provider requires an OS keyring to securely store passwords. "+
                "Please use a different secrets provider (e.g., 1password) "+
                "or ensure your system has a keyring service available"),
        http.StatusBadRequest,
)

// IsKeyringAvailable tests if any keyring backend is available
func IsKeyringAvailable() bool {
        provider := getKeyringProvider()
        return provider.IsAvailable()
}

// CreateSecretProvider creates the specified type of secrets provider.
// TODO CREATE function does not actually create anything, refactor or rename
func CreateSecretProvider(managerType ProviderType) (Provider, error) {
        return CreateSecretProviderWithPassword(managerType, "")
}

// CreateSecretProviderWithPassword creates the specified type of secrets provider with an optional password.
// If password is empty, it uses the current functionality (read from keyring or stdin).
// If password is provided, it uses that password and stores it in the keyring if not already setup.
func CreateSecretProviderWithPassword(managerType ProviderType, password string) (Provider, error) {
        // Create the primary provider
        var primary Provider
        var err error

        switch managerType {
        case EncryptedType:
                secretsPath, pathErr := xdg.DataFile("toolhive/secrets_encrypted")
                if pathErr != nil {
                        return nil, fmt.Errorf("unable to access secrets file path %w", pathErr)
                }
                manager, openErr := openEncryptedManager(secretsPath, password)
                if openErr != nil {
                        return nil, openErr
                }
                primary = manager
        case OnePasswordType:
                primary, err = NewOnePasswordManager()
        case EnvironmentType:
                // Direct environment provider - no fallback needed
                return NewEnvironmentProvider(), nil
        default:
                return nil, ErrUnknownManagerType
        }

        if err != nil {
                return nil, err
        }

        // Wrap with fallback provider if enabled
        if shouldEnableFallback() {
                return NewFallbackProvider(primary), nil
        }

        return primary, nil
}

// UpgradeEncryptedProtection upgrades the encrypted secrets store to the current
// password-protection format. It bypasses provider wrappers because upgrading is
// a store-specific operation, not a Provider capability.
func UpgradeEncryptedProtection() error {
        secretsPath, err := xdg.DataFile("toolhive/secrets_encrypted")
        if err != nil {
                return fmt.Errorf("unable to access secrets file path %w", err)
        }
        // Check before prompting: with no store and no keyring password, opening one
        // would silently run first-time setup instead of an upgrade.
        stat, err := os.Stat(secretsPath)
        if errors.Is(err, os.ErrNotExist) || (err == nil && stat.Size() == 0) {
                return errors.New("no encrypted secrets to upgrade; run 'thv secret setup' to configure the store")
        }
        if err != nil {
                return fmt.Errorf("failed to stat secrets file: %w", err)
        }
        manager, err := openEncryptedManager(secretsPath, "")
        if err != nil {
                return err
        }
        if err := manager.UpgradeProtection(); err != nil {
                return fmt.Errorf("upgrading encrypted protection: %w", err)
        }
        return nil
}

// openEncryptedManager opens the encrypted store at secretsPath. If password is
// empty it is read from the keyring or prompted for; a newly entered password is
// stored in the keyring only after it has successfully decrypted the store.
func openEncryptedManager(secretsPath, password string) (*EncryptedManager, error) {
        // Enforce keyring availability for encrypted provider
        if !IsKeyringAvailable() {
                return nil, ErrKeyringNotAvailable
        }

        secretsPassword, isNew, err := GetSecretsPassword(password)
        if err != nil {
                return nil, fmt.Errorf("failed to get secrets password: %w", err)
        }
        // The AES-256-GCM key is derived from the password inside the
        // manager, using Argon2id and a salt stored in the secrets file.
        manager, err := newEncryptedManager(secretsPath, secretsPassword)
        if err != nil {
                // Decryption failed - don't store the password in keyring
                // This allows the user to retry with the correct password
                return nil, fmt.Errorf("failed to create provider: %w", err)
        }

        // Only store password in keyring after successful validation (decryption)
        if isNew {
                if err := StoreSecretsPassword(secretsPassword); err != nil {
                        return nil, fmt.Errorf("failed to store password in keyring: %w", err)
                }
        }
        return manager, nil
}

// ProviderOption configures how CreateProvider wraps the underlying provider.
type ProviderOption func(*providerOptions) error

// providerOptions holds the accumulated state of all ProviderOptions passed to CreateProvider.
type providerOptions struct {
        scope      *SecretScope // nil = no scoping
        userFacing bool
}

// WithScope returns a ProviderOption that wraps the provider with a ScopedProvider
// for the given scope, keeping system secrets isolated under "__thv_<scope>_".
// Mutually exclusive with WithUserFacing.
func WithScope(scope SecretScope) ProviderOption {
        return func(o *providerOptions) error {
                if o.userFacing {
                        return errors.New("WithScope and WithUserFacing are mutually exclusive")
                }
                o.scope = &scope
                return nil
        }
}

// WithUserFacing returns a ProviderOption that wraps the provider with a UserProvider,
// blocking access to any key that starts with the system prefix "__thv_".
// Suitable for CLI, API, and MCP tool server callers.
// Mutually exclusive with WithScope.
func WithUserFacing() ProviderOption {
        return func(o *providerOptions) error {
                if o.scope != nil {
                        return errors.New("WithUserFacing and WithScope are mutually exclusive")
                }
                o.userFacing = true
                return nil
        }
}

// CreateProvider creates a secret Provider for the given manager type, optionally
// wrapped according to the supplied options.
//
// Without options it behaves identically to CreateSecretProvider.
// Pass WithUserFacing() for CLI/API callers or WithScope(scope) for internal callers.
func CreateProvider(managerType ProviderType, opts ...ProviderOption) (Provider, error) {
        inner, err := CreateSecretProvider(managerType)
        if err != nil {
                return nil, err
        }

        options := &providerOptions{}
        for _, opt := range opts {
                if err := opt(options); err != nil {
                        return nil, err
                }
        }

        if options.userFacing {
                return NewUserProvider(inner), nil
        }
        if options.scope != nil {
                return NewScopedProvider(inner, *options.scope), nil
        }
        return inner, nil
}

// shouldEnableFallback determines if environment variable fallback should be enabled
func shouldEnableFallback() bool {
        // Check for explicit opt-out
        if os.Getenv("TOOLHIVE_DISABLE_ENV_FALLBACK") == "true" {
                return false
        }

        // Enable by default for non-environment providers
        return true
}

// GetSecretsPassword returns the password to use for encrypting and decrypting secrets.
// It returns (password, isNew, error) where isNew indicates if the password was not found
// in the keyring and needs to be stored after successful validation.
// If optionalPassword is provided and keyring is not yet setup, it uses that password.
// Otherwise, it reads from keyring or prompts via stdin.
// IMPORTANT: When isNew is true, the caller MUST call StoreSecretsPassword after successfully
// validating the password (e.g., after successful decryption) to persist it in the keyring.
func GetSecretsPassword(optionalPassword string) ([]byte, bool, error) {
        provider := getKeyringProvider()

        // Attempt to load the password from the keyring
        keyringSecret, err := provider.Get(keyringService, keyringService)
        if err == nil {
                return []byte(keyringSecret), false, nil
        }

        // Handle key not found
        if errors.Is(err, keyring.ErrNotFound) {
                var password []byte

                // If optional password is provided, use it
                if optionalPassword != "" {
                        password = []byte(optionalPassword)
                } else {
                        // Keyring is available but no password stored - this should only happen during setup
                        if process.IsDetached() {
                                return nil, false, fmt.Errorf("detached process detected, cannot ask for password")
                        }

                        // Prompt for password during setup
                        var err error
                        password, err = readPasswordStdin()
                        if err != nil {
                                return nil, false, fmt.Errorf("failed to read password: %w", err)
                        }
                }

                // Return the password with isNew=true, caller must store after validation
                return password, true, nil
        }

        // Assume any other keyring error means keyring is not available
        return nil, false, fmt.Errorf("keyring is not available: %w", err)
}

// StoreSecretsPassword stores the password in the keyring.
// This should only be called after the password has been successfully validated
// (e.g., after successful decryption of the secrets file).
func StoreSecretsPassword(password []byte) error {
        provider := getKeyringProvider()
        //nolint:gosec // G706: provider name is from internal keyring provider
        slog.Debug("Writing password to keyring", "provider", provider.Name())
        err := provider.Set(keyringService, keyringService, string(password))
        if err != nil {
                return fmt.Errorf("failed to store password in keyring: %w", err)
        }
        return nil
}

func readPasswordStdin() ([]byte, error) {
        printPasswordPrompt()
        password, err := term.ReadPassword(int(os.Stdin.Fd())) //nolint:gosec // G115: stdin fd is always small
        // Start new line after receiving password to ensure errors are printed correctly.
        fmt.Println()
        if err != nil {
                return nil, fmt.Errorf("failed to read password: %w", err)
        }

        // Ensure the password is non-empty.
        if len(password) == 0 {
                return nil, errors.New("password cannot be empty")
        }
        return password, nil
}

// ResetKeyringSecret clears out the secret from the keystore (if present).
func ResetKeyringSecret() error {
        provider := getKeyringProvider()
        return provider.DeleteAll(keyringService)
}

// GenerateSecurePassword generates a cryptographically secure random password
func GenerateSecurePassword() (string, error) {
        // Generate 32 random bytes (256 bits)
        bytes := make([]byte, 32)
        if _, err := rand.Read(bytes); err != nil {
                return "", fmt.Errorf("failed to generate random bytes: %w", err)
        }

        // Encode as base64 to make it a readable string
        // This gives us a 44-character password with good entropy
        password := base64.URLEncoding.EncodeToString(bytes)
        return password, nil
}

func printPasswordPrompt() {
        fmt.Print("ToolHive needs a password to secure your credentials in the OS keyring.\n" +
                "This password will be used to encrypt and decrypt API tokens and other secrets\n" +
                "that need to be accessed by MCP servers. It will be securely stored in your OS keyring\n" +
                "so you won't need to enter it each time.\n" +
                "Please enter your keyring password: ")
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "context"
        "errors"
        "log/slog"
        "strings"
)

// FallbackProvider wraps a primary provider with environment variable fallback
type FallbackProvider struct {
        primary     Provider
        envProvider Provider
}

// NewFallbackProvider creates a new provider with environment variable fallback
func NewFallbackProvider(primary Provider) Provider {
        return &FallbackProvider{
                primary: primary,
                envProvider: &EnvironmentProvider{
                        prefix: EnvVarPrefix,
                },
        }
}

// GetSecret attempts to get a secret from the primary provider,
// falling back to environment variables if not found
func (f *FallbackProvider) GetSecret(ctx context.Context, name string) (string, error) {
        // First, try the primary provider
        value, err := f.primary.GetSecret(ctx, name)
        if err == nil {
                return value, nil
        }

        // Check if it's a "not found" error
        if !IsNotFoundError(err) {
                return "", err
        }

        // Try environment variable fallback
        envValue, envErr := f.envProvider.GetSecret(ctx, name)
        if envErr == nil {
                //nolint:gosec // G706: secret name is user-provided input used for diagnostics
                slog.Debug("Secret retrieved from environment variable fallback", "name", name)
                return envValue, nil
        }

        // Return the original error if no fallback found
        return "", err
}

// SetSecret always uses the primary provider (no env var writes)
func (f *FallbackProvider) SetSecret(ctx context.Context, name, value string) error {
        return f.primary.SetSecret(ctx, name, value)
}

// DeleteSecret always uses the primary provider (no env var deletes)
func (f *FallbackProvider) DeleteSecret(ctx context.Context, name string) error {
        return f.primary.DeleteSecret(ctx, name)
}

// ListSecrets only lists from the primary provider
// (env vars not listed in fallback mode for security)
func (f *FallbackProvider) ListSecrets(ctx context.Context) ([]SecretDescription, error) {
        return f.primary.ListSecrets(ctx)
}

// DeleteSecrets delegates to the primary provider.
func (f *FallbackProvider) DeleteSecrets(ctx context.Context, keys []string) error {
        return f.primary.DeleteSecrets(ctx, keys)
}

// Cleanup delegates to the primary provider
func (f *FallbackProvider) Cleanup() error {
        return f.primary.Cleanup()
}

// Capabilities returns the primary provider's capabilities
func (f *FallbackProvider) Capabilities() ProviderCapabilities {
        return f.primary.Capabilities()
}

// ErrSecretNotFound is the sentinel error returned by built-in Provider
// implementations when a requested secret does not exist. Callers should
// use IsNotFoundError rather than comparing directly, so that third-party
// backends whose errors cannot wrap this sentinel are still handled.
var ErrSecretNotFound = errors.New("secret not found")

// IsNotFoundError reports whether err indicates that a secret was not found.
// It first checks for the ErrSecretNotFound sentinel (used by all built-in
// backends) via errors.Is, then falls back to substring matching for
// third-party backends (e.g. 1Password SDK) that cannot wrap the sentinel.
func IsNotFoundError(err error) bool {
        if err == nil {
                return false
        }
        if errors.Is(err, ErrSecretNotFound) {
                return true
        }
        // Legacy fallback for third-party backends that don't wrap ErrSecretNotFound.
        errStr := err.Error()
        return strings.Contains(errStr, "not found") ||
                strings.Contains(errStr, "does not exist")
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "bytes"
        "crypto/rand"
        "crypto/subtle"
        "fmt"
        "sync"

        "golang.org/x/crypto/argon2"
)

// Argon2id cost parameters for secrets file format version 1.
//
// These live in code rather than in the file. A secrets file records only its
// format version, and the version determines the cost, which keeps attacker
// controllable values out of Argon2id's memory allocation and makes raising the
// cost an explicit new format version rather than a silent per-file property.
const (
        argon2Time        uint32 = 2
        argon2MemoryKiB   uint32 = 19456 // 19 MiB, the OWASP minimum for Argon2id
        argon2Parallelism uint8  = 1

        // derivedKeyLength is the key size required by AES-256-GCM.
        derivedKeyLength = 32
        // saltLength is the size of the per-file salt.
        saltLength = 16
)

// generateSalt returns a fresh cryptographically random salt for a secrets file.
func generateSalt() ([]byte, error) {
        salt := make([]byte, saltLength)
        if _, err := rand.Read(salt); err != nil {
                return nil, fmt.Errorf("failed to generate salt: %w", err)
        }
        return salt, nil
}

// deriveKey derives the AES-256-GCM key for a secrets file from the user's
// password and that file's salt.
func deriveKey(password, salt []byte) []byte {
        return argon2.IDKey(password, salt, argon2Time, argon2MemoryKiB, argon2Parallelism, derivedKeyLength)
}

// lastDerived memoises the most recently derived key for the lifetime of the
// process.
//
// Deriving a key costs roughly 16ms and 19MiB by design, and a new
// EncryptedManager is built for nearly every operation that touches a secret,
// including once per request in the API server. Without this memo those costs
// are paid per request, which would turn the memory hardness that protects the
// file at rest into a denial of service vector against the server.
//
// A single slot is enough: a process works with one secrets file under one
// password, so a larger cache would only add an eviction policy to maintain.
// Holding the lock across the derivation is deliberate — concurrent callers
// then wait for the first derivation and reuse its result, rather than each
// allocating 19MiB at once.
var (
        lastDerivedMu sync.Mutex
        lastDerived   *derivedKey
)

type derivedKey struct {
        path     string
        salt     []byte
        password []byte
        key      []byte
}

// deriveKeyCached returns the key for the given file and salt, deriving it only
// if it is not the one already memoised.
//
// The memo is matched on the password as well as the file and salt, so a second
// password for the same file derives its own key rather than being handed the
// first one. The password is compared, never hashed: a fast hash of a password
// is the very thing this file exists to avoid.
func deriveKeyCached(filePath string, password, salt []byte) []byte {
        lastDerivedMu.Lock()
        defer lastDerivedMu.Unlock()

        if d := lastDerived; d != nil && d.path == filePath && bytes.Equal(d.salt, salt) &&
                subtle.ConstantTimeCompare(d.password, password) == 1 {
                return d.key
        }

        key := deriveKey(password, salt)
        lastDerived = &derivedKey{
                path:     filePath,
                salt:     bytes.Clone(salt),
                password: bytes.Clone(password),
                key:      key,
        }
        return key
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package keyring provides a composite keyring provider that supports multiple backends.
// It supports macOS Keychain, Windows Credential Manager, and Linux D-Bus Secret Service,
// with keyctl as a fallback on Linux systems.
package keyring

import (
        "fmt"
        "log/slog"
        "runtime"
        "sync"
)

const linuxOS = "linux"

type compositeProvider struct {
        providers []Provider
        active    Provider
        mu        sync.RWMutex
}

// NewCompositeProvider creates a new composite keyring provider that combines multiple backends.
// It uses zalando/go-keyring as the primary provider and keyctl as a fallback on Linux.
//
//nolint:staticcheck
func NewCompositeProvider() Provider {
        var providers []Provider

        // Add zalando/go-keyring as primary provider
        // Handles macOS, Windows, and Linux D-Bus natively
        zkProvider := NewZalandoKeyringProvider()
        providers = append(providers, zkProvider)

        // Add keyctl provider as fallback ONLY on Linux
        if runtime.GOOS == linuxOS {
                if keyctl, err := NewKeyctlProvider(); err == nil {
                        providers = append(providers, keyctl)
                }
        }

        return &compositeProvider{
                providers: providers,
        }
}

func (c *compositeProvider) getActiveProvider() Provider {
        // First, try to read the active provider with a read lock
        c.mu.RLock()
        if c.active != nil && c.active.IsAvailable() {
                active := c.active
                c.mu.RUnlock()
                return active
        }
        c.mu.RUnlock()

        // If no active provider or it's not available, find a new one with write lock
        c.mu.Lock()
        defer c.mu.Unlock()

        // Double-check pattern: another goroutine might have set c.active while we were waiting
        if c.active != nil && c.active.IsAvailable() {
                return c.active
        }

        // Find the first available provider
        for _, provider := range c.providers {
                if provider.IsAvailable() {
                        if c.active == nil || c.active.Name() != provider.Name() {
                                // Log provider selection if logger is available
                                // Use fmt.Printf as fallback since logger might not be initialized in tests
                                c.logProviderSelection(provider.Name())
                        }
                        c.active = provider
                        return provider
                }
        }

        return nil
}

// logProviderSelection safely logs the provider selection
func (*compositeProvider) logProviderSelection(providerName string) {
        // Try to use the logger, but don't panic if it's not available
        defer func() {
                if r := recover(); r != nil {
                        // Logger not available, use fallback
                        fmt.Printf("Using keyring provider: %s\n", providerName)
                }
        }()

        //nolint:gosec // G706: provider name is from internal provider implementations
        slog.Debug("Using keyring provider", "provider", providerName)
}

func (c *compositeProvider) Set(service, key, value string) error {
        provider := c.getActiveProvider()
        if provider == nil {
                return fmt.Errorf("no keyring provider available")
        }
        return provider.Set(service, key, value)
}

func (c *compositeProvider) Get(service, key string) (string, error) {
        provider := c.getActiveProvider()
        if provider == nil {
                return "", fmt.Errorf("no keyring provider available")
        }
        return provider.Get(service, key)
}

func (c *compositeProvider) Delete(service, key string) error {
        provider := c.getActiveProvider()
        if provider == nil {
                return fmt.Errorf("no keyring provider available")
        }
        return provider.Delete(service, key)
}

func (c *compositeProvider) DeleteAll(service string) error {
        provider := c.getActiveProvider()
        if provider == nil {
                return fmt.Errorf("no keyring provider available")
        }
        return provider.DeleteAll(service)
}

func (c *compositeProvider) IsAvailable() bool {
        return c.getActiveProvider() != nil
}

func (c *compositeProvider) Name() string {
        if provider := c.getActiveProvider(); provider != nil {
                return provider.Name()
        }
        return "None Available"
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package keyring

import (
        "errors"
        "runtime"

        "github.com/zalando/go-keyring"
)

type dbusWrapperProvider struct{}

// NewZalandoKeyringProvider creates a new provider that wraps zalando/go-keyring.
// This provider supports macOS Keychain, Windows Credential Manager, and Linux D-Bus Secret Service.
func NewZalandoKeyringProvider() Provider {
        return &dbusWrapperProvider{}
}

func (*dbusWrapperProvider) Set(service, key, value string) error {
        return keyring.Set(service, key, value)
}

func (*dbusWrapperProvider) Get(service, key string) (string, error) {
        value, err := keyring.Get(service, key)
        if errors.Is(err, keyring.ErrNotFound) {
                return "", ErrNotFound
        }
        return value, err
}

func (*dbusWrapperProvider) Delete(service, key string) error {
        return keyring.Delete(service, key)
}

func (*dbusWrapperProvider) DeleteAll(service string) error {
        return keyring.DeleteAll(service)
}

func (*dbusWrapperProvider) IsAvailable() bool {
        // Test keyring availability with a unique test key to avoid conflicts
        testKey := GenerateUniqueTestKey()
        testValue := "test"

        if err := keyring.Set("toolhive-test", testKey, testValue); err != nil {
                return false
        }

        // Clean up the test key
        _ = keyring.Delete("toolhive-test", testKey)
        return true
}

func (*dbusWrapperProvider) Name() string {
        switch runtime.GOOS {
        case "darwin":
                return "macOS Keychain"
        case "windows":
                return "Windows Credential Manager"
        case linuxOS:
                return "D-Bus Secret Service"
        default:
                return "Platform Keyring"
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

//go:build linux

package keyring

import (
        "fmt"
        "sync"

        "golang.org/x/sys/unix"
)

type keyctlProvider struct {
        ringID int
        mu     sync.RWMutex
        keys   map[string]map[string]int // service -> key -> keyid mapping
}

// NewKeyctlProvider creates a new keyring provider using Linux keyctl.
// It initializes access to the user keyring for persistence across process invocations.
func NewKeyctlProvider() (Provider, error) {
        // Use user keyring for persistence across process invocations
        ringID, err := unix.KeyctlGetKeyringID(unix.KEY_SPEC_USER_KEYRING, false)
        if err != nil {
                return nil, fmt.Errorf("could not get user keyring: %w", err)
        }

        // Link to thread keyring for reads
        _, err = unix.KeyctlInt(unix.KEYCTL_LINK, ringID, unix.KEY_SPEC_THREAD_KEYRING, 0, 0)
        if err != nil {
                return nil, fmt.Errorf("unable to link user keyring to thread keyring: %w", err)
        }

        return &keyctlProvider{
                ringID: ringID,
                keys:   make(map[string]map[string]int),
        }, nil
}

func (k *keyctlProvider) Set(service, key, value string) error {
        k.mu.Lock()
        defer k.mu.Unlock()

        keyName := fmt.Sprintf("%s:%s", service, key)
        keyID, err := unix.AddKey("user", keyName, []byte(value), k.ringID)
        if err != nil {
                return fmt.Errorf("failed to set key '%s' in user keyring: %w", keyName, err)
        }

        if k.keys[service] == nil {
                k.keys[service] = make(map[string]int)
        }
        k.keys[service][key] = keyID

        return nil
}

func (k *keyctlProvider) Get(service, key string) (string, error) {
        k.mu.RLock()
        defer k.mu.RUnlock()

        keyName := fmt.Sprintf("%s:%s", service, key)
        keyID, err := unix.KeyctlSearch(k.ringID, "user", keyName, 0)
        if err != nil {
                return "", ErrNotFound
        }

        bufSize := 2048
        buf := make([]byte, bufSize)
        readBytes, err := unix.KeyctlBuffer(unix.KEYCTL_READ, keyID, buf, bufSize)
        if err != nil {
                return "", fmt.Errorf("read of key '%s' failed: %w", keyName, err)
        }

        if readBytes > bufSize {
                return "", fmt.Errorf("buffer too small for keyring payload")
        }

        return string(buf[:readBytes]), nil
}

func (k *keyctlProvider) Delete(service, key string) error {
        k.mu.Lock()
        defer k.mu.Unlock()
        return k.deleteKeyUnlocked(service, key)
}

// deleteKeyUnlocked performs the actual key deletion without acquiring the lock.
// This is used internally by Delete and DeleteAll to avoid deadlocks.
func (k *keyctlProvider) deleteKeyUnlocked(service, key string) error {
        keyName := fmt.Sprintf("%s:%s", service, key)
        keyID, err := unix.KeyctlSearch(k.ringID, "user", keyName, 0)
        if err != nil {
                return nil
        }

        // Unlink the key from the keyring first so it's no longer searchable,
        // then revoke it to invalidate any remaining references.
        _, err = unix.KeyctlInt(unix.KEYCTL_UNLINK, keyID, k.ringID, 0, 0)
        if err != nil {
                return fmt.Errorf("failed to unlink key '%s': %w", keyName, err)
        }
        // Best-effort revoke — key is already removed from keyring
        _, _ = unix.KeyctlInt(unix.KEYCTL_REVOKE, keyID, 0, 0, 0)

        // Remove from tracking
        if serviceKeys, exists := k.keys[service]; exists {
                delete(serviceKeys, key)
                if len(serviceKeys) == 0 {
                        delete(k.keys, service)
                }
        }

        return nil
}

func (k *keyctlProvider) DeleteAll(service string) error {
        k.mu.Lock()
        defer k.mu.Unlock()

        // Always try to delete the known key pattern from kernel keyring
        // regardless of what's in the in-memory map. This ensures cross-process
        // deletion works since the in-memory map is not persisted.
        if err := k.deleteKeyUnlocked(service, service); err != nil {
                return err
        }

        // Also delete any keys tracked in the in-memory map (for keys added in this process)
        if serviceKeys, exists := k.keys[service]; exists {
                var lastErr error
                for key := range serviceKeys {
                        if key == service {
                                // Already deleted above
                                continue
                        }
                        if err := k.deleteKeyUnlocked(service, key); err != nil {
                                lastErr = err
                        }
                }
                delete(k.keys, service)
                if lastErr != nil {
                        return lastErr
                }
        }

        return nil
}

func (k *keyctlProvider) IsAvailable() bool {
        // Use a unique test key to avoid conflicts with concurrent calls
        testKey := GenerateUniqueTestKey()
        testValue := "test"

        if err := k.Set("toolhive-test", testKey, testValue); err != nil {
                return false
        }

        _ = k.Delete("toolhive-test", testKey)
        return true
}

func (*keyctlProvider) Name() string {
        return "Linux Keyctl"
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package keyring

import (
        "crypto/rand"
        "fmt"
        "time"
)

// GenerateUniqueTestKey creates a unique key name used for keyring availability checks.
// It combines a timestamp and random bytes to prevent naming collisions
// when multiple checks run concurrently.
func GenerateUniqueTestKey() string {
        randomBytes := make([]byte, 4)
        if _, err := rand.Read(randomBytes); err != nil {
                return fmt.Sprintf("toolhive-keyring-test-%d", time.Now().UnixNano())
        }

        return fmt.Sprintf("toolhive-keyring-test-%d-%x", time.Now().UnixNano(), randomBytes)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "context"
        "fmt"
        "log/slog"
        "strings"
)

// KeyMigration describes a single key rename operation from OldKey to NewKey.
type KeyMigration struct {
        OldKey string
        NewKey string
}

// SystemKeyPrefixMappings maps known bare key prefixes to their target scope.
// Used by DiscoverMigrations to find keys that need migrating.
var SystemKeyPrefixMappings = []struct {
        Prefix string
        Scope  SecretScope
}{
        {"BEARER_TOKEN_", ScopeWorkloads},
        {"OAUTH_CLIENT_SECRET_", ScopeWorkloads},
        {"OAUTH_REFRESH_TOKEN_", ScopeWorkloads},
        {"registry-user-", ScopeWorkloads},
        {"registry-default-", ScopeWorkloads},
        {"BUILD_AUTH_FILE_", ScopeWorkloads},
        {"REGISTRY_OAUTH_", ScopeRegistry},
}

// MigrateSystemKeys renames keys from OldKey to NewKey in provider.
// The migration is idempotent: if the scoped key already exists the bare key
// is deleted without overwriting the scoped value, so a repeated or partial run
// never clobbers data that was already written under the new name.
// Write-before-delete ordering ensures that a crash between the two operations
// leaves the secret reachable under the new key. Keys that do not exist in
// provider are silently skipped, making the function safe to retry.
func MigrateSystemKeys(ctx context.Context, provider Provider, keyMigrations []KeyMigration) error {
        for _, m := range keyMigrations {
                // If the scoped key already exists (e.g. from a partial prior run),
                // skip the write and just clean up the bare key.
                _, err := provider.GetSecret(ctx, m.NewKey)
                if err == nil {
                        slog.Debug("migration: scoped key already exists, skipping write",
                                "old_key", m.OldKey, "new_key", m.NewKey)
                        if delErr := provider.DeleteSecret(ctx, m.OldKey); delErr != nil && !IsNotFoundError(delErr) {
                                return fmt.Errorf("migration: deleting stale bare key %q: %w", m.OldKey, delErr)
                        }
                        continue
                }
                if !IsNotFoundError(err) {
                        return fmt.Errorf("migration: checking scoped key %q: %w", m.NewKey, err)
                }

                val, err := provider.GetSecret(ctx, m.OldKey)
                if IsNotFoundError(err) {
                        continue
                }
                if err != nil {
                        return fmt.Errorf("migration: reading %q: %w", m.OldKey, err)
                }

                if err := provider.SetSecret(ctx, m.NewKey, val); err != nil {
                        return fmt.Errorf("migration: writing %q: %w", m.NewKey, err)
                }

                if err := provider.DeleteSecret(ctx, m.OldKey); err != nil {
                        return fmt.Errorf("migration: deleting %q: %w", m.OldKey, err)
                }
        }
        return nil
}

// DiscoverMigrations lists all secrets in provider and returns the set of
// migrations needed to move system-owned keys into their scoped namespaces.
// Keys that already start with SystemKeyPrefix are skipped (already migrated).
func DiscoverMigrations(ctx context.Context, provider Provider) ([]KeyMigration, error) {
        all, err := provider.ListSecrets(ctx)
        if err != nil {
                return nil, fmt.Errorf("migration discovery: listing secrets: %w", err)
        }

        var keyMigrations []KeyMigration
        for _, desc := range all {
                key := desc.Key
                // Skip already-migrated keys.
                if IsSystemKey(key) {
                        continue
                }
                for _, mapping := range SystemKeyPrefixMappings {
                        if strings.HasPrefix(key, mapping.Prefix) {
                                keyMigrations = append(keyMigrations, KeyMigration{
                                        OldKey: key,
                                        NewKey: SystemKeyPrefix + string(mapping.Scope) + "_" + key,
                                })
                                break
                        }
                }
        }
        return keyMigrations, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package secrets

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "strings"
)

// SecretScope is the type for system-managed secret scope identifiers.
//
// Invariants that every SecretScope value MUST satisfy:
//   - Non-empty: an empty scope would produce the prefix "__thv__", which is
//     ambiguous and cannot be reliably stripped.
//   - No underscores: the key format is "__thv_<scope>_<name>"; an underscore
//     inside the scope would make it impossible to determine where the scope
//     ends and the name begins.
//
// All constants declared in this package (ScopeRegistry, ScopeWorkloads,
// ScopeAuth, ScopeLLM) satisfy these invariants. Custom scopes introduced in
// the future must be validated against them.
type SecretScope string

const (
        // SystemKeyPrefix is the prefix used for all system-managed secret keys.
        // Any key starting with this prefix is reserved for internal use.
        // The double-underscore and trailing underscore make it visually distinct
        // and avoid conflicts with backends (e.g. 1Password) that treat "/" as a
        // path separator.
        SystemKeyPrefix = "__thv_"

        // ScopeRegistry is the scope for registry OAuth refresh tokens.
        ScopeRegistry SecretScope = "registry"

        // ScopeWorkloads is the scope for remote workload authentication tokens
        // (OAuth client secrets, bearer tokens, OAuth refresh tokens).
        ScopeWorkloads SecretScope = "workloads"

        // ScopeAuth is reserved for enterprise CLI/Desktop login tokens.
        ScopeAuth SecretScope = "auth"

        // ScopeLLM is the scope for LLM gateway OIDC refresh tokens.
        ScopeLLM SecretScope = "llm"
)

// ErrReservedKeyName is returned when a user command attempts to manage a
// secret whose name is reserved for system use.
var ErrReservedKeyName = errors.New("secret name is reserved for system use and cannot be managed via user commands")

// ScopedProvider wraps a Provider and namespaces all operations under a
// system-managed scope prefix ("__thv_<scope>_"). It is intended for
// internal callers such as registry auth and workload auth that need isolated
// key spaces inside the shared secrets store.
type ScopedProvider struct {
        provider Provider
        scope    SecretScope
}

// NewScopedProvider creates a Provider that transparently prefixes every key
// with "__thv_<scope>_", keeping system secrets isolated from user secrets.
func NewScopedProvider(inner Provider, scope SecretScope) Provider {
        return &ScopedProvider{
                provider: inner,
                scope:    scope,
        }
}

// GetSecret retrieves the secret identified by name under this provider's scope.
// If the scoped key is not found, it falls back to the bare (pre-migration) key.
// This makes the provider safe to use before or during secret scope migration:
// once migration completes and bare keys are deleted, the fallback is a no-op.
func (s *ScopedProvider) GetSecret(ctx context.Context, name string) (string, error) {
        val, err := s.provider.GetSecret(ctx, s.getScopedKey(name))
        if err == nil {
                return val, nil
        }
        if IsNotFoundError(err) {
                // Migration window: the scoped key does not exist yet. Try the bare key
                // that was used before secret scope migration ran. After migration
                // completes and the bare key is deleted, this lookup returns not-found
                // and we fall through to return the original scoped-key error.
                bareVal, bareErr := s.provider.GetSecret(ctx, name)
                if bareErr == nil {
                        slog.Debug("secret scope migration fallback: returning bare key",
                                "scope", s.scope, "name", name)
                        return bareVal, nil
                }
                if !IsNotFoundError(bareErr) {
                        // Bare-key lookup hit a real backend error (e.g. connection failure,
                        // auth error). Surface it so the caller doesn't misdiagnose a backend
                        // problem as "secret not found".
                        return "", fmt.Errorf("scoped key not found and bare-key fallback failed: %w", bareErr)
                }
        }
        return "", err
}

// SetSecret stores value under the scoped key for name.
func (s *ScopedProvider) SetSecret(ctx context.Context, name, value string) error {
        return s.provider.SetSecret(ctx, s.getScopedKey(name), value)
}

// DeleteSecret removes the scoped key for name from the underlying store.
func (s *ScopedProvider) DeleteSecret(ctx context.Context, name string) error {
        return s.provider.DeleteSecret(ctx, s.getScopedKey(name))
}

// ListSecrets returns only the entries that belong to this provider's scope,
// with the "__thv_<scope>_" prefix stripped from each Key so callers receive bare names.
func (s *ScopedProvider) ListSecrets(ctx context.Context) ([]SecretDescription, error) {
        all, err := s.provider.ListSecrets(ctx)
        if err != nil {
                return nil, err
        }

        prefix := s.getScopePrefix()
        var result []SecretDescription
        for _, desc := range all {
                if strings.HasPrefix(desc.Key, prefix) {
                        result = append(result, SecretDescription{
                                Key:         strings.TrimPrefix(desc.Key, prefix),
                                Description: desc.Description,
                        })
                }
        }
        return result, nil
}

// DeleteSecrets removes all named keys under this scope by delegating to the inner provider.
func (s *ScopedProvider) DeleteSecrets(ctx context.Context, names []string) error {
        keys := make([]string, len(names))
        for i, name := range names {
                keys[i] = s.getScopedKey(name)
        }
        return s.provider.DeleteSecrets(ctx, keys)
}

// Cleanup removes only the secrets that belong to this scope, leaving all
// other secrets untouched.
func (s *ScopedProvider) Cleanup() error {
        ctx := context.Background()

        all, err := s.provider.ListSecrets(ctx)
        if err != nil {
                return err
        }

        prefix := s.getScopePrefix()
        var toDelete []string
        for _, desc := range all {
                if strings.HasPrefix(desc.Key, prefix) {
                        toDelete = append(toDelete, desc.Key)
                }
        }
        if len(toDelete) == 0 {
                return nil
        }
        return s.provider.DeleteSecrets(ctx, toDelete)
}

// Capabilities delegates to the underlying provider.
func (s *ScopedProvider) Capabilities() ProviderCapabilities {
        return s.provider.Capabilities()
}

// getScopedKey builds the internal storage key in the form "__thv_<scope>_<name>".
func (s *ScopedProvider) getScopedKey(name string) string {
        return SystemKeyPrefix + string(s.scope) + "_" + name
}

// getScopePrefix returns the key prefix for this scope, i.e. "__thv_<scope>_".
func (s *ScopedProvider) getScopePrefix() string {
        return SystemKeyPrefix + string(s.scope) + "_"
}

// UserProvider wraps a Provider and hides all system-reserved keys from
// user-facing callers (CLI, API, MCP tool server). Any attempt to read or
// modify a key that starts with the system prefix is rejected with
// ErrReservedKeyName.
type UserProvider struct {
        provider Provider
}

// NewUserProvider creates a Provider that filters out system-reserved keys so
// that user-facing callers cannot accidentally read or overwrite internal
// secrets managed by ScopedProvider.
func NewUserProvider(inner Provider) Provider {
        return &UserProvider{provider: inner}
}

// GetSecret returns the secret for name, or ErrReservedKeyName if the name is
// a system-reserved key.
func (u *UserProvider) GetSecret(ctx context.Context, name string) (string, error) {
        if IsSystemKey(name) {
                return "", fmt.Errorf("%w: cannot get %q", ErrReservedKeyName, name)
        }
        return u.provider.GetSecret(ctx, name)
}

// SetSecret stores value under name, or returns ErrReservedKeyName if the name
// is system-reserved.
func (u *UserProvider) SetSecret(ctx context.Context, name, value string) error {
        if IsSystemKey(name) {
                return fmt.Errorf("%w: cannot set %q", ErrReservedKeyName, name)
        }
        return u.provider.SetSecret(ctx, name, value)
}

// DeleteSecret removes name from the underlying store, or returns
// ErrReservedKeyName if the name is system-reserved.
func (u *UserProvider) DeleteSecret(ctx context.Context, name string) error {
        if IsSystemKey(name) {
                return fmt.Errorf("%w: cannot delete %q", ErrReservedKeyName, name)
        }
        return u.provider.DeleteSecret(ctx, name)
}

// ListSecrets returns all non-system secrets from the underlying store.
// Entries whose Key starts with the system prefix are silently omitted.
func (u *UserProvider) ListSecrets(ctx context.Context) ([]SecretDescription, error) {
        all, err := u.provider.ListSecrets(ctx)
        if err != nil {
                return nil, err
        }

        var result []SecretDescription
        for _, desc := range all {
                if !IsSystemKey(desc.Key) {
                        result = append(result, desc)
                }
        }
        return result, nil
}

// DeleteSecrets removes all named keys with all-or-nothing semantics: it
// validates every name in the list before issuing any delete to the underlying
// store. If any name is system-reserved the entire operation is aborted and
// ErrReservedKeyName is returned without deleting anything.
func (u *UserProvider) DeleteSecrets(ctx context.Context, names []string) error {
        for _, name := range names {
                if IsSystemKey(name) {
                        return fmt.Errorf("%w: cannot delete %q", ErrReservedKeyName, name)
                }
        }
        return u.provider.DeleteSecrets(ctx, names)
}

// Cleanup removes only user-owned secrets (those that do not start with the
// system prefix). System secrets are managed independently through their own
// ScopedProvider.Cleanup calls and must not be touched here.
func (u *UserProvider) Cleanup() error {
        ctx := context.Background()

        all, err := u.provider.ListSecrets(ctx)
        if err != nil {
                return err
        }

        var toDelete []string
        for _, desc := range all {
                if !IsSystemKey(desc.Key) {
                        toDelete = append(toDelete, desc.Key)
                }
        }
        if len(toDelete) == 0 {
                return nil
        }
        return u.provider.DeleteSecrets(ctx, toDelete)
}

// Capabilities delegates to the underlying provider.
func (u *UserProvider) Capabilities() ProviderCapabilities {
        return u.provider.Capabilities()
}

// IsSystemKey reports whether name is reserved for system use, i.e. whether it
// starts with the system key prefix "__thv_".
func IsSystemKey(name string) bool {
        return strings.HasPrefix(name, SystemKeyPrefix)
}

// ParseSystemKey parses a system-managed key of the form "__thv_<scope>_<name>"
// and returns its scope and name components. ok is false if key does not start
// with SystemKeyPrefix or contains no scope separator.
func ParseSystemKey(key string) (scope, name string, ok bool) {
        rest, found := strings.CutPrefix(key, SystemKeyPrefix)
        if !found {
                return "", "", false
        }
        scope, name, ok = strings.Cut(rest, "_")
        return scope, name, ok
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package secrets contains the secrets management logic for ToolHive.
package secrets

import (
        "context"
        "fmt"
        "regexp"
)

const (
        // EnvVarPrefix is the prefix used for environment variable secrets
        EnvVarPrefix = "TOOLHIVE_SECRET_"
)

// regex to extract name and target from secret parameter, e.g. "name,target=target"
var secretParamRegex = regexp.MustCompile(`^([^,]+),target=(.+)$`)

// ProviderCapabilities represents what operations a secrets provider supports.
type ProviderCapabilities struct {
        CanRead    bool
        CanWrite   bool
        CanDelete  bool
        CanList    bool
        CanCleanup bool
}

// IsReadOnly returns true if the provider only supports read operations.
func (pc ProviderCapabilities) IsReadOnly() bool {
        return pc.CanRead && !pc.CanWrite && !pc.CanDelete && !pc.CanCleanup
}

// IsReadWrite returns true if the provider supports both read and write operations.
func (pc ProviderCapabilities) IsReadWrite() bool {
        return pc.CanRead && pc.CanWrite
}

// String returns a human-readable description of the capabilities.
func (pc ProviderCapabilities) String() string {
        if pc.IsReadWrite() {
                return "read-write"
        }
        if pc.IsReadOnly() {
                return "read-only"
        }
        return "custom"
}

// Provider describes a type which can manage secrets.
type Provider interface {
        GetSecret(ctx context.Context, name string) (string, error)
        SetSecret(ctx context.Context, name, value string) error
        DeleteSecret(ctx context.Context, name string) error
        ListSecrets(ctx context.Context) ([]SecretDescription, error)
        // DeleteSecrets removes all named keys. Read-only providers treat this as a no-op.
        DeleteSecrets(ctx context.Context, keys []string) error
        Cleanup() error
        // Capabilities returns what operations this provider supports
        Capabilities() ProviderCapabilities
}

// SecretParameter represents a parsed `--secret` parameter.
type SecretParameter struct {
        Name   string `json:"name"`
        Target string `json:"target"`
}

// ParseSecretParameter creates an instance of SecretParameter from a string.
// Expected format: `<Name>,target=<Target>`.
func ParseSecretParameter(parameter string) (SecretParameter, error) {
        if parameter == "" {
                return SecretParameter{}, fmt.Errorf("secret parameter cannot be empty")
        }

        // extract name and target using secretParamRegex
        matches := secretParamRegex.FindStringSubmatch(parameter)
        if len(matches) != 3 { // The first element is the full match, followed by capture groups
                return SecretParameter{}, fmt.Errorf("invalid secret parameter format: %s", parameter)
        }

        name := matches[1]
        target := matches[2]

        return SecretParameter{
                Name:   name,
                Target: target,
        }, nil
}

// ToCLIString converts a SecretParameter to CLI format string
func (sp SecretParameter) ToCLIString() string {
        return fmt.Sprintf("%s,target=%s", sp.Name, sp.Target)
}

// SecretParametersToCLI does the reverse of `ParseSecretParameter`
// TODO: It may be possible to get rid of this with refactoring.
func SecretParametersToCLI(params []SecretParameter) []string {
        result := make([]string, len(params))
        for i, p := range params {
                result[i] = p.ToCLIString()
        }
        return result
}

// SecretDescription is returned by `ListSecrets`.
type SecretDescription struct {
        // Key is the unique identifier for the secret, used when retrieving it.
        Key string `json:"key"`
        // Description provides a human-readable description of the secret
        // Particularly useful for 1password.
        // May be empty if no description is available.
        Description string `json:"description"`
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package sentry provides Sentry error tracking and distributed tracing for the ToolHive API server.
package sentry

import (
        "context"
        "fmt"
        "log/slog"
        "net/http"
        "sync"
        "sync/atomic"
        "time"

        "github.com/getsentry/sentry-go"
        sentryotel "github.com/getsentry/sentry-go/otel"
        sentryotlp "github.com/getsentry/sentry-go/otel/otlp"
        sdktrace "go.opentelemetry.io/otel/sdk/trace"

        "github.com/stacklok/toolhive/pkg/telemetry"
        "github.com/stacklok/toolhive/pkg/updates"
        "github.com/stacklok/toolhive/pkg/versions"
)

const flushTimeout = 2 * time.Second

const (
        // environmentKey and releaseKey are the attribute names Sentry itself uses
        // to carry Environment and Release on OTLP payloads (see sentry-go's
        // log.go/metrics.go), so exported spans must use the same names to be
        // grouped alongside Issues.
        environmentKey = "sentry.environment"
        releaseKey     = "sentry.release"
        // instanceIDKey carries the anonymous instance ID on both Sentry events
        // (as a scope tag) and exported spans (as a resource attribute), so Issues
        // and Traces can be correlated with toolhive-studio by the same value.
        instanceIDKey = "custom.instance_id"
)

// initialized tracks whether Sentry was successfully initialized.
var initialized atomic.Bool

// spanProcessor is the single Sentry OTLP span processor for this process,
// created on the first Init and reused by every subsequent one. Guarded by
// spanProcessorMu.
var (
        spanProcessorMu sync.Mutex
        spanProcessor   sdktrace.SpanProcessor
)

// Config holds the configuration for Sentry integration.
type Config struct {
        // DSN is the Sentry Data Source Name. When empty, Sentry is disabled.
        DSN string
        // Environment identifies the deployment environment (e.g. "production", "development").
        Environment string
        // TracesSampleRate controls the percentage of transactions captured for
        // performance monitoring (0.0–1.0).
        TracesSampleRate float64
        // Debug enables Sentry SDK debug logging.
        Debug bool
}

// Init initializes the Sentry SDK with the given configuration.
// If the DSN is empty, initialization is skipped and all Sentry operations become no-ops.
func Init(cfg Config) error {
        if cfg.DSN == "" {
                slog.Debug("sentry disabled (no DSN configured)")
                return nil
        }

        vi := versions.GetVersionInfo()
        // Reused verbatim as a span resource attribute below so Issues and Traces
        // report the same release string.
        release := fmt.Sprintf("toolhive@%s", vi.Version)

        err := sentry.Init(sentry.ClientOptions{
                Dsn:              cfg.DSN,
                Environment:      cfg.Environment,
                Release:          release,
                TracesSampleRate: cfg.TracesSampleRate,
                Debug:            cfg.Debug,
                EnableTracing:    true,
                AttachStacktrace: true,
                DataCollection:   noPIIDataCollection(),
                Integrations: func(integrations []sentry.Integration) []sentry.Integration {
                        return append(integrations, sentryotel.NewOtelIntegration())
                },
        })
        if err != nil {
                return fmt.Errorf("sentry init: %w", err)
        }

        if err := registerTraceExporter(cfg); err != nil {
                return err
        }
        initialized.Store(true)
        slog.Debug("sentry initialized", "environment", cfg.Environment)

        // Tag every event and transaction with the anonymous instance ID so that
        // Sentry events from the API server can be correlated with those from
        // toolhive-studio. Note: toolhive-studio currently uses "custom.user_id"
        // for the same value; these should be aligned to "custom.instance_id" in
        // both repos in a follow-up to avoid misleading PII detection heuristics.
        instanceID := ""
        if id, err := updates.TryGetAnonymousID(); err == nil && id != "" {
                instanceID = id
                sentry.ConfigureScope(func(scope *sentry.Scope) {
                        scope.SetTag(instanceIDKey, id)
                })
                slog.Debug("sentry anonymous instance ID tagged", "id", id)
        }

        // Spans are exported straight to Sentry's OTLP endpoint and never pass
        // through the Sentry client, so neither ClientOptions nor the scope
        // configured above reach them. Environment, release and instance ID have to
        // travel as OTEL resource attributes instead, or Traces would lose the
        // grouping that Issues keep and the two would disagree.
        telemetry.RegisterResourceAttributes(resourceAttributes(cfg.Environment, release, instanceID))

        return nil
}

// registerTraceExporter registers the Sentry OTLP span processor with the global
// OTEL registry, creating it on first use.
//
// The processor is cached because the registry deduplicates by pointer identity:
// sdktrace.NewBatchSpanProcessor allocates a fresh processor on every call, so
// without this a second Init (config reload, or a test that does not reset the
// registry) would register a second processor and double-export every span while
// leaking the first exporter's goroutine.
//
// Caching means a second Init keeps the first call's DSN. thv serve calls Init
// exactly once per process, and the registry already only feeds processors to
// providers created after registration, so re-initialising is not supported
// either way.
func registerTraceExporter(cfg Config) error {
        spanProcessorMu.Lock()
        defer spanProcessorMu.Unlock()

        if spanProcessor == nil {
                exporter, err := sentryotlp.NewTraceExporter(context.Background(), cfg.DSN)
                if err != nil {
                        return fmt.Errorf("create Sentry trace exporter: %w", err)
                }
                spanProcessor = sdktrace.NewBatchSpanProcessor(exporter)
        }

        telemetry.RegisterSpanProcessor(spanProcessor)
        // Spans no longer pass through the Sentry client, so TracesSampleRate has to
        // reach the OTEL sampler or --sentry-traces-sample-rate would be ignored and
        // every span would be exported.
        telemetry.RegisterSamplingRate(cfg.TracesSampleRate)
        slog.Debug("sentry trace exporter registered with OTEL registry",
                "traces_sample_rate", cfg.TracesSampleRate)
        return nil
}

// piiSensitiveTerms mirrors the deny-list that sentry-go applies internally for
// SendDefaultPII=false (its unexported extendedSensitiveTerms). These cover
// client-identifying data that an API server behind a proxy routinely sees:
// forwarding headers, remote addresses and user identifiers.
//
// The list has to be repeated here because sentry-go reaches it only through
// the deprecated SendDefaultPII path; the DataCollection API exposes no way to
// set it. In CollectionDenyList mode a behaviour's Terms are OR-ed with the
// SDK's built-in terms, so passing them per behaviour is equivalent.
//
// Re-check this against sentry-go's extendedSensitiveTerms on SDK upgrades — a
// term added upstream will not reach us automatically.
var piiSensitiveTerms = []string{
        "forwarded",
        "-ip",
        "remote-",
        "via",
        "-user",
}

// noPIIDataCollection returns the DataCollection that replaces the deprecated
// SendDefaultPII=false. It is deliberately equivalent to what sentry-go's
// legacyDataCollection built for that flag: no auto-populated user info, no
// HTTP bodies, no cookies, and headers and query params scrubbed against both
// the built-in and the extended deny-lists.
func noPIIDataCollection() *sentry.DataCollection {
        denyList := func() *sentry.KeyValueCollectionBehavior {
                return &sentry.KeyValueCollectionBehavior{
                        Mode:  sentry.CollectionDenyList,
                        Terms: piiSensitiveTerms,
                }
        }
        return &sentry.DataCollection{
                UserInfo:   sentry.Set(false),
                HTTPBodies: []sentry.BodyType{},
                Cookies:    &sentry.KeyValueCollectionBehavior{Mode: sentry.CollectionOff},
                HTTPHeaders: &sentry.HeaderCollectionConfig{
                        Request:  denyList(),
                        Response: denyList(),
                },
                QueryParams: denyList(),
        }
}

// resourceAttributes returns the OTEL resource attributes Sentry needs to group
// OTLP-ingested traces the same way it groups Issues. Empty values are omitted
// so they do not show up as blank attributes on other OTLP backends.
func resourceAttributes(environment, release, instanceID string) map[string]string {
        attrs := make(map[string]string, 3)
        for key, value := range map[string]string{
                environmentKey: environment,
                releaseKey:     release,
                instanceIDKey:  instanceID,
        } {
                if value != "" {
                        attrs[key] = value
                }
        }
        return attrs
}

// Close flushes buffered Sentry events and shuts down the SDK.
// Safe to call even when Sentry was not initialized.
func Close() {
        if !initialized.Load() {
                return
        }
        sentry.Flush(flushTimeout)
        initialized.Store(false)
        slog.Debug("sentry flushed and closed")
}

// Enabled reports whether the Sentry SDK was successfully initialized.
func Enabled() bool {
        return initialized.Load()
}

// CaptureException reports an error to Sentry using the hub from the request context.
// Falls back to the current hub if no hub is attached to the context.
// No-op when Sentry is not initialized.
//
// The API server's error handler calls this alongside span.RecordError so that
// 5xx errors appear as both OTEL span errors (distributed tracing) and
// standalone Sentry Issues (error tracking). The Sentry OTEL integration links
// those issues to the active OTEL trace; explicit hub calls are required to
// create Issues.
func CaptureException(r *http.Request, err error) {
        if !initialized.Load() || err == nil {
                return
        }
        hub := sentry.GetHubFromContext(r.Context())
        if hub == nil {
                hub = sentry.CurrentHub().Clone()
        }
        client := hub.Client()
        if client == nil {
                return
        }
        event := client.EventFromException(err, sentry.LevelError)
        hub.CaptureEventWithHint(event, &sentry.EventHint{
                OriginalException: err,
                Context:           r.Context(),
        })
}

// RecoverPanic reports a recovered panic value to Sentry.
// No-op when Sentry is not initialized.
func RecoverPanic(r *http.Request, recovered interface{}) {
        if !initialized.Load() || recovered == nil {
                return
        }
        hub := sentry.GetHubFromContext(r.Context())
        if hub == nil {
                hub = sentry.CurrentHub().Clone()
        }
        hub.RecoverWithContext(r.Context(), recovered)
        hub.Flush(flushTimeout)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package gitresolver

import (
        "log/slog"
        "net/url"
        "os"
        "strings"

        "github.com/go-git/go-git/v5/plumbing/transport"
        githttp "github.com/go-git/go-git/v5/plumbing/transport/http"
)

// tokenMapping maps environment variable names to the git hosts they are scoped to.
// Tokens are only sent to their matching host to prevent credential exfiltration.
var tokenMapping = []struct {
        envVar string
        hosts  []string // empty means the token is sent to any host (user opt-in)
}{
        {envVar: "GITHUB_TOKEN", hosts: []string{"github.com"}},
        {envVar: "GITLAB_TOKEN", hosts: []string{"gitlab.com"}},
        {envVar: "GIT_TOKEN", hosts: nil}, // fallback: sent to any host
}

// EnvFunc is a function that looks up an environment variable.
// The default is os.Getenv; tests can inject a custom implementation.
type EnvFunc func(string) string

// ResolveAuth attempts to find authentication credentials from the environment
// scoped to the given clone URL. Returns nil if no credentials match.
//
// Security: tokens are only sent to their designated hosts. GITHUB_TOKEN is
// only sent to github.com, GITLAB_TOKEN only to gitlab.com. GIT_TOKEN is a
// fallback sent to any host.
func ResolveAuth(cloneURL string) transport.AuthMethod {
        return ResolveAuthWith(os.Getenv, cloneURL)
}

// ResolveAuthWith is like ResolveAuth but uses the provided function to look up
// environment variables, making it testable without modifying process state.
func ResolveAuthWith(getenv EnvFunc, cloneURL string) transport.AuthMethod {
        host := extractHost(cloneURL)

        for _, mapping := range tokenMapping {
                token := getenv(mapping.envVar)
                if token == "" {
                        continue
                }
                // If hosts are specified, only send the token to matching hosts.
                if len(mapping.hosts) > 0 && !hostMatches(host, mapping.hosts) {
                        continue
                }
                // Log when the fallback GIT_TOKEN is used for non-standard hosts so
                // users can audit credential usage.
                if len(mapping.hosts) == 0 {
                        slog.Debug("Using fallback GIT_TOKEN for non-standard host — verify this is intended",
                                "env_var", mapping.envVar, "host", host)
                } else {
                        slog.Debug("Using git authentication from environment", "env_var", mapping.envVar)
                }
                return &githttp.BasicAuth{
                        Username: "x-access-token",
                        Password: token,
                }
        }

        return nil
}

// extractHost returns the lowercase hostname from a URL, or empty string on failure.
func extractHost(rawURL string) string {
        parsed, err := url.Parse(rawURL)
        if err != nil {
                return ""
        }
        return strings.ToLower(parsed.Hostname())
}

// hostMatches checks if host matches any of the allowed hosts (case-insensitive).
func hostMatches(host string, allowed []string) bool {
        for _, h := range allowed {
                if strings.EqualFold(host, h) {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package gitresolver

import (
        "fmt"
        "net"
        "os"
        "path"
        "strings"

        "github.com/stacklok/toolhive/pkg/networking"
)

const gitScheme = "git://"

// GitReference represents a parsed git:// skill reference.
type GitReference struct {
        // URL is the HTTPS clone URL (e.g., https://github.com/org/repo)
        URL string
        // Path is the subdirectory within repo (e.g., "path/to/skill"), empty = repo root
        Path string
        // Ref is the git ref: branch, tag, or commit (e.g., "v1.0.0"), empty = default branch
        Ref string
}

// IsGitReference returns true if name starts with "git://".
func IsGitReference(name string) bool {
        return strings.HasPrefix(name, gitScheme)
}

// ParseGitReference parses a git:// skill reference.
//
// Format: git://host/owner/repo[@ref][#path/to/skill]
//
// Examples:
//   - git://github.com/org/repo
//   - git://github.com/org/repo@v1.0.0
//   - git://github.com/org/repo#skills/my-skill
//   - git://github.com/org/repo@main#skills/my-skill
func ParseGitReference(raw string) (*GitReference, error) {
        if !IsGitReference(raw) {
                return nil, fmt.Errorf("not a git reference: must start with %q", gitScheme)
        }

        // Strip scheme
        rest := raw[len(gitScheme):]

        // Split off fragment (#path)
        var skillPath string
        if idx := strings.Index(rest, "#"); idx >= 0 {
                skillPath = rest[idx+1:]
                rest = rest[:idx]
        }

        // Split off ref (@ref)
        var ref string
        if idx := strings.Index(rest, "@"); idx >= 0 {
                ref = rest[idx+1:]
                rest = rest[:idx]
        }

        // rest is now "host/owner/repo" (or "host/owner/repo/...")
        if rest == "" {
                return nil, fmt.Errorf("invalid git reference: empty host/path")
        }

        // Extract host
        slashIdx := strings.Index(rest, "/")
        if slashIdx < 0 {
                return nil, fmt.Errorf("invalid git reference: no repository path after host")
        }
        host := rest[:slashIdx]
        repoPath := rest[slashIdx+1:]

        // Validate host
        if err := validateHost(host); err != nil {
                return nil, fmt.Errorf("invalid git reference: %w", err)
        }

        // Validate repo path has at least owner/repo
        if repoPath == "" || !strings.Contains(repoPath, "/") {
                return nil, fmt.Errorf("invalid git reference: repository path must be at least owner/repo")
        }

        // Validate ref
        if err := validateRef(ref); err != nil {
                return nil, fmt.Errorf("invalid git reference: %w", err)
        }

        // Validate skill path
        if err := validateSkillPath(skillPath); err != nil {
                return nil, fmt.Errorf("invalid git reference: %w", err)
        }

        // Build clone URL. In dev mode, use HTTP to support local test servers.
        scheme := "https"
        if isDevMode() {
                scheme = "http"
        }
        cloneURL := scheme + "://" + host + "/" + repoPath

        return &GitReference{
                URL:  cloneURL,
                Path: skillPath,
                Ref:  ref,
        }, nil
}

// SkillName extracts the expected skill name from the reference.
// Uses the last component of Path if set, otherwise the last component of the repo URL.
func (r *GitReference) SkillName() string {
        if r.Path != "" {
                return path.Base(r.Path)
        }
        // Extract from URL: "https://github.com/org/repo" -> "repo"
        trimmed := strings.TrimSuffix(r.URL, ".git")
        return path.Base(trimmed)
}

// validateHost checks the host is not localhost, a private IP, or empty.
// Reuses pkg/networking SSRF utilities as the single source of truth.
//
// NOTE: This check only validates literal IPs and known localhost strings.
// Hostnames that DNS-resolve to private IPs (DNS rebinding) are NOT caught here
// because go-git does not expose a DialContext hook. A pre-clone DNS resolution
// check could be added as defense-in-depth.
func validateHost(host string) error {
        if host == "" {
                return fmt.Errorf("host must not be empty")
        }

        // Strip port if present
        hostname := host
        if h, _, err := net.SplitHostPort(host); err == nil {
                hostname = h
        }

        // In dev mode, allow localhost/private IPs for testing. This mirrors
        // the OCI plain-HTTP dev mode enabled by TOOLHIVE_DEV=true.
        if !isDevMode() {
                // Reject localhost variants using the shared networking utility.
                if networking.IsLocalhost(hostname) {
                        return fmt.Errorf("host %q is not allowed: localhost is rejected for SSRF prevention", host)
                }

                // Reject private/loopback IPs using the shared networking utility.
                ip := net.ParseIP(hostname)
                if ip != nil && networking.IsPrivateIP(ip) {
                        return fmt.Errorf("host %q is not allowed: private/loopback IPs are rejected for SSRF prevention", host)
                }
        }

        return nil
}

// validateRef checks that the ref doesn't contain shell metacharacters.
func validateRef(ref string) error {
        if ref == "" {
                return nil
        }
        // Reject characters that could be used in shell injection or path traversal
        for _, c := range ref {
                switch {
                case c >= 'a' && c <= 'z',
                        c >= 'A' && c <= 'Z',
                        c >= '0' && c <= '9',
                        c == '.', c == '-', c == '_', c == '/':
                        continue
                default:
                        return fmt.Errorf("ref %q contains invalid character %q", ref, c)
                }
        }
        if strings.Contains(ref, "..") {
                return fmt.Errorf("ref %q must not contain '..' segments", ref)
        }
        return nil
}

// validateSkillPath checks that the path doesn't contain traversal, null bytes,
// absolute paths, or backslashes.
func validateSkillPath(p string) error {
        if p == "" {
                return nil
        }
        if strings.ContainsRune(p, 0) {
                return fmt.Errorf("path contains null bytes")
        }
        if strings.HasPrefix(p, "/") || strings.HasPrefix(p, "\\") {
                return fmt.Errorf("path %q must be relative", p)
        }
        if strings.Contains(p, "\\") {
                return fmt.Errorf("path %q must not contain backslashes", p)
        }
        for _, segment := range strings.Split(p, "/") {
                if segment == ".." {
                        return fmt.Errorf("path %q must not contain '..' traversal segments", p)
                }
        }
        return nil
}

// isDevMode returns true when TOOLHIVE_DEV=true. In dev mode, SSRF checks
// for localhost and private IPs are relaxed to enable E2E testing with local
// git HTTP servers.
func isDevMode() bool {
        return strings.EqualFold(os.Getenv("TOOLHIVE_DEV"), "true")
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package gitresolver resolves skill installations from git repositories.
package gitresolver

//go:generate mockgen -destination=mocks/mock_resolver.go -package=mocks -source=resolver.go Resolver

import (
        "context"
        "fmt"
        "io/fs"
        "path"
        "regexp"
        "time"

        "github.com/go-git/go-git/v5/plumbing/object"

        "github.com/stacklok/toolhive/pkg/git"
        "github.com/stacklok/toolhive/pkg/skills"
)

// CloneTimeout is the maximum time allowed for cloning a git repository.
// Exported so callers (e.g. pluginsvc.installFromGit) can reuse the same
// timeout without duplicating the constant.
const CloneTimeout = 2 * time.Minute

// semverLike matches refs that look like semantic version tags (v1.0, v1.2.3, v1.2.3-rc1, etc.).
// Requires at least one dot-separated numeric segment after the major version to avoid matching
// branch names like "v1-beta-branch".
var semverLike = regexp.MustCompile(`^v\d+\.\d+(\.\d+)*(-[a-zA-Z0-9._-]+)?$`)

// Resolver clones a git repository and extracts skill files.
type Resolver interface {
        // Resolve clones the repo, validates the skill, and returns the skill
        // directory contents as files ready for installation.
        Resolve(ctx context.Context, ref *GitReference) (*ResolveResult, error)
}

// ResolveResult contains the outcome of resolving a git skill reference.
type ResolveResult struct {
        // SkillConfig is the parsed SKILL.md
        SkillConfig *skills.ParseResult
        // Files is all files in the skill directory
        Files []FileEntry
        // CommitHash is the git commit hash (for digest/upgrade detection)
        CommitHash string
        // CommitSignature is the armored signature attached to the resolved
        // commit, empty when the commit is unsigned. It is UNVERIFIED here —
        // the install-time verifier checks it cryptographically before it is
        // trusted or recorded as provenance.
        CommitSignature string
        // CommitPayload is the encoded commit object without its signature —
        // the bytes CommitSignature signs.
        CommitPayload []byte
}

// FileEntry represents a single file from the cloned repository.
type FileEntry struct {
        Path    string
        Content []byte
        Mode    fs.FileMode
}

// ResolverOption configures a defaultResolver.
type ResolverOption func(*defaultResolver)

// WithGitClient sets a fixed git client, bypassing per-clone auth resolution.
// Primarily used for testing with mock clients.
func WithGitClient(client git.Client) ResolverOption {
        return func(r *defaultResolver) {
                r.fixedClient = client
        }
}

// NewResolver creates a new git skill resolver.
func NewResolver(opts ...ResolverOption) Resolver {
        r := &defaultResolver{}
        for _, o := range opts {
                o(r)
        }
        return r
}

type defaultResolver struct {
        // fixedClient, when set, is used for all clones (testing).
        // When nil, a new client is created per-clone with host-scoped auth.
        fixedClient git.Client
}

// ClientForURL returns a git client for the given clone URL. When fixedClient
// is non-nil it is returned as-is (testing); otherwise a new client is created
// with host-scoped auth from the environment. Exported so
// pluginsvc.installFromGit can reuse the same auth-resolution path without
// duplicating clientForURL.
func ClientForURL(cloneURL string, fixedClient git.Client) git.Client {
        if fixedClient != nil {
                return fixedClient
        }
        auth := ResolveAuth(cloneURL)
        var opts []git.ClientOption
        if auth != nil {
                opts = append(opts, git.WithAuth(auth))
        }
        return git.NewDefaultGitClient(opts...)
}

// CloneConfigForRef classifies ref as commit/tag/branch and returns the
// matching git.CloneConfig (with URL set from ref.URL). Exported so
// pluginsvc.installFromGit can reuse the same ref-classification logic
// without duplicating semverLike/isHex.
func CloneConfigForRef(ref *GitReference) *git.CloneConfig {
        cfg := &git.CloneConfig{URL: ref.URL}
        if ref.Ref != "" {
                switch {
                case len(ref.Ref) == 40 && isHex(ref.Ref):
                        // Full commit hash → checkout specific commit
                        cfg.Commit = ref.Ref
                case semverLike.MatchString(ref.Ref):
                        // Semver-like pattern (v1.0.0) → clone as tag
                        cfg.Tag = ref.Ref
                default:
                        // Everything else → treat as branch
                        cfg.Branch = ref.Ref
                }
        }
        return cfg
}

// Resolve clones a git repository and extracts skill files from it.
func (r *defaultResolver) Resolve(ctx context.Context, ref *GitReference) (*ResolveResult, error) {
        // Enforce a clone timeout to prevent indefinite hangs from slow/malicious servers.
        ctx, cancel := context.WithTimeout(ctx, CloneTimeout)
        defer cancel()

        // Build clone config from the git reference
        cloneConfig := CloneConfigForRef(ref)

        client := ClientForURL(ref.URL, r.fixedClient)

        repoInfo, err := client.Clone(ctx, cloneConfig)
        if err != nil {
                return nil, fmt.Errorf("cloning repository: %w", err)
        }
        defer client.Cleanup(ctx, repoInfo) //nolint:errcheck // best-effort cleanup

        // Get the commit hash (for digest tracking) and its signature in one
        // lookup, so both describe the same commit.
        head, err := client.HeadCommit(repoInfo)
        if err != nil {
                return nil, fmt.Errorf("getting HEAD commit: %w", err)
        }

        // Read SKILL.md from the skill path
        skillMDPath := path.Join(ref.Path, "SKILL.md")
        if ref.Path == "" {
                skillMDPath = "SKILL.md"
        }

        skillContent, err := client.GetFileContent(repoInfo, skillMDPath)
        if err != nil {
                return nil, fmt.Errorf("reading SKILL.md at %q: %w", skillMDPath, err)
        }

        // Parse the skill definition
        parsed, err := skills.ParseSkillMD(skillContent)
        if err != nil {
                return nil, fmt.Errorf("parsing SKILL.md: %w", err)
        }

        // Validate skill name
        if err := skills.ValidateSkillName(parsed.Name); err != nil {
                return nil, fmt.Errorf("invalid skill name in SKILL.md: %w", err)
        }

        // Collect all files in the skill directory, recursively walking nested
        // subtrees so that companion files in subdirectories (e.g. references/,
        // scripts/) are included in the resolved skill bundle.
        files, err := r.collectFiles(repoInfo, ref.Path)
        if err != nil {
                return nil, fmt.Errorf("collecting skill files: %w", err)
        }

        return &ResolveResult{
                SkillConfig:     parsed,
                Files:           files,
                CommitHash:      head.Hash,
                CommitSignature: head.Signature,
                CommitPayload:   head.Payload,
        }, nil
}

// collectFiles reads all files from the given path in the repository,
// walking nested subtrees recursively. Returned paths are forward-slash
// relative to basePath. WriteContainedFile creates parent directories and
// guards against path traversal; the in-memory clone is bounded by
// LimitedFs in pkg/git, and the OCI packager re-asserts file count and
// total size limits independently.
func (*defaultResolver) collectFiles(repoInfo *git.RepositoryInfo, basePath string) ([]FileEntry, error) {
        ref, err := repoInfo.Repository.Head()
        if err != nil {
                return nil, fmt.Errorf("getting HEAD: %w", err)
        }

        commit, err := repoInfo.Repository.CommitObject(ref.Hash())
        if err != nil {
                return nil, fmt.Errorf("getting commit: %w", err)
        }

        tree, err := commit.Tree()
        if err != nil {
                return nil, fmt.Errorf("getting tree: %w", err)
        }

        if basePath != "" {
                tree, err = tree.Tree(basePath)
                if err != nil {
                        return nil, fmt.Errorf("navigating to path %q: %w", basePath, err)
                }
        }

        var files []FileEntry
        err = tree.Files().ForEach(func(f *object.File) error {
                content, contentErr := f.Contents()
                if contentErr != nil {
                        return fmt.Errorf("reading content of %q: %w", f.Name, contentErr)
                }

                // All files are capped to 0644 by the writer; set a uniform mode here.
                files = append(files, FileEntry{
                        Path:    f.Name,
                        Content: []byte(content),
                        Mode:    fs.FileMode(0644),
                })
                return nil
        })
        if err != nil {
                return nil, fmt.Errorf("iterating tree: %w", err)
        }

        return files, nil
}

// isHex checks if a string is a valid non-empty hexadecimal string.
func isHex(s string) bool {
        if s == "" {
                return false
        }
        for _, c := range s {
                switch {
                case c >= '0' && c <= '9',
                        c >= 'a' && c <= 'f',
                        c >= 'A' && c <= 'F':
                        continue
                default:
                        return false
                }
        }
        return true
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package gitresolver

import (
        "fmt"
        "os"
        "path/filepath"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/skills"
)

// WriteFiles writes resolved skill files to the target directory.
// If force is true, any existing directory is removed before writing.
//
// Security: targetDir is produced by PathResolver.GetSkillPath (a trusted
// internal source that builds paths from known base directories). File paths
// within the archive are validated via containment check against targetDir.
func WriteFiles(files []FileEntry, targetDir string, force bool) error {
        // Sanitize targetDir early so all downstream os calls use the clean path.
        targetDir = filepath.Clean(targetDir)

        // Ensure the parent directory exists before acquiring the lock.
        // WithFileLock opens targetDir+".lock", which requires the parent to exist.
        if err := os.MkdirAll(filepath.Dir(targetDir), skills.DirPermissions); err != nil {
                return fmt.Errorf("creating parent directory: %w", err)
        }

        return fileutils.WithFileLock(targetDir, func() error {
                // Handle existing directory
                if _, statErr := os.Stat(targetDir); statErr == nil { // lgtm[go/path-injection] #nosec G304
                        if !force {
                                return fmt.Errorf("target directory %q already exists; use force to overwrite", targetDir)
                        }
                        if err := os.RemoveAll(targetDir); err != nil { // lgtm[go/path-injection] #nosec G304 -- targetDir is cleaned above
                                return fmt.Errorf("removing existing directory: %w", err)
                        }
                }

                // Pre-extraction: validate that no existing path components are symlinks.
                // Reuses the same check as the OCI installer (pkg/skills/installer.go).
                if err := skills.ValidatePathNoSymlinks(targetDir); err != nil {
                        return fmt.Errorf("target path validation: %w", err)
                }

                if err := os.MkdirAll(targetDir, skills.DirPermissions); err != nil { // lgtm[go/path-injection] #nosec G304
                        return fmt.Errorf("creating target directory: %w", err)
                }

                for _, f := range files {
                        mode := (f.Mode & 0o777) & skills.FilePermissionMask
                        if err := fileutils.WriteContainedFile(targetDir, f.Path, f.Content, skills.DirPermissions, mode); err != nil {
                                return err
                        }
                }

                return nil
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package skills

import (
        "fmt"
        "os"
        "path/filepath"
        "strings"

        ociskills "github.com/stacklok/toolhive-core/oci/skills"
        "github.com/stacklok/toolhive/pkg/fileutils"
)

const (
        // MaxTotalExtractSize is the maximum total decompressed size (500MB).
        MaxTotalExtractSize int64 = 500 * 1024 * 1024
        // MaxFileExtractSize is the maximum size per file in the tar archive (100MB).
        // Matches the toolhive-core default.
        MaxFileExtractSize int64 = 100 * 1024 * 1024
        // MaxExtractFileCount is the maximum number of files allowed in an archive.
        MaxExtractFileCount = 1000

        // DirPermissions is the permission mode for created directories.
        DirPermissions os.FileMode = 0750
        // FilePermissionMask strips setuid, setgid, sticky bits and caps at 0644.
        FilePermissionMask os.FileMode = 0644
        // PluginFilePermissionMask strips setuid, setgid, sticky bits and caps at
        // 0755, preserving the executable bit so plugin hook scripts keep +x.
        // For non-executable files (mode 0644), 0644 & 0755 = 0644 — no change.
        PluginFilePermissionMask os.FileMode = 0755
)

// ExtractResult contains the outcome of an Extract operation.
type ExtractResult struct {
        // SkillDir is the absolute path where the skill was extracted.
        SkillDir string
        // Files is the number of files written.
        Files int
}

// defaultInstaller is the production implementation of Installer.
type defaultInstaller struct{}

// NewInstaller returns a production Installer that delegates to the package-level
// Extract and Remove functions.
func NewInstaller() Installer {
        return &defaultInstaller{}
}

func (*defaultInstaller) Extract(layerData []byte, targetDir string, force bool) (*ExtractResult, error) {
        return Extract(layerData, targetDir, force)
}

func (*defaultInstaller) ExtractPlugin(layerData []byte, targetDir string, force bool) (*ExtractResult, error) {
        return ExtractPlugin(layerData, targetDir, force)
}

func (*defaultInstaller) Remove(skillDir string) error {
        return Remove(skillDir)
}

// Extract decompresses a tar.gz OCI layer and writes files to targetDir.
// If targetDir exists and force is false, an error is returned.
// If force is true, the existing directory is removed before extraction.
func Extract(layerData []byte, targetDir string, force bool) (*ExtractResult, error) {
        // Decompress gzip with total size limit
        tarData, err := ociskills.DecompressWithLimit(layerData, MaxTotalExtractSize)
        if err != nil {
                return nil, fmt.Errorf("decompressing layer: %w", err)
        }

        // Extract tar with per-file size limit (rejects symlinks, hardlinks, path traversal)
        files, err := ociskills.ExtractTarWithLimit(tarData, MaxFileExtractSize)
        if err != nil {
                return nil, fmt.Errorf("extracting tar: %w", err)
        }

        if len(files) > MaxExtractFileCount {
                return nil, fmt.Errorf("archive contains %d files, exceeding limit of %d", len(files), MaxExtractFileCount)
        }

        // Handle existing directory
        if _, statErr := os.Stat(targetDir); statErr == nil {
                if !force {
                        return nil, fmt.Errorf("target directory %q already exists; use force to overwrite", targetDir)
                }
                if err := Remove(targetDir); err != nil {
                        return nil, fmt.Errorf("removing existing directory: %w", err)
                }
        }

        // Pre-extraction: validate that no existing path components are symlinks.
        // This prevents an attacker from placing a symlink at a parent directory
        // that would cause MkdirAll/writes to follow through to an unintended location.
        if err := ValidatePathNoSymlinks(targetDir); err != nil {
                return nil, fmt.Errorf("target path validation: %w", err)
        }

        if err := os.MkdirAll(targetDir, DirPermissions); err != nil {
                return nil, fmt.Errorf("creating target directory: %w", err)
        }

        if err := writeFiles(files, targetDir, FilePermissionMask); err != nil {
                return nil, err
        }

        // Defense in depth: verify the extracted directory post-extraction
        if err := CheckFilesystem(targetDir); err != nil {
                _ = os.RemoveAll(targetDir) // clean up on verification failure
                return nil, fmt.Errorf("post-extraction verification failed: %w", err)
        }

        return &ExtractResult{
                SkillDir: targetDir,
                Files:    len(files),
        }, nil
}

// ExtractPlugin is like Extract but uses PluginFilePermissionMask (0755) so
// executable files (hook scripts, entry points) keep their +x bit. Used by
// plugin adapters which install multi-component trees that may include
// executable hooks — unlike skills, which are single markdown files.
func ExtractPlugin(layerData []byte, targetDir string, force bool) (*ExtractResult, error) {
        // Decompress gzip with total size limit
        tarData, err := ociskills.DecompressWithLimit(layerData, MaxTotalExtractSize)
        if err != nil {
                return nil, fmt.Errorf("decompressing layer: %w", err)
        }

        // Extract tar with per-file size limit (rejects symlinks, hardlinks, path traversal)
        files, err := ociskills.ExtractTarWithLimit(tarData, MaxFileExtractSize)
        if err != nil {
                return nil, fmt.Errorf("extracting tar: %w", err)
        }

        if len(files) > MaxExtractFileCount {
                return nil, fmt.Errorf("archive contains %d files, exceeding limit of %d", len(files), MaxExtractFileCount)
        }

        // Handle existing directory
        if _, statErr := os.Stat(targetDir); statErr == nil {
                if !force {
                        return nil, fmt.Errorf("target directory %q already exists; use force to overwrite", targetDir)
                }
                if err := Remove(targetDir); err != nil {
                        return nil, fmt.Errorf("removing existing directory: %w", err)
                }
        }

        // Pre-extraction: validate that no existing path components are symlinks.
        if err := ValidatePathNoSymlinks(targetDir); err != nil {
                return nil, fmt.Errorf("target path validation: %w", err)
        }

        if err := os.MkdirAll(targetDir, DirPermissions); err != nil {
                return nil, fmt.Errorf("creating target directory: %w", err)
        }

        if err := writeFiles(files, targetDir, PluginFilePermissionMask); err != nil {
                return nil, err
        }

        // Defense in depth: verify the extracted directory post-extraction
        if err := CheckFilesystem(targetDir); err != nil {
                _ = os.RemoveAll(targetDir) // clean up on verification failure
                return nil, fmt.Errorf("post-extraction verification failed: %w", err)
        }

        return &ExtractResult{
                SkillDir: targetDir,
                Files:    len(files),
        }, nil
}

// writeFiles writes extracted file entries to targetDir with containment checks
// and sanitized permissions. The mask controls how much of the original mode
// is preserved (FilePermissionMask for skills, PluginFilePermissionMask for
// plugins).
func writeFiles(files []ociskills.FileEntry, targetDir string, mask os.FileMode) error {
        cleanTarget := filepath.Clean(targetDir)

        for _, f := range files {
                // Sanitize file permissions: strip setuid/setgid/sticky, apply mask.
                // Pre-write containment check is handled by WriteContainedFile.
                mode := os.FileMode(f.Mode&0o777) & mask //nolint:gosec // mode is masked to 9 bits before conversion

                if err := fileutils.WriteContainedFile(cleanTarget, f.Path, f.Content, DirPermissions, mode); err != nil {
                        return err
                }
        }
        return nil
}

// Remove safely removes a skill directory. Returns nil if the directory does not exist.
func Remove(skillDir string) error {
        if skillDir == "" {
                return fmt.Errorf("skill directory path must not be empty")
        }

        // Resolve to absolute path for safety checks
        absPath, err := filepath.Abs(skillDir)
        if err != nil {
                return fmt.Errorf("resolving absolute path: %w", err)
        }

        // Use Lstat (not Stat) to detect symlinks without following them.
        info, err := os.Lstat(absPath)
        if err != nil {
                if os.IsNotExist(err) {
                        return nil
                }
                return fmt.Errorf("checking path %q: %w", absPath, err)
        }

        // Refuse to remove if the path itself is a symlink — prevents an attacker
        // from replacing the skill directory with a symlink to trick us into
        // deleting an arbitrary location.
        if info.Mode()&os.ModeSymlink != 0 {
                return fmt.Errorf("refusing to remove symlink at %q: expected a directory", absPath)
        }

        // Resolve any symlinks in parent components to get the real path for
        // the dangerous-path checks below.
        realPath, err := filepath.EvalSymlinks(absPath)
        if err != nil {
                return fmt.Errorf("resolving symlinks in path: %w", err)
        }

        // Guard against removing dangerous paths (checked against resolved path).
        // filepath.Dir(p) == p is true at filesystem roots on all platforms
        // (e.g. "/" on Unix, "C:\" on Windows).
        homeDir, homeErr := os.UserHomeDir()
        if filepath.Dir(realPath) == realPath {
                return fmt.Errorf("refusing to remove dangerous path %q", realPath)
        }
        if homeErr == nil && realPath == homeDir {
                return fmt.Errorf("refusing to remove dangerous path %q", realPath)
        }
        // If we couldn't determine the home directory, refuse shallow paths as a safety net.
        // Count path depth by splitting on separator (e.g., "/var/home/user" → 4 components).
        if homeErr != nil && pathDepth(realPath) < 4 {
                return fmt.Errorf("refusing to remove shallow path %q (could not determine home directory)", realPath)
        }

        return os.RemoveAll(absPath)
}

// ValidatePathNoSymlinks walks up from the target path checking each existing
// path component for symlinks. This prevents symlink attacks where an attacker
// places a symlink at a parent directory before extraction.
func ValidatePathNoSymlinks(targetDir string) error {
        absTarget, err := filepath.Abs(targetDir)
        if err != nil {
                return fmt.Errorf("resolving absolute path: %w", err)
        }

        // Walk each component from the root down, checking existing segments.
        // Use filepath.VolumeName to determine the root correctly on all platforms
        // (e.g. "/" on Unix, "C:\" on Windows).
        current := func() string {
                if vol := filepath.VolumeName(absTarget); vol != "" {
                        return vol + string(os.PathSeparator)
                }
                return string(os.PathSeparator)
        }()
        for _, component := range strings.Split(absTarget, string(os.PathSeparator)) {
                if component == "" {
                        continue
                }
                current = filepath.Join(current, component)

                info, err := os.Lstat(current)
                if err != nil {
                        // Path doesn't exist yet — remaining components will be created by MkdirAll.
                        break
                }
                if info.Mode()&os.ModeSymlink != 0 {
                        return fmt.Errorf("symlink found at %q: refusing to extract through symlinks", current)
                }
        }
        return nil
}

// RemoveEmptyParents walks up from dir, removing each directory that is empty,
// stopping at stopAt (which is never removed) or when a non-empty directory is
// encountered. Errors are silently ignored — this is best-effort cleanup.
func RemoveEmptyParents(dir, stopAt string) {
        dir = filepath.Clean(dir)
        stopAt = filepath.Clean(stopAt)
        for dir != stopAt && filepath.Dir(dir) != dir {
                if err := os.Remove(dir); err != nil {
                        // Directory is not empty, doesn't exist, or we lack permission — stop.
                        return
                }
                dir = filepath.Dir(dir)
        }
}

// pathDepth counts the number of non-empty components in an absolute path.
// For example, "/var/home/user/skills" returns 4.
func pathDepth(absPath string) int {
        count := 0
        for _, part := range strings.Split(absPath, string(os.PathSeparator)) {
                if part != "" {
                        count++
                }
        }
        return count
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package lockfile

import (
        "crypto/sha256"
        "encoding/hex"
        "fmt"
        "io"
        "io/fs"
        "os"
        "path/filepath"
        "slices"
        "strings"
)

// ContentFile is a single file used for contentDigest computation.
type ContentFile struct {
        // Path is the relative path within the skill directory, using either
        // native or slash separators.
        Path string
        // Content is the raw file bytes.
        Content []byte
}

// ContentDigest computes a deterministic SHA-256 dirhash over files.
//
// Algorithm (content-only; file modes and timestamps are ignored):
//  1. Normalize each path to slash-separated form and sort files by path.
//  2. For each file, feed path + "\x00" + sha256(content) + "\n" into a
//     running SHA-256.
//  3. Return "sha256:" + hex(aggregate).
//
// The algorithm is frozen: lock files in the wild pin its output, so any
// change would report every installed skill as drifted. Golden vectors in
// the package tests guard against accidental changes.
//
// Known limitations (accepted, like go.sum's): the digest covers file
// content only — file modes (e.g. an exec-bit flip) are invisible to it,
// and [ContentDigestFromDir] skips symlinks and other non-regular files
// entirely. Trust in the file set's provenance is the Sigstore
// verification layer's job, not this integrity pin's.
func ContentDigest(files []ContentFile) (string, error) {
        if len(files) == 0 {
                return "", fmt.Errorf("content digest requires at least one file")
        }

        normalized := make([]ContentFile, 0, len(files))
        for _, f := range files {
                path := filepath.ToSlash(f.Path)
                for strings.HasPrefix(path, "./") {
                        path = strings.TrimPrefix(path, "./")
                }
                if path == "" || path == ".." || strings.HasPrefix(path, "../") || strings.Contains(path, "/../") ||
                        strings.HasSuffix(path, "/..") || strings.HasPrefix(path, "/") || hasControlChar(path) {
                        return "", fmt.Errorf("invalid content file path %q", f.Path)
                }
                normalized = append(normalized, ContentFile{Path: path, Content: f.Content})
        }
        slices.SortFunc(normalized, func(a, b ContentFile) int { return strings.Compare(a.Path, b.Path) })

        h := sha256.New()
        for _, f := range normalized {
                fileHash := sha256.Sum256(f.Content)
                // Writes to a hash.Hash never fail.
                _, _ = io.WriteString(h, f.Path)
                _, _ = h.Write([]byte{0})
                _, _ = h.Write(fileHash[:])
                _, _ = h.Write([]byte{'\n'})
        }
        return ContentDigestPrefix + hex.EncodeToString(h.Sum(nil)), nil
}

// hasControlChar reports whether s contains a NUL, newline, or other C0/DEL
// control byte. The serialization format used by [ContentDigest] delimits
// fields with NUL and newline bytes; without this check, a Path containing
// those bytes could reproduce the exact byte stream of a different,
// unrelated multi-file tree and collide with its digest.
func hasControlChar(s string) bool {
        for i := 0; i < len(s); i++ {
                if s[i] < 0x20 || s[i] == 0x7f {
                        return true
                }
        }
        return false
}

// ContentDigestFromDir walks skillDir and computes the content digest from
// the on-disk file set. The walk is confined to skillDir via [os.Root], so
// symlinks cannot escape it.
func ContentDigestFromDir(skillDir string) (string, error) {
        root, err := os.OpenRoot(skillDir)
        if err != nil {
                return "", fmt.Errorf("opening skill directory: %w", err)
        }
        defer func() { _ = root.Close() }()

        var files []ContentFile
        err = fs.WalkDir(root.FS(), ".", func(path string, d fs.DirEntry, walkErr error) error {
                if walkErr != nil {
                        return walkErr
                }
                if d.IsDir() || !d.Type().IsRegular() {
                        return nil
                }
                data, err := fs.ReadFile(root.FS(), path)
                if err != nil {
                        return fmt.Errorf("reading %q: %w", path, err)
                }
                files = append(files, ContentFile{Path: path, Content: data})
                return nil
        })
        if err != nil {
                return "", err
        }
        return ContentDigest(files)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package lockfile manages the project-level lock file (toolhive.lock.yaml).
// The lock file pins the exact name, version, source, and digests of every
// project-scoped skill and plugin install so a team can restore
// ("thv skill sync" / "thv ai-plugin sync") or refresh
// ("thv skill upgrade" / "thv ai-plugin upgrade") the pinned state on any
// machine. RFC THV-0080 owns the version and skills: keys; plugins: is the
// sibling key ratified by the AI-plugin lock track.
//
// All filesystem access goes through [Root], a capability type that can only
// be constructed from a validated project root. Root confines every read and
// write to that directory using [os.Root] (OS-enforced containment, not just
// string validation), so no function in this package can be made to open a
// path outside the project root regardless of what name is requested.
package lockfile

import (
        "crypto/rand"
        "encoding/hex"
        "errors"
        "fmt"
        "io/fs"
        "os"
        "path/filepath"
        "reflect"
        "slices"
        "strings"

        "go.yaml.in/yaml/v3"

        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/skills"
)

// FileName is the name of the project-level lock file, written to the
// project root alongside .git.
const FileName = "toolhive.lock.yaml"

// CurrentVersion is the schema version written to new lock files. Loading a
// lock file with a different version is a hard error, never a silent partial
// parse.
const CurrentVersion = 1

// ErrEntryChanged reports that a conditional lock-entry update no longer
// matches the entry it was planned against.
var ErrEntryChanged = errors.New("lockfile: entry changed")

// Entry represents a single pinned skill or plugin installation in the lock file.
type Entry struct {
        // Name is the artifact's unique name within its key (skills: or plugins:).
        // The same name may appear once under each key.
        Name string `yaml:"name"`
        // Version is the artifact's declared version, if any (SKILL.md frontmatter
        // or .claude-plugin/plugin.json).
        Version string `yaml:"version,omitempty"`
        // Source is exactly what the user (or the registry resolver) originally
        // requested — a plain registry name, an OCI reference, or a git://
        // reference. It is never rewritten; upgrade re-resolves this value to
        // check for newer content.
        Source string `yaml:"source"`
        // ResolvedReference is the concrete OCI reference or git:// URL that
        // Source resolved to at install time.
        ResolvedReference string `yaml:"resolvedReference,omitempty"`
        // Digest pins the exact content installed: an OCI "sha256:..." manifest
        // digest or a full git commit hash.
        Digest string `yaml:"digest"`
        // ContentDigest is a deterministic SHA-256 dirhash of the materialized
        // file set (skill directory, or the canonical plugin tree ExtractPlugin
        // writes), used for on-disk integrity verification.
        ContentDigest string `yaml:"contentDigest,omitempty"`
        // Provenance records the verified signer identity of the installed
        // artifact — the trust-on-first-use anchor later installs, syncs, and
        // upgrades are checked against. Nil for entries recorded before
        // verification existed and for unsigned entries (see Unsigned); an
        // entry never carries both.
        Provenance *Provenance `yaml:"provenance,omitempty"`
        // Unsigned is true when the artifact carried no signature and the user
        // explicitly accepted that with --allow-unsigned. It is a recorded
        // trust decision, not an omission: replacing an unsigned entry with a
        // signed one clears it, and a signed entry never becomes unsigned
        // without the same explicit flag.
        Unsigned bool `yaml:"unsigned,omitempty"`
        // RequiredBy lists parent names in the same key (skills: or plugins:)
        // for transitively materialized dependencies. Plugin v1 does not
        // populate this; the field is reserved for when plugin requires are
        // materialized.
        RequiredBy []string `yaml:"requiredBy,omitempty"`
        // Explicit is true when the user directly installed this artifact;
        // explicit entries are exempt from cascade removal when RequiredBy
        // becomes empty.
        Explicit bool `yaml:"explicit,omitempty"`
        // Extra round-trips fields this binary does not know about, so a
        // Load→modify→Save cycle by an older binary preserves rather than strips
        // them. It applies per entry: an entry this binary rewrites (reinstall,
        // upgrade) is built fresh, so its unknown fields — which described the
        // previous install — are intentionally dropped along with it.
        Extra map[string]any `yaml:",inline"`
}

// Provenance is the Sigstore trust anchor recorded for a verified entry.
//
// Exactly one anchor is recorded. Keyless (Fulcio) entries record a
// certificate identity — SignerIdentity plus CertIssuer, optionally narrowed
// by the certificate-derived fields below. Key-pair (cosign) entries record
// PublicKey instead and leave every certificate field empty, because a
// key-pair signature carries no certificate to derive them from.
type Provenance struct {
        // SignerIdentity is the certificate subject identity: for GitHub
        // Actions certificates, the workflow path relative to the repository;
        // otherwise the certificate SAN verbatim (a URI, email, or SPIFFE ID).
        SignerIdentity string `yaml:"signerIdentity"`
        // CertIssuer is the OIDC issuer that authenticated the signer.
        CertIssuer string `yaml:"certIssuer"`
        // RepositoryURI is the source repository from the Fulcio certificate
        // extensions, when present.
        RepositoryURI string `yaml:"repositoryUri,omitempty"`
        // RepositoryRef is the git ref the signing workflow ran on, from Fulcio
        // certificate extension 1.3.6.1.4.1.57264.1.14, when present. Empty
        // means unconstrained: every lock file written before this field existed
        // omits it, and such entries must keep verifying against any ref.
        RepositoryRef string `yaml:"repositoryRef,omitempty"`
        // RunnerEnvironment is the runner class the signing workflow executed in
        // (e.g. "github-hosted"), from Fulcio certificate extension
        // 1.3.6.1.4.1.57264.1.11, when present. Empty means unconstrained, for
        // the same backward-compatibility reason as RepositoryRef.
        RunnerEnvironment string `yaml:"runnerEnvironment,omitempty"`
        // SigstoreURL is the Sigstore instance the signature chains to.
        SigstoreURL string `yaml:"sigstoreUrl,omitempty"`
        // PublicKey is the base64-encoded DER SPKI form of the cosign public key
        // a key-pair-signed entry is pinned to — the PEM body with its armor and
        // line breaks removed, since a lock value may not contain whitespace
        // (see validateProvenance).
        //
        // The full key is stored, not a digest of it, because the key is
        // recoverable from neither the artifact nor the stored bundle: cosign's
        // signature manifest defines no annotation carrying it, so there would
        // be nothing to hash at verification time. This value is therefore the
        // entry's only trust anchor, and re-verification depends on it.
        //
        // Mutually exclusive with SignerIdentity/CertIssuer. Whether the value
        // parses as a usable key is the verifier's concern; validation here is
        // syntactic, as for every other field.
        PublicKey string `yaml:"publicKey,omitempty"`
        // Provisional marks provenance whose verification has a documented
        // gap — currently git-commit signatures, verified for signature and
        // certificate chain but without transparency-log proof of signing
        // time. Visible in lock diffs so reduced assurance is never silent;
        // removed when the gap closes.
        Provisional bool `yaml:"provisional,omitempty"`
}

// Lockfile is the parsed contents of a project's toolhive.lock.yaml.
type Lockfile struct {
        // Version is the lock file schema version.
        Version int `yaml:"version"`
        // Skills is the set of pinned skill installations, sorted by name for
        // stable diffs.
        Skills []Entry `yaml:"skills,omitempty"`
        // Plugins is the set of pinned plugin installations, sorted by name
        // for stable diffs. Validated as its own name/requiredBy graph, so a
        // skill and a plugin may share a name.
        Plugins []Entry `yaml:"plugins,omitempty"`
        // Extra round-trips unknown top-level fields, mirroring Entry.Extra.
        Extra map[string]any `yaml:",inline"`
}

// Root is a validated project root directory. It is the only way to address
// a lock file on disk: constructing one runs full project-root validation
// (absolute, NUL-free, no traversal segments, symlink-canonical, git-rooted).
// Every read and write additionally goes through a freshly opened [os.Root]
// scoped to the validated directory, so access is confined at the OS level.
type Root struct {
        dir string
}

// OpenRoot validates projectRoot and returns a Root for it. The zero Root is
// unusable; all lock file operations require a Root produced here.
func OpenRoot(projectRoot string) (Root, error) {
        dir, err := skills.ValidateProjectRoot(projectRoot)
        if err != nil {
                return Root{}, err
        }
        return Root{dir: dir}, nil
}

// Dir returns the validated project root directory.
func (r Root) Dir() string {
        return r.dir
}

// Path returns the absolute path of the lock file inside the project root,
// for display purposes (e.g. error messages, CLI output). It is not used by
// this package to open the file; use [Load] and [Save] for that.
func (r Root) Path() (string, error) {
        if r.dir == "" {
                return "", errors.New("lockfile: uninitialized Root, use OpenRoot")
        }
        return filepath.Join(r.dir, FileName), nil
}

// osRoot opens an OS-level containment handle for r. Every name passed to
// its methods is resolved beneath r.dir; a name that would escape it (via
// "..", an absolute path, or a symlink) is rejected by the OS, not by string
// inspection.
func (r Root) osRoot() (*os.Root, error) {
        if r.dir == "" {
                return nil, errors.New("lockfile: uninitialized Root, use OpenRoot")
        }
        return os.OpenRoot(r.dir)
}

// Load reads and parses the lock file for root. A missing lock file is not
// an error — it returns an empty Lockfile ready to be populated. A lock file
// with an unsupported schema version is a hard error.
func Load(root Root) (*Lockfile, error) {
        osRoot, err := root.osRoot()
        if err != nil {
                return nil, err
        }
        defer func() { _ = osRoot.Close() }()

        data, err := osRoot.ReadFile(FileName)
        if errors.Is(err, fs.ErrNotExist) {
                return &Lockfile{Version: CurrentVersion}, nil
        }
        if err != nil {
                return nil, fmt.Errorf("reading lock file: %w", err)
        }

        var lf Lockfile
        if err := yaml.Unmarshal(data, &lf); err != nil {
                return nil, fmt.Errorf("parsing lock file: %w", err)
        }
        if lf.Version == 0 {
                // Tolerate a hand-written lock file that omits the version key.
                lf.Version = CurrentVersion
        }
        sortEntries(lf.Skills)
        sortEntries(lf.Plugins)
        if err := validateLockfile(&lf); err != nil {
                return nil, err
        }
        return &lf, nil
}

// Get returns the skills: entry for name, if present.
func (l *Lockfile) Get(name string) (Entry, bool) {
        return getNamed(l.Skills, name)
}

// GetPlugin returns the plugins: entry for name, if present.
func (l *Lockfile) GetPlugin(name string) (Entry, bool) {
        return getNamed(l.Plugins, name)
}

// Upsert inserts or replaces the skills: entry with a matching name, keeping
// the slice sorted by name for stable diffs. It does not touch plugins:.
func (l *Lockfile) Upsert(entry Entry) {
        upsertNamed(&l.Skills, entry)
}

// UpsertPlugin inserts or replaces the plugins: entry with a matching name,
// keeping the slice sorted by name for stable diffs. It does not touch skills:.
func (l *Lockfile) UpsertPlugin(entry Entry) {
        upsertNamed(&l.Plugins, entry)
}

// Remove deletes the skills: entry with the given name, if present. Reports
// whether an entry was removed.
func (l *Lockfile) Remove(name string) bool {
        return removeNamed(&l.Skills, name)
}

// RemovePlugin deletes the plugins: entry with the given name, if present.
// Reports whether an entry was removed.
func (l *Lockfile) RemovePlugin(name string) bool {
        return removeNamed(&l.Plugins, name)
}

// RemoveParentFromRequiredBy removes parent from every skills: entry's
// RequiredBy list and returns the names of entries that lost their last
// parent and are not explicit — the candidates for cascade removal. It
// does not touch plugins:.
func (l *Lockfile) RemoveParentFromRequiredBy(parent string) []string {
        var cascadeCandidates []string
        for i := range l.Skills {
                entry := &l.Skills[i]
                if !slices.Contains(entry.RequiredBy, parent) {
                        continue
                }
                entry.RequiredBy = slices.DeleteFunc(entry.RequiredBy, func(s string) bool { return s == parent })
                if len(entry.RequiredBy) == 0 {
                        entry.RequiredBy = nil
                        if !entry.Explicit {
                                cascadeCandidates = append(cascadeCandidates, entry.Name)
                        }
                }
        }
        return cascadeCandidates
}

// tmpFileName is the temporary file Save writes before renaming it into
// tmpFileName returns a per-call temporary file name for Save's atomic
// write. The random suffix means two concurrent direct Save calls (which,
// unlike Update/UpsertEntry/RemoveEntry, hold no file lock) cannot write
// to, rename, or delete each other's temp file — last rename still wins,
// but neither can promote or destroy a half-written file.
func tmpFileName() (string, error) {
        var buf [8]byte
        if _, err := rand.Read(buf[:]); err != nil {
                return "", fmt.Errorf("generating temp file suffix: %w", err)
        }
        return ".toolhive.lock." + hex.EncodeToString(buf[:]) + ".tmp", nil
}

// Save writes the lock file into root atomically (temp file + rename), after
// validating it with the same rules [Load] enforces on read. This prevents a
// caller bug from writing a lock file that every subsequent Load/Update call
// would then hard-fail on, with no recovery path through this package.
// Callers that need read-modify-write atomicity across processes must use
// [UpsertEntry], [RemoveEntry], or [Update] instead of Load+Save — Save alone
// is write-atomic but does nothing to serialize concurrent read-modify-write
// cycles.
func (l *Lockfile) Save(root Root) error {
        osRoot, err := root.osRoot()
        if err != nil {
                return err
        }
        defer func() { _ = osRoot.Close() }()

        if l.Version == 0 {
                l.Version = CurrentVersion
        }
        sortEntries(l.Skills)
        sortEntries(l.Plugins)

        if err := validateLockfile(l); err != nil {
                return fmt.Errorf("refusing to save invalid lock file: %w", err)
        }

        data, err := yaml.Marshal(l)
        if err != nil {
                return fmt.Errorf("marshaling lock file: %w", err)
        }

        tmpName, err := tmpFileName()
        if err != nil {
                return err
        }
        // The lock file is committed to git and not sensitive; 0o644 matches any
        // other source file.
        if err := osRoot.WriteFile(tmpName, data, 0o644); err != nil {
                return fmt.Errorf("writing lock file: %w", err)
        }
        if err := osRoot.Rename(tmpName, FileName); err != nil {
                _ = osRoot.Remove(tmpName)
                return fmt.Errorf("saving lock file: %w", err)
        }
        return nil
}

// UpsertEntry loads the lock file, upserts a skills: entry, and saves it
// back, all under a single file lock so concurrent installs cannot race on
// read-modify-write.
func UpsertEntry(root Root, entry Entry) error {
        return Update(root, func(lf *Lockfile) error {
                lf.Upsert(entry)
                return nil
        })
}

// CompareAndSwapEntry replaces expected with replacement only when the
// complete current skills: entry still equals expected. The comparison and
// replacement run under Update's advisory file lock, so a mutation made by a
// non-participating process between planning and persistence cannot be
// overwritten by a stale plan.
func CompareAndSwapEntry(root Root, expected, replacement Entry) error {
        if expected.Name == "" || replacement.Name != expected.Name {
                return errors.New("lockfile: conditional replacement requires matching non-empty names")
        }
        return Update(root, func(lf *Lockfile) error {
                current, ok := lf.Get(expected.Name)
                if !ok || !reflect.DeepEqual(current, expected) {
                        return fmt.Errorf("%w: skill %q no longer matches the expected entry", ErrEntryChanged, expected.Name)
                }
                lf.Upsert(replacement)
                return nil
        })
}

// CompareAndSwapPluginEntry is the plugins: counterpart to
// CompareAndSwapEntry. It performs the complete-entry comparison and
// replacement under the same advisory file lock.
func CompareAndSwapPluginEntry(root Root, expected, replacement Entry) error {
        if expected.Name == "" || replacement.Name != expected.Name {
                return errors.New("lockfile: conditional plugin replacement requires matching non-empty names")
        }
        return Update(root, func(lf *Lockfile) error {
                current, ok := lf.GetPlugin(expected.Name)
                if !ok || !reflect.DeepEqual(current, expected) {
                        return fmt.Errorf("%w: plugin %q no longer matches the expected entry", ErrEntryChanged, expected.Name)
                }
                lf.UpsertPlugin(replacement)
                return nil
        })
}

// UpsertPluginEntry loads the lock file, upserts a plugins: entry, and
// saves it back, all under a single file lock so concurrent installs cannot
// race on read-modify-write.
func UpsertPluginEntry(root Root, entry Entry) error {
        return Update(root, func(lf *Lockfile) error {
                lf.UpsertPlugin(entry)
                return nil
        })
}

// RemoveEntry loads the lock file, removes the named skills: entry if
// present, and saves it back, all under a single file lock. Removing an
// entry that does not exist is a no-op, not an error.
func RemoveEntry(root Root, name string) error {
        return Update(root, func(lf *Lockfile) error {
                if !lf.Remove(name) {
                        return errSkipSave
                }
                return nil
        })
}

// RemovePluginEntry loads the lock file, removes the named plugins: entry
// if present, and saves it back, all under a single file lock. Removing an
// entry that does not exist is a no-op, not an error.
func RemovePluginEntry(root Root, name string) error {
        return Update(root, func(lf *Lockfile) error {
                if !lf.RemovePlugin(name) {
                        return errSkipSave
                }
                return nil
        })
}

// Update loads the lock file, applies fn, and saves the result, all under a
// single file lock. If fn returns an error the lock file is left unchanged.
func Update(root Root, fn func(*Lockfile) error) error {
        path, err := root.Path()
        if err != nil {
                return err
        }
        return fileutils.WithFileLock(path, func() error {
                lf, err := Load(root)
                if err != nil {
                        return err
                }
                if err := fn(lf); err != nil {
                        if errors.Is(err, errSkipSave) {
                                return nil
                        }
                        return err
                }
                return lf.Save(root)
        })
}

// errSkipSave signals from an Update callback that nothing changed and the
// save should be skipped without reporting an error.
var errSkipSave = errors.New("lockfile: no changes to save")

func sortEntries(entries []Entry) {
        slices.SortFunc(entries, func(a, b Entry) int { return strings.Compare(a.Name, b.Name) })
}

func getNamed(entries []Entry, name string) (Entry, bool) {
        for _, e := range entries {
                if e.Name == name {
                        return e, true
                }
        }
        return Entry{}, false
}

func upsertNamed(entries *[]Entry, entry Entry) {
        for i := range *entries {
                if (*entries)[i].Name == entry.Name {
                        (*entries)[i] = entry
                        return
                }
        }
        *entries = append(*entries, entry)
        sortEntries(*entries)
}

func removeNamed(entries *[]Entry, name string) bool {
        for i, e := range *entries {
                if e.Name == name {
                        *entries = slices.Delete(*entries, i, i+1)
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package lockfile

import (
        "crypto/x509"
        "encoding/base64"
        "errors"
        "fmt"
        "strings"
        "unicode"

        nameref "github.com/google/go-containerregistry/pkg/name"

        "github.com/stacklok/toolhive/pkg/skills"
        "github.com/stacklok/toolhive/pkg/skills/gitresolver"
)

// ErrUnsupportedVersion indicates the lock file schema version is not
// supported by this build.
var ErrUnsupportedVersion = errors.New("unsupported lock file version")

const (
        // ContentDigestPrefix is the prefix for contentDigest values in the lock file.
        ContentDigestPrefix = "sha256:"

        sha256HexLength = 64
        sha1HexLength   = 40
)

// validateLockfile checks the schema version and every entry. skills: and
// plugins: are validated as independent name/requiredBy graphs, so a skill
// and a plugin may share a name, and a requiredBy parent must live in the
// same key as the entry that names it.
func validateLockfile(lf *Lockfile) error {
        if lf.Version != CurrentVersion {
                return fmt.Errorf("%w: version %d (upgrade thv to read this lock file)", ErrUnsupportedVersion, lf.Version)
        }
        if err := validateEntryGraph(lf.Skills); err != nil {
                return err
        }
        if err := validateEntryGraph(lf.Plugins); err != nil {
                return fmt.Errorf("plugins: %w", err)
        }
        return nil
}

// validateEntryGraph checks uniqueness, per-entry fields, requiredBy
// cross-references, and cycles within one lock-file key.
func validateEntryGraph(entries []Entry) error {
        names := make(map[string]struct{}, len(entries))
        for _, entry := range entries {
                if err := validateEntry(entry); err != nil {
                        return err
                }
                if _, dup := names[entry.Name]; dup {
                        return fmt.Errorf("duplicate entry %q", entry.Name)
                }
                names[entry.Name] = struct{}{}
        }

        for _, entry := range entries {
                for _, parent := range entry.RequiredBy {
                        if parent == entry.Name {
                                return fmt.Errorf("entry %q: requiredBy references itself", entry.Name)
                        }
                        if _, ok := names[parent]; !ok {
                                return fmt.Errorf("entry %q: requiredBy references unknown parent %q", entry.Name, parent)
                        }
                }
        }

        if cycle := findRequiredByCycle(entries); len(cycle) > 0 {
                return fmt.Errorf("requiredBy cycle: %s", strings.Join(cycle, " -> "))
        }
        return nil
}

// findRequiredByCycle detects a cycle in the requiredBy graph and returns
// one such cycle path, or nil. Normal installs can never produce a cycle
// (the install-time depState active stack rejects them), but a hand-edited or badly
// merge-resolved lock file can — and a ring of mutually-required,
// non-explicit entries would then be impossible to ever cascade-remove, so
// it is rejected at validation instead of persisting silently.
func findRequiredByCycle(entries []Entry) []string {
        requiredBy := make(map[string][]string, len(entries))
        for _, e := range entries {
                requiredBy[e.Name] = e.RequiredBy
        }

        const (
                unvisited = 0
                inStack   = 1
                done      = 2
        )
        state := make(map[string]int, len(entries))

        var visit func(name string, path []string) []string
        visit = func(name string, path []string) []string {
                state[name] = inStack
                path = append(path, name)
                for _, parent := range requiredBy[name] {
                        switch state[parent] {
                        case inStack:
                                // Trim the path to start at the cycle entry point.
                                for i, n := range path {
                                        if n == parent {
                                                return append(path[i:], parent)
                                        }
                                }
                        case unvisited:
                                if cycle := visit(parent, path); cycle != nil {
                                        return cycle
                                }
                        }
                }
                state[name] = done
                return nil
        }

        for _, e := range entries {
                if state[e.Name] == unvisited {
                        if cycle := visit(e.Name, nil); cycle != nil {
                                return cycle
                        }
                }
        }
        return nil
}

// maxReferenceLength bounds ResolvedReference: longer than any legitimate
// OCI reference or git URL, short enough to stop megabyte-scale garbage
// from a corrupted or hostile lock file reaching the fetch path.
const maxReferenceLength = 512

// MaxEncodedPublicKeyLength bounds the base64 DER SPKI in a provenance block.
// Key material needs more room than an identifier: the largest key cosign can
// import, RSA-4096, encodes to 736 characters, so the reference bound would
// reject a legitimate anchor rather than the oversized garbage these limits
// exist to stop. Exported because the same value must bound the field before
// it is ever written — see verifier.EncodePublicKey.
const MaxEncodedPublicKeyLength = 1024

func validateEntry(entry Entry) error {
        if err := skills.ValidateSkillName(entry.Name); err != nil {
                return fmt.Errorf("entry name: %w", err)
        }
        if entry.Source == "" {
                return fmt.Errorf("entry %q: source is required", entry.Name)
        }
        if entry.Digest == "" {
                return fmt.Errorf("entry %q: digest is required", entry.Name)
        }
        if err := validateDigest(entry.Digest); err != nil {
                return fmt.Errorf("entry %q: digest: %w", entry.Name, err)
        }
        if entry.ContentDigest != "" {
                if err := validateContentDigest(entry.ContentDigest); err != nil {
                        return fmt.Errorf("entry %q: contentDigest: %w", entry.Name, err)
                }
        }
        if entry.ResolvedReference != "" {
                if err := validateResolvedReference(entry.ResolvedReference); err != nil {
                        return fmt.Errorf("entry %q: resolvedReference: %w", entry.Name, err)
                }
                if err := validateDigestKind(entry); err != nil {
                        return fmt.Errorf("entry %q: %w", entry.Name, err)
                }
        }
        if err := validateEntryTrust(entry); err != nil {
                return fmt.Errorf("entry %q: %w", entry.Name, err)
        }
        return nil
}

// validateEntryTrust checks the entry's trust fields: that it records at most
// one of a signature and an unsigned exception, that the provenance block is
// well-formed, and that its anchor suits the kind of artifact the entry
// restores. Separated from the rest of validateEntry because these checks
// need the whole entry — the anchor's fitness depends on the source — where
// validateProvenance below sees only the provenance block.
func validateEntryTrust(entry Entry) error {
        if entry.Provenance != nil && entry.Unsigned {
                return errors.New("provenance and unsigned are mutually exclusive" +
                        " — an entry is either a verified signature or a recorded unsigned exception")
        }
        if entry.Provenance == nil {
                return nil
        }
        if err := validateProvenance(entry.Provenance); err != nil {
                return fmt.Errorf("provenance: %w", err)
        }
        // A cosign key pair signs an OCI artifact, while a git entry's signature
        // lives on the commit and is always certificate-based. A key-pinned git
        // entry pins an anchor no verification of that entry could ever use.
        if entry.Provenance.PublicKey != "" && entryIsGitSource(entry) {
                return errors.New("provenance: publicKey is only valid for an OCI artifact;" +
                        " a git commit signature is verified against a certificate, not a key")
        }
        return nil
}

// entryIsGitSource classifies an entry the way the restore path does: from
// resolvedReference, which is the field buildPinnedReference dispatches on.
// Reading the deciding field — rather than inferring the kind from the digest
// beside it — is what keeps this classification from drifting away from the
// code that acts on it. An entry recording no resolved reference is
// classified by its digest form, the only signal left.
func entryIsGitSource(entry Entry) bool {
        if entry.ResolvedReference != "" {
                return gitresolver.IsGitReference(entry.ResolvedReference)
        }
        return !strings.HasPrefix(entry.Digest, ContentDigestPrefix)
}

// validateDigestKind rejects an entry whose digest form contradicts the
// source it is restored from. Restore dispatches on resolvedReference but
// pins from digest, so a git reference paired with an OCI digest yields
// "git://host/repo@sha256:..." — a reference no fetch can satisfy, and one
// whose malformedness surfaces only once the fetch is attempted. The install
// path cannot write such a pair (a git install records a bare commit hash, an
// OCI install a prefixed manifest digest); a hand edit or a botched merge
// resolution can, which is why the lock boundary is where it belongs.
//
// Callers apply this only to an entry that records a resolved reference:
// without one there is no second field to disagree with.
func validateDigestKind(entry Entry) error {
        ociDigest := strings.HasPrefix(entry.Digest, ContentDigestPrefix)
        if entryIsGitSource(entry) {
                if ociDigest {
                        return errors.New("digest is an OCI manifest digest but resolvedReference is a git reference;" +
                                " a git entry pins a full commit hash")
                }
                return nil
        }
        if !ociDigest {
                return fmt.Errorf("digest is a git commit hash but resolvedReference is an OCI reference;"+
                        " an OCI entry pins %q + 64 hex chars", ContentDigestPrefix)
        }
        return nil
}

// validateProvenance syntactically constrains a provenance block. Like
// resolvedReference, these values are hand-editable and feed the identity
// policy that future verifications are checked against, so they must be
// well-formed graphic strings of bounded length. Validation is purely
// syntactic — whether the identity is trustworthy is the verifier's job.
func validateProvenance(p *Provenance) error {
        fields := map[string]string{
                "signerIdentity":    p.SignerIdentity,
                "certIssuer":        p.CertIssuer,
                "repositoryUri":     p.RepositoryURI,
                "repositoryRef":     p.RepositoryRef,
                "runnerEnvironment": p.RunnerEnvironment,
                "sigstoreUrl":       p.SigstoreURL,
                "publicKey":         p.PublicKey,
        }
        for name, value := range fields {
                if value == "" {
                        continue
                }
                if limit := provenanceFieldLimit(name); len(value) > limit {
                        return fmt.Errorf("%s exceeds %d characters", name, limit)
                }
                if strings.TrimSpace(value) != value {
                        return fmt.Errorf("%s has leading or trailing whitespace", name)
                }
                for _, r := range value {
                        if !unicode.IsGraphic(r) || unicode.IsSpace(r) {
                                return fmt.Errorf("%s contains non-graphic or whitespace character %q", name, r)
                        }
                }
        }
        // Ordered after the syntactic checks so publicKey is length-bounded before
        // it is base64-decoded: this file is hand-editable, and decoding first
        // would size the allocation off a value that has not been checked yet.
        return validateProvenanceAnchor(p)
}

// provenanceFieldLimit returns the length bound for one provenance field.
// Every field is a short identifier except publicKey, which carries encoded
// key material and needs the wider bound.
func provenanceFieldLimit(name string) int {
        if name == "publicKey" {
                return MaxEncodedPublicKeyLength
        }
        return maxReferenceLength
}

// validateProvenanceAnchor enforces that an entry records exactly one trust
// anchor: a keyless certificate identity, or a cosign public key. The two are
// verified by different policies against different trust roots, so an entry
// carrying both would not say which applies, and an entry carrying neither
// pins nothing at all.
//
// Rejecting a key-pinned entry outright is also what makes this field safe to
// add without a schema version bump: a build that predates PublicKey sees an
// entry with no signerIdentity, reports it as required, and fails the whole
// lock file closed rather than silently treating the entry as unpinned.
func validateProvenanceAnchor(p *Provenance) error {
        keyed := p.PublicKey != ""
        identified := p.SignerIdentity != "" || p.CertIssuer != ""
        switch {
        case keyed && identified:
                return errors.New("publicKey and signerIdentity/certIssuer are mutually exclusive" +
                        " — an entry is pinned to either a cosign key or a keyless certificate identity")
        case keyed:
                // The remaining fields are all read off a Fulcio certificate, which a
                // key-pair signature does not have. Populated here they would pin
                // constraints that no verification could ever check.
                for name, value := range map[string]string{
                        "repositoryUri":     p.RepositoryURI,
                        "repositoryRef":     p.RepositoryRef,
                        "runnerEnvironment": p.RunnerEnvironment,
                        "sigstoreUrl":       p.SigstoreURL,
                } {
                        if value != "" {
                                return fmt.Errorf("%s cannot be set on a publicKey-pinned entry"+
                                        " — it is read from a certificate, and a key-pair signature has none", name)
                        }
                }
                der, err := base64.StdEncoding.DecodeString(p.PublicKey)
                if err != nil {
                        return fmt.Errorf("publicKey is not valid base64: %w", err)
                }
                // Decoding proves the encoding, not the content. This value is the
                // entry's only trust anchor, so a blob that is merely well-encoded
                // would be accepted here and then fail deep inside verification, long
                // after the lock file stopped being the obvious suspect.
                if _, err := x509.ParsePKIXPublicKey(der); err != nil {
                        return fmt.Errorf("publicKey is not a DER SPKI public key: %w", err)
                }
                return nil
        case p.SignerIdentity == "":
                return errors.New("signerIdentity is required")
        case p.CertIssuer == "":
                return errors.New("certIssuer is required")
        }
        return nil
}

// validateResolvedReference syntactically constrains the resolvedReference
// field. Sync fetches from this value without re-resolving Source, and the
// lock file is hand-editable, so a value that is not a plausible git:// or
// OCI reference must never reach the fetch path. Validation is purely
// syntactic — no network access, no allow-list policy.
func validateResolvedReference(ref string) error {
        if len(ref) > maxReferenceLength {
                return fmt.Errorf("exceeds %d characters", maxReferenceLength)
        }
        if strings.TrimSpace(ref) != ref {
                return errors.New("has leading or trailing whitespace")
        }
        for _, r := range ref {
                if !unicode.IsGraphic(r) || unicode.IsSpace(r) {
                        return fmt.Errorf("contains non-graphic or whitespace character %q", r)
                }
        }
        if gitresolver.IsGitReference(ref) {
                if _, err := gitresolver.ParseGitReference(ref); err != nil {
                        return fmt.Errorf("invalid git reference: %w", err)
                }
                return nil
        }
        // StrictValidation requires an explicit tag or digest, matching what the
        // install path records (qualifiedOCIRef always includes one) — and,
        // unlike weak validation, rejects URL-shaped strings such as
        // "http://169.254.169.254/…" that must never reach the fetch path.
        if _, err := nameref.ParseReference(ref, nameref.StrictValidation); err != nil {
                return fmt.Errorf("not a valid git:// or OCI reference: %w", err)
        }
        return nil
}

// validateDigest accepts the two pin formats the lock file records: an OCI
// manifest digest ("sha256:" + 64 hex chars) or a full git commit hash
// (40 hex chars for SHA-1 repositories, 64 for SHA-256 repositories).
// Abbreviated digests are rejected: a truncated pin weakens the guarantee
// the lock file exists to provide.
func validateDigest(d string) error {
        if hexPart, ok := strings.CutPrefix(d, ContentDigestPrefix); ok {
                if err := validateHex(hexPart, sha256HexLength); err != nil {
                        return fmt.Errorf("OCI digest: %w", err)
                }
                return nil
        }
        if len(d) != sha1HexLength && len(d) != sha256HexLength {
                return fmt.Errorf("expected %q + 64 hex chars or a full git commit hash, got %q", ContentDigestPrefix, d)
        }
        if err := validateHex(d, len(d)); err != nil {
                return fmt.Errorf("git commit hash: %w", err)
        }
        return nil
}

func validateContentDigest(d string) error {
        hexPart, ok := strings.CutPrefix(d, ContentDigestPrefix)
        if !ok {
                return fmt.Errorf("must start with %q", ContentDigestPrefix)
        }
        return validateHex(hexPart, sha256HexLength)
}

func validateHex(s string, wantLen int) error {
        if len(s) != wantLen {
                return fmt.Errorf("expected %d hex characters, got %d", wantLen, len(s))
        }
        for _, c := range s {
                if (c < '0' || c > '9') && (c < 'a' || c > 'f') {
                        return fmt.Errorf("invalid hex character %q", c)
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package skills

import (
        "bytes"
        "errors"
        "fmt"

        "go.yaml.in/yaml/v3"
)

// frontmatterDelimiter is the YAML frontmatter delimiter.
var frontmatterDelimiter = []byte("---")

// MaxFrontmatterSize is the maximum size of frontmatter content in bytes.
// This prevents YAML parsing attacks (e.g., billion laughs).
const MaxFrontmatterSize = 64 * 1024 // 64KB

// MaxDependencies is the maximum number of dependencies allowed per skill.
// This prevents resource exhaustion from malicious or misconfigured skills.
const MaxDependencies = 100

// ErrInvalidFrontmatter indicates that the SKILL.md frontmatter is malformed.
var ErrInvalidFrontmatter = errors.New("invalid frontmatter")

// ParseSkillMD parses a SKILL.md file and extracts frontmatter and body.
func ParseSkillMD(content []byte) (*ParseResult, error) {
        fm, body, err := extractFrontmatter(content)
        if err != nil {
                return nil, err
        }

        requires, err := toDependencies(fm.Requires)
        if err != nil {
                return nil, fmt.Errorf("parsing requires: %w", err)
        }

        return &ParseResult{
                Name:          fm.Name,
                Description:   fm.Description,
                Version:       fm.Version,
                AllowedTools:  fm.AllowedTools,
                License:       fm.License,
                Compatibility: fm.Compatibility,
                Metadata:      fm.Metadata,
                Requires:      requires,
                Body:          body,
        }, nil
}

// extractFrontmatter extracts YAML frontmatter from content.
// Returns the parsed frontmatter and the remaining body content.
func extractFrontmatter(content []byte) (*SkillFrontmatter, []byte, error) {
        content = bytes.TrimSpace(content)

        if !bytes.HasPrefix(content, frontmatterDelimiter) {
                return nil, nil, ErrInvalidFrontmatter
        }

        // Skip opening delimiter and optional newline
        rest := content[len(frontmatterDelimiter):]
        rest = bytes.TrimPrefix(rest, []byte("\n"))

        // Limit the search scope for the closing delimiter to avoid scanning
        // arbitrarily large inputs. Search within MaxFrontmatterSize + delimiter
        // bytes for the closing "---"; if not found, the frontmatter is too large.
        searchLimit := rest
        maxSearch := MaxFrontmatterSize + len(frontmatterDelimiter)
        if len(searchLimit) > maxSearch {
                searchLimit = searchLimit[:maxSearch]
        }

        endIdx := bytes.Index(searchLimit, frontmatterDelimiter)
        if endIdx == -1 {
                if len(rest) > MaxFrontmatterSize {
                        return nil, nil, fmt.Errorf("%w: frontmatter exceeds maximum size of %d bytes",
                                ErrInvalidFrontmatter, MaxFrontmatterSize)
                }
                return nil, nil, ErrInvalidFrontmatter
        }

        frontmatterBytes := rest[:endIdx]
        body := rest[endIdx+len(frontmatterDelimiter):]
        body = bytes.TrimPrefix(body, []byte("\n"))

        var fm SkillFrontmatter
        if err := yaml.Unmarshal(frontmatterBytes, &fm); err != nil {
                return nil, nil, fmt.Errorf("%w: %v", ErrInvalidFrontmatter, err)
        }

        return &fm, body, nil
}

// toDependencies converts a list of OCI reference strings to Dependencies.
func toDependencies(refs []string) ([]Dependency, error) {
        if len(refs) == 0 {
                return nil, nil
        }

        if len(refs) > MaxDependencies {
                return nil, fmt.Errorf("too many dependencies: %d (max %d)", len(refs), MaxDependencies)
        }

        deps := make([]Dependency, 0, len(refs))
        for _, ref := range refs {
                if ref != "" {
                        deps = append(deps, Dependency{Reference: ref})
                }
        }

        return deps, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package skills

import (
        "errors"
        "fmt"
        "io/fs"
        "os"
        "path/filepath"
        "strings"
)

// ValidateProjectRoot validates a project root path and returns its cleaned
// form. Symlinks are rejected to avoid unexpected filesystem traversal.
func ValidateProjectRoot(projectRoot string) (string, error) {
        if err := validateProjectRootInput(projectRoot); err != nil {
                return "", err
        }

        projectRoot = filepath.Clean(projectRoot)
        resolvedRoot, err := resolveProjectRoot(projectRoot)
        if err != nil {
                return "", err
        }
        if err := validateProjectRootDir(resolvedRoot); err != nil {
                return "", err
        }
        if err := validateProjectRootGitDir(resolvedRoot); err != nil {
                return "", err
        }

        return resolvedRoot, nil
}

// NormalizeScopeAndProjectRoot validates scope and project_root and returns
// normalized values.
func NormalizeScopeAndProjectRoot(scope Scope, projectRoot string) (Scope, string, error) {
        if projectRoot != "" && scope == "" {
                scope = ScopeProject
        }
        if err := ValidateScope(scope); err != nil {
                return scope, projectRoot, err
        }
        if projectRoot != "" && scope != ScopeProject {
                return scope, projectRoot, errors.New("project_root is only valid with project scope")
        }
        if scope == ScopeProject {
                cleaned, err := ValidateProjectRoot(projectRoot)
                if err != nil {
                        return scope, projectRoot, err
                }
                return scope, cleaned, nil
        }
        return scope, projectRoot, nil
}

func validateProjectRootInput(projectRoot string) error {
        if projectRoot == "" {
                return errors.New("project_root is required for project scope")
        }
        if strings.ContainsRune(projectRoot, 0) {
                return errors.New("project_root contains null bytes")
        }
        if !filepath.IsAbs(projectRoot) {
                return fmt.Errorf("project_root must be absolute, got %q", projectRoot)
        }
        return validateNoTraversal(projectRoot)
}

func resolveProjectRoot(projectRoot string) (string, error) {
        resolvedRoot, err := filepath.EvalSymlinks(projectRoot)
        if err != nil {
                if errors.Is(err, fs.ErrNotExist) {
                        return "", errors.New("project_root does not exist")
                }
                return "", fmt.Errorf("resolving project_root: %w", err)
        }
        resolvedRoot = filepath.Clean(resolvedRoot)
        cleanedRoot := filepath.Clean(projectRoot)
        if resolvedRoot != cleanedRoot {
                return "", errors.New("project_root must not contain symlinks")
        }
        if !filepath.IsAbs(resolvedRoot) {
                return "", fmt.Errorf("project_root must be absolute, got %q", resolvedRoot)
        }
        if err := validateNoTraversal(resolvedRoot); err != nil {
                return "", err
        }
        return resolvedRoot, nil
}

func validateProjectRootDir(projectRoot string) error {
        // project_root is user-provided, but already validated for absolute paths,
        // traversal, and symlink usage before any filesystem access.
        info, err := os.Stat(projectRoot) // #nosec G304 -- path is validated and resolved above
        if err != nil {
                if errors.Is(err, fs.ErrNotExist) {
                        return errors.New("project_root does not exist")
                }
                return fmt.Errorf("checking project_root: %w", err)
        }
        if !info.IsDir() {
                return errors.New("project_root must be a directory")
        }
        return nil
}

func validateProjectRootGitDir(projectRoot string) error {
        gitPath := filepath.Join(projectRoot, ".git")
        gitInfo, err := os.Stat(gitPath) // #nosec G304 -- path is validated and resolved above
        if err != nil {
                if errors.Is(err, fs.ErrNotExist) {
                        return errors.New("project_root must be a git repository")
                }
                return fmt.Errorf("checking project_root .git: %w", err)
        }
        if !gitInfo.IsDir() && !gitInfo.Mode().IsRegular() {
                return errors.New("project_root must contain a .git directory or file")
        }
        return nil
}

func validateNoTraversal(path string) error {
        for _, segment := range strings.Split(filepath.ToSlash(path), "/") {
                if segment == ".." {
                        return errors.New("project_root must not contain '..' traversal segments")
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package skills provides types and interfaces for managing ToolHive skills.
package skills

import (
        "fmt"
        "strings"
        "time"

        "go.yaml.in/yaml/v3"
)

// Scope represents the scope at which a skill is installed.
type Scope string

const (
        // ScopeUser indicates a skill installed for the current user (user-wide, ~/).
        ScopeUser Scope = "user"
        // ScopeProject indicates a skill installed for a specific project (project-local).
        ScopeProject Scope = "project"
)

// ValidateScope checks that a scope value is valid. An empty scope is accepted
// (meaning "unscoped" / "all"). Otherwise only "user" and "project" are allowed.
func ValidateScope(s Scope) error {
        switch s {
        case "", ScopeUser, ScopeProject:
                return nil
        default:
                return fmt.Errorf("invalid scope %q: must be empty, %q, or %q", s, ScopeUser, ScopeProject)
        }
}

// InstallStatus represents the current status of a skill installation.
type InstallStatus string

const (
        // InstallStatusInstalled indicates a skill is fully installed and ready.
        InstallStatusInstalled InstallStatus = "installed"
        // InstallStatusPending indicates a skill installation is in progress.
        InstallStatusPending InstallStatus = "pending"
        // InstallStatusFailed indicates a skill installation has failed.
        InstallStatusFailed InstallStatus = "failed"
)

// StringOrSlice is a custom type that can unmarshal from either a YAML string
// (space-delimited per spec, or comma-delimited for compatibility) or a YAML array.
type StringOrSlice []string

// UnmarshalYAML implements yaml.Unmarshaler for StringOrSlice.
// Per the Agent Skills spec, allowed-tools is space-delimited, but we also
// support comma-delimited for compatibility with existing skills.
func (s *StringOrSlice) UnmarshalYAML(value *yaml.Node) error {
        switch value.Kind {
        case yaml.ScalarNode:
                str := value.Value
                if str == "" {
                        *s = nil
                        return nil
                }

                // Delimiter precedence: if any comma is present, split on commas;
                // otherwise split on whitespace (space-delimited is the canonical
                // format per the Agent Skills spec). This means a mixed-delimiter
                // string like "Read,Glob Grep" splits on comma, yielding
                // ["Read", "Glob Grep"] — comma takes priority.
                var parts []string
                if strings.Contains(str, ",") {
                        parts = strings.Split(str, ",")
                } else {
                        parts = strings.Fields(str)
                }

                result := make([]string, 0, len(parts))
                for _, part := range parts {
                        trimmed := strings.TrimSpace(part)
                        if trimmed != "" {
                                result = append(result, trimmed)
                        }
                }
                *s = result
                return nil
        case yaml.SequenceNode:
                var arr []string
                if err := value.Decode(&arr); err != nil {
                        return fmt.Errorf("decoding allowed-tools array: %w", err)
                }
                *s = arr
                return nil
        case yaml.DocumentNode, yaml.MappingNode, yaml.AliasNode:
                return fmt.Errorf("allowed-tools: expected string or array, got unsupported YAML node type")
        }
        return fmt.Errorf("allowed-tools: unexpected YAML node kind %d", value.Kind)
}

// SkillFrontmatter represents the raw YAML frontmatter from a SKILL.md file.
type SkillFrontmatter struct {
        Name          string            `yaml:"name"`
        Description   string            `yaml:"description"`
        Version       string            `yaml:"version,omitempty"`
        AllowedTools  StringOrSlice     `yaml:"allowed-tools,omitempty"`
        Requires      StringOrSlice     `yaml:"toolhive.requires,omitempty"`
        License       string            `yaml:"license,omitempty"`
        Compatibility string            `yaml:"compatibility,omitempty"`
        Metadata      map[string]string `yaml:"metadata,omitempty"`
}

// Dependency represents an external skill dependency (OCI reference).
type Dependency struct {
        // Name is the dependency name.
        Name string `json:"name,omitempty"`
        // Reference is the OCI reference for the dependency.
        Reference string `json:"reference"`
        // Digest is the OCI digest for upgrade detection.
        Digest string `json:"digest,omitempty"`
}

// ParseResult contains the parsed contents of a SKILL.md file.
type ParseResult struct {
        Name          string
        Description   string
        Version       string
        AllowedTools  []string
        License       string
        Compatibility string
        Metadata      map[string]string
        Requires      []Dependency
        Body          []byte
}

// SkillMetadata contains metadata about a skill.
type SkillMetadata struct {
        // Name is the unique name of the skill.
        Name string `json:"name"`
        // Version is the semantic version of the skill.
        Version string `json:"version"`
        // Description is a human-readable description of the skill.
        Description string `json:"description"`
        // Author is the skill author or maintainer.
        Author string `json:"author"`
        // Tags is a list of tags for categorization.
        Tags []string `json:"tags,omitempty"`
}

// InstalledSkill represents a skill that has been installed locally.
type InstalledSkill struct {
        // Metadata contains the skill's metadata.
        Metadata SkillMetadata `json:"metadata"`
        // Scope is the installation scope (user or project).
        Scope Scope `json:"scope"`
        // ProjectRoot is the project root path for project-scoped skills. Empty for user-scoped.
        ProjectRoot string `json:"project_root,omitempty"`
        // Reference is the full OCI reference (e.g. ghcr.io/org/skill:v1).
        Reference string `json:"reference,omitempty"`
        // Tag is the OCI tag (e.g. v1.0.0).
        Tag string `json:"tag,omitempty"`
        // Digest is the OCI digest (sha256:...) for upgrade detection.
        Digest string `json:"digest,omitempty"`
        // Status is the current installation status.
        Status InstallStatus `json:"status"`
        // InstalledAt is the timestamp when the skill was installed.
        InstalledAt time.Time `json:"installed_at"`
        // Clients is the list of client identifiers the skill is installed for.
        // TODO: Refactor client.ClientApp to a shared package so it can be used here instead of []string.
        Clients []string `json:"clients,omitempty"`
        // Dependencies is the list of external skill dependencies.
        Dependencies []Dependency `json:"dependencies,omitempty"`
        // Managed indicates this install is tracked in the project's
        // toolhive.lock.yaml. Only ever true for project-scoped installs. No
        // omitempty: false is an observable state (unmanaged), not an absence.
        Managed bool `json:"managed"`
        // SigstoreBundle is the signature material captured at install time,
        // used for offline re-verification during sync. It is core's durable
        // form (see verifier.Result.Bundle), not necessarily bare Sigstore
        // bundle JSON, and is re-verified against the ARTIFACT digest. Nil for
        // unsigned installs. Never serialized to API responses.
        SigstoreBundle []byte `json:"-"`
}

// SkillIndexEntry represents a single skill entry in a remote skill index.
type SkillIndexEntry struct {
        // Metadata contains the skill's metadata.
        Metadata SkillMetadata `json:"metadata"`
        // Repository is the OCI repository reference for the skill.
        Repository string `json:"repository"`
        // Group is the optional group this skill belongs to.
        Group string `json:"group,omitempty"`
}

// SkillIndex represents a collection of available skills from a remote index.
type SkillIndex struct {
        // Skills is the list of available skills.
        Skills []SkillIndexEntry `json:"skills"`
}

//go:generate mockgen -destination=mocks/mock_path_resolver.go -package=mocks -source=types.go PathResolver Installer

// PathResolver resolves filesystem paths for skill installations.
// It uses string (not client.ClientApp) to avoid importing pkg/client from pkg/skills.
type PathResolver interface {
        // GetSkillPath returns the filesystem path where a skill should be installed.
        GetSkillPath(clientType, skillName string, scope Scope, projectRoot string) (string, error)
        // ListSkillSupportingClients returns all client identifiers that support skills.
        ListSkillSupportingClients() []string
}

// Installer handles filesystem operations for skill installation and removal.
type Installer interface {
        // Extract decompresses a tar.gz OCI layer and writes files to targetDir.
        Extract(layerData []byte, targetDir string, force bool) (*ExtractResult, error)
        // ExtractPlugin is like Extract but preserves the executable bit (cap 0755)
        // so plugin hook scripts keep +x. Used by plugin adapters.
        ExtractPlugin(layerData []byte, targetDir string, force bool) (*ExtractResult, error)
        // Remove safely removes a skill directory.
        Remove(skillDir string) error
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package skills

import (
        "bytes"
        "fmt"
        "io/fs"
        "os"
        "path/filepath"
        "regexp"
        "strings"
)

// skillNameRegex validates skill names: 2-64 chars, lowercase alphanumeric and hyphens,
// must start and end with alphanumeric.
var skillNameRegex = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,62}[a-z0-9]$`)

// MaxCompatibilityLength is the maximum allowed length for the compatibility field.
const MaxCompatibilityLength = 500

// MaxDescriptionLength is the maximum allowed length for the description field.
const MaxDescriptionLength = 1024

// RecommendedMaxSkillMDLines is the recommended maximum number of lines in SKILL.md.
// Exceeding this generates a warning, not an error.
const RecommendedMaxSkillMDLines = 500

// ValidateSkillDir validates a skill directory at the given path.
// I/O errors are returned as error; validation issues are returned in ValidationResult.
func ValidateSkillDir(path string) (*ValidationResult, error) {
        // Defense-in-depth: sanitize and validate the path before any filesystem access.
        // The caller (skillsvc.Validate) also validates via validateLocalPath, but we
        // re-check here because ValidateSkillDir is exported and may be called directly.
        path = filepath.Clean(path)
        if !filepath.IsAbs(path) {
                return nil, fmt.Errorf("path must be absolute, got %q", path)
        }

        var errs []string
        var warnings []string

        // Check SKILL.md exists
        skillMDPath := filepath.Join(path, "SKILL.md")
        content, err := os.ReadFile(skillMDPath) //#nosec G304 -- path is cleaned and validated as absolute above
        if err != nil {
                if os.IsNotExist(err) {
                        return &ValidationResult{
                                Valid:  false,
                                Errors: []string{"SKILL.md not found in skill directory"},
                        }, nil
                }
                return nil, fmt.Errorf("reading SKILL.md: %w", err)
        }

        // Check for symlinks and path traversal in a single walk
        if err := CheckFilesystem(path); err != nil {
                errs = append(errs, err.Error())
        }

        // Parse frontmatter
        result, err := ParseSkillMD(content)
        if err != nil {
                errs = append(errs, fmt.Sprintf("invalid SKILL.md: %v", err))
                return &ValidationResult{
                        Valid:  false,
                        Errors: errs,
                }, nil
        }

        // Validate parsed fields
        errs = append(errs, validateFields(result, filepath.Base(path))...)

        // Collect warnings
        warnings = append(warnings, collectWarnings(result, content)...)

        return &ValidationResult{
                Valid:    len(errs) == 0,
                Errors:   errs,
                Warnings: warnings,
        }, nil
}

// validateFields checks parsed frontmatter fields against spec constraints.
func validateFields(result *ParseResult, dirName string) []string {
        var errs []string

        if result.Name == "" {
                errs = append(errs, "name is required")
        } else {
                if err := ValidateSkillName(result.Name); err != nil {
                        errs = append(errs, err.Error())
                }
                if result.Name != dirName {
                        errs = append(errs,
                                fmt.Sprintf("skill name %q must match directory name %q", result.Name, dirName))
                }
        }
        if result.Description == "" {
                errs = append(errs, "description is required")
        }
        if len(result.Description) > MaxDescriptionLength {
                errs = append(errs,
                        fmt.Sprintf("description exceeds maximum length of %d characters", MaxDescriptionLength))
        }
        if len(result.Compatibility) > MaxCompatibilityLength {
                errs = append(errs,
                        fmt.Sprintf("compatibility field exceeds maximum length of %d characters", MaxCompatibilityLength))
        }

        return errs
}

// collectWarnings generates non-blocking warnings for spec compliance.
func collectWarnings(result *ParseResult, content []byte) []string {
        var warnings []string

        if len(result.AllowedTools) > 0 && bytes.Contains(content, []byte(",")) &&
                bytes.Contains(content, []byte("allowed-tools:")) {
                warnings = append(warnings,
                        "allowed-tools uses comma-delimited format, which is a ToolHive extension; "+
                                "the Agent Skills spec defines space-delimited as the canonical format")
        }
        lineCount := bytes.Count(content, []byte("\n")) + 1
        if lineCount > RecommendedMaxSkillMDLines {
                warnings = append(warnings,
                        fmt.Sprintf("SKILL.md has %d lines (recommended max: %d)", lineCount, RecommendedMaxSkillMDLines))
        }

        return warnings
}

// ValidateSkillName checks that a skill name conforms to the Agent Skills specification.
// Names must be 2-64 lowercase alphanumeric characters or hyphens, starting and ending
// with alphanumeric, with no consecutive hyphens.
// See: https://agentskills.io/specification
func ValidateSkillName(name string) error {
        if name == "" {
                return fmt.Errorf("invalid skill name: must not be empty")
        }
        if !skillNameRegex.MatchString(name) {
                return fmt.Errorf("invalid skill name %q: must be 2-64 lowercase alphanumeric characters or hyphens, "+
                        "starting and ending with alphanumeric", name)
        }
        if strings.Contains(name, "--") {
                return fmt.Errorf("invalid skill name %q: must not contain consecutive hyphens", name)
        }
        return nil
}

// CheckFilesystem walks the directory once, checking for symlinks and path traversal.
func CheckFilesystem(path string) error {
        return filepath.WalkDir(path, func(p string, _ fs.DirEntry, err error) error {
                if err != nil {
                        return nil // Skip inaccessible paths
                }

                rel, err := filepath.Rel(path, p)
                if err != nil {
                        return nil
                }

                // Check for path traversal
                for _, component := range strings.Split(filepath.ToSlash(rel), "/") {
                        if component == ".." {
                                return fmt.Errorf("path traversal detected in %q: '..' components are not allowed", rel)
                        }
                }

                // Check for symlinks (WalkDir doesn't stat, so use Lstat)
                info, err := os.Lstat(p)
                if err != nil {
                        return nil
                }
                if info.Mode()&os.ModeSymlink != 0 {
                        return fmt.Errorf("symlink found at %q: symlinks are not allowed in skill directories", rel)
                }

                return nil
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package state

import (
        "github.com/stacklok/toolhive/pkg/container/runtime"
)

const (
        // RunConfigsDir is the directory name for storing run configurations
        RunConfigsDir = "runconfigs"

        // GroupConfigsDir is the directory name for storing group configurations
        GroupConfigsDir = "groups"
)

// NewRunConfigStore creates a store for run configuration state
func NewRunConfigStore(appName string) (Store, error) {
        return NewRunConfigStoreWithDetector(appName)
}

// NewRunConfigStoreWithDetector creates a store
func NewRunConfigStoreWithDetector(appName string) (Store, error) {
        if runtime.IsKubernetesRuntime() {
                return NewKubernetesStore(), nil
        }
        return NewLocalStore(appName, RunConfigsDir)
}

// NewGroupConfigStore creates a store for group configurations
func NewGroupConfigStore(appName string) (Store, error) {
        return NewGroupConfigStoreWithDetector(appName)
}

// NewGroupConfigStoreWithDetector creates a store
func NewGroupConfigStoreWithDetector(appName string) (Store, error) {
        if runtime.IsKubernetesRuntime() {
                return NewKubernetesStore(), nil
        }
        return NewLocalStore(appName, GroupConfigsDir)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package state provides functionality for storing and retrieving runner state
// across different environments (local filesystem, Kubernetes, etc.)
package state

import (
        "context"
        "fmt"
        "io"
)

// Aborter is implemented by state writers that can discard uncommitted data.
// Abort must release the writer's resources without publishing its contents.
type Aborter interface {
        Abort() error
}

// AbortWriter discards uncommitted data written to writer.
//
// Writers returned by Store.GetWriter and Store.CreateExclusive must implement
// Aborter. Callers must use AbortWriter, rather than Close, when abandoning a
// write because Close may publish the data.
func AbortWriter(writer io.WriteCloser) error {
        aborter, ok := writer.(Aborter)
        if !ok {
                return fmt.Errorf("state writer does not support abort")
        }
        return aborter.Abort()
}

//go:generate mockgen -destination=mocks/mock_store.go -package=mocks -source=interface.go Store

// Store defines the interface for runner state storage operations. Writers returned by
// GetWriter and CreateExclusive must support Aborter; callers must call AbortWriter
// rather than Close when abandoning a write, since Close may publish it.
type Store interface {
        // GetReader returns a reader for the state data
        // This is useful for streaming large state data
        GetReader(ctx context.Context, name string) (io.ReadCloser, error)

        // GetWriter returns a writer for the state data. For stores that publish state on Close,
        // callers must return Close errors because the write may not be visible until Close succeeds.
        // Callers must use AbortWriter instead of Close when abandoning a write, because Close may publish it.
        // This is useful for streaming large state data.
        GetWriter(ctx context.Context, name string) (io.WriteCloser, error)

        // CreateExclusive creates a new state entry exclusively, returning an error if it already exists.
        // For stores that publish state on Close, exclusivity is enforced at Close and Close can return
        // an error with http.StatusConflict if another writer published the entry first.
        // Callers must return Close errors and use AbortWriter rather than Close when abandoning a write,
        // because the write may not be visible until Close succeeds.
        // This provides atomic check-and-create semantics to prevent race conditions.
        CreateExclusive(ctx context.Context, name string) (io.WriteCloser, error)

        // Delete removes the data for the given name
        Delete(ctx context.Context, name string) error

        // List returns all available state names
        List(ctx context.Context) ([]string, error)

        // Exists checks if data exists for the given name
        Exists(ctx context.Context, name string) (bool, error)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package state

import (
        "context"
        "io"
        "strings"
)

// KubernetesStore is a no-op implementation of Store for Kubernetes environments.
// In Kubernetes, workload state is managed by the cluster, not by local files.
type KubernetesStore struct{}

// NewKubernetesStore creates a new no-op store for Kubernetes environments.
func NewKubernetesStore() Store {
        return &KubernetesStore{}
}

// Exists always returns false for Kubernetes stores since state is not persisted locally.
func (*KubernetesStore) Exists(_ context.Context, _ string) (bool, error) {
        return false, nil
}

// List always returns an empty slice for Kubernetes stores.
func (*KubernetesStore) List(_ context.Context) ([]string, error) {
        return []string{}, nil
}

// GetReader returns a no-op reader for Kubernetes stores.
func (*KubernetesStore) GetReader(_ context.Context, _ string) (io.ReadCloser, error) {
        return io.NopCloser(strings.NewReader("")), nil
}

// GetWriter returns a no-op writer for Kubernetes stores.
func (*KubernetesStore) GetWriter(_ context.Context, _ string) (io.WriteCloser, error) {
        return &noopWriteCloser{}, nil
}

// CreateExclusive returns a no-op writer for Kubernetes stores.
// In Kubernetes, state management is handled by the cluster, not local files.
func (*KubernetesStore) CreateExclusive(_ context.Context, _ string) (io.WriteCloser, error) {
        return &noopWriteCloser{}, nil
}

// Delete is a no-op for Kubernetes stores.
func (*KubernetesStore) Delete(_ context.Context, _ string) error {
        return nil
}

// noopWriteCloser is a writer that discards all writes.
type noopWriteCloser struct{}

// Write discards all data and returns success.
func (*noopWriteCloser) Write(p []byte) (n int, err error) {
        return len(p), nil
}

// Close is a no-op.
func (*noopWriteCloser) Close() error {
        return nil
}

// Abort is a no-op.
func (*noopWriteCloser) Abort() error {
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package state

import (
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "os"
        "path/filepath"
        "strings"

        "github.com/adrg/xdg"

        "github.com/stacklok/toolhive-core/httperr"
)

const (
        // DefaultAppName is the default application name used for XDG paths
        DefaultAppName = "toolhive"

        // FileExtension is the file extension for stored configurations
        FileExtension = ".json"
)

// LocalStore implements the Store interface using the local filesystem
// following the XDG Base Directory Specification
type LocalStore struct {
        // basePath is the base directory path for storing configurations
        basePath string
}

// NewLocalStore creates a new LocalStore with the given application name and store type
// If appName is empty, DefaultAppName will be used
func NewLocalStore(appName string, storeName string) (*LocalStore, error) {
        if appName == "" {
                appName = DefaultAppName
        }

        // Create the base directory path following XDG spec
        basePath := filepath.Join(xdg.StateHome, appName, storeName)

        // Ensure the directory exists
        if err := os.MkdirAll(basePath, 0750); err != nil {
                return nil, fmt.Errorf("failed to create state directory: %w", err)
        }

        return &LocalStore{
                basePath: basePath,
        }, nil
}

// getFilePath returns the full file path for a configuration, validating that
// the resolved path stays within the store's base directory.
func (s *LocalStore) getFilePath(name string) (string, error) {
        if !strings.HasSuffix(name, FileExtension) {
                name = name + FileExtension
        }
        filePath := filepath.Join(s.basePath, name)
        // filepath.Join calls filepath.Clean, which resolves ".." components.
        // Verify the result is still within basePath to prevent path traversal.
        if !strings.HasPrefix(filePath, s.basePath+string(filepath.Separator)) {
                return "", fmt.Errorf("path traversal detected: name escapes state directory")
        }
        return filePath, nil
}

// GetReader returns a reader for the state data
func (s *LocalStore) GetReader(_ context.Context, name string) (io.ReadCloser, error) {
        filePath, err := s.getFilePath(name)
        if err != nil {
                return nil, err
        }
        // #nosec G304 - path traversal is prevented by the containment check in getFilePath
        file, err := os.Open(filePath)
        if err != nil {
                if os.IsNotExist(err) {
                        return nil, httperr.WithCode(fmt.Errorf("state '%s' not found", name), http.StatusNotFound)
                }
                return nil, fmt.Errorf("failed to open state file: %w", err)
        }

        return file, nil
}

// localWriter writes state to a temporary file and publishes it when closed.
type localWriter struct {
        file       *os.File
        tempPath   string
        targetPath string
        exclusive  bool
}

// Write writes state data to the temporary file.
func (w *localWriter) Write(p []byte) (int, error) {
        return w.file.Write(p)
}

// Sync flushes the temporary file to durable storage.
func (w *localWriter) Sync() error {
        return w.file.Sync()
}

// Close syncs and closes the temporary file before atomically publishing it.
func (w *localWriter) Close() error {
        if err := w.file.Sync(); err != nil {
                return w.cleanupAfterError(err)
        }
        if err := w.file.Close(); err != nil {
                return w.removeTemp(err)
        }

        if w.exclusive {
                if err := os.Link(w.tempPath, w.targetPath); err != nil {
                        if os.IsExist(err) {
                                return w.removeTemp(httperr.WithCode(
                                        fmt.Errorf("state '%s' already exists", strings.TrimSuffix(filepath.Base(w.targetPath), FileExtension)),
                                        http.StatusConflict,
                                ))
                        }
                        return w.removeTemp(fmt.Errorf("failed to publish state file: %w", err))
                }
                if err := os.Remove(w.tempPath); err != nil && !os.IsNotExist(err) {
                        slog.Warn("failed to remove published temporary state file", "path", w.tempPath, "error", err)
                }
                return nil
        }

        if err := os.Rename(w.tempPath, w.targetPath); err != nil {
                return w.removeTemp(fmt.Errorf("failed to publish state file: %w", err))
        }
        return nil
}

// Abort closes and removes the temporary file without publishing it.
func (w *localWriter) Abort() error {
        return w.cleanupAfterError(nil)
}

func (w *localWriter) cleanupAfterError(err error) error {
        return w.removeTemp(errors.Join(err, w.file.Close()))
}

func (w *localWriter) removeTemp(err error) error {
        if removeErr := os.Remove(w.tempPath); removeErr != nil && !os.IsNotExist(removeErr) {
                return errors.Join(err, fmt.Errorf("failed to remove temporary state file: %w", removeErr))
        }
        return err
}

// GetWriter returns a writer for the state data. Data is atomically published when the writer is closed.
func (s *LocalStore) GetWriter(_ context.Context, name string) (io.WriteCloser, error) {
        filePath, err := s.getFilePath(name)
        if err != nil {
                return nil, err
        }

        return s.newWriter(filePath, false)
}

// CreateExclusive creates a new state entry exclusively. Data is atomically published when the writer is closed.
func (s *LocalStore) CreateExclusive(_ context.Context, name string) (io.WriteCloser, error) {
        filePath, err := s.getFilePath(name)
        if err != nil {
                return nil, err
        }
        if _, err := os.Lstat(filePath); err == nil {
                return nil, httperr.WithCode(
                        fmt.Errorf("state '%s' already exists", name),
                        http.StatusConflict,
                )
        } else if !os.IsNotExist(err) {
                return nil, fmt.Errorf("failed to check state file: %w", err)
        }

        return s.newWriter(filePath, true)
}

// Delete removes the data for the given name
func (s *LocalStore) Delete(_ context.Context, name string) error {
        filePath, err := s.getFilePath(name)
        if err != nil {
                return err
        }
        if err := os.Remove(filePath); err != nil {
                if os.IsNotExist(err) {
                        return fmt.Errorf("state '%s' not found", name)
                }
                return fmt.Errorf("failed to delete state file: %w", err)
        }
        return nil
}

// List returns all available state names
func (s *LocalStore) List(_ context.Context) ([]string, error) {
        // Read the directory
        entries, err := os.ReadDir(s.basePath)
        if err != nil {
                if os.IsNotExist(err) {
                        return []string{}, nil
                }
                return nil, fmt.Errorf("failed to read state directory: %w", err)
        }

        // Filter and process the file names
        var names []string
        for _, entry := range entries {
                if entry.IsDir() {
                        continue
                }

                name := entry.Name()
                if strings.HasSuffix(name, FileExtension) {
                        // Remove the file extension
                        name = strings.TrimSuffix(name, FileExtension)
                        names = append(names, name)
                }
        }

        return names, nil
}

// Exists checks if data exists for the given name
func (s *LocalStore) Exists(_ context.Context, name string) (bool, error) {
        filePath, err := s.getFilePath(name)
        if err != nil {
                return false, err
        }
        _, err = os.Stat(filePath)
        if err != nil {
                if os.IsNotExist(err) {
                        return false, nil
                }
                return false, fmt.Errorf("failed to check if state exists: %w", err)
        }
        return true, nil
}

func (*LocalStore) newWriter(filePath string, exclusive bool) (io.WriteCloser, error) {
        file, err := os.CreateTemp(filepath.Dir(filePath), "."+filepath.Base(filePath)+".tmp-*")
        if err != nil {
                return nil, fmt.Errorf("failed to create temporary state file: %w", err)
        }

        return &localWriter{file: file, tempPath: file.Name(), targetPath: filePath, exclusive: exclusive}, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package state

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"

        "github.com/stacklok/toolhive/pkg/workloads/types/errors"
)

// LoadRunConfigJSON loads a run configuration from the state store and returns the raw reader
func LoadRunConfigJSON(ctx context.Context, name string) (io.ReadCloser, error) {
        // Create a state store
        store, err := NewRunConfigStore(DefaultAppName)
        if err != nil {
                return nil, fmt.Errorf("failed to create state store: %w", err)
        }

        // Check if the configuration exists
        exists, err := store.Exists(ctx, name)
        if err != nil {
                return nil, fmt.Errorf("failed to check if run configuration exists: %w", err)
        }
        if !exists {
                return nil, fmt.Errorf("%w: %s", errors.ErrRunConfigNotFound, name)
        }

        // Get a reader for the state
        reader, err := store.GetReader(ctx, name)
        if err != nil {
                return nil, fmt.Errorf("failed to get reader for state: %w", err)
        }

        return reader, nil
}

// DeleteSavedRunConfig deletes a saved run configuration
func DeleteSavedRunConfig(ctx context.Context, name string) error {
        // Create a state store
        store, err := NewRunConfigStore(DefaultAppName)
        if err != nil {
                return fmt.Errorf("failed to create state store: %w", err)
        }

        // Check if the configuration exists
        exists, err := store.Exists(ctx, name)
        if err != nil {
                return fmt.Errorf("failed to check if run configuration exists: %w", err)
        }
        if !exists {
                return fmt.Errorf("run configuration for %s not found", name)
        }

        // Delete the configuration
        if err := store.Delete(ctx, name); err != nil {
                return fmt.Errorf("failed to delete run configuration: %w", err)
        }

        slog.Debug("Deleted run configuration", "name", name)
        return nil
}

// RunConfigPersister defines an interface for objects that can be persisted and loaded as JSON
type RunConfigPersister interface {
        // WriteJSON serializes the object to JSON and writes it to the provided writer
        WriteJSON(w io.Writer) error
        // GetBaseName returns the base name used for persistence
        GetBaseName() string
}

// ReadJSONFunc defines a function type for reading JSON into an object
type ReadJSONFunc[T any] func(r io.Reader) (T, error)

// SaveRunConfig saves a run configuration to the state store
func SaveRunConfig[T RunConfigPersister](ctx context.Context, config T) error {
        // Create a state store
        store, err := NewRunConfigStore(DefaultAppName)
        if err != nil {
                return fmt.Errorf("failed to create state store: %w", err)
        }

        // Get a writer for the state
        writer, err := store.GetWriter(ctx, config.GetBaseName())
        if err != nil {
                return fmt.Errorf("failed to get writer for state: %w", err)
        }
        closed := false
        defer func() {
                if !closed {
                        if err := AbortWriter(writer); err != nil {
                                slog.Warn("failed to abort run configuration writer", "name", config.GetBaseName(), "error", err)
                        }
                }
        }()

        // Serialize the configuration to JSON and write it directly to the state store
        if err := config.WriteJSON(writer); err != nil {
                return fmt.Errorf("failed to write run configuration: %w", err)
        }
        if err := writer.Close(); err != nil {
                closed = true
                return fmt.Errorf("failed to close run configuration writer: %w", err)
        }
        closed = true

        slog.Debug("Saved run configuration", "name", config.GetBaseName())
        return nil
}

// LoadRunConfig loads a run configuration from the state store using the provided reader function
func LoadRunConfig[T any](ctx context.Context, name string, readJSONFunc ReadJSONFunc[T]) (T, error) {
        var zero T
        reader, err := LoadRunConfigJSON(ctx, name)
        if err != nil {
                return zero, err
        }
        defer func() {
                if err := reader.Close(); err != nil {
                        slog.Warn("Failed to close reader", "error", err)
                }
        }()

        // Deserialize the configuration using the provided function
        return readJSONFunc(reader)
}

// ReadJSON deserializes JSON from the provided reader into a generic interface
// This function is moved from the runner package to avoid circular dependencies
func ReadJSON(r io.Reader, target interface{}) error {
        decoder := json.NewDecoder(r)
        return decoder.Decode(target)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package syncutil provides concurrency utilities.
package syncutil

import (
        "sync/atomic"
        "time"
)

// AtMost executes a function at most once per the configured interval.
// Subsequent calls within the interval are silently skipped.
// Safe for concurrent use from multiple goroutines.
type AtMost struct {
        interval time.Duration
        last     atomic.Int64
        now      func() time.Time
}

// NewAtMost creates an AtMost that allows execution at most once per interval.
// The first call to Do always executes.
func NewAtMost(interval time.Duration) *AtMost {
        return &AtMost{
                interval: interval,
                now:      time.Now,
        }
}

// Do calls fn if at least interval has elapsed since the last successful
// invocation. Otherwise the call is silently skipped.
// The first call always executes because last is initialized to zero,
// which is treated as "never called".
func (a *AtMost) Do(fn func()) {
        now := a.now().UnixNano()
        last := a.last.Load()
        if last != 0 && now-last < a.interval.Nanoseconds() {
                return
        }
        if a.last.CompareAndSwap(last, now) {
                fn()
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package telemetry provides OpenTelemetry instrumentation for ToolHive MCP server proxies.
package telemetry

import (
        "fmt"
        "strings"

        "go.opentelemetry.io/otel/attribute"
)

// ParseCustomAttributes parses a comma-separated list of key=value pairs into a map.
// Example input: "server_type=production,region=us-east-1,team=platform"
// Returns a map[string]string that can be converted to resource attributes.
func ParseCustomAttributes(input string) (map[string]string, error) {
        if input == "" {
                return map[string]string{}, nil
        }

        attributes := make(map[string]string)
        pairs := strings.Split(input, ",")

        for _, pair := range pairs {
                trimmedPair := strings.TrimSpace(pair)
                if trimmedPair == "" {
                        continue
                }

                parts := strings.SplitN(trimmedPair, "=", 2)
                if len(parts) != 2 {
                        return nil, fmt.Errorf("invalid attribute format '%s': expected key=value", trimmedPair)
                }

                key := strings.TrimSpace(parts[0])
                value := strings.TrimSpace(parts[1])

                if key == "" {
                        return nil, fmt.Errorf("empty attribute key in '%s'", trimmedPair)
                }

                // Store as key-value pair in map
                attributes[key] = value
        }

        return attributes, nil
}

// ConvertMapToAttributes converts a map[string]string to OpenTelemetry attributes
func ConvertMapToAttributes(attrs map[string]string) []attribute.KeyValue {
        var result []attribute.KeyValue
        for k, v := range attrs {
                result = append(result, attribute.String(k, v))
        }
        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package telemetry provides OpenTelemetry instrumentation for ToolHive MCP server proxies.
package telemetry

import (
        "context"
        "fmt"
        "maps"
        "net/http"
        "strconv"
        "strings"

        "go.opentelemetry.io/otel"
        "go.opentelemetry.io/otel/metric"
        "go.opentelemetry.io/otel/propagation"
        sdktrace "go.opentelemetry.io/otel/sdk/trace"
        "go.opentelemetry.io/otel/trace"

        "github.com/stacklok/toolhive-core/telemetry/providers"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/versions"
)

// Config holds the configuration for OpenTelemetry instrumentation.
// +kubebuilder:object:generate=true
// +gendoc
type Config struct {
        // Endpoint is the OTLP endpoint URL
        // +optional
        Endpoint string `json:"endpoint,omitempty" yaml:"endpoint,omitempty"`

        // ServiceName is the service name for telemetry.
        // When omitted, defaults to the server name (e.g., VirtualMCPServer name).
        // +optional
        ServiceName string `json:"serviceName,omitempty" yaml:"serviceName,omitempty"`

        // ServiceVersion is the service version for telemetry.
        // When omitted, defaults to the ToolHive version.
        // +optional
        ServiceVersion string `json:"serviceVersion,omitempty" yaml:"serviceVersion,omitempty"`

        // TracingEnabled controls whether distributed tracing is enabled.
        // When false, no tracer provider is created even if an endpoint is configured.
        // +kubebuilder:default=false
        // +optional
        TracingEnabled bool `json:"tracingEnabled" yaml:"tracingEnabled"`

        // MetricsEnabled controls whether OTLP metrics are enabled.
        // When false, OTLP metrics are not sent even if an endpoint is configured.
        // This is independent of EnablePrometheusMetricsPath.
        // +kubebuilder:default=false
        // +optional
        MetricsEnabled bool `json:"metricsEnabled" yaml:"metricsEnabled"`

        // SamplingRate is the trace sampling rate (0.0-1.0) as a string.
        // Only used when TracingEnabled is true.
        // Example: "0.05" for 5% sampling.
        // +kubebuilder:default="0.05"
        // +optional
        SamplingRate string `json:"samplingRate,omitempty" yaml:"samplingRate,omitempty"`

        // Headers contains authentication headers for the OTLP endpoint.
        // +optional
        Headers map[string]string `json:"headers,omitempty" yaml:"headers,omitempty"`

        // Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint.
        // +kubebuilder:default=false
        // +optional
        Insecure bool `json:"insecure,omitempty" yaml:"insecure,omitempty"`

        // EnablePrometheusMetricsPath controls whether to expose Prometheus-style /metrics endpoint.
        // The metrics are served at /metrics on a dedicated diagnostics port rather than on the
        // main transport port, so the endpoint can be restricted by port and is not routed
        // alongside application traffic. The endpoint is unauthenticated either way.
        // See PrometheusPort and pkg/diagnostics.
        // This is separate from OTLP metrics which are sent to the Endpoint.
        // +kubebuilder:default=false
        // +optional
        EnablePrometheusMetricsPath bool `json:"enablePrometheusMetricsPath,omitempty" yaml:"enablePrometheusMetricsPath,omitempty"`

        // PrometheusPort is the port the Prometheus /metrics endpoint is served on when
        // EnablePrometheusMetricsPath is true. It is deliberately not the main transport port,
        // so that access can be restricted with a NetworkPolicy: NetworkPolicy matches on port,
        // not on HTTP path, so a shared port makes "allow MCP traffic, deny metrics scraping"
        // impossible to express. The endpoint itself is unauthenticated, so restricting who can
        // reach this port is how it is protected.
        //
        // Zero selects the default diagnostics port (9464, the OpenTelemetry specification's
        // Prometheus exporter default). If that port is taken the listener falls back to an
        // available one and logs the resolved address. Do not route this port publicly.
        // +optional
        PrometheusPort int `json:"prometheusPort,omitempty" yaml:"prometheusPort,omitempty"`

        // MetricsOnTransportPort controls whether /metrics is ALSO served on the main
        // transport port, in addition to the diagnostics port. It exists to give
        // deployments a migration window: while true, an existing scrape configuration
        // aimed at the transport port keeps working, and a new one aimed at
        // PrometheusPort works too, so a scraper can be moved and verified before the
        // old location goes away. See https://github.com/stacklok/toolhive/issues/6384 for
        // the removal timeline.
        //
        // +optional
        MetricsOnTransportPort *bool `json:"metricsOnTransportPort,omitempty" yaml:"metricsOnTransportPort,omitempty"`

        // EnvironmentVariables is a list of environment variable names that should be
        // included in telemetry spans as attributes. Only variables in this list will
        // be read from the host machine and included in spans for observability.
        // Example: ["NODE_ENV", "DEPLOYMENT_ENV", "SERVICE_VERSION"]
        // +optional
        EnvironmentVariables []string `json:"environmentVariables,omitempty" yaml:"environmentVariables,omitempty"`

        // CustomAttributes contains custom resource attributes to be added to all telemetry signals.
        // These are parsed from CLI flags (--otel-custom-attributes) or environment variables
        // (OTEL_RESOURCE_ATTRIBUTES) as key=value pairs.
        // +optional
        CustomAttributes map[string]string `json:"customAttributes,omitempty" yaml:"customAttributes,omitempty"`

        // UseLegacyAttributes controls whether legacy (pre-MCP OTEL semconv) attribute names
        // are emitted alongside the new standard attribute names. When true, spans include both
        // old and new attribute names for backward compatibility with existing dashboards.
        // Currently defaults to true; this will change to false in a future release.
        // +kubebuilder:default=true
        // +optional
        UseLegacyAttributes bool `json:"useLegacyAttributes" yaml:"useLegacyAttributes"`

        // CACertPath is the file path to a CA certificate bundle for the OTLP endpoint.
        // When set, the OTLP exporters use this CA to verify the collector's TLS certificate
        // instead of relying solely on the system CA pool.
        // +optional
        CACertPath string `json:"caCertPath,omitempty" yaml:"caCertPath,omitempty"`
}

// Ensure Config implements fmt.Stringer and fmt.GoStringer
var _ fmt.Stringer = Config{}
var _ fmt.GoStringer = Config{}

// GoString returns the same redacted representation as String().
// This prevents credential leakage via the %#v format verb, which calls GoString() instead of String().
func (c Config) GoString() string {
        return c.String()
}

// String returns a human-readable representation of the Config with sensitive header values
// and environment variable entries redacted.
func (c Config) String() string {
        // Redact header values to prevent credential leakage
        redactedHeaders := make(map[string]string, len(c.Headers))
        for key := range c.Headers {
                redactedHeaders[key] = "[REDACTED]"
        }

        return fmt.Sprintf("Config{Endpoint: %q, ServiceName: %q, ServiceVersion: %q, TracingEnabled: %t, "+
                "MetricsEnabled: %t, SamplingRate: %q, Headers: %v, Insecure: %t, "+
                "EnablePrometheusMetricsPath: %t, PrometheusPort: %d, MetricsOnTransportPort: %t, "+
                "EnvironmentVariables: [%d redacted], CustomAttributes: %v, "+
                "UseLegacyAttributes: %t, CACertPath: %q}",
                c.Endpoint, c.ServiceName, c.ServiceVersion, c.TracingEnabled,
                c.MetricsEnabled, c.SamplingRate, redactedHeaders, c.Insecure,
                c.EnablePrometheusMetricsPath, c.PrometheusPort, c.ServeMetricsOnTransportPort(),
                len(c.EnvironmentVariables), c.CustomAttributes,
                c.UseLegacyAttributes, c.CACertPath)
}

// DefaultMetricsOnTransportPort is the value used when Config.MetricsOnTransportPort
// is unset.
//
// True during the deprecation window: /metrics is served on both the transport port
// and the diagnostics port, so no existing scrape configuration breaks. Flipping this
// to false is the entire cutover — it stops serving the transport-port copy and leaves
// only the diagnostics port. Deployments that set MetricsOnTransportPort explicitly
// are unaffected by the flip, by design.
//
// Flipping it is not the whole job: #6384 carries the timeline and the cleanup this
// field, its CLI flag, and its CRD entry all need afterwards.
const DefaultMetricsOnTransportPort = true

// ServeMetricsOnTransportPort reports whether /metrics should also be served on the
// main transport port. It resolves the unset case at the point of use, which is what
// makes changing DefaultMetricsOnTransportPort take effect for workloads that already
// exist. A nil receiver reports the default so callers need no nil check.
func (c *Config) ServeMetricsOnTransportPort() bool {
        if c == nil || c.MetricsOnTransportPort == nil {
                return DefaultMetricsOnTransportPort
        }
        return *c.MetricsOnTransportPort
}

// GetSamplingRateFloat parses the SamplingRate string and returns it as float64.
// Returns 0.0 if the string is empty or cannot be parsed.
func (c *Config) GetSamplingRateFloat() float64 {
        if c.SamplingRate == "" {
                return 0.0
        }
        rate, err := strconv.ParseFloat(c.SamplingRate, 64)
        if err != nil {
                return 0.0
        }
        return rate
}

// SetSamplingRateFromFloat sets the SamplingRate from a float64 value.
func (c *Config) SetSamplingRateFromFloat(rate float64) {
        c.SamplingRate = strconv.FormatFloat(rate, 'f', -1, 64)
}

// DefaultServiceNamePrefix is prepended to the workload name when deriving the
// OTel service name automatically (e.g. "thv-fetch", "thv-github").
const DefaultServiceNamePrefix = "thv-"

// DefaultConfig returns a default telemetry configuration.
func DefaultConfig() Config {
        return Config{
                ServiceName:                 "",     // empty — resolved at runtime from the workload name
                ServiceVersion:              "",     // resolved at runtime in NewProvider()
                TracingEnabled:              true,   // Enable tracing by default if endpoint is configured
                MetricsEnabled:              true,   // Enable metrics by default if endpoint is configured
                SamplingRate:                "0.05", // 5% sampling by default
                Headers:                     make(map[string]string),
                Insecure:                    false,
                EnablePrometheusMetricsPath: false,      // No metrics endpoint by default
                EnvironmentVariables:        []string{}, // No environment variables by default
                UseLegacyAttributes:         true,       // Dual emission for backward compat
        }
}

// MaybeMakeConfig creates a new telemetry configuration from the given values.
// It may return nil if no telemetry is configured.
func MaybeMakeConfig(
        otelEndpoint string,
        otelEnablePrometheusMetricsPath bool,
        otelTracingEnabled bool,
        otelMetricsEnabled bool,
        otelServiceName string,
        otelSamplingRate float64,
        otelHeaders []string,
        otelInsecure bool,
        otelEnvironmentVariables []string,
        otelUseLegacyAttributes bool,
) *Config {
        if otelEndpoint == "" && !otelEnablePrometheusMetricsPath {
                return nil
        }
        // Parse headers from key=value format
        headers := make(map[string]string)
        for _, header := range otelHeaders {
                parts := strings.SplitN(header, "=", 2)
                if len(parts) == 2 {
                        headers[parts[0]] = parts[1]
                }
        }

        // Process environment variables - split comma-separated values
        var processedEnvVars []string
        for _, envVarEntry := range otelEnvironmentVariables {
                // Split by comma and trim whitespace
                envVars := strings.Split(envVarEntry, ",")
                for _, envVar := range envVars {
                        trimmed := strings.TrimSpace(envVar)
                        if trimmed != "" {
                                processedEnvVars = append(processedEnvVars, trimmed)
                        }
                }
        }
        return &Config{
                Endpoint:                    otelEndpoint,
                ServiceName:                 otelServiceName,
                ServiceVersion:              "", // resolved at runtime in NewProvider()
                TracingEnabled:              otelTracingEnabled,
                MetricsEnabled:              otelMetricsEnabled,
                SamplingRate:                strconv.FormatFloat(otelSamplingRate, 'f', -1, 64),
                Headers:                     headers,
                Insecure:                    otelInsecure,
                EnablePrometheusMetricsPath: otelEnablePrometheusMetricsPath,
                EnvironmentVariables:        processedEnvVars,
                UseLegacyAttributes:         otelUseLegacyAttributes,
        }
}

// ResolveServiceName sets the telemetry service name on the config if it has
// not been explicitly provided. When empty, it derives the name from the
// workload/server name with the "thv-" prefix (e.g. "thv-fetch").
func ResolveServiceName(config *Config, serverName string) {
        if config == nil || config.ServiceName != "" {
                return
        }
        if serverName != "" {
                config.ServiceName = DefaultServiceNamePrefix + serverName
        }
}

// Provider encapsulates OpenTelemetry providers and configuration.
type Provider struct {
        config            Config
        tracerProvider    trace.TracerProvider
        meterProvider     metric.MeterProvider
        prometheusHandler http.Handler
        shutdown          func(context.Context) error
}

// NewProvider creates a new OpenTelemetry provider with the given configuration.
// Optional extra span processors (e.g. a Sentry bridge) can be registered via extraProcessors.
func NewProvider(ctx context.Context, config Config, extraProcessors ...sdktrace.SpanProcessor) (*Provider, error) {
        // Validate configuration
        if err := validateOtelConfig(config); err != nil {
                return nil, err
        }

        // Always use the current binary version so that restarts and exports
        // report the version actually running, not the version that originally
        // created the config. See https://github.com/stacklok/toolhive/issues/2296
        serviceVersion := config.ServiceVersion
        if serviceVersion == "" {
                serviceVersion = versions.GetVersionInfo().Version
        }

        telemetryOptions := []providers.ProviderOption{
                providers.WithServiceName(config.ServiceName),
                providers.WithServiceVersion(serviceVersion),
                providers.WithOTLPEndpoint(config.Endpoint),
                providers.WithHeaders(config.Headers),
                providers.WithInsecure(config.Insecure),
                providers.WithCACertPath(config.CACertPath),
                providers.WithTracingEnabled(config.TracingEnabled),
                providers.WithMetricsEnabled(config.MetricsEnabled),
                providers.WithSamplingRate(config.GetSamplingRateFloat()),
                providers.WithEnablePrometheusMetricsPath(config.EnablePrometheusMetricsPath),
                providers.WithCustomAttributes(mergeRegisteredResourceAttributes(config.CustomAttributes)),
        }

        // Merge globally registered processors (self-registered by integrations such
        // as a Sentry bridge) with any explicitly passed ones.
        allProcessors := append(registeredSpanProcessors(), extraProcessors...)
        if len(allProcessors) > 0 {
                telemetryOptions = append(telemetryOptions, providers.WithExtraSpanProcessors(allProcessors...))
        }

        telemetryProviders, err := providers.NewCompositeProvider(ctx, telemetryOptions...)
        if err != nil {
                return nil, fmt.Errorf("failed to build telemetry providers: %w", err)
        }

        return setGlobalProvidersAndReturn(telemetryProviders, config)
}

// mergeRegisteredResourceAttributes overlays the caller's custom attributes on
// top of any registered via RegisterResourceAttributes, so an explicitly
// configured attribute always wins over one a self-registered integration
// supplied as a default. The caller's map is never mutated.
func mergeRegisteredResourceAttributes(configured map[string]string) map[string]string {
        registered := RegisteredResourceAttributes()
        if len(registered) == 0 {
                return configured
        }
        merged := make(map[string]string, len(registered)+len(configured))
        maps.Copy(merged, registered)
        maps.Copy(merged, configured)
        return merged
}

// setGlobalProvidersAndReturn sets the global providers for OTEL and returns the providers
func setGlobalProvidersAndReturn(telemetryProviders *providers.CompositeProvider, config Config) (*Provider, error) {
        tracingProvider := telemetryProviders.TracerProvider()
        meterProvider := telemetryProviders.MeterProvider()

        // set the global providers for OTEL
        otel.SetTracerProvider(tracingProvider)
        otel.SetMeterProvider(meterProvider)
        otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
                propagation.TraceContext{},
                propagation.Baggage{},
        ))

        return &Provider{
                config:            config,
                tracerProvider:    tracingProvider,
                meterProvider:     meterProvider,
                prometheusHandler: telemetryProviders.PrometheusHandler(),
                shutdown:          telemetryProviders.Shutdown,
        }, nil
}

// Middleware returns an HTTP middleware that instruments requests with OpenTelemetry.
// serverName is the name of the MCP server (e.g., "github", "fetch")
// transport is the backend transport type ("stdio", "sse", or "streamable-http").
func (p *Provider) Middleware(serverName, transport string) types.MiddlewareFunction {
        return NewHTTPMiddleware(p.config, p.tracerProvider, p.meterProvider, serverName, transport)
}

// Shutdown gracefully shuts down the telemetry provider.
func (p *Provider) Shutdown(ctx context.Context) error {
        if p.shutdown != nil {
                return p.shutdown(ctx)
        }
        return nil
}

// TracerProvider returns the configured tracer provider.
func (p *Provider) TracerProvider() trace.TracerProvider {
        return p.tracerProvider
}

// MeterProvider returns the configured meter provider.
func (p *Provider) MeterProvider() metric.MeterProvider {
        return p.meterProvider
}

// PrometheusHandler returns the Prometheus metrics handler if configured.
// Returns nil if no metrics port is configured.
func (p *Provider) PrometheusHandler() http.Handler {
        return p.prometheusHandler
}

// PrometheusPort returns the configured diagnostics port for the Prometheus
// /metrics endpoint. Zero means unset; see diagnostics.ResolvePort.
func (p *Provider) PrometheusPort() int {
        return p.config.PrometheusPort
}

// ServeMetricsOnTransportPort reports whether this provider's configuration also
// wants /metrics on the main transport port. Resolves the unset case against the
// current default; see Config.ServeMetricsOnTransportPort.
func (p *Provider) ServeMetricsOnTransportPort() bool {
        return p.config.ServeMetricsOnTransportPort()
}

// validateOtelConfig validates the otel configuration
func validateOtelConfig(config Config) error {
        // If OTLP endpoint is configured but both tracing and metrics are disabled, that's an error
        if config.Endpoint != "" && !config.TracingEnabled && !config.MetricsEnabled {
                return fmt.Errorf("OTLP endpoint is configured but both tracing and metrics are disabled; " +
                        "either enable tracing or metrics, or remove the endpoint")
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package telemetry

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "net"
        "net/http"
        "os"
        "strconv"
        "strings"
        "time"
        "unicode/utf8"

        "go.opentelemetry.io/otel"
        "go.opentelemetry.io/otel/attribute"
        "go.opentelemetry.io/otel/codes"
        "go.opentelemetry.io/otel/metric"
        "go.opentelemetry.io/otel/propagation"
        "go.opentelemetry.io/otel/trace"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive-core/telemetry/metrics"
        mcpparser "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // instrumentationName is the name of this instrumentation package
        instrumentationName = "github.com/stacklok/toolhive/pkg/telemetry"
        // methodPromptsGet is the MCP method name for prompts/get
        methodPromptsGet = "prompts/get"
        // networkTransportTCP is the OTEL value for TCP transport
        networkTransportTCP = "tcp"
        // networkProtocolHTTP is the OTEL value for HTTP protocol
        networkProtocolHTTP = "http"
        // maxLabelValueBytes bounds the length of client-controlled metric label
        // values. Cumulative metric readers (Prometheus, OTLP PeriodicReader) keep
        // every distinct attribute set resident for the process lifetime, so an
        // unbounded label value (e.g. an 8 MB MCP method name) stays in memory
        // permanently and can exhaust the pod. This caps the byte length; the OTEL
        // SDK's 2000-series limit caps the count. Spans are left untruncated on
        // purpose — they are sampled and ephemeral, and the OTEL MCP semconv wants
        // the real value.
        maxLabelValueBytes = 128
        // labelTruncationMarker is appended to truncated label values so consumers
        // can tell a value was clamped.
        labelTruncationMarker = "..."
)

// MCPHistogramBuckets are the bucket boundaries defined by the MCP OTEL semantic conventions
// for MCP server histograms (e.g. mcp.server.operation.duration).
// See https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/mcp.md
var MCPHistogramBuckets = metrics.BucketsMCPSemconv()

// HTTPMiddleware provides OpenTelemetry instrumentation for HTTP requests.
type HTTPMiddleware struct {
        config         Config
        tracerProvider trace.TracerProvider
        tracer         trace.Tracer
        meterProvider  metric.MeterProvider
        meter          metric.Meter
        serverName     string
        transport      string

        // Metrics
        requestCounter    metric.Int64Counter
        requestDuration   metric.Float64Histogram
        operationDuration metric.Float64Histogram
        activeConnections metric.Int64UpDownCounter
        toolCallCounter   metric.Int64Counter
}

// NewHTTPMiddleware creates a new HTTP middleware for OpenTelemetry instrumentation.
// serverName is the name of the MCP server (e.g., "github", "fetch")
// transport is the backend transport type ("stdio", "sse", or "streamable-http").
func NewHTTPMiddleware(
        config Config,
        tracerProvider trace.TracerProvider,
        meterProvider metric.MeterProvider,
        serverName, transport string,
) types.MiddlewareFunction {
        meter := meterProvider.Meter(instrumentationName)

        // Initialize metrics
        requestCounter, err := meter.Int64Counter(
                "toolhive_mcp_requests", // The exporter adds the _total suffix automatically
                metric.WithDescription("Total number of MCP requests"),
        )
        if err != nil {
                slog.Debug("failed to create request counter metric", "error", err)
        }

        requestDuration, err := meter.Float64Histogram(
                "toolhive_mcp_request_duration", // The exporter adds the _seconds suffix automatically
                metric.WithDescription("Duration of MCP requests in seconds"),
                metric.WithUnit("s"),
                metric.WithExplicitBucketBoundaries(MCPHistogramBuckets...),
        )
        if err != nil {
                slog.Debug("failed to create request duration metric", "error", err)
        }

        activeConnections, err := meter.Int64UpDownCounter(
                "toolhive_mcp_active_connections",
                metric.WithDescription("Number of active MCP connections"),
        )
        if err != nil {
                slog.Debug("failed to create active connections metric", "error", err)
        }

        operationDuration, err := meter.Float64Histogram(
                "mcp.server.operation.duration",
                metric.WithDescription("Duration of MCP server operations"),
                metric.WithUnit("s"),
                metric.WithExplicitBucketBoundaries(MCPHistogramBuckets...),
        )
        if err != nil {
                slog.Debug("failed to create operation duration metric", "error", err)
        }

        toolCallCounter, err := meter.Int64Counter(
                "toolhive_mcp_tool_calls",
                metric.WithDescription("Total number of MCP tool calls"),
        )
        if err != nil {
                slog.Debug("failed to create tool call counter metric", "error", err)
        }

        middleware := &HTTPMiddleware{
                config:            config,
                tracerProvider:    tracerProvider,
                tracer:            tracerProvider.Tracer(instrumentationName),
                meterProvider:     meterProvider,
                meter:             meter,
                serverName:        serverName,
                transport:         transport,
                requestCounter:    requestCounter,
                requestDuration:   requestDuration,
                operationDuration: operationDuration,
                activeConnections: activeConnections,
                toolCallCounter:   toolCallCounter,
        }

        return middleware.Handler
}

// Handler implements the middleware function that wraps HTTP handlers.
// This middleware should be placed after the MCP parsing middleware in the chain
// to leverage the parsed MCP data.
// Note: Panic recovery is handled by the dedicated recovery middleware.
func (m *HTTPMiddleware) Handler(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                ctx := r.Context()

                // Handle SSE endpoints specially - they are long-lived connections
                // that don't follow the normal request/response pattern
                if strings.HasSuffix(r.URL.Path, "/sse") {
                        // Record SSE connection establishment immediately
                        m.recordSSEConnection(ctx, r)

                        // Track active SSE connections with defer to ensure decrement on close
                        sseAttrs := metric.WithAttributes(
                                attribute.String("server", m.serverName),
                                attribute.String("transport", m.transport),
                                attribute.String("connection_type", "sse"),
                        )
                        m.activeConnections.Add(ctx, 1, sseAttrs)
                        defer m.activeConnections.Add(ctx, -1, sseAttrs)

                        // Pass through to SSE handler - blocks for the lifetime of the SSE connection
                        next.ServeHTTP(w, r)
                        return
                }

                // Normal HTTP request handling
                // Reuse an outer holder (for example, audit's) so webhook mutations are
                // visible to every wrapper. When there is no outer holder, install one for
                // telemetry itself and seed it with the parse published by the parser.
                parsedHolder, ok := mcpparser.ParsedRequestHolderFromContext(ctx)
                if !ok {
                        parsedHolder = &mcpparser.ParsedRequestHolder{
                                Parsed: mcpparser.GetParsedMCPRequest(ctx),
                        }
                        ctx = mcpparser.WithParsedRequestHolder(ctx, parsedHolder)
                } else if parsedHolder.Parsed == nil {
                        parsedHolder.Parsed = mcpparser.GetParsedMCPRequest(ctx)
                }

                // Extract trace context from incoming request headers
                ctx = otel.GetTextMapPropagator().Extract(ctx, propagation.HeaderCarrier(r.Header))

                // Extract trace context from MCP _meta field if present.
                // Per the MCP OTEL spec, servers should use traceparent/tracestate from
                // params._meta as the parent span context. This takes priority over HTTP
                // headers since _meta is the MCP-specified propagation mechanism.
                if parsedMCP := mcpparser.GetParsedMCPRequest(ctx); parsedMCP != nil && parsedMCP.Meta != nil {
                        carrier := NewMetaCarrier(parsedMCP.Meta)
                        ctx = otel.GetTextMapPropagator().Extract(ctx, carrier)
                }

                // Increment active connections
                m.activeConnections.Add(ctx, 1, metric.WithAttributes(
                        attribute.String("server", m.serverName),
                        attribute.String("transport", m.transport),
                ))
                defer m.activeConnections.Add(ctx, -1, metric.WithAttributes(
                        attribute.String("server", m.serverName),
                        attribute.String("transport", m.transport),
                ))

                // Start the request span before rate limiting. MCP identity is populated
                // after the inner chain returns because a mutating webhook may republish a
                // different method, resource, or argument set while the request is in flight.
                ctx, span := m.tracer.Start(
                        ctx,
                        fmt.Sprintf("%s %s", r.Method, r.URL.Path),
                        trace.WithSpanKind(trace.SpanKindServer),
                )
                defer span.End()

                // Create a response writer wrapper to capture response details
                rw := &responseWriter{
                        ResponseWriter: w,
                        statusCode:     http.StatusOK,
                        bytesWritten:   0,
                }

                // Add HTTP attributes
                m.addHTTPAttributes(span, r)

                // Add environment variables as attributes
                m.addEnvironmentAttributes(span)

                // Record request start time
                startTime := time.Now()

                // Call the next handler with the instrumented context
                next.ServeHTTP(rw, r.WithContext(ctx))

                // A mutating webhook publishes its replacement parse through the shared
                // holder. Use that final parse for all MCP identity exported by telemetry.
                finalCtx := contextWithFinalParsedMCPRequest(ctx, parsedHolder)
                if spanName := m.createSpanName(finalCtx); spanName != "" {
                        span.SetName(spanName)
                }
                m.addMCPAttributes(finalCtx, span, r)

                // Record completion metrics and finalize span.
                duration := time.Since(startTime)
                m.finalizeSpan(span, rw, duration)
                m.recordMetrics(finalCtx, r, rw, duration)
        })
}

func contextWithFinalParsedMCPRequest(
        ctx context.Context,
        holder *mcpparser.ParsedRequestHolder,
) context.Context {
        if holder == nil || holder.Parsed == nil || holder.Parsed == mcpparser.GetParsedMCPRequest(ctx) {
                return ctx
        }
        return context.WithValue(ctx, mcpparser.MCPRequestContextKey, holder.Parsed)
}

func (*HTTPMiddleware) createSpanName(ctx context.Context) string {
        parsedMCP := mcpparser.GetParsedMCPRequest(ctx)
        if parsedMCP == nil || parsedMCP.Method == "" {
                return ""
        }
        // OTEL MCP semconv: span name should be "{mcp.method.name} {target}"
        // where target is the tool/prompt/resource name when available.
        if parsedMCP.ResourceID != "" {
                return parsedMCP.Method + " " + parsedMCP.ResourceID
        }
        return parsedMCP.Method
}

// addHTTPAttributes adds standard HTTP attributes to the span.
func (m *HTTPMiddleware) addHTTPAttributes(span trace.Span, r *http.Request) {
        // New OTEL HTTP semantic convention attributes (always emitted)
        span.SetAttributes(
                attribute.String("http.request.method", r.Method),
                attribute.String("url.full", r.URL.String()),
                attribute.String("url.scheme", r.URL.Scheme),
                attribute.String("server.address", r.Host),
                attribute.String("url.path", r.URL.Path),
                attribute.String("user_agent.original", r.UserAgent()),
        )

        if r.ContentLength > 0 {
                span.SetAttributes(attribute.Int64("http.request.body.size", r.ContentLength))
        }

        if r.URL.RawQuery != "" {
                span.SetAttributes(attribute.String("url.query", r.URL.RawQuery))
        }

        // Legacy attribute names (emitted only when UseLegacyAttributes is true)
        if m.config.UseLegacyAttributes {
                span.SetAttributes(
                        attribute.String("http.method", r.Method),
                        attribute.String("http.url", r.URL.String()),
                        attribute.String("http.scheme", r.URL.Scheme),
                        attribute.String("http.host", r.Host),
                        attribute.String("http.target", r.URL.Path),
                        attribute.String("http.user_agent", r.UserAgent()),
                )

                if contentLength := r.Header.Get("Content-Length"); contentLength != "" {
                        span.SetAttributes(attribute.String("http.request_content_length", contentLength))
                }

                if r.URL.RawQuery != "" {
                        span.SetAttributes(attribute.String("http.query", r.URL.RawQuery))
                }
        }
}

func (m *HTTPMiddleware) addEnvironmentAttributes(span trace.Span) {
        // Include environment variables from host machine as configured
        // Only environment variables specified in the config will be read and included
        for _, envVar := range m.config.EnvironmentVariables {
                // A "NAME=value" assignment must not become part of the span attribute key.
                if envVar == "" || strings.Contains(envVar, "=") {
                        continue // Skip empty or assignment-style environment variable names
                }

                value := os.Getenv(envVar)
                // Always set the attribute, even if the environment variable is empty
                // This helps distinguish between unset variables and empty string values
                span.SetAttributes(
                        attribute.String(fmt.Sprintf("environment.%s", envVar), value),
                )
        }
}

// addMCPAttributes adds MCP-specific attributes using the parsed MCP data from context.
func (m *HTTPMiddleware) addMCPAttributes(ctx context.Context, span trace.Span, r *http.Request) {
        // Get parsed MCP request from context (set by MCP parsing middleware)
        parsedMCP := mcpparser.GetParsedMCPRequest(ctx)
        if parsedMCP == nil {
                // No MCP data available, this might be a non-MCP request (e.g., health check)
                return
        }

        // New OTEL MCP semantic convention attributes (always emitted)
        span.SetAttributes(
                attribute.String("mcp.method.name", parsedMCP.Method),
                attribute.String("rpc.system.name", "jsonrpc"),
                attribute.String("jsonrpc.protocol.version", "2.0"),
        )

        if parsedMCP.ID != nil {
                span.SetAttributes(attribute.String("jsonrpc.request.id", formatRequestID(parsedMCP.ID)))
        }

        // Resource URI: only set for resource-related methods
        if parsedMCP.ResourceID != "" {
                switch parsedMCP.Method {
                case "resources/read", "resources/subscribe", "resources/unsubscribe", "notifications/resources/updated":
                        span.SetAttributes(attribute.String("mcp.resource.uri", parsedMCP.ResourceID))
                }
        }

        // Legacy attribute names (emitted only when UseLegacyAttributes is true)
        if m.config.UseLegacyAttributes {
                span.SetAttributes(
                        attribute.String("mcp.method", parsedMCP.Method),
                        attribute.String("rpc.system", "jsonrpc"),
                        attribute.String("rpc.service", "mcp"),
                )

                if parsedMCP.ID != nil {
                        span.SetAttributes(attribute.String("mcp.request.id", formatRequestID(parsedMCP.ID)))
                }

                if parsedMCP.ResourceID != "" {
                        span.SetAttributes(attribute.String("mcp.resource.id", parsedMCP.ResourceID))
                }
        }

        // Add method-specific attributes
        m.addMethodSpecificAttributes(span, parsedMCP)

        // Extract server name from the request
        serverName := m.extractServerName(r)
        span.SetAttributes(attribute.String("mcp.server.name", serverName))

        // Add network, client, and session attributes
        backendTransport := m.extractBackendTransport(r)
        m.addNetworkAttributes(span, r, backendTransport)

        if m.config.UseLegacyAttributes {
                span.SetAttributes(attribute.String("mcp.transport", backendTransport))
        }

        // Add batch indicator
        if parsedMCP.IsBatch {
                span.SetAttributes(attribute.Bool("mcp.is_batch", true))
        }
}

// addNetworkAttributes adds network, client, and session attributes to the span.
func (*HTTPMiddleware) addNetworkAttributes(span trace.Span, r *http.Request, backendTransport string) {
        networkTransport, protocolName, backendProtocolVersion := mapTransport(backendTransport)
        span.SetAttributes(attribute.String("network.transport", networkTransport))
        if protocolName != "" {
                span.SetAttributes(attribute.String("network.protocol.name", protocolName))
        }
        if backendProtocolVersion != "" {
                span.SetAttributes(attribute.String("mcp.backend.protocol.version", backendProtocolVersion))
        }

        // HTTP protocol version from the incoming request
        if protocolVer := httpProtocolVersion(r); protocolVer != "" {
                span.SetAttributes(attribute.String("network.protocol.version", protocolVer))
        }

        // Client address and port
        if clientAddr, clientPort := parseRemoteAddr(r.RemoteAddr); clientAddr != "" {
                span.SetAttributes(attribute.String("client.address", clientAddr))
                if clientPort > 0 {
                        span.SetAttributes(attribute.Int("client.port", clientPort))
                }
        }

        // Session ID if available
        if sessionID := r.Header.Get("Mcp-Session-Id"); sessionID != "" {
                span.SetAttributes(attribute.String("mcp.session.id", sessionID))
        }

        // MCP protocol version from the streamable HTTP transport header
        if mcpVersion := r.Header.Get("MCP-Protocol-Version"); mcpVersion != "" {
                span.SetAttributes(attribute.String("mcp.protocol.version", mcpVersion))
        }
}

// addMethodSpecificAttributes adds attributes specific to certain MCP methods.
// Despite the name, the mcp.client.name block below runs for every method: under
// the Modern (stateless) MCP revision there is no initialize request, so per-request
// _meta.clientInfo is the only source of client attribution. It is a no-op for
// Legacy requests, which carry no _meta.clientInfo and still get mcp.client.name
// from the "initialize" case below.
func (m *HTTPMiddleware) addMethodSpecificAttributes(span trace.Span, parsedMCP *mcpparser.ParsedMCPRequest) {
        if name, ok := parsedMCP.ClientInfo["name"].(string); ok && name != "" {
                span.SetAttributes(attribute.String("mcp.client.name", name))
        }

        switch parsedMCP.Method {
        case string(mcp.MethodToolsCall):
                // New gen_ai namespace attributes (always emitted)
                if parsedMCP.ResourceID != "" {
                        span.SetAttributes(attribute.String("gen_ai.tool.name", parsedMCP.ResourceID))
                }
                span.SetAttributes(attribute.String("gen_ai.operation.name", "execute_tool"))

                sanitizedArgs := m.sanitizeArguments(parsedMCP.Arguments)
                if sanitizedArgs != "" {
                        span.SetAttributes(attribute.String("gen_ai.tool.call.arguments", sanitizedArgs))
                }

                // Legacy names
                if m.config.UseLegacyAttributes {
                        if parsedMCP.ResourceID != "" {
                                span.SetAttributes(attribute.String("mcp.tool.name", parsedMCP.ResourceID))
                        }
                        if sanitizedArgs != "" {
                                span.SetAttributes(attribute.String("mcp.tool.arguments", sanitizedArgs))
                        }
                }

        case methodPromptsGet:
                // New gen_ai namespace attribute (always emitted)
                if parsedMCP.ResourceID != "" {
                        span.SetAttributes(attribute.String("gen_ai.prompt.name", parsedMCP.ResourceID))
                }

                // Legacy name
                if m.config.UseLegacyAttributes {
                        if parsedMCP.ResourceID != "" {
                                span.SetAttributes(attribute.String("mcp.prompt.name", parsedMCP.ResourceID))
                        }
                }

        case "initialize":
                if parsedMCP.ResourceID != "" {
                        span.SetAttributes(attribute.String("mcp.client.name", parsedMCP.ResourceID))
                }
        }
}

// extractServerName extracts the MCP server name from the HTTP request.
// It checks for an explicit X-MCP-Server-Name header first, then falls back to the
// configured server name. This approach is more reliable than parsing URL paths since
// the server name is already known during middleware construction.
func (m *HTTPMiddleware) extractServerName(r *http.Request) string {
        // Check for explicit server name header (for advanced routing scenarios)
        if serverName := r.Header.Get("X-MCP-Server-Name"); serverName != "" {
                return serverName
        }

        // Always use the configured server name - this is the correct server name
        // that was passed during middleware construction and doesn't depend on URL structure
        //
        // NOTE: Previously this function attempted to parse server names from URL paths by
        // splitting r.URL.Path and filtering out known endpoint segments like "sse", "messages",
        // "api", "v1", etc. This approach was fundamentally flawed because:
        // 1. It incorrectly treated endpoint names like "message" as server names
        // 2. It made assumptions about URL structure that don't always hold
        // 3. The actual server name is already available via m.serverName
        // Adding more exclusions (like "message") would just be treating symptoms, not the root cause.
        return m.serverName
}

// extractBackendTransport determines the backend transport type.
// ToolHive supports multiple transport types: stdio, sse, streamable-http.
func (m *HTTPMiddleware) extractBackendTransport(r *http.Request) string {
        // Try to get transport info from custom headers (if set by proxy)
        if transport := r.Header.Get("X-MCP-Transport"); transport != "" {
                return transport
        }

        return m.transport
}

func mapTransport(mcpTransport string) (networkTransport, protocolName, protocolVersion string) {
        switch mcpTransport {
        case "stdio":
                return "pipe", "", ""
        case "sse":
                return networkTransportTCP, networkProtocolHTTP, "1.1"
        case "streamable-http":
                return networkTransportTCP, networkProtocolHTTP, ""
        default:
                return networkTransportTCP, networkProtocolHTTP, ""
        }
}

// httpProtocolVersion extracts the HTTP protocol version from the request.
func httpProtocolVersion(r *http.Request) string {
        if r.ProtoMajor == 0 {
                return ""
        }
        if r.ProtoMajor >= 2 && r.ProtoMinor == 0 {
                return strconv.Itoa(r.ProtoMajor)
        }
        return fmt.Sprintf("%d.%d", r.ProtoMajor, r.ProtoMinor)
}

// parseRemoteAddr parses the remote address into host and port.
func parseRemoteAddr(remoteAddr string) (string, int) {
        if remoteAddr == "" {
                return "", 0
        }
        host, portStr, err := net.SplitHostPort(remoteAddr)
        if err != nil {
                return remoteAddr, 0
        }
        port, err := strconv.Atoi(portStr)
        if err != nil {
                return host, 0
        }
        return host, port
}

// sanitizeArguments converts arguments to a safe string representation.
func (m *HTTPMiddleware) sanitizeArguments(arguments map[string]interface{}) string {
        if len(arguments) == 0 {
                return ""
        }

        // Create a sanitized representation
        var parts []string
        for key, value := range arguments {
                // Check for sensitive keys
                if m.isSensitiveKey(key) {
                        parts = append(parts, fmt.Sprintf("%s=[REDACTED]", key))
                        continue
                }

                // Limit value length and convert to string
                valueStr := fmt.Sprintf("%v", value)
                if len(valueStr) > 100 {
                        valueStr = valueStr[:100] + "..."
                }

                parts = append(parts, fmt.Sprintf("%s=%s", key, valueStr))
        }

        result := strings.Join(parts, ", ")
        if len(result) > 200 {
                result = result[:200] + "..."
        }

        return result
}

// isSensitiveKey checks if a key might contain sensitive information.
func (*HTTPMiddleware) isSensitiveKey(key string) bool {
        sensitivePatterns := []string{
                "password", "token", "secret", "key", "auth", "credential",
                "api_key", "access_token", "refresh_token", "private",
        }

        keyLower := strings.ToLower(key)
        for _, pattern := range sensitivePatterns {
                if strings.Contains(keyLower, pattern) {
                        return true
                }
        }
        return false
}

// formatRequestID converts the request ID to a string representation.
func formatRequestID(id interface{}) string {
        switch v := id.(type) {
        case string:
                return v
        case float64:
                return strconv.FormatFloat(v, 'f', -1, 64)
        case int:
                return strconv.Itoa(v)
        case int64:
                return strconv.FormatInt(v, 10)
        default:
                return fmt.Sprintf("%v", v)
        }
}

// finalizeSpan adds response attributes and sets the span status.
func (m *HTTPMiddleware) finalizeSpan(span trace.Span, rw *responseWriter, duration time.Duration) {
        // New OTEL HTTP semantic convention response attributes (always emitted)
        span.SetAttributes(
                attribute.Int("http.response.status_code", rw.statusCode),
                attribute.Int64("http.response.body.size", rw.bytesWritten),
        )

        // Legacy response attributes
        if m.config.UseLegacyAttributes {
                span.SetAttributes(
                        attribute.Int("http.status_code", rw.statusCode),
                        attribute.Int64("http.response_content_length", rw.bytesWritten),
                        attribute.Float64("http.duration_ms", float64(duration.Nanoseconds())/1e6),
                )
        }

        // Set span status based on HTTP status code per OTEL semconv
        if rw.statusCode >= 500 {
                // 5xx: Server errors set span status to Error with error.type
                span.SetStatus(codes.Error, fmt.Sprintf("HTTP %d", rw.statusCode))
                span.SetAttributes(attribute.String("error.type", strconv.Itoa(rw.statusCode)))
        } else if rw.statusCode >= 400 {
                // 4xx: Client errors leave span status Unset (not server errors per OTEL semconv)
        } else {
                // 2xx/3xx: Success
                span.SetStatus(codes.Ok, "")
        }
}

// responseWriter wraps http.ResponseWriter to capture response details.
type responseWriter struct {
        http.ResponseWriter
        statusCode    int
        bytesWritten  int64
        headerWritten bool // Guard against double WriteHeader calls
}

// WriteHeader captures the status code. Guards against duplicate calls which
// can cause panics in Go's reverse proxy (http: superfluous response.WriteHeader call).
func (rw *responseWriter) WriteHeader(statusCode int) {
        if rw.headerWritten {
                return // Silently ignore duplicate WriteHeader calls
        }
        rw.headerWritten = true
        rw.statusCode = statusCode
        rw.ResponseWriter.WriteHeader(statusCode)
}

// Write captures the number of bytes written.
// Note: Write() implicitly calls WriteHeader(200) on the underlying ResponseWriter
// if headers haven't been written yet. This is standard HTTP behavior - once headers
// are written, the status code cannot be changed. We track this to accurately record
// what actually happened and to prevent subsequent WriteHeader() calls from panicking.
//
// Important: If a non-200 status code is needed, WriteHeader() MUST be called BEFORE Write().
// Once Write() is called first, the status code is fixed at 200 and cannot be changed.
func (rw *responseWriter) Write(data []byte) (int, error) {
        // If headers haven't been written yet, Write() will implicitly write them with status 200.
        // This is what the underlying ResponseWriter actually does - we're tracking what happened,
        // not forcing a status code. Mark headers as written to prevent subsequent WriteHeader()
        // calls from panicking.
        if !rw.headerWritten {
                rw.headerWritten = true
                rw.statusCode = http.StatusOK // Write() implicitly uses 200 - this is what actually happened
        }

        n, err := rw.ResponseWriter.Write(data)
        rw.bytesWritten += int64(n)
        return n, err
}

// Flush implements http.Flusher if the underlying ResponseWriter supports it.
func (rw *responseWriter) Flush() {
        if flusher, ok := rw.ResponseWriter.(http.Flusher); ok {
                flusher.Flush()
        }
}

// truncateLabelValue bounds a client-controlled metric label value to
// maxLabelValueBytes, clamping on a UTF-8 rune boundary so the result is never
// invalid UTF-8, and appending labelTruncationMarker to signal truncation.
// Values within the cap are returned unchanged. Use this only for metric
// labels; span attributes keep the full value.
func truncateLabelValue(s string) string {
        if len(s) <= maxLabelValueBytes {
                return s
        }
        // Reserve room for the marker so the returned value stays within the cap.
        limit := maxLabelValueBytes - len(labelTruncationMarker)
        // Back up to a rune boundary so we never split a multi-byte character.
        for limit > 0 && !utf8.RuneStart(s[limit]) {
                limit--
        }
        return s[:limit] + labelTruncationMarker
}

// recordMetrics records request metrics.
func (m *HTTPMiddleware) recordMetrics(ctx context.Context, r *http.Request, rw *responseWriter, duration time.Duration) {
        // Get MCP method from context if available
        mcpMethod := mcpparser.GetMCPMethod(ctx)
        if mcpMethod == "" {
                mcpMethod = "unknown"
        }

        // Determine status (success/error)
        status := "success"
        if rw.statusCode >= 400 {
                status = "error"
        }

        // Get the resource ID from the parsed MCP request if available.
        // For tools/call this is the tool name, for resources/read the URI,
        // and for prompts/get the prompt name.
        mcpResourceID := ""
        if parsedMCP := mcpparser.GetParsedMCPRequest(ctx); parsedMCP != nil {
                mcpResourceID = parsedMCP.ResourceID
        }

        // Common attributes for all metrics
        attrs := metric.WithAttributes(
                attribute.String("method", r.Method),
                attribute.String("status_code", strconv.Itoa(rw.statusCode)),
                attribute.String("status", status),
                attribute.String("mcp_method", truncateLabelValue(mcpMethod)),
                attribute.String("mcp_resource_id", truncateLabelValue(mcpResourceID)),
                attribute.String("server", m.serverName),
                attribute.String("transport", m.transport),
        )

        // Record request count
        m.requestCounter.Add(ctx, 1, attrs)

        // Record request duration
        m.requestDuration.Record(ctx, duration.Seconds(), attrs)

        // Record OTEL MCP spec mcp.server.operation.duration for actual MCP requests.
        // Only POST requests carry a JSON-RPC body; GET (SSE stream open) and DELETE
        // (session termination) are valid Streamable HTTP lifecycle requests with no
        // MCP method to record. An unknown method on a POST indicates either a
        // misconfigured middleware chain (see #3687) or a parse failure.
        if mcpMethod != "unknown" {
                m.recordOperationDuration(ctx, r, mcpMethod, mcpResourceID, rw.statusCode, duration)
        } else if r.Method == http.MethodPost {
                //nolint:gosec // G706: HTTP method and URL path from request
                slog.Warn("mcp method could not be determined, middleware may be misconfigured",
                        "http_method", r.Method, "path", r.URL.Path)
        }

        // For tools/call, record tool-specific metrics
        if mcpMethod == string(mcp.MethodToolsCall) {
                if parsedMCP := mcpparser.GetParsedMCPRequest(ctx); parsedMCP != nil && parsedMCP.ResourceID != "" {
                        toolAttrs := metric.WithAttributes(
                                attribute.String("server", m.serverName),
                                attribute.String("tool", truncateLabelValue(parsedMCP.ResourceID)),
                                attribute.String("status", status),
                        )
                        m.toolCallCounter.Add(ctx, 1, toolAttrs)
                }
        }
}

// recordOperationDuration records the mcp.server.operation.duration metric
// per the OTEL MCP semantic conventions.
func (m *HTTPMiddleware) recordOperationDuration(
        ctx context.Context, r *http.Request, mcpMethod, resourceID string, statusCode int, duration time.Duration,
) {
        networkTransport, protocolName, _ := mapTransport(m.transport)

        specAttrs := []attribute.KeyValue{
                attribute.String("mcp.method.name", truncateLabelValue(mcpMethod)),
                attribute.String("jsonrpc.protocol.version", "2.0"),
                attribute.String("network.transport", networkTransport),
        }
        if protocolName != "" {
                specAttrs = append(specAttrs, attribute.String("network.protocol.name", protocolName))
        }
        if pv := httpProtocolVersion(r); pv != "" {
                specAttrs = append(specAttrs, attribute.String("network.protocol.version", pv))
        }

        // error.type: Conditionally required on error.
        // NOTE: This only captures HTTP-level errors (5xx). JSON-RPC errors returned
        // with HTTP 200 are not yet captured here — that requires response body parsing
        // which is tracked as future work (rpc.response.status_code, error.type for
        // JSON-RPC error codes like -32601).
        if statusCode >= 500 {
                specAttrs = append(specAttrs, attribute.String("error.type", strconv.Itoa(statusCode)))
        }

        // Method-specific attributes
        switch mcpMethod {
        case string(mcp.MethodToolsCall):
                specAttrs = append(specAttrs, attribute.String("gen_ai.operation.name", "execute_tool"))
                if resourceID != "" {
                        specAttrs = append(specAttrs, attribute.String("gen_ai.tool.name", truncateLabelValue(resourceID)))
                }
        case methodPromptsGet:
                if resourceID != "" {
                        specAttrs = append(specAttrs, attribute.String("gen_ai.prompt.name", truncateLabelValue(resourceID)))
                }
        }

        m.operationDuration.Record(ctx, duration.Seconds(), metric.WithAttributes(specAttrs...))
}

// recordSSEConnection records telemetry for SSE connection establishment.
// SSE connections are long-lived and don't follow the normal request/response pattern,
// so we record the connection establishment event immediately.
func (m *HTTPMiddleware) recordSSEConnection(ctx context.Context, r *http.Request) {
        // Create a short-lived span for SSE connection establishment
        spanName := "sse.connection_established"
        _, span := m.tracer.Start(ctx, spanName, trace.WithSpanKind(trace.SpanKindServer))

        // Add HTTP attributes for the connection
        m.addHTTPAttributes(span, r)

        // Add SSE-specific attributes
        networkTransport, protocolName, backendProtocolVersion := mapTransport(m.transport)
        span.SetAttributes(
                attribute.String("sse.event_type", "connection_established"),
                attribute.String("mcp.server.name", m.serverName),
                attribute.String("network.transport", networkTransport),
        )
        if protocolName != "" {
                span.SetAttributes(attribute.String("network.protocol.name", protocolName))
        }
        if backendProtocolVersion != "" {
                span.SetAttributes(attribute.String("mcp.backend.protocol.version", backendProtocolVersion))
        }
        if protocolVer := httpProtocolVersion(r); protocolVer != "" {
                span.SetAttributes(attribute.String("network.protocol.version", protocolVer))
        }
        if m.config.UseLegacyAttributes {
                span.SetAttributes(attribute.String("mcp.transport", m.transport))
        }

        // End the span immediately since this is just the connection establishment
        span.SetStatus(codes.Ok, "SSE connection established")
        span.End()

        // Record SSE connection metrics
        attrs := metric.WithAttributes(
                attribute.String("method", r.Method),
                attribute.String("status_code", "200"), // SSE connections start with 200
                attribute.String("status", "success"),
                attribute.String("mcp_method", "sse_connection"),
                attribute.String("server", m.serverName),
                attribute.String("transport", m.transport),
        )

        // Record the connection establishment
        m.requestCounter.Add(ctx, 1, attrs)
}

// Factory middleware type constant
const (
        MiddlewareType = "telemetry"
)

// FactoryMiddlewareParams represents the parameters for telemetry middleware
type FactoryMiddlewareParams struct {
        Config     *Config `json:"config"`
        ServerName string  `json:"server_name"`
        Transport  string  `json:"transport"`
}

// FactoryMiddleware wraps telemetry middleware functionality for factory pattern
type FactoryMiddleware struct {
        provider          *Provider
        middleware        types.MiddlewareFunction
        prometheusHandler http.Handler
}

// Handler returns the middleware function used by the proxy.
func (m *FactoryMiddleware) Handler() types.MiddlewareFunction {
        return m.middleware
}

// Close cleans up any resources used by the middleware.
func (m *FactoryMiddleware) Close() error {
        if m.provider != nil {
                return m.provider.Shutdown(context.Background())
        }
        return nil
}

// PrometheusHandler returns the Prometheus metrics handler.
func (m *FactoryMiddleware) PrometheusHandler() http.Handler {
        return m.prometheusHandler
}

// CreateMiddleware factory function for telemetry middleware
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params FactoryMiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal telemetry middleware parameters: %w", err)
        }

        if params.Config == nil {
                return fmt.Errorf("telemetry config is required")
        }

        provider, err := NewProvider(context.Background(), *params.Config)
        if err != nil {
                return fmt.Errorf("failed to create telemetry provider: %w", err)
        }

        middleware := provider.Middleware(params.ServerName, params.Transport)

        var prometheusHandler http.Handler
        if params.Config.EnablePrometheusMetricsPath {
                prometheusHandler = provider.PrometheusHandler()
        }

        telemetryMw := &FactoryMiddleware{
                provider:          provider,
                middleware:        middleware,
                prometheusHandler: prometheusHandler,
        }

        // Add middleware to runner
        runner.AddMiddleware(config.Type, telemetryMw)

        // Set Prometheus handler if enabled. The runner serves it on a dedicated
        // diagnostics listener and logs the resolved address there, so no port is
        // logged here — it is not the application port.
        if prometheusHandler != nil {
                runner.SetPrometheusHandler(prometheusHandler)
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package telemetry

import (
        "context"
        "maps"

        "go.opentelemetry.io/otel"
        "go.opentelemetry.io/otel/propagation"
)

// Compile-time assertion that MetaCarrier implements propagation.TextMapCarrier
var _ propagation.TextMapCarrier = (*MetaCarrier)(nil)

// MetaCarrier implements propagation.TextMapCarrier for MCP _meta fields.
// This enables W3C Trace Context propagation through MCP request params._meta,
// as recommended by the MCP OpenTelemetry specification.
//
// The carrier wraps a map[string]interface{} (the _meta field from MCP params)
// and allows the OpenTelemetry propagator to inject/extract traceparent and
// tracestate headers into/from the map.
type MetaCarrier struct {
        meta map[string]interface{}
}

// NewMetaCarrier creates a new MetaCarrier wrapping the given meta map.
// If meta is nil, a new empty map is created.
func NewMetaCarrier(meta map[string]interface{}) *MetaCarrier {
        if meta == nil {
                meta = make(map[string]interface{})
        }
        return &MetaCarrier{meta: meta}
}

// Get returns the value associated with the passed key.
func (c *MetaCarrier) Get(key string) string {
        if v, ok := c.meta[key]; ok {
                if s, ok := v.(string); ok {
                        return s
                }
        }
        return ""
}

// Set stores the key-value pair.
func (c *MetaCarrier) Set(key string, value string) {
        c.meta[key] = value
}

// Keys lists the keys stored in this carrier.
func (c *MetaCarrier) Keys() []string {
        keys := make([]string, 0, len(c.meta))
        for k := range c.meta {
                keys = append(keys, k)
        }
        return keys
}

// Meta returns the underlying meta map. Use this after injection to retrieve
// the enriched map containing trace context fields.
func (c *MetaCarrier) Meta() map[string]interface{} {
        return c.meta
}

// InjectMetaTraceContext injects the current trace context from ctx directly into
// the given meta map using W3C Trace Context format (traceparent, tracestate)
// plus W3C Baggage, subject to the globally configured OTEL propagator.
//
// This function operates directly on the meta map contents. Use this when you
// already have the _meta map and want to inject trace context fields into it.
func InjectMetaTraceContext(ctx context.Context, meta map[string]interface{}) {
        if meta == nil {
                return
        }
        carrier := NewMetaCarrier(meta)
        otel.GetTextMapPropagator().Inject(ctx, carrier)
}

// MetaWithTraceContext returns a COPY of meta enriched with the current W3C
// trace context (traceparent/tracestate/baggage) from ctx, per SEP-414.
// The caller's map is never mutated (copy-before-mutate). Returns nil when the
// result is empty (no caller meta and no active trace context) so callers can
// omit _meta entirely for 2025-11-25 peers that don't expect it.
//
// Overwrite semantics: injecting the CURRENT span's context overwrites any
// inbound traceparent already present in the cloned caller map. This is the
// correct behavior for an intermediary — the outbound request must be a child
// of the span active here, not of whatever the caller inherited upstream.
//
// SECURITY: the injected trace context AND baggage cross the trust boundary into
// backend requests, and backends may be untrusted. Callers must never place
// secrets or PII in baggage, since it is serialized verbatim into the outbound
// _meta. Baggage is intentionally NOT stripped: SEP-414 explicitly reserves it.
func MetaWithTraceContext(ctx context.Context, meta map[string]interface{}) map[string]interface{} {
        enriched := maps.Clone(meta) // nil-safe: returns nil for nil input
        if enriched == nil {
                enriched = make(map[string]interface{})
        }
        InjectMetaTraceContext(ctx, enriched)
        if len(enriched) == 0 {
                return nil
        }
        return enriched
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package telemetry

import (
        "maps"
        "sync"

        sdktrace "go.opentelemetry.io/otel/sdk/trace"
)

// DefaultRegisteredSamplingRate is the rate used when a registered integration
// did not specify one, preserving the "sample everything" behaviour that
// processor-only mode had before rates were configurable.
const DefaultRegisteredSamplingRate = 1.0

var (
        globalProcessors    []sdktrace.SpanProcessor
        globalResourceAttrs map[string]string
        globalSamplingRate  *float64
        globalProcessorsMu  sync.Mutex
)

// RegisterSpanProcessor registers an extra OTEL span processor to be included
// in any provider created via NewProvider. This allows optional integrations
// (e.g. a Sentry bridge, Datadog exporter) to self-register during their own
// Init without coupling to the caller that creates the OTEL provider.
//
// Registration must happen before NewProvider is called; processors registered
// after provider creation will not be included in the already-created provider.
//
// Duplicate registrations of the same processor pointer are silently ignored
// to prevent OnStart/OnEnd from firing twice on a single span when Init is
// called more than once (e.g. during tests or config reload).
func RegisterSpanProcessor(p sdktrace.SpanProcessor) {
        if p == nil {
                return
        }
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        for _, existing := range globalProcessors {
                if existing == p {
                        return
                }
        }
        globalProcessors = append(globalProcessors, p)
}

// RegisterResourceAttributes merges attributes into the OTEL resource of any
// provider created via NewProvider. Integrations whose exporter bypasses their
// own SDK — such as the Sentry OTLP exporter, which never passes spans through
// the Sentry client — use this to attach the grouping keys their backend needs.
//
// As with RegisterSpanProcessor, registration must happen before NewProvider is
// called. Repeated registrations of the same key overwrite earlier values.
//
// Note that resource attributes apply to the whole provider, so they are also
// exported to any configured OTLP collector, not just to the integration that
// registered them.
func RegisterResourceAttributes(attrs map[string]string) {
        if len(attrs) == 0 {
                return
        }
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        if globalResourceAttrs == nil {
                globalResourceAttrs = make(map[string]string, len(attrs))
        }
        maps.Copy(globalResourceAttrs, attrs)
}

// RegisterSamplingRate records the trace sampling rate an integration wants
// applied to the spans it receives. In processor-only mode (no OTLP endpoint)
// NewServeProvider hands this to the SDK sampler, so unsampled spans are never
// constructed at all.
//
// The SDK sampler is shared by the whole provider, so this rate is not
// per-processor: when an OTLP endpoint is also configured its own sampling rate
// wins and this value is ignored. Repeated registrations overwrite the previous
// value; only one integration is expected to register a rate.
func RegisterSamplingRate(rate float64) {
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        globalSamplingRate = &rate
}

// RegisteredSamplingRate returns the rate registered via RegisterSamplingRate,
// or DefaultRegisteredSamplingRate when no integration registered one.
func RegisteredSamplingRate() float64 {
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        if globalSamplingRate == nil {
                return DefaultRegisteredSamplingRate
        }
        return *globalSamplingRate
}

// HasRegisteredSpanProcessors returns true if any extra span processors have
// been registered. Callers can use this to decide whether to initialise an
// OTEL provider even when no OTLP endpoint is configured.
func HasRegisteredSpanProcessors() bool {
        return RegisteredSpanProcessorCount() > 0
}

// RegisteredSpanProcessorCount returns how many extra span processors are
// currently registered.
func RegisteredSpanProcessorCount() int {
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        return len(globalProcessors)
}

// RegisteredResourceAttributes returns a copy of every resource attribute
// registered via RegisterResourceAttributes, or nil when there are none.
func RegisteredResourceAttributes() map[string]string {
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        if len(globalResourceAttrs) == 0 {
                return nil
        }
        return maps.Clone(globalResourceAttrs)
}

// ResetSpanProcessorsForTesting clears all registered span processors, resource
// attributes and the sampling rate. For use in tests only.
func ResetSpanProcessorsForTesting() {
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        globalProcessors = nil
        globalResourceAttrs = nil
        globalSamplingRate = nil
}

// registeredSpanProcessors returns a snapshot of all registered processors.
func registeredSpanProcessors() []sdktrace.SpanProcessor {
        globalProcessorsMu.Lock()
        defer globalProcessorsMu.Unlock()
        if len(globalProcessors) == 0 {
                return nil
        }
        result := make([]sdktrace.SpanProcessor, len(globalProcessors))
        copy(result, globalProcessors)
        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package telemetry

import (
        "context"
        "fmt"
        "log/slog"

        "github.com/stacklok/toolhive/pkg/config"
)

// NewServeProvider initialises the OTEL provider for thv serve using the global
// config (set via `thv config otel set-endpoint`). No new CLI flags are
// introduced; serve reuses the same OTEL config as thv run. Any span processors
// registered via RegisterSpanProcessor (e.g. by sentrypkg.Init) are
// automatically included.
//
// The caller is responsible for calling provider.Shutdown when done. The
// returned shutdown function is always safe to call even when otelEnabled is
// false (it is a no-op in that case).
//
// Registration of span processors (e.g. sentrypkg.Init) must happen before
// calling NewServeProvider so that the processors are picked up by NewProvider.
func NewServeProvider(ctx context.Context) (provider *Provider, otelEnabled bool, err error) {
        configProvider := config.NewDefaultProvider()
        appConfig := configProvider.GetConfig()

        otelCfg := appConfig.OTEL
        hasRegisteredProcessors := HasRegisteredSpanProcessors()

        handleUnusedEndpoint(&otelCfg)

        if otelCfg.Endpoint == "" && !otelCfg.EnablePrometheusMetricsPath && !hasRegisteredProcessors {
                return nil, false, nil
        }

        telemetryCfg := Config{
                ServiceName:                 "thv-api",
                Endpoint:                    otelCfg.Endpoint,
                TracingEnabled:              otelCfg.TracingEnabled != nil && *otelCfg.TracingEnabled,
                MetricsEnabled:              otelCfg.MetricsEnabled != nil && *otelCfg.MetricsEnabled,
                Insecure:                    otelCfg.Insecure,
                EnablePrometheusMetricsPath: otelCfg.EnablePrometheusMetricsPath,
                EnvironmentVariables:        otelCfg.EnvVars,
        }
        if otelCfg.SamplingRate != 0.0 {
                telemetryCfg.SetSamplingRateFromFloat(otelCfg.SamplingRate)
        }
        if telemetryCfg.SamplingRate == "" {
                telemetryCfg.SamplingRate = "0.05"
        }

        if otelCfg.Endpoint == "" && hasRegisteredProcessors {
                applyProcessorOnlySampling(&telemetryCfg)
        } else if rate, ignored := ignoredRegisteredSamplingRate(hasRegisteredProcessors); ignored {
                slog.Warn("integration sampling rate is ignored because an OTLP endpoint is configured; "+
                        "the endpoint's sampling rate applies to every backend, so the integration receives "+
                        "more traces than it requested",
                        "ignored_sampling_rate", rate,
                        "effective_sampling_rate", telemetryCfg.GetSamplingRateFloat(),
                        "endpoint", otelCfg.Endpoint)
        }

        p, err := NewProvider(ctx, telemetryCfg)
        if err != nil {
                return nil, false, fmt.Errorf("failed to initialize telemetry: %w", err)
        }

        slog.Debug("OTEL provider initialized for thv serve",
                "endpoint", otelCfg.Endpoint,
                "tracing", telemetryCfg.TracingEnabled,
                "metrics", telemetryCfg.MetricsEnabled)

        return p, true, nil
}

// applyProcessorOnlySampling configures tracing for the case where no OTLP
// endpoint is set but registered processors are active (e.g. a Sentry bridge).
// Tracing has to be forced on because the registered processors are the only
// consumers, and the SDK sampler is given their requested rate directly so that
// unsampled spans are never constructed. A processor that registered no rate
// gets DefaultRegisteredSamplingRate.
func applyProcessorOnlySampling(cfg *Config) {
        cfg.TracingEnabled = true
        cfg.SetSamplingRateFromFloat(RegisteredSamplingRate())
}

// ignoredRegisteredSamplingRate returns the rate a registered integration asked
// for, and whether that rate is being ignored.
//
// It is ignored as soon as an OTLP endpoint is configured: the SDK sampler is
// shared by the whole provider, so the endpoint's rate applies to every backend
// and the integration receives everything that sampler passes instead of its own
// share of it. An integration that wants everything (the default rate) is never
// short-changed, so only a rate below the default is reported.
//
// This is a warning rather than an error because exporting more traces than
// requested is not worth refusing to start the server over.
func ignoredRegisteredSamplingRate(hasRegisteredProcessors bool) (float64, bool) {
        if !hasRegisteredProcessors {
                return 0, false
        }
        rate := RegisteredSamplingRate()
        return rate, rate < DefaultRegisteredSamplingRate
}

// handleUnusedEndpoint enables tracing by default when an OTLP endpoint is
// configured but both tracing and metrics are disabled, so the server can start
// normally instead of crashing with a fatal validation error.
func handleUnusedEndpoint(otelCfg *config.OpenTelemetryConfig) {
        if otelCfg.Endpoint == "" {
                return
        }
        tracingOff := otelCfg.TracingEnabled == nil || !*otelCfg.TracingEnabled
        metricsOff := otelCfg.MetricsEnabled == nil || !*otelCfg.MetricsEnabled
        if tracingOff && metricsOff {
                slog.Warn("OTLP endpoint is configured but tracing and metrics are both disabled; enabling tracing by default",
                        "endpoint", otelCfg.Endpoint)
                enabled := true
                otelCfg.TracingEnabled = &enabled
        }
}

//go:build !ignore_autogenerated

/*
Copyright 2025 Stacklok

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Code generated by controller-gen. DO NOT EDIT.

package telemetry

import ()

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *Config) DeepCopyInto(out *Config) {
        *out = *in
        if in.Headers != nil {
                in, out := &in.Headers, &out.Headers
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.MetricsOnTransportPort != nil {
                in, out := &in.MetricsOnTransportPort, &out.MetricsOnTransportPort
                *out = new(bool)
                **out = **in
        }
        if in.EnvironmentVariables != nil {
                in, out := &in.EnvironmentVariables, &out.EnvironmentVariables
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.CustomAttributes != nil {
                in, out := &in.CustomAttributes, &out.CustomAttributes
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Config.
func (in *Config) DeepCopy() *Config {
        if in == nil {
                return nil
        }
        out := new(Config)
        in.DeepCopyInto(out)
        return out
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package templates

import (
        "encoding/json"
        "fmt"
        "text/template"
)

// FuncMap returns the standard template functions available in composite tools.
// These functions are used both at validation time and runtime to ensure
// templates using json, quote, or fromJson are valid.
//
// Smell: It is odd to expose this low-level detail of which functions are available to templating to other packages.
// Ideally, this would be encapsulated and the higher-level functions like validation would be implemented within this package.
// However, validation is currently implemented against the config types and not the composer types.
// We should make this function internal and consolidate all validation capabilities here once we've
// replaced the composer types for the config types. They are semantically the same and the config types represent the
// documented API types.
func FuncMap() template.FuncMap {
        return template.FuncMap{
                "json":     jsonEncode,
                "quote":    quote,
                "fromJson": fromJson,
        }
}

// jsonEncode is a template function that encodes a value as JSON.
func jsonEncode(v any) (string, error) {
        b, err := json.Marshal(v)
        if err != nil {
                return "", fmt.Errorf("failed to encode JSON: %w", err)
        }
        return string(b), nil
}

// quote is a template function that quotes a string.
func quote(s string) string {
        return fmt.Sprintf("%q", s)
}

// fromJson is a template function that parses a JSON string into a value.
// It is useful when the underlying MCP server does not support structured content.
func fromJson(s string) (any, error) {
        var v any
        if err := json.Unmarshal([]byte(s), &v); err != nil {
                return nil, fmt.Errorf("failed to unmarshal JSON: %w", err)
        }
        return v, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package templates provides utilities for parsing and analyzing Go templates.
package templates

import (
        "strings"
        "text/template"
        "text/template/parse"
)

// ExtractReferences finds all field references in a template string.
// Returns the full field paths (e.g., ".steps.step1.output", ".params.message").
// It uses FuncMap() to ensure templates with custom functions can be parsed.
func ExtractReferences(tmplStr string) ([]string, error) {
        tmpl, err := template.New("extract").Funcs(FuncMap()).Parse(tmplStr)
        if err != nil {
                return nil, err
        }
        return ExtractReferencesFromTemplate(tmpl), nil
}

// ExtractReferencesFromTemplate finds all field references in a parsed template.
func ExtractReferencesFromTemplate(tmpl *template.Template) []string {
        references := make(map[string]bool)

        for _, t := range tmpl.Templates() {
                if t.Root != nil {
                        walkNode(t.Root, references)
                }
        }

        // Convert map to slice
        result := make([]string, 0, len(references))
        for ref := range references {
                result = append(result, ref)
        }
        return result
}

//nolint:gocyclo // walkNode has to reason about many potential node types from the underlying template package.
func walkNode(node parse.Node, refs map[string]bool) {
        if node == nil {
                return
        }

        switch n := node.(type) {
        case *parse.ListNode:
                if n != nil {
                        for _, child := range n.Nodes {
                                walkNode(child, refs)
                        }
                }
        case *parse.ActionNode:
                walkNode(n.Pipe, refs)
        case *parse.PipeNode:
                for _, cmd := range n.Cmds {
                        walkNode(cmd, refs)
                }
                // Also walk variable declarations
                for _, variable := range n.Decl {
                        walkNode(variable, refs)
                }
        case *parse.CommandNode:
                for _, arg := range n.Args {
                        walkNode(arg, refs)
                }
        case *parse.FieldNode:
                // This is something like .foo or .data.foo
                // FieldNode.Ident contains the path segments
                ref := "." + strings.Join(n.Ident, ".")
                refs[ref] = true
        case *parse.ChainNode:
                // ChainNode is for chained field access like (.foo).bar
                walkNode(n.Node, refs)
                if len(n.Field) > 0 {
                        ref := "." + strings.Join(n.Field, ".")
                        refs[ref] = true
                }
        case *parse.IfNode:
                walkNode(n.Pipe, refs)
                walkNode(n.List, refs)
                if n.ElseList != nil {
                        walkNode(n.ElseList, refs)
                }
        case *parse.RangeNode:
                walkNode(n.Pipe, refs)
                walkNode(n.List, refs)
                if n.ElseList != nil {
                        walkNode(n.ElseList, refs)
                }
        case *parse.WithNode:
                walkNode(n.Pipe, refs)
                walkNode(n.List, refs)
                if n.ElseList != nil {
                        walkNode(n.ElseList, refs)
                }
        case *parse.TemplateNode:
                walkNode(n.Pipe, refs)
        default:
                // The other nodes do not potentially contain field references.
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package transport

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "strings"
        "sync"

        "github.com/stacklok/toolhive-core/mcpcompat/client"
        "github.com/stacklok/toolhive-core/mcpcompat/client/transport"
        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive-core/mcpcompat/server"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/versions"
)

// StdioBridge connects stdin/stdout to a target MCP server using the specified transport type.
type StdioBridge struct {
        name      string
        mode      types.TransportType
        rawTarget string // upstream base URL

        up  *client.Client
        srv *server.MCPServer

        wg     sync.WaitGroup
        cancel context.CancelFunc
}

// NewStdioBridge creates a new StdioBridge instance for the given target URL and transport type.
func NewStdioBridge(name, rawURL string, mode types.TransportType) (*StdioBridge, error) {
        return &StdioBridge{
                name:      name,
                mode:      mode,
                rawTarget: rawURL,
        }, nil
}

// Start initializes the bridge and connects to the upstream MCP server.
func (b *StdioBridge) Start(ctx context.Context) {
        ctx, b.cancel = context.WithCancel(ctx)
        b.wg.Add(1)
        go b.run(ctx)
}

// Shutdown gracefully stops the bridge, closing connections and waiting for cleanup.
func (b *StdioBridge) Shutdown() {
        if b.cancel != nil {
                b.cancel()
        }
        if b.up != nil {
                _ = b.up.Close()
        }
        b.wg.Wait()
}

func (b *StdioBridge) run(ctx context.Context) {
        //nolint:gosec // G706: logging target URL and mode from config
        slog.Debug("starting StdioBridge", "target", b.rawTarget, "mode", b.mode)
        defer b.wg.Done()

        up, err := b.connectUpstream(ctx)
        if err != nil {
                slog.Error("upstream connect failed", "error", err)
                return
        }
        b.up = up
        //nolint:gosec // G706: logging target URL from config
        slog.Debug("connected to upstream", "target", b.rawTarget)

        if err := b.initializeUpstream(ctx); err != nil {
                slog.Error("upstream initialize failed", "error", err)
                return
        }
        slog.Debug("upstream initialized successfully")

        // Tiny local stdio server
        b.srv = server.NewMCPServer(
                fmt.Sprintf("thv-%s", b.name),
                versions.Version,
                server.WithToolCapabilities(true),
                server.WithResourceCapabilities(true, true),
                server.WithPromptCapabilities(true),
        )
        slog.Debug("starting local stdio server")

        b.up.OnConnectionLost(func(err error) { slog.Warn("upstream lost", "error", err) })

        // Handle upstream notifications
        b.up.OnNotification(func(n mcp.JSONRPCNotification) {
                slog.Info("upstream notification received", "method", n.Method)
                // Convert the Params struct to JSON and back to a generic map
                var params map[string]any
                if buf, err := json.Marshal(n.Params); err != nil {
                        slog.Warn("failed to marshal params", "error", err)
                        params = map[string]any{}
                } else if err := json.Unmarshal(buf, &params); err != nil {
                        slog.Warn("failed to unmarshal to map", "error", err)
                        params = map[string]any{}
                }

                // On a *_list_changed notification, re-fetch the upstream capability set
                // before forwarding the notification. forwardAll is additive only:
                // AddTool/AddResource/AddPrompt upsert by name/URI, so re-running it adds
                // or updates capabilities but does NOT prune ones the upstream removed —
                // a stale capability stays advertised locally after an upstream removal.
                // (SetTools exists if pruning is needed later.)
                switch n.Method {
                case "notifications/tools/list_changed",
                        "notifications/resources/list_changed",
                        "notifications/prompts/list_changed":
                        b.forwardAll(context.Background())
                }

                b.srv.SendNotificationToAllClients(n.Method, params)
        })

        // Forwarders (register once; no pagination/refresh to keep it simple)
        b.forwardAll(ctx)

        // Serve stdio (blocks)
        if err := server.ServeStdio(b.srv); err != nil {
                slog.Error("stdio server error", "error", err)
        }
}

func (b *StdioBridge) connectUpstream(_ context.Context) (*client.Client, error) {
        //nolint:gosec // G706: logging target URL and mode from config
        slog.Debug("connecting to upstream", "target", b.rawTarget, "mode", b.mode)

        switch b.mode {
        case types.TransportTypeStreamableHTTP:
                c, err := client.NewStreamableHttpClient(
                        b.rawTarget,
                        transport.WithHTTPTimeout(0),
                        transport.WithContinuousListening(),
                )
                if err != nil {
                        return nil, err
                }
                // use separate, never-ending context for the client
                if err := c.Start(context.Background()); err != nil {
                        return nil, err
                }
                return c, nil
        case types.TransportTypeSSE:
                c, err := client.NewSSEMCPClient(
                        b.rawTarget,
                )
                if err != nil {
                        return nil, err
                }
                if err := c.Start(context.Background()); err != nil {
                        return nil, err
                }
                return c, nil
        case types.TransportTypeStdio:
                // if url contains sse it's sse else streamable-http
                var c *client.Client
                var err error
                if strings.Contains(b.rawTarget, "sse") {
                        c, err = client.NewSSEMCPClient(
                                b.rawTarget,
                        )
                        if err != nil {
                                return nil, err
                        }
                } else {
                        c, err = client.NewStreamableHttpClient(
                                b.rawTarget,
                        )
                        if err != nil {
                                return nil, err
                        }
                }
                if err := c.Start(context.Background()); err != nil {
                        return nil, err
                }
                return c, nil
        case types.TransportTypeInspector:
                fallthrough
        default:
                return nil, fmt.Errorf("unsupported mode %q", b.mode)
        }
}

func (b *StdioBridge) initializeUpstream(ctx context.Context) error {
        //nolint:gosec // G706: logging target URL from config
        slog.Debug("initializing upstream", "target", b.rawTarget)
        _, err := b.up.Initialize(ctx, mcp.InitializeRequest{
                Params: mcp.InitializeParams{
                        ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
                        ClientInfo:      mcp.Implementation{Name: "toolhive-bridge", Version: "0.1.0"},
                        Capabilities:    mcp.ClientCapabilities{},
                },
        })
        if err != nil {
                return err
        }
        return nil
}

func (b *StdioBridge) forwardAll(ctx context.Context) {
        slog.Debug("forwarding all upstream data to local stdio server")
        // Tools -> straight passthrough
        slog.Debug("forwarding tools from upstream to local stdio server")
        if lt, err := b.up.ListTools(ctx, mcp.ListToolsRequest{}); err == nil {
                for _, tool := range lt.Tools {
                        toolCopy := tool
                        b.srv.AddTool(toolCopy, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
                                return b.up.CallTool(ctx, req)
                        })
                }
        }

        // Resources -> full result passthrough so upstream _meta reaches the client
        slog.Debug("forwarding resources from upstream to local stdio server")
        if lr, err := b.up.ListResources(ctx, mcp.ListResourcesRequest{}); err == nil {
                for _, res := range lr.Resources {
                        resCopy := res
                        b.srv.AddResourceWithResult(resCopy,
                                func(ctx context.Context, req mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) {
                                        return b.up.ReadResource(ctx, req)
                                })
                }
        }

        // Resource templates -> same result-returning shape as resources
        slog.Debug("forwarding resource templates from upstream to local stdio server")
        if lt, err := b.up.ListResourceTemplates(ctx, mcp.ListResourceTemplatesRequest{}); err == nil {
                for _, tpl := range lt.ResourceTemplates {
                        tplCopy := tpl
                        b.srv.AddResourceTemplateWithResult(tplCopy,
                                func(ctx context.Context, req mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) {
                                        return b.up.ReadResource(ctx, req)
                                })
                }
        }

        // Prompts -> straight passthrough
        slog.Debug("forwarding prompts from upstream to local stdio server")
        if lp, err := b.up.ListPrompts(ctx, mcp.ListPromptsRequest{}); err == nil {
                for _, p := range lp.Prompts {
                        pCopy := p
                        b.srv.AddPrompt(pCopy, func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
                                return b.up.GetPrompt(ctx, req)
                        })
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package transport provides utilities for handling different transport modes
// for communication between the client and MCP server.
package transport

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/transport/errors"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// Factory creates transports
type Factory struct{}

// NewFactory creates a new transport factory
func NewFactory() *Factory {
        return &Factory{}
}

// Option is a function that configures a transport
type Option func(types.Transport) error

// WithContainerName returns an option that sets the container name on a transport
func WithContainerName(containerName string) Option {
        return func(t types.Transport) error {
                if setter, ok := t.(interface{ setContainerName(string) }); ok {
                        setter.setContainerName(containerName)
                }
                return nil
        }
}

// WithTargetURI returns an option that sets the target URI on a transport
func WithTargetURI(targetURI string) Option {
        return func(t types.Transport) error {
                if setter, ok := t.(interface{ setTargetURI(string) }); ok {
                        setter.setTargetURI(targetURI)
                }
                return nil
        }
}

// Create creates a transport based on the provided configuration
func (*Factory) Create(config types.Config, opts ...Option) (types.Transport, error) {
        if config.ReadTimeout < 0 {
                return nil, fmt.Errorf("read timeout must be non-negative, got %s", config.ReadTimeout)
        }

        var tr types.Transport

        switch config.Type {
        case types.TransportTypeStdio:
                stdio := NewStdioTransport(
                        config.Host, config.ProxyPort, config.Deployer, config.Debug, config.TrustProxyHeaders,
                        config.PrometheusHandler, config.Middlewares...,
                )
                stdio.SetProxyMode(config.ProxyMode)
                stdio.SetStrictProtocolValidation(config.StrictProtocolValidation)
                if config.SessionStorage != nil {
                        stdio.SetSessionStorage(config.SessionStorage)
                }
                stdio.SetSessionTTL(config.SessionTTL)
                stdio.SetReadTimeout(config.ReadTimeout)
                if config.AuthInfoHandler != nil {
                        stdio.SetAuthInfoHandler(config.AuthInfoHandler)
                }
                if len(config.PrefixHandlers) > 0 {
                        stdio.SetPrefixHandlers(config.PrefixHandlers)
                }
                tr = stdio
        case types.TransportTypeSSE:
                httpTransport := NewHTTPTransport(
                        types.TransportTypeSSE,
                        config.Host,
                        config.ProxyPort,
                        config.TargetPort,
                        config.Deployer,
                        config.Debug,
                        config.TargetHost,
                        config.AuthInfoHandler,
                        config.PrometheusHandler,
                        config.PrefixHandlers,
                        config.EndpointPrefix,
                        config.TrustProxyHeaders,
                        config.Middlewares...,
                )
                httpTransport.sessionStorage = config.SessionStorage
                httpTransport.sessionTTL = config.SessionTTL
                httpTransport.readTimeout = config.ReadTimeout
                httpTransport.stripConsentCookie = config.StripConsentCookie
                tr = httpTransport
        case types.TransportTypeStreamableHTTP:
                httpTransport := NewHTTPTransport(
                        types.TransportTypeStreamableHTTP,
                        config.Host,
                        config.ProxyPort,
                        config.TargetPort,
                        config.Deployer,
                        config.Debug,
                        config.TargetHost,
                        config.AuthInfoHandler,
                        config.PrometheusHandler,
                        config.PrefixHandlers,
                        config.EndpointPrefix,
                        config.TrustProxyHeaders,
                        config.Middlewares...,
                )
                httpTransport.sessionStorage = config.SessionStorage
                httpTransport.sessionTTL = config.SessionTTL
                httpTransport.readTimeout = config.ReadTimeout
                httpTransport.stripConsentCookie = config.StripConsentCookie
                tr = httpTransport
        case types.TransportTypeInspector:
                // HTTP transport is not implemented yet
                return nil, errors.ErrUnsupportedTransport
        default:
                return nil, errors.ErrUnsupportedTransport
        }

        // Apply options to the transport
        for _, opt := range opts {
                if err := opt(tr); err != nil {
                        return nil, err
                }
        }

        return tr, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package transport

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "net/http"
        "net/url"
        "os"
        "strings"
        "sync"
        "time"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/container"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/oauthproto/tokenexchange"
        transporterrors "github.com/stacklok/toolhive/pkg/transport/errors"
        "github.com/stacklok/toolhive/pkg/transport/middleware"
        "github.com/stacklok/toolhive/pkg/transport/proxy/transparent"
        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // LocalhostName is the standard hostname for localhost
        LocalhostName = "localhost"
        // LocalhostIPv4 is the standard IPv4 address for localhost
        LocalhostIPv4 = "127.0.0.1"
)

// HTTPTransport implements the Transport interface using Server-Sent/Streamable Events.
type HTTPTransport struct {
        transportType      types.TransportType
        host               string
        proxyPort          int
        targetPort         int
        targetHost         string
        containerName      string
        targetURI          string
        deployer           rt.Deployer
        debug              bool
        middlewares        []types.NamedMiddleware
        prometheusHandler  http.Handler
        authInfoHandler    http.Handler
        prefixHandlers     map[string]http.Handler
        stripConsentCookie bool

        // endpointPrefix is an explicit prefix to prepend to SSE endpoint URLs
        endpointPrefix string

        // trustProxyHeaders indicates whether to trust X-Forwarded-* headers
        trustProxyHeaders bool

        // Remote MCP server support
        remoteURL string

        // stateless indicates the server is POST-only (no SSE/GET support)
        stateless bool

        // tokenSource is the OAuth token source for remote authentication
        tokenSource oauth2.TokenSource

        // onHealthCheckFailed is called when a health check fails for remote servers
        onHealthCheckFailed types.HealthCheckFailedCallback

        // onUnauthorizedResponse is called when a 401 Unauthorized response is received
        onUnauthorizedResponse types.UnauthorizedResponseCallback

        // isMarkedUnauthorized tracks if we've already marked the workload as unauthenticated
        // This prevents repeated status updates on every 401 response
        isMarkedUnauthorized bool

        // Mutex for protecting shared state
        mutex sync.Mutex

        // sessionStorage overrides the default in-memory session store when set.
        // Used for Redis-backed session sharing across replicas.
        sessionStorage session.Storage

        // sessionTTL overrides the inactivity timeout for sessions managed by the
        // underlying proxy. Zero uses the proxy's default.
        sessionTTL time.Duration

        // readTimeout overrides http.Server.ReadTimeout on the underlying transparent
        // proxy. Zero uses the proxy's default.
        readTimeout time.Duration

        // Transparent proxy
        proxy types.Proxy

        // Shutdown channel
        shutdownCh chan struct{}

        // Container monitor
        monitor        rt.Monitor
        monitorRuntime rt.Runtime // Stored for monitor reconnection on container restart
        errorCh        <-chan error

        // Container exit error (for determining if restart is needed)
        containerExitErr error
        exitErrMutex     sync.Mutex
}

// NewHTTPTransport creates a new HTTP transport.
func NewHTTPTransport(
        transportType types.TransportType,
        host string,
        proxyPort int,
        targetPort int,
        deployer rt.Deployer,
        debug bool,
        targetHost string,
        authInfoHandler http.Handler,
        prometheusHandler http.Handler,
        prefixHandlers map[string]http.Handler,
        endpointPrefix string,
        trustProxyHeaders bool,
        middlewares ...types.NamedMiddleware,
) *HTTPTransport {
        if host == "" {
                host = LocalhostIPv4
        }

        // If targetHost is not specified, default to localhost
        if targetHost == "" {
                targetHost = LocalhostIPv4
        }

        return &HTTPTransport{
                transportType:     transportType,
                host:              host,
                proxyPort:         proxyPort,
                middlewares:       middlewares,
                targetPort:        targetPort,
                targetHost:        targetHost,
                deployer:          deployer,
                debug:             debug,
                prometheusHandler: prometheusHandler,
                authInfoHandler:   authInfoHandler,
                prefixHandlers:    prefixHandlers,
                endpointPrefix:    endpointPrefix,
                trustProxyHeaders: trustProxyHeaders,
                shutdownCh:        make(chan struct{}),
        }
}

// SetRemoteURL sets the remote URL for the MCP server
func (t *HTTPTransport) SetRemoteURL(remoteURL string) {
        t.remoteURL = remoteURL
}

// SetTokenSource sets the OAuth token source for remote authentication
func (t *HTTPTransport) SetTokenSource(tokenSource oauth2.TokenSource) {
        t.tokenSource = tokenSource
}

// SetOnHealthCheckFailed sets the callback for health check failures
func (t *HTTPTransport) SetOnHealthCheckFailed(callback types.HealthCheckFailedCallback) {
        t.onHealthCheckFailed = callback
}

// SetStateless configures the transport for a stateless server.
func (t *HTTPTransport) SetStateless(stateless bool) {
        t.stateless = stateless
}

// SetOnUnauthorizedResponse sets the callback for 401 Unauthorized responses
// The callback is wrapped to check the unauthorized flag to prevent repeated status updates
func (t *HTTPTransport) SetOnUnauthorizedResponse(callback types.UnauthorizedResponseCallback) {
        if callback == nil {
                t.onUnauthorizedResponse = nil
                return
        }
        // Wrap the callback to check the flag before calling it
        t.onUnauthorizedResponse = func() {
                // Check if we've already marked as unauthenticated (skip if already marked)
                if t.checkAndMarkUnauthorized() {
                        return // Already marked, skip update
                }
                // Call the original callback
                callback()
        }
}

// checkAndMarkUnauthorized checks if we've already marked as unauthenticated and marks it if not
// Returns true if we should skip the status update (already marked)
func (t *HTTPTransport) checkAndMarkUnauthorized() bool {
        t.mutex.Lock()
        defer t.mutex.Unlock()
        if t.isMarkedUnauthorized {
                return true // Already marked, skip update
        }
        t.isMarkedUnauthorized = true
        return false // Not marked yet, proceed with update
}

// createTokenInjectionMiddleware creates a middleware that injects the OAuth token into requests
func (t *HTTPTransport) createTokenInjectionMiddleware() types.MiddlewareFunction {
        return middleware.CreateTokenInjectionMiddleware(t.tokenSource)
}

// hasTokenExchangeMiddleware checks if any middleware in the slice is a token exchange middleware.
// When token exchange is configured, it handles its own Authorization header injection,
// so the oauth-token-injection middleware should be skipped to avoid overwriting the exchanged token.
func hasTokenExchangeMiddleware(middlewares []types.NamedMiddleware) bool {
        for _, mw := range middlewares {
                if mw.Name == tokenexchange.MiddlewareType {
                        return true
                }
        }
        return false
}

// shouldEnableHealthCheck determines whether health checks should be enabled based on workload type.
// For local workloads, health checks are always enabled.
// For remote workloads, health checks are only enabled if explicitly opted in via the
// TOOLHIVE_REMOTE_HEALTHCHECKS environment variable (set to "true" or "1").
func shouldEnableHealthCheck(isRemote bool) bool {
        if !isRemote {
                // Always enable health checks for local workloads
                return true
        }
        // For remote workloads, only enable if explicitly opted in via environment variable
        envVal := os.Getenv("TOOLHIVE_REMOTE_HEALTHCHECKS")
        return strings.ToLower(envVal) == "true" || envVal == "1"
}

// Mode returns the transport mode.
func (t *HTTPTransport) Mode() types.TransportType {
        return t.transportType
}

// ProxyPort returns the proxy port used by the transport.
func (t *HTTPTransport) ProxyPort() int {
        return t.proxyPort
}

// setContainerName configures the transport with the container name.
// This is an unexported method used by the option pattern.
func (t *HTTPTransport) setContainerName(containerName string) {
        t.mutex.Lock()
        defer t.mutex.Unlock()
        t.containerName = containerName
}

// setTargetURI configures the transport with the target URI for proxying.
// This is an unexported method used by the option pattern.
func (t *HTTPTransport) setTargetURI(targetURI string) {
        t.mutex.Lock()
        defer t.mutex.Unlock()
        t.targetURI = targetURI
}

// Start initializes the transport and begins processing messages.
// The transport is responsible for starting the container.
//
//nolint:gocyclo // Start is a candidate for refactoring; keeping this PR focused on the Redis wiring
func (t *HTTPTransport) Start(ctx context.Context) error {
        t.mutex.Lock()
        defer t.mutex.Unlock()

        if t.deployer == nil && t.remoteURL == "" {
                return fmt.Errorf("container deployer not set")
        }

        // Determine target URI
        var targetURI string

        // remoteBasePath holds the path component from the remote URL (e.g., "/v2" from
        // "https://mcp.asana.com/v2/mcp"). This must be prepended to incoming request
        // paths so they reach the correct endpoint on the remote server.
        var remoteBasePath string

        // remoteRawQuery holds the raw query string from the remote URL (e.g.,
        // "toolsets=core,alerting" from "https://mcp.example.com/mcp?toolsets=core,alerting").
        // This must be forwarded on every outbound request or it is silently dropped.
        var remoteRawQuery string

        if t.remoteURL != "" {
                // For remote MCP servers, construct target URI from remote URL
                remoteURL, err := url.Parse(t.remoteURL)
                if err != nil {
                        return fmt.Errorf("failed to parse remote URL: %w", err)
                }
                targetURI = (&url.URL{
                        Scheme: remoteURL.Scheme,
                        Host:   remoteURL.Host,
                }).String()

                // Extract the path prefix that needs to be prepended to incoming requests.
                // The target URI only has scheme+host, so without this the remote path is lost.
                remoteBasePath = remoteURL.Path

                remoteRawQuery = remoteURL.RawQuery

                //nolint:gosec // G706: logging proxy port and remote URL from config
                slog.Debug("setting up transparent proxy to forward to remote URL",
                        "port", t.proxyPort, "target", targetURI, "base_path", remoteBasePath, "raw_query", remoteRawQuery)
        } else {
                if t.containerName == "" {
                        return transporterrors.ErrContainerNameNotSet
                }

                // For local containers, use the configured target URI
                if t.targetURI == "" {
                        return fmt.Errorf("target URI not set for HTTP transport")
                }
                targetURI = t.targetURI
                //nolint:gosec // G706: logging proxy port and target URI from config
                slog.Debug("setting up transparent proxy to forward to target",
                        "port", t.proxyPort, "target", targetURI)
        }

        // Create middlewares slice
        var middlewares []types.NamedMiddleware

        // Add the transport's existing middlewares
        middlewares = append(middlewares, t.middlewares...)

        isRemote := t.remoteURL != ""

        // Add OAuth token injection middleware for remote authentication if we have a token source.
        // Skip if token exchange is configured (it handles its own Authorization header injection).
        if isRemote && t.tokenSource != nil && !hasTokenExchangeMiddleware(t.middlewares) {
                tokenMiddleware := t.createTokenInjectionMiddleware()
                middlewares = append(middlewares, types.NamedMiddleware{
                        Name:     "oauth-token-injection",
                        Function: tokenMiddleware,
                })
        }

        // Determine whether to enable health checks based on workload type
        enableHealthCheck := shouldEnableHealthCheck(isRemote)

        // Build proxy options
        proxyOptions := t.buildProxyOptions(remoteBasePath, remoteRawQuery)

        // Create the transparent proxy
        t.proxy = transparent.NewTransparentProxyWithOptions(
                t.host,
                t.proxyPort,
                targetURI,
                t.prometheusHandler,
                t.authInfoHandler,
                t.prefixHandlers,
                enableHealthCheck,
                isRemote,
                string(t.transportType),
                t.onHealthCheckFailed,
                t.onUnauthorizedResponse,
                t.endpointPrefix,
                t.trustProxyHeaders,
                middlewares,
                proxyOptions...)
        if err := t.proxy.Start(ctx); err != nil {
                return err
        }

        //nolint:gosec // G706: logging container name and port from config
        slog.Debug("http transport started",
                "container", t.containerName, "port", t.proxyPort)

        // For remote MCP servers, we don't need container monitoring
        if isRemote {
                return nil
        }

        // Create a container monitor
        monitorRuntime, err := container.NewFactory().Create(ctx)
        if err != nil {
                return fmt.Errorf("failed to create container monitor: %w", err)
        }
        t.monitorRuntime = monitorRuntime // Store for reconnection
        t.monitor = container.NewMonitor(monitorRuntime, t.containerName)

        // Start monitoring the container
        t.errorCh, err = t.monitor.StartMonitoring(ctx)
        if err != nil {
                return fmt.Errorf("failed to start container monitoring: %w", err)
        }

        // Start a goroutine to handle container exit
        go t.handleContainerExit(ctx)

        return nil
}

// Stop gracefully shuts down the transport and the container.
func (t *HTTPTransport) Stop(ctx context.Context) error {
        t.mutex.Lock()
        defer t.mutex.Unlock()

        // Signal shutdown (guard against double-close if Stop is called
        // both from handleContainerExit and externally by the runner)
        select {
        case <-t.shutdownCh:
                // Already closed/stopping
                return nil
        default:
                close(t.shutdownCh)
        }

        // For remote MCP servers, we don't need container monitoring
        if t.remoteURL == "" {
                // Stop the monitor if it's running
                if t.monitor != nil {
                        t.monitor.StopMonitoring()
                        t.monitor = nil
                }

                // Stop the container if deployer is available
                if t.deployer != nil && t.containerName != "" {
                        if err := t.deployer.StopWorkload(ctx, t.containerName); err != nil {
                                return fmt.Errorf("failed to stop workload: %w", err)
                        }
                }
        }

        // Stop the transparent proxy
        if t.proxy != nil {
                if err := t.proxy.Stop(ctx); err != nil {
                        slog.Warn("failed to stop proxy", "error", err)
                }
        }

        return nil
}

// buildProxyOptions constructs the transparent proxy options for this transport.
func (t *HTTPTransport) buildProxyOptions(remoteBasePath, remoteRawQuery string) []transparent.Option {
        var opts []transparent.Option
        if remoteBasePath != "" {
                opts = append(opts, transparent.WithRemoteBasePath(remoteBasePath))
        }
        opts = append(opts, transparent.WithRemoteRawQuery(remoteRawQuery))
        if t.stateless {
                opts = append(opts, transparent.WithStateless())
        }
        if t.sessionTTL > 0 {
                opts = append(opts, transparent.WithSessionTTL(t.sessionTTL))
        }
        if t.readTimeout > 0 {
                opts = append(opts, transparent.WithReadTimeout(t.readTimeout))
        }
        if t.sessionStorage != nil {
                opts = append(opts, transparent.WithSessionStorage(t.sessionStorage))
        }
        if t.stripConsentCookie {
                opts = append(opts, transparent.WithStripConsentCookie())
        }
        return opts
}

// handleContainerExit handles container exit events.
// It loops to support reconnecting the monitor when a container is restarted
// by Docker (e.g., via restart policy) rather than truly exiting.
func (t *HTTPTransport) handleContainerExit(ctx context.Context) {
        for {
                select {
                case <-ctx.Done():
                        return
                case <-t.shutdownCh:
                        return
                case err := <-t.errorCh:
                        t.exitErrMutex.Lock()
                        t.containerExitErr = err
                        t.exitErrMutex.Unlock()

                        if errors.Is(err, rt.ErrContainerRestarted) {
                                //nolint:gosec // G706: logging container name from config
                                slog.Debug("container was restarted by Docker, reconnecting monitor",
                                        "container", t.containerName)
                                if reconnectErr := t.reconnectMonitor(ctx); reconnectErr != nil {
                                        //nolint:gosec // G706: logging container name from config
                                        slog.Error("failed to reconnect monitor, stopping transport",
                                                "container", t.containerName, "error", reconnectErr)
                                } else {
                                        t.exitErrMutex.Lock()
                                        t.containerExitErr = nil
                                        t.exitErrMutex.Unlock()
                                        continue
                                }
                        }

                        //nolint:gosec // G706: logging container name from config
                        slog.Warn("container exited", "container", t.containerName, "error", err)

                        // Check if container was removed (not just exited) using typed error
                        if errors.Is(err, rt.ErrContainerRemoved) {
                                //nolint:gosec // G706: logging container name from config
                                slog.Debug("container was removed, stopping proxy and cleaning up",
                                        "container", t.containerName)
                        } else {
                                //nolint:gosec // G706: logging container name from config
                                slog.Debug("container exited, will attempt automatic restart",
                                        "container", t.containerName)
                        }

                        // Stop the transport when the container exits/removed
                        if stopErr := t.Stop(ctx); stopErr != nil {
                                slog.Error("error stopping transport after container exit", "error", stopErr)
                        }
                        return
                }
        }
}

// reconnectMonitor stops the current monitor and starts a new one.
// This is used when a container is restarted by Docker -- the proxy keeps running
// but the monitor needs to track the new container start time.
func (t *HTTPTransport) reconnectMonitor(ctx context.Context) error {
        t.mutex.Lock()
        defer t.mutex.Unlock()

        // Stop the old monitor (safe even if goroutine already returned)
        if t.monitor != nil {
                t.monitor.StopMonitoring()
        }

        // Create a new monitor that records the current (post-restart) start time
        t.monitor = container.NewMonitor(t.monitorRuntime, t.containerName)

        // Start monitoring — errorCh is reassigned here, which is safe because
        // handleContainerExit (the only reader) runs on the same goroutine and
        // will see the new channel when it re-enters the select after continue.
        var err error
        t.errorCh, err = t.monitor.StartMonitoring(ctx)
        if err != nil {
                return fmt.Errorf("failed to restart container monitoring: %w", err)
        }

        return nil
}

// ShouldRestart returns true if the container exited and should be restarted.
// Returns false if the container was removed (intentionally deleted) or
// restarted by Docker (already running, no ToolHive restart needed).
func (t *HTTPTransport) ShouldRestart() bool {
        t.exitErrMutex.Lock()
        defer t.exitErrMutex.Unlock()

        if t.containerExitErr == nil {
                return false // No exit error, normal shutdown
        }

        // Don't restart if container was removed or restarted by Docker (use typed error check)
        return !errors.Is(t.containerExitErr, rt.ErrContainerRemoved) &&
                !errors.Is(t.containerExitErr, rt.ErrContainerRestarted)
}

// IsRunning checks if the transport is currently running.
// It checks both the transport's shutdown channel and the proxy's running state.
// This ensures that if the proxy stops (e.g., due to health check failure),
// the transport is also reported as not running, allowing the runner to exit cleanly.
func (t *HTTPTransport) IsRunning() (bool, error) {
        t.mutex.Lock()
        defer t.mutex.Unlock()

        // Check if the shutdown channel is closed
        select {
        case <-t.shutdownCh:
                return false, nil
        default:
                // Also check if the proxy is still running (handles health check failure case)
                // When health checks fail, the proxy stops itself but the transport's
                // shutdownCh may not be closed, causing the runner to hang as a zombie process.
                proxyRunning := true
                var err error
                if t.proxy != nil {
                        proxyRunning, err = t.proxy.IsRunning()
                        if err != nil {
                                return false, err
                        }
                }
                return proxyRunning, nil
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package middleware

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "maps"
        "net/http"
        "slices"
        "strings"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// HeaderForwardMiddlewareName is the type constant for the header forward middleware.
const HeaderForwardMiddlewareName = "header-forward"

// RestrictedHeaders is the set of headers that cannot be configured for forwarding.
// Keys are in canonical form (http.CanonicalHeaderKey).
var RestrictedHeaders = map[string]bool{
        // Routing manipulation
        "Host": true,
        // Hop-by-hop headers (RFC 7230, RFC 7540)
        "Connection":     true,
        "Keep-Alive":     true,
        "Te":             true,
        "Trailer":        true,
        "Upgrade":        true,
        "Http2-Settings": true, // RFC 7540 Section 3.2.1
        // Hop-by-hop proxy headers
        "Proxy-Authorization": true,
        "Proxy-Authenticate":  true,
        "Proxy-Connection":    true,
        // Request smuggling vectors
        "Transfer-Encoding": true,
        "Content-Length":    true,
        // Identity spoofing
        "Forwarded":         true, // RFC 7239 (standardized X-Forwarded-*)
        "X-Forwarded-For":   true,
        "X-Forwarded-Host":  true,
        "X-Forwarded-Proto": true,
        "X-Real-Ip":         true,
}

// HeaderForwardMiddlewareParams holds the parameters for the header forward middleware factory.
type HeaderForwardMiddlewareParams struct {
        // AddHeaders is a map of header names to values to inject into requests.
        AddHeaders map[string]string `json:"add_headers"`
}

// HeaderForwardFactoryMiddleware wraps header forward functionality for the factory pattern.
type HeaderForwardFactoryMiddleware struct {
        handler types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *HeaderForwardFactoryMiddleware) Handler() types.MiddlewareFunction {
        return m.handler
}

// Close cleans up any resources used by the middleware.
func (*HeaderForwardFactoryMiddleware) Close() error {
        return nil
}

// CreateMiddleware is the factory function for header forward middleware.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params HeaderForwardMiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal header forward middleware parameters: %w", err)
        }

        handler, err := createHeaderForwardHandler(params.AddHeaders)
        if err != nil {
                return err
        }

        mw := &HeaderForwardFactoryMiddleware{
                handler: handler,
        }
        runner.AddMiddleware(HeaderForwardMiddlewareName, mw)
        return nil
}

// CreateHeaderForwardMiddleware returns a middleware function that injects configured headers
// into requests before they are forwarded to remote MCP servers.
// This is a convenience function for use outside the factory pattern (e.g., thv proxy).
// It returns an error if any header name is in the restricted set.
func CreateHeaderForwardMiddleware(addHeaders map[string]string) (types.MiddlewareFunction, error) {
        return createHeaderForwardHandler(addHeaders)
}

// createHeaderForwardHandler returns a middleware that injects configured headers
// into requests before they are forwarded to remote MCP servers.
// Header names are pre-canonicalized at creation time.
// Returns an error if any configured header is in the RestrictedHeaders blocklist.
func createHeaderForwardHandler(addHeaders map[string]string) (types.MiddlewareFunction, error) {
        // Return no-op middleware if no headers configured
        if len(addHeaders) == 0 {
                return func(next http.Handler) http.Handler {
                        return next
                }, nil
        }

        // Pre-canonicalize header names and validate against blocklist
        canonicalHeaders := make(map[string]string, len(addHeaders))
        for name, value := range addHeaders {
                canonical := http.CanonicalHeaderKey(name)

                if RestrictedHeaders[canonical] {
                        return nil, fmt.Errorf("header %q is restricted and cannot be configured for forwarding", canonical)
                }

                if canonical == "Authorization" {
                        slog.Warn("authorization header is configured for forwarding; ensure the value is appropriate for the target server")
                }

                canonicalHeaders[canonical] = value
        }

        // Log configured header names once at startup (never log values)
        headerNames := slices.Sorted(maps.Keys(canonicalHeaders))
        slog.Debug("header forward middleware configured",
                "headers", strings.Join(headerNames, ", "))

        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        for name, value := range canonicalHeaders {
                                r.Header.Set(name, value)
                        }
                        next.ServeHTTP(w, r)
                })
        }, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package origin provides HTTP middleware that enforces MCP Origin header
// validation (DNS-rebinding protection) per MCP 2025-11-25 §"Security Warning"
// (https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#security-warning).
//
// When the Origin header is present on an inbound request, it MUST exactly
// match one of the configured allowed origins. Otherwise the middleware
// responds with HTTP 403 and a JSON-RPC error body. Requests without an
// Origin header (typical for non-browser clients) are permitted through.
package origin

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "maps"
        "net"
        "net/http"
        "net/url"
        "slices"
        "strings"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // MiddlewareType is the type identifier registered in the middleware factory map.
        MiddlewareType = "origin"

        // jsonRPCCodeInvalidRequest is the JSON-RPC 2.0 error code for an invalid
        // request. We reuse it for rejected Origin values because the request is
        // not well-formed from the server's security policy perspective.
        jsonRPCCodeInvalidRequest int64 = -32600

        // forbiddenBodyFallback is returned if JSON marshalling of the error body
        // fails (should never happen with simple map types).
        forbiddenBodyFallback = `{"jsonrpc":"2.0","error":{"code":-32600,"message":"Origin not allowed"}}`
)

// MiddlewareParams holds the parameters for the origin middleware factory.
type MiddlewareParams struct {
        // AllowedOrigins is the exact-match allowlist of acceptable Origin values.
        // An empty list disables the middleware (requests pass through unchanged).
        AllowedOrigins []string `json:"allowed_origins"`
}

// FactoryMiddleware wraps origin-validation as a factory-pattern middleware.
type FactoryMiddleware struct {
        handler types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *FactoryMiddleware) Handler() types.MiddlewareFunction {
        return m.handler
}

// Close releases any resources held by the middleware.
func (*FactoryMiddleware) Close() error {
        return nil
}

// CreateMiddleware is the factory function registered in
// runner.GetSupportedMiddlewareFactories.
//
// If params.AllowedOrigins is empty the factory still registers a pass-through
// handler so the middleware slot is occupied, but logs at Warn level to make
// the security-disabled state visible in operator logs. Callers that want to
// avoid registration entirely should skip calling this factory (see
// pkg/runner.prependOriginMiddleware).
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params MiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal origin middleware parameters: %w", err)
        }

        if len(params.AllowedOrigins) == 0 {
                slog.Warn("origin middleware registered with empty allowlist; Origin validation disabled")
        }

        handler := NewHandler(params.AllowedOrigins)
        runner.AddMiddleware(MiddlewareType, &FactoryMiddleware{handler: handler})
        return nil
}

// NewHandler returns a middleware function that enforces Origin header
// validation against the provided allowlist. It is the single entry point used
// by both the factory path (CreateMiddleware) and callers that build their
// middleware chain directly (e.g. `thv proxy`).
//
// What this solves: DNS-rebinding protection per MCP 2025-11-25 §"Security
// Warning" — requests whose Origin header is present and not in allowedOrigins
// receive HTTP 403 with a JSON-RPC error body.
//
// What this does NOT solve: CORS, CSRF token validation, authentication, or
// Origin-header injection via trusted reverse proxies (the caller's reverse
// proxy must deduplicate Origin headers upstream).
//
// An empty allowedOrigins slice produces a pass-through handler — the caller
// is responsible for deciding whether that is acceptable (e.g. when bind is
// loopback-only and the caller derived an allowlist via ResolveAllowedOrigins).
//
// Matching rules: exact match on byte representation except that the scheme
// and host portions of the Origin value are lowercased (RFC 6454 §4: scheme
// and host are ASCII-case-insensitive). Configured allowlist entries are
// canonicalized once at construction time.
func NewHandler(allowedOrigins []string) types.MiddlewareFunction {
        if len(allowedOrigins) == 0 {
                return func(next http.Handler) http.Handler { return next }
        }

        // Build a set for O(1) lookups. Entries are canonicalized so that
        // case-variant Origin values (RFC 6454 §4 makes scheme + host case-
        // insensitive) match predictably. Preserve the sorted list for logging.
        allowedSet := make(map[string]struct{}, len(allowedOrigins))
        for _, o := range allowedOrigins {
                allowedSet[canonicalizeOrigin(o)] = struct{}{}
        }
        slog.Debug("origin middleware configured",
                "allowed_origin_count", len(allowedSet),
                "allowed_origins", slices.Sorted(maps.Keys(allowedSet)),
        )

        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // Reject requests with multiple Origin headers outright — the
                        // Fetch spec defines Origin as a single-value header and browsers
                        // never legitimately send more than one. Splitting / merging at an
                        // upstream proxy is the only way this fires.
                        if values := r.Header.Values("Origin"); len(values) > 1 {
                                slog.Warn("rejecting request with multiple Origin headers",
                                        "count", len(values),
                                        "method", r.Method,
                                        "path", r.URL.Path,
                                        "remote", r.RemoteAddr,
                                )
                                writeForbidden(w)
                                return
                        }

                        origin := r.Header.Get("Origin")
                        if origin == "" {
                                // MCP spec §"Security Warning" only mandates validation when
                                // the header is present. Non-browser clients (stdio bridges,
                                // SDK clients) typically omit Origin entirely.
                                next.ServeHTTP(w, r)
                                return
                        }
                        if _, ok := allowedSet[canonicalizeOrigin(origin)]; !ok {
                                slog.Warn("rejecting request with disallowed Origin",
                                        "origin", origin,
                                        "method", r.Method,
                                        "path", r.URL.Path,
                                        "remote", r.RemoteAddr,
                                )
                                writeForbidden(w)
                                return
                        }
                        next.ServeHTTP(w, r)
                })
        }
}

// canonicalizeOrigin normalizes an Origin value for exact-match comparison.
// It parses the value with net/url.Parse and rebuilds it as
// "scheme://host[:port]" with the scheme and host lowercased (RFC 6454 §4
// makes both ASCII-case-insensitive) and the port preserved verbatim. Using
// the standard parser handles IPv6 bracket literals, percent-encoding, and the
// host/port split correctly instead of hand-rolling them.
//
// Per RFC 6454 §6, a serialized origin has no userinfo, path, query, or
// fragment. Any input that carries userinfo (e.g. "https://user:pass@host"),
// or that net/url cannot parse, or that lacks a scheme or host, is returned
// with a "\x00invalid:" sentinel prefix so it can never collide with a
// legitimate allowlist entry. This makes such values fail closed: a malformed
// configured entry will not match anything, and a malformed request Origin
// will not match the allowlist.
func canonicalizeOrigin(raw string) string {
        if raw == "" {
                return raw
        }
        const invalid = "\x00invalid:" // sentinel that no real Origin can produce
        u, err := url.Parse(raw)
        if err != nil {
                return invalid + raw
        }
        // A serialized origin (RFC 6454 §6) carries no userinfo and must have both
        // a scheme and a host. Reject anything else to avoid ambiguous matches.
        if u.Scheme == "" || u.Host == "" || u.User != nil {
                return invalid + raw
        }
        scheme := strings.ToLower(u.Scheme)
        // u.Hostname() strips IPv6 brackets and the port; u.Port() returns the port
        // (possibly empty). Re-add brackets for IPv6 literals (those containing ":").
        host := strings.ToLower(u.Hostname())
        if strings.Contains(host, ":") {
                host = "[" + host + "]"
        }
        if port := u.Port(); port != "" {
                return scheme + "://" + host + ":" + port
        }
        return scheme + "://" + host
}

// ResolveAllowedOrigins picks the effective Origin allowlist for a proxy
// listener. Resolution order:
//  1. If explicit is non-empty, use it verbatim.
//  2. Otherwise, if host is a loopback IP or the string "localhost", and port
//     is valid, return loopback-only defaults
//     (http://localhost:PORT, http://127.0.0.1:PORT, http://[::1]:PORT).
//  3. Otherwise, return nil — operators exposing the proxy publicly must
//     configure an explicit allowlist.
//
// Shared by the runner middleware-config helper (pkg/runner) and the
// standalone `thv proxy` command to keep the default-derivation logic in one
// place; exported because the `thv proxy` call site is outside the runner
// package and cannot reach an internal helper.
//
// What this does NOT solve: it does not validate that `explicit` entries are
// well-formed Origin values. Callers that pass operator-supplied slices must
// rely on the middleware's canonical matching to either accept or reject
// malformed entries at request time (they will simply fail to match).
func ResolveAllowedOrigins(host string, port int, explicit []string) []string {
        if len(explicit) > 0 {
                return explicit
        }
        if port <= 0 {
                return nil
        }
        if !isLoopbackHost(host) {
                return nil
        }
        return []string{
                fmt.Sprintf("http://localhost:%d", port),
                fmt.Sprintf("http://127.0.0.1:%d", port),
                fmt.Sprintf("http://[::1]:%d", port),
        }
}

// isLoopbackHost reports whether host refers to a loopback address. Accepts
// the literal string "localhost" plus any IP literal that net.ParseIP
// classifies as loopback (e.g. 127.0.0.0/8, ::1). IPv6 is currently rejected
// by cmd/thv/app/run.go:ValidateAndNormaliseHostFlag; this helper nevertheless
// handles it so future IPv6 support does not silently lose default Origin
// protection.
func isLoopbackHost(host string) bool {
        if host == "localhost" {
                return true
        }
        // Strip bracket form for IPv6 literals: "[::1]" → "::1".
        trimmed := strings.TrimSuffix(strings.TrimPrefix(host, "["), "]")
        if ip := net.ParseIP(trimmed); ip != nil {
                return ip.IsLoopback()
        }
        return false
}

// writeForbidden emits a 403 response with a JSON-RPC error body. The "id"
// key is omitted: this rejection happens on the Origin header alone, before
// any JSON-RPC body is parsed, so there is no incoming id to echo. MCP
// narrows base JSON-RPC 2.0 here -- schema/2025-11-25 types the error
// response id as optional, string|number (no null), so an absent id must be
// encoded by omitting the key, never as "id":null.
func writeForbidden(w http.ResponseWriter) {
        body, err := json.Marshal(map[string]any{
                "jsonrpc": "2.0",
                "error": map[string]any{
                        "code":    jsonRPCCodeInvalidRequest,
                        "message": "Origin not allowed",
                },
        })
        if err != nil {
                // Marshal of a static map should never fail; fall back to a literal.
                body = []byte(forbiddenBodyFallback)
        }
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusForbidden)
        //nolint:gosec // G104: writing a static JSON error response to an HTTP client
        _, _ = w.Write(body)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package middleware

import (
        "net/http"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// StripAuthMiddlewareName is the type constant for the credential stripping middleware.
const StripAuthMiddlewareName = "strip-auth"

// clientCredentialHeaders are the request headers removed by the strip-auth
// middleware before a request is forwarded to the backend. Authorization
// carries the ToolHive JWT; Cookie and Proxy-Authorization are the other
// headers a client can use to carry credentials. The set mirrors the headers
// net/http refuses to copy across cross-host redirects.
var clientCredentialHeaders = []string{"Authorization", "Cookie", "Proxy-Authorization"}

// StripAuthMiddleware removes client credential headers from requests so they
// never reach the backend. It is used when clients are authenticated by the
// proxy but the backend itself is public (DisableUpstreamTokenInjection): by
// the time this middleware runs, the auth middleware has already validated
// the client JWT and stored the identity in the request context, so the
// backend receives an unauthenticated request. Credentials injected by
// middlewares that run closer to the backend (e.g. header-forward) are
// unaffected.
type StripAuthMiddleware struct{}

// Handler returns the middleware function.
func (*StripAuthMiddleware) Handler() types.MiddlewareFunction {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        for _, h := range clientCredentialHeaders {
                                r.Header.Del(h)
                        }
                        next.ServeHTTP(w, r)
                })
        }
}

// Close cleans up resources. The middleware holds none.
func (*StripAuthMiddleware) Close() error { return nil }

// CreateStripAuthMiddleware is the factory function for the strip-auth
// middleware. It takes no parameters.
func CreateStripAuthMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        runner.AddMiddleware(config.Type, &StripAuthMiddleware{})
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package middleware provides middleware functions for the transport package.
package middleware

import (
        "fmt"
        "log/slog"
        "net/http"

        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// retryAfterSecs tells MCP clients how long to wait before retrying.
// Matches the initial MonitoredTokenSource backoff interval so that clients
// retry around the same time the next token refresh attempt happens.
const retryAfterSecs = "10"

// CreateTokenInjectionMiddleware returns a middleware that injects a Bearer token
// from the provided oauth2.TokenSource. It returns 503 Service Unavailable with a
// Retry-After header when the token cannot be retrieved, so that MCP clients treat
// the failure as transient rather than initiating an OAuth discovery flow.
func CreateTokenInjectionMiddleware(tokenSource oauth2.TokenSource) types.MiddlewareFunction {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        if tokenSource != nil {
                                token, err := tokenSource.Token()
                                if err != nil {
                                        slog.Warn("unable to retrieve OAuth token", "error", err)
                                        // The token source (AuthenticatedTokenSource) handles marking
                                        // the workload as unauthenticated in its Token() method.
                                        // Return 503 instead of 401 so MCP clients do not mistake this
                                        // for a server that requires client-side OAuth authentication.
                                        w.Header().Set("Retry-After", retryAfterSecs)
                                        http.Error(w, "Token temporarily unavailable", http.StatusServiceUnavailable)
                                        return
                                }

                                r.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
                        }
                        next.ServeHTTP(w, r)
                })
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package middleware

import (
        "log/slog"
        "net/http"
        "strings"
        "sync"
        "time"
)

const defaultResponseWriteTimeout = 30 * time.Second

type postResponseWriter struct {
        http.ResponseWriter
        writeTimeout time.Duration
        armOnce      sync.Once
}

func (w *postResponseWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }

func (w *postResponseWriter) armWriteDeadline() {
        w.armOnce.Do(func() {
                // A streamable-HTTP POST may legitimately negotiate an SSE response.
                // Like a qualifying GET stream, its lifetime is governed by request
                // cancellation rather than a socket write deadline.
                if strings.Contains(w.Header().Get("Content-Type"), "text/event-stream") {
                        return
                }
                if err := http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(w.writeTimeout)); err != nil {
                        slog.Warn("failed to arm MCP response write deadline", "error", err)
                }
        })
}

func (w *postResponseWriter) WriteHeader(statusCode int) {
        w.armWriteDeadline()
        w.ResponseWriter.WriteHeader(statusCode)
}

func (w *postResponseWriter) Write(p []byte) (int, error) {
        w.armWriteDeadline()
        return w.ResponseWriter.Write(p)
}

func (w *postResponseWriter) Flush() {
        w.armWriteDeadline()
        if err := http.NewResponseController(w.ResponseWriter).Flush(); err != nil {
                slog.Debug("failed to flush MCP response", "error", err)
        }
}

// WriteTimeout clears the server-level write deadline while an MCP POST is
// computing, then arms a fresh deadline when a non-streaming response begins.
// Qualifying SSE requests remain unbounded and are canceled through request
// contexts. Other requests are left untouched. The optional duration controls
// the response-write allowance and defaults to 30 seconds.
func WriteTimeout(endpointPath string, responseWriteTimeout ...time.Duration) func(http.Handler) http.Handler {
        writeTimeout := defaultResponseWriteTimeout
        if len(responseWriteTimeout) > 0 && responseWriteTimeout[0] > 0 {
                writeTimeout = responseWriteTimeout[0]
        }

        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        isMCPPost := r.Method == http.MethodPost && r.URL.Path == endpointPath
                        isMCPStream := r.Method == http.MethodGet &&
                                strings.Contains(r.Header.Get("Accept"), "text/event-stream") &&
                                r.URL.Path == endpointPath
                        if isMCPPost || isMCPStream {
                                rc := http.NewResponseController(w)
                                if err := rc.SetWriteDeadline(time.Time{}); err != nil {
                                        slog.Warn("failed to clear write deadline for MCP request; request may be killed by server WriteTimeout",
                                                "error", err,
                                                "method", r.Method,
                                                "path", r.URL.Path,
                                                "remote", r.RemoteAddr,
                                        )
                                }
                        }
                        if isMCPPost {
                                w = &postResponseWriter{ResponseWriter: w, writeTimeout: writeTimeout}
                        }
                        next.ServeHTTP(w, r)
                })
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package httpsse provides an HTTP proxy implementation for Server-Sent Events (SSE)
// used in communication between the client and MCP server.
package httpsse

import (
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "maps"
        "net"
        "net/http"
        "net/url"
        "path"
        "strconv"
        "sync"
        "time"

        "github.com/google/uuid"
        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/sessionbinding"
        "github.com/stacklok/toolhive/pkg/bodylimit"
        "github.com/stacklok/toolhive/pkg/diagnostics"
        "github.com/stacklok/toolhive/pkg/healthcheck"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/proxy/socket"
        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/ssecommon"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// defaultReadTimeout bounds reading the entire request (headers + body) on the
// proxy http.Server, mitigating slow-upload connection exhaustion. It does not
// affect SSE responses, which stream on the response side. Matches the vMCP
// server default. WriteTimeout is intentionally not set: it would sever
// long-lived SSE response streams, and the request-read side (the DoS vector) is
// already bounded by ReadTimeout.
const defaultReadTimeout = 30 * time.Second

// Proxy defines the interface for proxying messages between clients and destinations.
type Proxy interface {
        // Start starts the proxy.
        Start(ctx context.Context) error

        // Stop stops the proxy.
        Stop(ctx context.Context) error

        // GetMessageChannel returns the channel for messages to/from the destination.
        GetMessageChannel() chan jsonrpc2.Message

        // GetResponseChannel returns the channel for receiving messages from the destination.
        GetResponseChannel() <-chan jsonrpc2.Message

        // SendMessageToDestination sends a message to the destination.
        SendMessageToDestination(msg jsonrpc2.Message) error

        // ForwardResponseToClients delivers a message from the destination to the
        // client session(s) it belongs to. Implementations must never deliver a
        // message to a session that did not originate it.
        ForwardResponseToClients(ctx context.Context, msg jsonrpc2.Message) error

        // SendResponseMessage sends a message to the response channel.
        SendResponseMessage(msg jsonrpc2.Message) error
}

// HTTPSSEProxy encapsulates the HTTP proxy functionality for SSE transports.
// It provides SSE endpoints and JSON-RPC message handling.
//
//nolint:revive // Intentionally named HTTPSSEProxy despite package name
type HTTPSSEProxy struct {
        // Basic configuration
        host              string
        port              int
        middlewares       []types.NamedMiddleware
        trustProxyHeaders bool

        // HTTP server
        server     *http.Server
        shutdownCh chan struct{}

        // Optional Prometheus metrics handler
        prometheusHandler http.Handler

        // Session manager for SSE clients
        sessionManager *session.Manager

        // sessionTTL is the resolved inactivity timeout for the session manager.
        // Defaults to session.DefaultSessionTTL; overridable via WithSessionTTL.
        sessionTTL time.Duration

        // readTimeout is the resolved http.Server.ReadTimeout for this proxy.
        // Defaults to defaultReadTimeout; overridable via WithReadTimeout.
        readTimeout time.Duration

        // sessionStorage is the optional custom storage backend for the session manager.
        // When nil, in-memory LocalStorage is used. Set via WithSessionStorage.
        sessionStorage session.Storage

        // authInfoHandler is the optional RFC 9728 OAuth protected resource discovery handler.
        // When nil, /.well-known/ returns a clean JSON 404. Set via WithAuthInfoHandler.
        authInfoHandler http.Handler

        // prefixHandlers contains additional HTTP handlers mounted outside the middleware chain.
        // Keys are URL path prefixes. Set via WithPrefixHandlers.
        prefixHandlers map[string]http.Handler

        // liveSSESessions tracks active SSE connections local to this instance.
        // Keys are clientID strings; values are *session.SSESession.
        // This is separate from sessionManager so that distributed storage backends
        // (e.g. Redis) can be used for session metadata without breaking response
        // routing and notification broadcast, which must reach the live in-memory
        // connection regardless of storage backend.
        liveSSESessions sync.Map

        // warnedKeys records which fail-closed drop or rejection events have
        // already been logged at Warn (see warnOnce). Guarded by warnedMutex.
        warnedKeys  map[string]struct{}
        warnedMutex sync.Mutex

        // Message channel
        messageCh      chan jsonrpc2.Message
        messageHandler http.Handler

        // Health checker
        healthChecker *healthcheck.HealthChecker

        // stopOnce ensures Stop is idempotent even when called concurrently.
        stopOnce sync.Once
}

// Option configures an HTTPSSEProxy.
type Option func(*HTTPSSEProxy)

// WithSessionStorage injects a custom storage backend into the session manager.
// When not provided, the proxy uses in-memory LocalStorage (single-replica default).
// Provide a Redis-backed storage for multi-replica deployments so all replicas
// share the same session store.
//
// Architectural note: HTTPSSEProxy is used by StdioTransport for stdio-backed MCP
// servers. Response routing (ForwardResponseToClients) and POST handling are both local
// to the instance holding the live SSE connection, so Redis storage enables
// cross-replica session metadata sharing but does NOT solve cross-replica message
// delivery — a POST accepted on replica B won't reach a client whose SSE connection
// is on replica A. Callers must ensure an external load balancer provides session
// affinity (sticky sessions) when using distributed storage with this proxy.
//
// Prefer Streamable HTTP (ProxyModeStreamableHTTP), also supported on StdioTransport,
// which does not have this affinity constraint.
//
// Note: response routing and graceful disconnect use a separate in-memory liveSSESessions
// registry, not the session manager, so any Storage implementation is safe to inject here.
func WithSessionStorage(storage session.Storage) Option {
        return func(p *HTTPSSEProxy) {
                if storage == nil {
                        return
                }
                p.sessionStorage = storage
        }
}

// WithSessionTTL overrides the session inactivity timeout used by this proxy.
// Zero or negative values are ignored so the constructor's default is preserved.
func WithSessionTTL(ttl time.Duration) Option {
        return func(p *HTTPSSEProxy) {
                if ttl <= 0 {
                        return
                }
                p.sessionTTL = ttl
        }
}

// WithReadTimeout overrides http.Server.ReadTimeout for this proxy, which bounds
// reading the entire request (headers + body). Zero or negative values are
// ignored so the constructor's default (defaultReadTimeout) is preserved.
func WithReadTimeout(d time.Duration) Option {
        return func(p *HTTPSSEProxy) {
                if d <= 0 {
                        return
                }
                p.readTimeout = d
        }
}

// WithAuthInfoHandler sets the RFC 9728 OAuth protected resource discovery handler.
// When nil (the default), requests to /.well-known/ return a clean JSON 404.
func WithAuthInfoHandler(h http.Handler) Option {
        return func(p *HTTPSSEProxy) {
                p.authInfoHandler = h
        }
}

// WithPrefixHandlers sets additional HTTP handlers that are mounted outside the
// middleware chain, keyed by URL path prefix.
func WithPrefixHandlers(handlers map[string]http.Handler) Option {
        return func(p *HTTPSSEProxy) {
                p.prefixHandlers = maps.Clone(handlers)
        }
}

// NewHTTPSSEProxy creates a new HTTP SSE proxy for transports.
func NewHTTPSSEProxy(
        host string,
        port int,
        trustProxyHeaders bool,
        prometheusHandler http.Handler,
        middlewares []types.NamedMiddleware,
        opts ...Option,
) *HTTPSSEProxy {
        proxy := &HTTPSSEProxy{
                middlewares:       middlewares,
                host:              host,
                port:              port,
                trustProxyHeaders: trustProxyHeaders,
                shutdownCh:        make(chan struct{}),
                messageCh:         make(chan jsonrpc2.Message, 100),
                sessionTTL:        session.DefaultSessionTTL,
                readTimeout:       defaultReadTimeout,
                warnedKeys:        make(map[string]struct{}),
                prometheusHandler: prometheusHandler,
        }

        for _, opt := range opts {
                opt(proxy)
        }

        // Construct the session manager once, after options have resolved sessionTTL and sessionStorage.
        sseFactory := func(id string) session.Session { return session.NewSSESession(id) }
        if proxy.sessionStorage != nil {
                proxy.sessionManager = session.NewManagerWithStorage(proxy.sessionTTL, sseFactory, proxy.sessionStorage)
        } else {
                proxy.sessionManager = session.NewManager(proxy.sessionTTL, sseFactory)
        }

        ownership, err := sessionbinding.NewMiddleware(
                func(r *http.Request) (string, error) { return sessionbinding.RequestID(r, "session_id") },
                proxy.sessionManager.LookupOwner,
                func(w http.ResponseWriter, _ *http.Request, err error) {
                        sessionbinding.WriteOwnershipError(w, nil, err)
                },
        )
        if err != nil {
                panic(err) // All required callbacks above are constructor invariants.
        }
        proxy.messageHandler = ownership(http.HandlerFunc(proxy.handleOwnedPostRequest))

        // Create MCP pinger and health checker
        mcpPinger := NewMCPPinger(proxy)
        proxy.healthChecker = healthcheck.NewHealthChecker("stdio", mcpPinger)

        return proxy
}

// applyMiddlewares applies a chain of middlewares to a handler
func applyMiddlewares(handler http.Handler, middlewares ...types.NamedMiddleware) http.Handler {
        // Apply middleware chain in reverse order (last middleware is applied first)
        for i := len(middlewares) - 1; i >= 0; i-- {
                handler = middlewares[i].Function(handler)
        }
        return handler
}

// Start starts the HTTP SSE proxy.
func (p *HTTPSSEProxy) Start(_ context.Context) error {
        // Create a new HTTP server
        mux := http.NewServeMux()

        // Mount prefix handlers (e.g. embedded auth server routes) outside the middleware chain.
        // RFC 9728 requires discovery endpoints to be reachable without authentication.
        for prefix, h := range p.prefixHandlers {
                mux.Handle(prefix, h)
                slog.Debug("mounted prefix handler", "prefix", prefix)
        }

        // Mount RFC 9728 OAuth protected resource discovery endpoint (no middlewares).
        // Always register so OAuth discovery gets a clean JSON 404 when auth is off.
        wellKnownHandler := auth.NewWellKnownHandler(p.authInfoHandler)
        mux.Handle("/.well-known/", wellKnownHandler)
        if p.authInfoHandler != nil {
                slog.Debug("rfc 9728 OAuth discovery endpoint enabled at /.well-known/oauth-protected-resource")
        }

        // Add handlers for SSE and JSON-RPC with middlewares
        // At some point we should add support for Streamable HTTP transport here
        // https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http
        // The SSE endpoint serves a long-lived GET stream. No http.Server.WriteTimeout
        // is set on this proxy, so the stream is not bounded on the write side; the
        // request-read phase is bounded by ReadTimeout.
        mux.Handle(ssecommon.HTTPSSEEndpoint, applyMiddlewares(
                http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        if r.Method != http.MethodGet {
                                http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
                                return
                        }
                        p.handleSSEConnection(w, r)
                }),
                p.middlewares...,
        ))

        mux.Handle(ssecommon.HTTPMessagesEndpoint, applyMiddlewares(http.HandlerFunc(p.handlePostRequest), p.middlewares...))

        // Add health check endpoint with MCP status (no middlewares)
        mux.Handle("/health", p.healthChecker)

        // Add Prometheus metrics endpoint if handler is provided (no middlewares),
        // otherwise return 404 so /metrics is not proxied to the backend. The runner
        // serves metrics on a separate diagnostics listener instead (see
        // pkg/diagnostics) so access can be restricted by port rather than by HTTP
        // path.
        if p.prometheusHandler != nil {
                mux.Handle("/metrics", p.prometheusHandler)
                slog.Debug("prometheus metrics endpoint enabled at /metrics")
        } else {
                mux.Handle(diagnostics.MetricsPath, diagnostics.NotServedHereHandler())
        }

        // Create a listener to get the actual port when using port 0
        // Use ListenConfig with SO_REUSEADDR to allow port reuse after unclean shutdown
        addr := fmt.Sprintf("%s:%d", p.host, p.port)
        lc := socket.ListenConfig()
        listener, err := lc.Listen(context.Background(), "tcp", addr)
        if err != nil {
                return fmt.Errorf("failed to create listener: %w", err)
        }

        // Update the server address with the actual address
        actualAddr := listener.Addr().String()

        // Create the server
        p.server = &http.Server{
                Handler:           mux,
                ReadHeaderTimeout: 10 * time.Second, // Prevent Slowloris attacks
                ReadTimeout:       p.readTimeout,    // Bound slow body uploads
        }

        // Store the actual address
        p.server.Addr = actualAddr

        // Start the server in a goroutine
        go func() {
                // Parse the actual port for logging
                _, portStr, _ := net.SplitHostPort(actualAddr)
                actualPort, _ := strconv.Atoi(portStr)

                slog.Debug("http proxy started", "port", actualPort)
                //nolint:gosec // G706: logging configured SSE and JSON-RPC endpoint addresses
                slog.Debug("sse endpoint",
                        "url", fmt.Sprintf("http://%s%s", actualAddr, ssecommon.HTTPSSEEndpoint))
                //nolint:gosec // G706: logging configured JSON-RPC endpoint address
                slog.Debug("json-RPC endpoint",
                        "url", fmt.Sprintf("http://%s%s", actualAddr, ssecommon.HTTPMessagesEndpoint))

                if err := p.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
                        slog.Error("http server error", "error", err)
                }
        }()

        // Give the server a moment to start
        time.Sleep(10 * time.Millisecond)

        return nil
}

// Stop stops the HTTP SSE proxy. It is safe to call Stop more than once or
// concurrently; only the first call performs the shutdown sequence.
func (p *HTTPSSEProxy) Stop(ctx context.Context) error {
        var stopErr error
        p.stopOnce.Do(func() {
                // Signal shutdown to SSE handlers waiting on shutdownCh.
                close(p.shutdownCh)

                // Disconnect all active SSE connections.
                p.liveSSESessions.Range(func(_, value interface{}) bool {
                        if sess, ok := value.(*session.SSESession); ok {
                                sess.Disconnect()
                        }
                        return true
                })

                // Stop the session manager last: terminates the cleanup goroutine and
                // closes any underlying storage connections (e.g. Redis client).
                // Deferred so it always runs even if server.Shutdown returns an error.
                defer func() {
                        if p.sessionManager != nil {
                                if err := p.sessionManager.Stop(); err != nil {
                                        slog.Error("failed to stop session manager", "error", err)
                                }
                        }
                }()

                // Drain active HTTP connections before tearing down storage. This ensures
                // that removeClient calls (triggered by SSE handler cancellation) can still
                // reach sessionManager.Delete without hitting a closed storage backend.
                if p.server != nil {
                        if err := p.server.Shutdown(ctx); err != nil {
                                stopErr = err
                                return
                        }
                }
        })
        return stopErr
}

// IsRunning checks if the proxy is running.
func (p *HTTPSSEProxy) IsRunning() (bool, error) {
        select {
        case <-p.shutdownCh:
                return false, nil
        default:
                return true, nil
        }
}

// GetMessageChannel returns the channel for messages to/from the destination.
func (p *HTTPSSEProxy) GetMessageChannel() chan jsonrpc2.Message {
        return p.messageCh
}

// SendMessageToDestination sends a message to the destination via the message channel.
func (p *HTTPSSEProxy) SendMessageToDestination(msg jsonrpc2.Message) error {
        select {
        case p.messageCh <- msg:
                // Message sent successfully
                return nil
        default:
                // Channel is full or closed
                return fmt.Errorf("failed to send message to destination")
        }
}

// ForwardResponseToClients routes one message from the shared backend to the
// SSE session(s) it belongs to. Every message the backend emits arrives here
// untagged, so delivery is decided by JSON-RPC shape (see routing.go):
//   - a *jsonrpc2.Response goes only to the session encoded in its routed ID,
//     with the client's original ID restored (routeResponse);
//   - a notification is broadcast if it is global, otherwise dropped
//     (routeNotification);
//   - a server-initiated ping is answered (answerBackendPing); any other
//     server-initiated request is rejected back to the backend
//     (rejectServerRequest).
//
// It never delivers a message to a session that did not originate it.
func (p *HTTPSSEProxy) ForwardResponseToClients(_ context.Context, msg jsonrpc2.Message) error {
        switch m := msg.(type) {
        case *jsonrpc2.Response:
                return p.routeResponse(m)
        case *jsonrpc2.Request:
                if !m.ID.IsValid() {
                        return p.routeNotification(m)
                }
                if m.Method == methodPing {
                        return p.answerBackendPing(m)
                }
                return p.rejectServerRequest(m)
        default:
                slog.Warn("dropping backend message of unknown JSON-RPC type", "type", fmt.Sprintf("%T", msg))
                return nil
        }
}

// handleSSEConnection handles an SSE connection.
func (p *HTTPSSEProxy) handleSSEConnection(w http.ResponseWriter, r *http.Request) {
        // Set headers for SSE
        w.Header().Set("Content-Type", "text/event-stream")
        w.Header().Set("Cache-Control", "no-cache")
        w.Header().Set("Connection", "keep-alive")
        // CORS headers deliberately omitted: the origin middleware
        // (pkg/transport/middleware/origin) enforces Origin validation per
        // MCP 2025-11-25 §"Security Warning". Reflecting Origin or emitting
        // `*` here would bypass that protection.

        // Create a unique client ID
        clientID := uuid.New().String()

        // Create a channel for sending messages to this client
        messageCh := make(chan string, 100)

        // Create SSE client info
        clientInfo := &ssecommon.SSEClient{
                MessageCh: messageCh,
                CreatedAt: time.Now(),
        }

        // Create and register the SSE session
        sseSession := session.NewSSESessionWithClient(clientID, clientInfo)
        if err := sessionbinding.BindOwner(r.Context(), sseSession); err != nil {
                http.Error(w, "Invalid identity", http.StatusUnauthorized)
                return
        }
        if err := p.sessionManager.AddSession(sseSession); err != nil {
                slog.Error("failed to add SSE session", "error", err)
                http.Error(w, "Failed to create session", http.StatusServiceUnavailable)
                return
        }
        p.liveSSESessions.Store(clientID, sseSession)

        // Create a flusher for SSE
        flusher, ok := w.(http.Flusher)
        if !ok {
                http.Error(w, "Streaming not supported", http.StatusInternalServerError)
                return
        }

        // Build and send the endpoint event
        endpointURL := p.buildEndpointURL(r, clientID)
        endpointMsg := ssecommon.NewSSEMessage("endpoint", endpointURL)

        // Send the initial event
        if _, err := fmt.Fprint(w, endpointMsg.ToSSEString()); err != nil { //nolint:gosec // G705: SSE data from internal MCP protocol
                slog.Debug("failed to write endpoint message", "error", err)
                return
        }
        flusher.Flush()

        // Create a context that is canceled when the client disconnects
        ctx, cancel := context.WithCancel(r.Context())
        defer cancel()

        // Start keep-alive ticker
        keepAliveTicker := time.NewTicker(30 * time.Second)
        defer keepAliveTicker.Stop()

        // Create a goroutine to monitor for client disconnection
        go func() {
                <-ctx.Done()
                p.removeClient(clientID)
                slog.Debug("client disconnected", "client_id", clientID)
        }()

        // Send messages to the client
        for {
                select {
                case <-ctx.Done():
                        return
                case msg, ok := <-messageCh:
                        if !ok {
                                return
                        }
                        if _, err := fmt.Fprint(w, msg); err != nil {
                                slog.Debug("failed to write message", "error", err)
                                return
                        }
                        flusher.Flush()
                case <-keepAliveTicker.C:
                        // Refresh session TTL while the SSE socket is open so the cleanup
                        // goroutine does not evict clients that haven't sent a POST recently.
                        p.sessionManager.Get(clientID)
                        if _, err := fmt.Fprint(w, ": keep-alive\n\n"); err != nil {
                                slog.Debug("failed to write keep-alive", "error", err)
                                return
                        }
                        flusher.Flush()
                }
        }
}

// handlePostRequest handles a POST request with a JSON-RPC message.
func (p *HTTPSSEProxy) handlePostRequest(w http.ResponseWriter, r *http.Request) {
        // Only accept POST requests
        if r.Method != http.MethodPost {
                http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
                return
        }

        // Extract session ID from query parameters
        query := r.URL.Query()
        sessionID := query.Get("session_id")
        if sessionID == "" {
                http.Error(w, "session_id is required", http.StatusBadRequest)
                return
        }

        p.messageHandler.ServeHTTP(w, r)
}

// handleOwnedPostRequest runs only after authentication and ownership middleware.
func (p *HTTPSSEProxy) handleOwnedPostRequest(w http.ResponseWriter, r *http.Request) {
        sessionID := r.URL.Query().Get("session_id")

        // Verify the live SSE connection for this session is held by this instance.
        // With a distributed storage backend (e.g. Redis), sessionManager.Get succeeds
        // on any replica, but response delivery only reaches the live connection held
        // by this instance. Rejecting
        // here with 503 surfaces the affinity failure explicitly instead of silently
        // dropping the response after forwarding to the backend.
        if _, local := p.liveSSESessions.Load(sessionID); !local {
                http.Error(w, "SSE connection not held by this instance", http.StatusServiceUnavailable)
                return
        }

        // Read the request body
        body, err := io.ReadAll(r.Body)
        if err != nil {
                // A body that exceeds the configured limit without a Content-Length
                // (e.g. chunked) trips http.MaxBytesReader here rather than at the
                // early Content-Length check. Surface it as 413, not 500.
                if bodylimit.IsRequestTooLarge(err) {
                        slog.Warn("rejected unreadable MCP request body", "reason", "body_too_large")
                        http.Error(w, "Request Entity Too Large", http.StatusRequestEntityTooLarge)
                        return
                }
                slog.Warn("rejected unreadable MCP request body", "reason", "read_error")
                http.Error(w, "Invalid request body", http.StatusBadRequest)
                return
        }

        // Parse the JSON-RPC message
        msg, err := mcp.DecodeMessage(body)
        if err != nil {
                mcp.WriteClassificationError(w, nil, err)
                return
        }

        slog.Debug("received JSON-RPC message", "type", fmt.Sprintf("%T", msg))

        // Tag each call's wire ID with the session that issued it, so the shared
        // backend's echoed response can be routed back to this session alone (see
        // routing.go). sessionID is the same value the ownership middleware just
        // validated for this request. Shared admission preserves exact numeric IDs.
        // A notifications/cancelled names its target request by id in params, so
        // that id is rewritten the same way; other notifications pass through
        // unchanged. A *jsonrpc2.Response from a client is refused: the proxy
        // answers or rejects every server-initiated request itself before any
        // client sees it, so a client response can only be an attempt to answer
        // backend work on behalf of some other session.
        switch m := msg.(type) {
        case *jsonrpc2.Response:
                http.Error(w, "JSON-RPC responses from clients are not accepted by the SSE proxy", http.StatusBadRequest)
                return
        case *jsonrpc2.Request:
                switch {
                case m.ID.IsValid():
                        routed, err := encodeRoutedID(sessionID, m.ID)
                        if err != nil {
                                http.Error(w, fmt.Sprintf("Unsupported JSON-RPC id: %v", err), http.StatusBadRequest)
                                return
                        }
                        msg = &jsonrpc2.Request{ID: jsonrpc2.StringID(routed), Method: m.Method, Params: m.Params}
                case m.Method == methodCancelled:
                        rewritten, err := rewriteCancelledRequestID(sessionID, m)
                        if err != nil {
                                http.Error(w, fmt.Sprintf("Unsupported JSON-RPC id: %v", err), http.StatusBadRequest)
                                return
                        }
                        msg = rewritten
                }
        }

        // Send the message to the destination
        if err := p.SendMessageToDestination(msg); err != nil {
                http.Error(w, "Failed to send message to destination", http.StatusInternalServerError)
                return
        }

        // Return a success response
        w.WriteHeader(http.StatusAccepted)
        if _, err := w.Write([]byte("Accepted")); err != nil {
                slog.Warn("failed to write response", "error", err)
        }
}

// sendSSEEvent broadcasts one SSE event to every live client. It is reserved
// for global notifications with no session-specific payload (see
// routeNotification); responses are delivered to a single session by
// routeResponse and must never pass through here.
func (p *HTTPSSEProxy) sendSSEEvent(msg *ssecommon.SSEMessage) error {
        // Convert the message to an SSE-formatted string
        sseString := msg.ToSSEString()

        // Iterate through all live SSE connections and deliver the event
        p.liveSSESessions.Range(func(key, value interface{}) bool {
                clientID, ok := key.(string)
                if !ok {
                        return true // Continue iteration
                }

                sseSession, ok := value.(*session.SSESession)
                if !ok {
                        return true // Continue iteration
                }

                // Try to send the message
                if err := sseSession.SendMessage(sseString); err != nil {
                        // Log the error but continue sending to other clients
                        switch {
                        case errors.Is(err, session.ErrSessionDisconnected):
                                slog.Debug("client is disconnected, skipping message", "client_id", clientID)
                        case errors.Is(err, session.ErrMessageChannelFull):
                                slog.Debug("client channel full, skipping message", "client_id", clientID)
                        }
                }

                return true // Continue iteration
        })

        return nil
}

// removeClient removes a client and closes its channel.
// liveSSESessions.LoadAndDelete is atomic, so concurrent calls for the same
// clientID are safe: only one will find the entry and call Disconnect.
func (p *HTTPSSEProxy) removeClient(clientID string) {
        // Disconnect the live session directly from liveSSESessions. With a
        // distributed storage backend (e.g. Redis), sessionManager.Get returns a
        // freshly-deserialized SSESession with a different MessageCh than the
        // actual live connection, so calling Disconnect() on it would close the
        // wrong channel and leave the real connection undrained.
        if val, ok := p.liveSSESessions.LoadAndDelete(clientID); ok {
                if sseSession, ok := val.(*session.SSESession); ok {
                        sseSession.Disconnect()
                }
        }

        // Remove the session from the manager
        if err := p.sessionManager.Delete(clientID); err != nil {
                slog.Debug("failed to delete session", "client_id", clientID, "error", err)
        }
}

// buildEndpointURL constructs the endpoint URL from request headers and proxy configuration.
func (p *HTTPSSEProxy) buildEndpointURL(r *http.Request, clientID string) string {
        host := r.Host
        prefix := ""

        scheme := "http"
        if r.TLS != nil {
                scheme = "https"
        }
        if forwardedProto := r.Header.Get("X-Forwarded-Proto"); forwardedProto != "" {
                scheme = forwardedProto
        }

        // Handle X-Forwarded headers from reverse proxies only if trusted
        if p.trustProxyHeaders {
                if forwardedHost := r.Header.Get("X-Forwarded-Host"); forwardedHost != "" {
                        host = forwardedHost
                        if forwardedPort := r.Header.Get("X-Forwarded-Port"); forwardedPort != "" {
                                // Strip any existing port from host before adding the forwarded port
                                if hostOnly, _, err := net.SplitHostPort(host); err == nil {
                                        host = hostOnly
                                }
                                host = net.JoinHostPort(host, forwardedPort)
                        }
                }

                prefix = r.Header.Get("X-Forwarded-Prefix")
        }

        // Strip the SSE endpoint suffix from prefix if present, since we'll add the full messages path
        prefix = stripSSEEndpointSuffix(prefix)

        u := &url.URL{
                Scheme: scheme,
                Host:   host,
                Path:   path.Join(prefix, ssecommon.HTTPMessagesEndpoint),
        }
        q := u.Query()
        q.Set("session_id", clientID)
        u.RawQuery = q.Encode()

        return u.String()
}

// stripSSEEndpointSuffix removes the SSE endpoint suffix from a path prefix if present.
func stripSSEEndpointSuffix(prefix string) string {
        sseEndpointLen := len(ssecommon.HTTPSSEEndpoint)
        if len(prefix) < sseEndpointLen {
                return prefix
        }

        // Check if the prefix ends with the SSE endpoint
        suffixStart := len(prefix) - sseEndpointLen
        if prefix[suffixStart:] == ssecommon.HTTPSSEEndpoint {
                return prefix[:suffixStart]
        }

        return prefix
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package httpsse provides MCP ping functionality for HTTP SSE proxies.
package httpsse

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "time"

        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/healthcheck"
)

// MCPPinger implements healthcheck.MCPPinger for HTTP SSE proxies
type MCPPinger struct {
        proxy *HTTPSSEProxy
}

// NewMCPPinger creates a new MCP pinger for HTTP SSE proxies
func NewMCPPinger(proxy *HTTPSSEProxy) healthcheck.MCPPinger {
        return &MCPPinger{
                proxy: proxy,
        }
}

// Ping sends a JSON-RPC ping request to the MCP server and waits for the response
// Following the MCP ping specification:
// https://modelcontextprotocol.io/specification/2025-03-26/basic/utilities/ping
func (p *MCPPinger) Ping(ctx context.Context) (time.Duration, error) {
        if p.proxy == nil {
                return 0, fmt.Errorf("proxy not available")
        }

        messageCh := p.proxy.GetMessageChannel()
        if messageCh == nil {
                return 0, fmt.Errorf("message channel not available")
        }

        // Create a ping request following MCP specification
        // {"jsonrpc": "2.0", "id": "123", "method": "ping"}
        pingID := fmt.Sprintf("ping_%d", time.Now().UnixNano())
        pingRequest, err := jsonrpc2.NewCall(jsonrpc2.StringID(pingID), "ping", json.RawMessage("{}"))
        if err != nil {
                return 0, fmt.Errorf("failed to create ping request: %w", err)
        }

        start := time.Now()

        // Send the ping request
        select {
        case messageCh <- pingRequest:
                slog.Debug("sent MCP ping request", "id", pingID)
        case <-ctx.Done():
                return 0, ctx.Err()
        default:
                return 0, fmt.Errorf("message channel is full or closed")
        }

        // The HTTP SSE proxy has no direct response channel. The backend's pong
        // carries this ping's own id, which no session route was minted for, so
        // the proxy drops it by design (see routeResponse) rather than pushing an
        // unsolicited response onto every client's stream. Only the time taken to
        // hand the request to the backend is measured.
        duration := time.Since(start)

        slog.Debug("mcp ping request sent", "duration", duration)
        return duration, nil
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package httpsse

import (
        "bytes"
        "encoding/json"
        "fmt"
        "log/slog"
        "strconv"
        "strings"

        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/ssecommon"
)

// Routed IDs carry the originating session inside each request's wire ID so
// the shared stdio backend's echoed response can be delivered to that session
// alone (GHSA-wm2j-ch74-276r). Wire format: "<sessionID>|n:<int64>" or
// "<sessionID>|s:<string>". Session IDs are proxy-minted UUIDs and never
// contain the separator, so cutting at the first "|" is unambiguous even when
// the client's own string ID contains one; the type tag lets the original ID
// be restored with its exact type. Routing is stateless because the
// destination, the session's SSE stream, is already indexed by session ID in
// liveSSESessions. Design and the full dispatch table:
// docs/arch/03-transport-architecture.md, "Server->Client Routing (Legacy SSE)".

const (
        routedIDSeparator = "|"
        routedIDNumberTag = "n:"
        routedIDStringTag = "s:"

        // methodCancelled is the client->server notification that names an
        // in-flight request by id in params.requestId. The backend only knows
        // that request by its routed id, so the proxy rewrites the parameter too
        // (see rewriteCancelledRequestID).
        methodCancelled = "notifications/cancelled"

        // methodPing is the liveness check either MCP party may issue. A backend
        // ping needs no session attribution, so the proxy answers it itself (see
        // answerBackendPing) instead of rejecting it like other server requests.
        methodPing = "ping"

        // maxWarnOnceKeys bounds the set of distinct events warnOnce remembers.
        // Real backends emit a handful of notification methods; the cap only
        // matters for a backend inventing method names.
        maxWarnOnceKeys = 64
)

// listChangedNotificationMethods is the set of server->client notification
// methods that describe a server-wide capability change with no
// session-specific payload. Delivering one copy to every connected session is
// correct and safe; every other notification the shared backend emits is
// unattributable and is dropped (see routeNotification). Mirrors the
// streamable proxy's set of the same name.
var listChangedNotificationMethods = map[string]bool{
        "notifications/tools/list_changed":     true,
        "notifications/resources/list_changed": true,
        "notifications/prompts/list_changed":   true,
}

// encodeRoutedID returns the wire ID to send upstream for a client request
// with id, issued on sessionID. id must be valid (a call, not a notification).
// It fails rather than guess if id carries a type jsonrpc2 does not construct
// today, so a library change surfaces as a rejected request instead of a
// silently retyped response the client would never match.
func encodeRoutedID(sessionID string, id jsonrpc2.ID) (string, error) {
        switch v := id.Raw().(type) {
        case int64:
                return sessionID + routedIDSeparator + routedIDNumberTag + strconv.FormatInt(v, 10), nil
        case string:
                return sessionID + routedIDSeparator + routedIDStringTag + v, nil
        default:
                return "", fmt.Errorf("unsupported JSON-RPC id type %T", v)
        }
}

// decodeRoutedID parses a wire ID the backend echoed back. It reports ok=false
// for any ID the proxy did not mint via encodeRoutedID (a non-string ID, a
// missing separator or type tag, or a malformed number), so callers drop such
// responses rather than guess a destination.
func decodeRoutedID(id jsonrpc2.ID) (sessionID string, original jsonrpc2.ID, ok bool) {
        raw, isString := id.Raw().(string)
        if !isString {
                return "", jsonrpc2.ID{}, false
        }
        sessionID, rest, found := strings.Cut(raw, routedIDSeparator)
        if !found || sessionID == "" {
                return "", jsonrpc2.ID{}, false
        }
        switch {
        case strings.HasPrefix(rest, routedIDNumberTag):
                n, err := strconv.ParseInt(strings.TrimPrefix(rest, routedIDNumberTag), 10, 64)
                if err != nil {
                        return "", jsonrpc2.ID{}, false
                }
                return sessionID, jsonrpc2.Int64ID(n), true
        case strings.HasPrefix(rest, routedIDStringTag):
                return sessionID, jsonrpc2.StringID(strings.TrimPrefix(rest, routedIDStringTag)), true
        default:
                return "", jsonrpc2.ID{}, false
        }
}

// rewriteCancelledRequestID returns a copy of the notifications/cancelled
// notification n whose params.requestId is the routed id the backend was sent
// for that request on sessionID. It never mutates n: params are freshly
// decoded from n.Params. A cancellation with no requestId, non-object params,
// or a requestId that is not a JSON-RPC id (string or integer) is returned as
// is for the backend to ignore; the proxy does not police MCP semantics.
func rewriteCancelledRequestID(sessionID string, n *jsonrpc2.Request) (*jsonrpc2.Request, error) {
        if len(n.Params) == 0 {
                return n, nil
        }
        params := map[string]any{}
        dec := json.NewDecoder(bytes.NewReader(n.Params))
        dec.UseNumber() // keep integer request ids exact; float64 would round above 2^53
        if err := dec.Decode(&params); err != nil {
                // Params are valid JSON (DecodeMessage already parsed the envelope)
                // but not an object, e.g. positional params; nothing to route on.
                return n, nil
        }
        id, ok := requestIDParam(params["requestId"])
        if !ok {
                return n, nil
        }
        routed, err := encodeRoutedID(sessionID, id)
        if err != nil {
                return nil, err
        }
        params["requestId"] = routed
        data, err := json.Marshal(params)
        if err != nil {
                return nil, fmt.Errorf("failed to encode cancellation params: %w", err)
        }
        return &jsonrpc2.Request{Method: n.Method, Params: data}, nil
}

// requestIDParam converts a decoded params.requestId into a jsonrpc2.ID. Only
// strings and integers are JSON-RPC ids; anything else reports ok=false.
func requestIDParam(v any) (jsonrpc2.ID, bool) {
        switch id := v.(type) {
        case string:
                return jsonrpc2.StringID(id), true
        case json.Number:
                n, err := id.Int64()
                if err != nil {
                        return jsonrpc2.ID{}, false
                }
                return jsonrpc2.Int64ID(n), true
        default:
                return jsonrpc2.ID{}, false
        }
}

// routeResponse delivers a backend response to the one session whose routed
// ID it carries, with the client's original ID restored. Anything else is
// dropped at Debug: an ID the proxy never minted (a health-check ping echo, or
// a backend inventing IDs), or a session that disconnected before its
// response arrived. Dropped responses are never queued or broadcast.
func (p *HTTPSSEProxy) routeResponse(resp *jsonrpc2.Response) error {
        sessionID, originalID, ok := decodeRoutedID(resp.ID)
        if !ok {
                slog.Debug("dropping response whose id carries no session route")
                return nil
        }
        sess, live := p.liveSession(sessionID)
        if !live {
                slog.Debug("dropping response for a session that is no longer connected", "session_id", sessionID)
                return nil
        }

        restored := &jsonrpc2.Response{ID: originalID, Result: resp.Result, Error: resp.Error}
        data, err := jsonrpc2.EncodeMessage(restored)
        if err != nil {
                return fmt.Errorf("failed to encode JSON-RPC response: %w", err)
        }
        // A disconnected session or a full channel drops the response, matching
        // sendSSEEvent's per-client policy.
        if err := sess.SendMessage(ssecommon.NewSSEMessage("message", string(data)).ToSSEString()); err != nil {
                slog.Debug("failed to deliver response to session", "session_id", sessionID, "error", err)
        }
        return nil
}

// routeNotification handles a server->client notification (a request with no
// ID). list_changed notifications (listChangedNotificationMethods) are
// broadcast to every live session as a payload-free notification rebuilt from
// the method alone: the spec allows params._meta on them, and anything the
// backend put there could belong to one session. With no session connected
// they reach nobody, as there is no replay buffer. Every other notification
// is dropped: the shared backend cannot say which session caused a notifications/progress,
// notifications/message, or notifications/resources/updated, and forwarding
// it to all sessions, or a guessed one, would leak one session's activity to
// another. Routing progress by a proxy-minted token, as the streamable proxy
// does, is a follow-up. The first drop of each method is logged at Warn so
// the missing feature is diagnosable; later drops at Debug.
func (p *HTTPSSEProxy) routeNotification(n *jsonrpc2.Request) error {
        if !listChangedNotificationMethods[n.Method] {
                p.warnOnce("drop:"+n.Method,
                        "dropping server notification the shared SSE proxy cannot attribute to a session; "+
                                "use streamable-http proxy mode for per-session delivery", "method", n.Method)
                return nil
        }

        stripped, err := jsonrpc2.NewNotification(n.Method, nil)
        if err != nil {
                return fmt.Errorf("failed to build JSON-RPC notification: %w", err)
        }
        data, err := jsonrpc2.EncodeMessage(stripped)
        if err != nil {
                return fmt.Errorf("failed to encode JSON-RPC notification: %w", err)
        }
        return p.sendSSEEvent(ssecommon.NewSSEMessage("message", string(data)))
}

// answerBackendPing replies to a backend's ping with an empty result, written
// back to the backend. A ping is a liveness check with no session-specific
// content, and MCP requires the receiver to answer promptly; rejecting it
// like other server-initiated requests would read as a failed check.
func (p *HTTPSSEProxy) answerBackendPing(req *jsonrpc2.Request) error {
        resp := &jsonrpc2.Response{ID: req.ID, Result: json.RawMessage("{}")}
        if err := p.SendMessageToDestination(resp); err != nil {
                return fmt.Errorf("failed to answer backend ping: %w", err)
        }
        return nil
}

// rejectServerRequest answers a server-initiated request other than ping (a
// request with a valid ID, e.g. sampling/createMessage or elicitation/create) with a JSON-RPC
// error written back to the backend, so its blocking call fails fast instead
// of hanging. The shared proxy has no way to pick which session should answer
// without guessing or broadcasting, and a request emitted by a session that
// has since disconnected would otherwise land on whoever remains connected.
// Same policy and error code as the streamable proxy's rejectServerRequestToBackend.
func (p *HTTPSSEProxy) rejectServerRequest(req *jsonrpc2.Request) error {
        p.warnOnce("reject:"+req.Method,
                "rejecting server-initiated request; the shared SSE proxy cannot attribute it to a session", "method", req.Method)
        resp := &jsonrpc2.Response{
                ID: req.ID,
                Error: jsonrpc2.NewError(-32601, fmt.Sprintf(
                        "server-initiated %s is not supported by the shared SSE proxy; requires a per-session backend deployment",
                        req.Method,
                )),
        }
        if err := p.SendMessageToDestination(resp); err != nil {
                return fmt.Errorf("failed to reject server-initiated %s: %w", req.Method, err)
        }
        return nil
}

// warnOnce logs the first occurrence of key at Warn, so an operator can see
// that a fail-closed drop or rejection is happening, and every later
// occurrence at Debug, so a chatty backend cannot flood the log.
func (p *HTTPSSEProxy) warnOnce(key, msg string, args ...any) {
        if p.firstOccurrence(key) {
                slog.Warn(msg, args...)
                return
        }
        slog.Debug(msg, args...)
}

// firstOccurrence reports whether key has not been seen before and records
// it. Once maxWarnOnceKeys distinct keys are recorded, every new key reports
// false, so a backend inventing method names cannot grow the set.
func (p *HTTPSSEProxy) firstOccurrence(key string) bool {
        p.warnedMutex.Lock()
        defer p.warnedMutex.Unlock()
        if _, seen := p.warnedKeys[key]; seen || len(p.warnedKeys) >= maxWarnOnceKeys {
                return false
        }
        p.warnedKeys[key] = struct{}{}
        return true
}

// liveSession returns the live SSE session for sessionID on this instance.
func (p *HTTPSSEProxy) liveSession(sessionID string) (*session.SSESession, bool) {
        val, ok := p.liveSSESessions.Load(sessionID)
        if !ok {
                return nil, false
        }
        sess, ok := val.(*session.SSESession)
        return sess, ok
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

//go:build !windows

// Package socket provides platform-specific socket configuration.
package socket

import (
        "net"
        "syscall"

        "golang.org/x/sys/unix"
)

// ListenConfig returns a net.ListenConfig with SO_REUSEADDR enabled
// This allows the server to restart immediately even if the port is in TIME_WAIT state
// or held by a zombie process (common after laptop sleep/wake cycles)
func ListenConfig() net.ListenConfig {
        return net.ListenConfig{
                Control: func(_, _ string, c syscall.RawConn) error {
                        var opErr error
                        if err := c.Control(func(fd uintptr) {
                                //nolint:gosec // G115: fd is a valid file descriptor
                                opErr = unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_REUSEADDR, 1)
                        }); err != nil {
                                return err
                        }
                        return opErr
                },
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package streamable

import (
        "fmt"
        "log/slog"

        "golang.org/x/exp/jsonrpc2"
)

// JSON-RPC method names used for server<->client routing decisions across
// this package (dispatcher.go and streamable_proxy.go). Collected here as
// the single source of truth so the two files can never drift on a wire
// string; none of these are exported by any SDK the proxy depends on (see
// #5744's review).
const (
        methodToolsListChanged     = "notifications/tools/list_changed"
        methodResourcesListChanged = "notifications/resources/list_changed"
        methodPromptsListChanged   = "notifications/prompts/list_changed"
        methodProgress             = "notifications/progress"
        methodResourcesUpdated     = "notifications/resources/updated"
        methodLoggingMessage       = "notifications/message"
        methodResourcesSubscribe   = "resources/subscribe"
        methodResourcesUnsubscribe = "resources/unsubscribe"
)

// listChangedNotificationMethods is the set of server->client notification
// methods that describe a server-wide capability change (not tied to any
// particular request or subscription), so delivering one copy to every
// connected session's standalone stream is correct and safe.
var listChangedNotificationMethods = map[string]bool{
        methodToolsListChanged:     true,
        methodResourcesListChanged: true,
        methodPromptsListChanged:   true,
}

// dispatchResponses routes container messages arriving on responseCh to the
// appropriate destination based on their concrete JSON-RPC type:
//   - *jsonrpc2.Response routes to the waiter correlated by its composite ID
//     (see routeResponseToWaiter).
//   - *jsonrpc2.Request with no valid ID is a server->client notification,
//     routed by method (see routeNotification).
//   - *jsonrpc2.Request with a valid ID is a server->client REQUEST (e.g.
//     sampling/createMessage, elicitation/create). It is rejected back to the
//     backend rather than forwarded to any client (see
//     rejectServerRequestToBackend).
func (p *HTTPProxy) dispatchResponses() {
        for {
                select {
                case <-p.shutdownCh:
                        return
                case msg := <-p.responseCh:
                        switch m := msg.(type) {
                        case *jsonrpc2.Response:
                                p.routeResponseToWaiter(m)
                        case *jsonrpc2.Request:
                                if !m.ID.IsValid() {
                                        p.routeNotification(m)
                                        continue
                                }
                                p.rejectServerRequestToBackend(m)
                        default:
                                slog.Warn("received invalid message that is not a valid response",
                                        "type", fmt.Sprintf("%T", msg))
                        }
                }
        }
}

// routeNotification delivers a server->client notification (m.ID invalid) to
// its correct destination based on m.Method, per the shared-backend routing
// design (see #5744 and docs/arch/03-transport-architecture.md):
//
//   - */list_changed (listChangedNotificationMethods): GLOBAL, no
//     session-specific payload -- broadcast to every connected session's
//     standalone GET stream.
//   - notifications/progress: request-scoped -- delivered ONLY to the
//     originating request's POST-SSE stream, via routeProgress's
//     progress-token correlation. Never broadcast or fanned out.
//   - notifications/resources/updated: subscription-scoped -- delivered ONLY
//     to sessions currently subscribed to the notification's uri, via
//     routeResourceUpdated.
//   - notifications/message (logging): SECURE-DROP. The shared backend is one
//     process serving every session; a log message it emits carries no
//     session attribution, so there is no way to deliver it to the session
//     that caused it without either guessing or leaking it to every other
//     connected session. Dropping it is the only safe choice until per-session
//     backends (see the vMCP architecture) are in the request path.
//   - anything else: dropped, logged at Debug.
func (p *HTTPProxy) routeNotification(m *jsonrpc2.Request) {
        switch {
        case listChangedNotificationMethods[m.Method]:
                p.serverStreams.broadcast(m)
        case m.Method == methodProgress:
                p.routeProgress(m)
        case m.Method == methodResourcesUpdated:
                p.routeResourceUpdated(m)
        case m.Method == methodLoggingMessage:
                // SECURE-DROP: shared-backend log content is unattributable to any
                // one session and forwarding it (to all sessions, or a guessed one)
                // would leak cross-session information. Deferred to per-session
                // backend work (#5744).
                slog.Debug("dropping logging notification; shared backend cannot attribute it to a session")
        default:
                //nolint:gosec // G706: method is from parsed JSON-RPC request
                slog.Debug("dropping unrecognized server->client notification", "method", m.Method)
        }
}

// routeProgress delivers a notifications/progress message to the single
// in-flight POST-SSE request that asked for it, correlated by the raw
// (proxy-minted) progressToken the backend echoes back (see
// rewriteMetaProgressToken and handleSingleRequestSSE's setupProgressRouting).
// If the token is missing or not currently registered (the request already
// completed, or the token was never ours), the notification is dropped: there
// is no session to deliver it to without guessing, and guessing risks
// cross-session leakage.
func (p *HTTPProxy) routeProgress(m *jsonrpc2.Request) {
        ptGlobal, ok := extractStringParam(m.Params, "progressToken")
        if !ok || ptGlobal == "" {
                slog.Debug("dropping progress notification with no progressToken")
                return
        }

        route, found := p.routing.lookupProgressToken(ptGlobal)
        if !found {
                slog.Debug("dropping progress notification for unknown or expired progress token")
                return
        }

        restored, err := rewriteRequestParam(m, "progressToken", route.originalToken)
        if err != nil {
                slog.Error("failed to restore client progressToken on progress notification", "error", err)
                return
        }

        select {
        case route.deliver <- restored:
        default:
                slog.Warn("progress delivery channel full; dropping progress notification")
        }
}

// routeResourceUpdated delivers a notifications/resources/updated message to
// every session currently subscribed to its uri (see
// sessionRouter.subscribersOf), and only those sessions. A notification for a
// uri with no recorded subscriber is dropped.
func (p *HTTPProxy) routeResourceUpdated(m *jsonrpc2.Request) {
        uri, ok := extractStringParam(m.Params, "uri")
        if !ok || uri == "" {
                slog.Debug("dropping resources/updated notification with no uri")
                return
        }

        subscribers := p.routing.subscribersOf(uri)
        if len(subscribers) == 0 {
                //nolint:gosec // G706: uri is from parsed JSON-RPC notification
                slog.Debug("dropping resources/updated notification with no subscribers", "uri", uri)
                return
        }

        p.serverStreams.deliverToMany(subscribers, m)
}

// rejectServerRequestToBackend responds to a server-initiated request (a
// *jsonrpc2.Request with a valid ID, e.g. sampling/createMessage or
// elicitation/create) with a JSON-RPC error, written back to the BACKEND (not
// to any client) via SendMessageToDestination, so the backend's own blocking
// call unblocks with an error instead of hanging until its own timeout.
//
// The shared-backend streamable proxy has no way to route a client's reply
// back to the correct originating backend call when multiple sessions share
// one backend process, so it cannot forward this to any client without either
// guessing (misdelivery) or broadcasting (cross-session leakage). Per-session
// backends (see the vMCP architecture, docs/arch/10-virtual-mcp-architecture.md)
// are the correct place to support server->client requests; see #5744.
func (p *HTTPProxy) rejectServerRequestToBackend(m *jsonrpc2.Request) {
        resp := &jsonrpc2.Response{
                ID: m.ID,
                Error: jsonrpc2.NewError(-32601, fmt.Sprintf(
                        "server-initiated %s is not supported by the shared streamable proxy; requires a per-session backend deployment",
                        m.Method,
                )),
        }
        if err := p.SendMessageToDestination(resp); err != nil {
                //nolint:gosec // G706: method is from parsed JSON-RPC request
                slog.Error("failed to send server-request-rejected error back to backend", "method", m.Method, "error", err)
        }
}

// routeResponseToWaiter delivers resp to the waiter channel correlated by its
// composite request ID (the composite key), if one is registered.
func (p *HTTPProxy) routeResponseToWaiter(resp *jsonrpc2.Response) {
        if !resp.ID.IsValid() {
                slog.Warn("received invalid message that is not a valid response",
                        "type", fmt.Sprintf("%T", resp))
                return
        }

        rawID := resp.ID.Raw()
        // Composite-only routing: responses must carry the composite key minted by createWaiter
        sID, ok := rawID.(string)
        if !ok {
                slog.Warn("non-string response id (expected composite string); dropping",
                        "raw_id", fmt.Sprintf("%v", rawID))
                return
        }

        chVal, ok := p.waiters.Load(sID)
        if !ok {
                slog.Warn("no waiter found for composite key; dropping", "composite_key", sID)
                return
        }
        ch, ok := chVal.(chan jsonrpc2.Message)
        if !ok {
                slog.Warn("no waiter found for composite key; dropping", "composite_key", sID)
                return
        }

        select {
        case ch <- resp:
        default:
                slog.Warn("waiter channel full; dropping response", "composite_key", sID)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package streamable

import (
        "sync"

        "golang.org/x/exp/jsonrpc2"
)

// progressRoute is the per-request delivery route for notifications/progress
// messages that echo back a proxy-minted ptGlobal token (see
// rewriteMetaProgressToken and handleSingleRequestSSE's setupProgressRouting).
type progressRoute struct {
        // deliver is the in-flight POST-SSE handler's per-request delivery
        // channel. It is buffered (see progressDeliverBufferSize) and owned by
        // that handler; routeProgress only ever sends to it, never closes it.
        deliver chan jsonrpc2.Message
        // originalToken is the client's own progressToken value (string or
        // number), restored into the notification before delivery so the client
        // sees the token it originally asked for, never ptGlobal.
        originalToken any
}

// sessionRouter owns three INDEPENDENT registries used to route server->client
// messages that arrive on the shared backend's response channel to the
// correct session(s) without leaking to any other session:
//
//   - progress tokens: request-scoped, correlate a notifications/progress
//     reply to the single in-flight POST-SSE request that asked for it.
//   - subscriptions: session-scoped, correlate a resources/updated
//     notification to the set of sessions currently subscribed to its uri,
//     and ref-count resources/subscribe|unsubscribe so the shared backend only
//     sees one upstream (un)subscribe per uri.
//   - log levels: session-scoped, track each session's requested logging
//     verbosity so the shared backend's single process-wide level can be
//     reconciled to the maximum (most verbose) requested by any session.
//
// Each registry has its own mutex (see the "one synchronization primitive per
// data structure" style rule): the three are never locked together, and no
// method here locks more than one of ptMu/subMu/logMu at a time.
type sessionRouter struct {
        ptMu   sync.Mutex
        ptToks map[string]progressRoute // key: ptGlobal (proxy-minted, unique per request)

        subMu sync.Mutex
        subs  map[string]map[string]struct{} // uri -> set(sessID)

        logMu  sync.Mutex
        logLvl map[string]int // sessID -> requested verbosity rank (higher = more verbose)
}

// newSessionRouter creates an empty sessionRouter.
func newSessionRouter() *sessionRouter {
        return &sessionRouter{
                ptToks: make(map[string]progressRoute),
                subs:   make(map[string]map[string]struct{}),
                logLvl: make(map[string]int),
        }
}

// recordProgressToken registers route under ptGlobal, so a later
// notifications/progress echoing ptGlobal is delivered via route.deliver (see
// lookupProgressToken). The caller (handleSingleRequestSSE) MUST call
// dropProgressToken(ptGlobal) once the request completes -- progress tokens
// are request-scoped, not session-scoped, so they are never reaped by reap.
func (r *sessionRouter) recordProgressToken(ptGlobal string, route progressRoute) {
        r.ptMu.Lock()
        defer r.ptMu.Unlock()
        r.ptToks[ptGlobal] = route
}

// lookupProgressToken returns the route registered for ptGlobal, if any.
func (r *sessionRouter) lookupProgressToken(ptGlobal string) (progressRoute, bool) {
        r.ptMu.Lock()
        defer r.ptMu.Unlock()
        route, ok := r.ptToks[ptGlobal]
        return route, ok
}

// dropProgressToken removes the route registered for ptGlobal, if any. It is
// always safe to call, including for a ptGlobal that was never recorded.
func (r *sessionRouter) dropProgressToken(ptGlobal string) {
        r.ptMu.Lock()
        defer r.ptMu.Unlock()
        delete(r.ptToks, ptGlobal)
}

// addSubscription records sess as a subscriber of uri and reports whether
// sess is the FIRST subscriber currently recorded for uri.
//
// The forward-vs-dedup decision is NOT driven by this return value: to record
// a subscription only after the upstream subscribe succeeds (see
// interceptSubscribe), the caller first peeks subscribersOf(uri) to decide
// whether an upstream call is needed, forwards, and only then calls
// addSubscription on success. Once any session is subscribed to uri the shared
// backend is already sending updates for it, so a later subscribe for the same
// uri is served locally without another upstream call. firstForURI is retained
// for tests/observability.
func (r *sessionRouter) addSubscription(uri, sess string) (firstForURI bool) {
        r.subMu.Lock()
        defer r.subMu.Unlock()

        set, ok := r.subs[uri]
        if !ok {
                set = make(map[string]struct{})
                r.subs[uri] = set
        }
        firstForURI = len(set) == 0
        set[sess] = struct{}{}
        return firstForURI
}

// removeSubscription removes sess as a subscriber of uri and reports whether
// sess was the LAST subscriber recorded for uri. The caller uses
// lastForURI to decide whether to forward resources/unsubscribe upstream: the
// backend should only stop sending updates for uri once no session wants them
// anymore.
func (r *sessionRouter) removeSubscription(uri, sess string) (lastForURI bool) {
        r.subMu.Lock()
        defer r.subMu.Unlock()

        set, ok := r.subs[uri]
        if !ok {
                return false
        }
        delete(set, sess)
        lastForURI = len(set) == 0
        if lastForURI {
                delete(r.subs, uri)
        }
        return lastForURI
}

// subscribersOf returns the session IDs currently subscribed to uri (nil if
// none). The caller (dispatcher.go's routeResourceUpdated) uses this as the
// routeKeys for serverStreamRegistry.deliverToMany, so a resources/updated
// notification reaches only sessions that actually subscribed to that uri.
func (r *sessionRouter) subscribersOf(uri string) []string {
        r.subMu.Lock()
        defer r.subMu.Unlock()

        set, ok := r.subs[uri]
        if !ok {
                return nil
        }
        out := make([]string, 0, len(set))
        for sess := range set {
                out = append(out, sess)
        }
        return out
}

// setLogLevel records sess's requested verbosity rank and reports whether the
// MAXIMUM rank across all currently-tracked sessions changed as a result,
// along with the new maximum. The caller (streamable_proxy.go's
// logging/setLevel handling) uses maxChanged to decide whether to reconcile
// the shared backend's single process-wide log level upstream: forwarding
// each session's own request directly would let the last writer silently
// under-serve an earlier, more-verbose requester, since every session shares
// one backend log level.
func (r *sessionRouter) setLogLevel(sess string, rank int) (maxChanged bool, newMax int) {
        r.logMu.Lock()
        defer r.logMu.Unlock()

        before := r.maxLogLevelLocked()
        r.logLvl[sess] = rank
        after := r.maxLogLevelLocked()
        return after != before, after
}

// dropLogLevel removes sess's tracked verbosity rank, if any. It is always
// safe to call, including for a sess that was never recorded.
func (r *sessionRouter) dropLogLevel(sess string) {
        r.logMu.Lock()
        defer r.logMu.Unlock()
        delete(r.logLvl, sess)
}

// maxLogLevelLocked returns the maximum rank across all tracked sessions, or
// -1 if none are tracked. Callers MUST hold logMu.
func (r *sessionRouter) maxLogLevelLocked() int {
        maxRank := -1
        for _, rank := range r.logLvl {
                if rank > maxRank {
                        maxRank = rank
                }
        }
        return maxRank
}

// purgeSession removes sess from every subscription set and from the
// log-level table. Progress tokens are request-scoped (see
// recordProgressToken's doc comment), not session-scoped, so they are
// untouched here -- an in-flight request's progress token outlives this call
// by design, and is dropped by its own handler's deferred dropProgressToken.
//
// It reports whether removing sess changed the maximum log-level rank, so the
// caller (handleDelete) can decide whether the shared backend's level should
// be reconciled down now that sess is gone.
func (r *sessionRouter) purgeSession(sess string) (maxChanged bool, newMax int) {
        r.subMu.Lock()
        for uri, set := range r.subs {
                if _, ok := set[sess]; ok {
                        delete(set, sess)
                        if len(set) == 0 {
                                delete(r.subs, uri)
                        }
                }
        }
        r.subMu.Unlock()

        r.logMu.Lock()
        defer r.logMu.Unlock()
        if _, ok := r.logLvl[sess]; !ok {
                return false, r.maxLogLevelLocked()
        }
        before := r.maxLogLevelLocked()
        delete(r.logLvl, sess)
        after := r.maxLogLevelLocked()
        return after != before, after
}

// reap removes subs/logLvl entries for every session isActive reports as no
// longer active, bounding the router's memory growth for sessions whose owning
// client disappeared without an explicit DELETE (see purgeSession for that
// path). Progress tokens are request-scoped and are never touched by reap
// (see recordProgressToken's doc comment).
//
// isActive is called OUTSIDE of ptMu/subMu/logMu: it may perform I/O (e.g. a
// Redis-backed session.Manager), so resolving liveness for every candidate
// session happens in a first pass with no lock held, and stale entries are
// then removed in a second pass that only ever does in-memory map work while
// holding a lock. This keeps reap from blocking concurrent
// addSubscription/removeSubscription/setLogLevel calls for the (possibly
// slow) duration of every liveness check.
//
// TOCTOU, the other direction: because isActive is snapshotted in the first
// pass and never re-checked once the second pass's locks are held, a session
// that (a) expires, (b) is re-created with the SAME session ID, and (c)
// re-subscribes or re-sets a log level -- all within the window between this
// tick's first and second pass -- could have its brand-new entry wrongly
// swept away by this same tick, since isActive's stale "false" result for the
// old incarnation is applied to the new one's state too. This is practically
// unreachable: session IDs are server-minted UUIDs (see resolveSessionForRequest),
// never reused by any real client in practice, so the "same ID reappears
// mid-tick" precondition does not occur outside a contrived test. Re-checking
// liveness a second time inside the phase-3 lock would close this gap, but is
// deliberately avoided: isActive's underlying session-store read (e.g. Redis
// GETEX) is I/O, and doing it while holding subMu/logMu would reintroduce the
// exact blocking-under-lock problem the two-pass design exists to avoid.
func (r *sessionRouter) reap(isActive func(sess string) bool) {
        seen := make(map[string]struct{})

        r.subMu.Lock()
        for _, set := range r.subs {
                for sess := range set {
                        seen[sess] = struct{}{}
                }
        }
        r.subMu.Unlock()

        r.logMu.Lock()
        for sess := range r.logLvl {
                seen[sess] = struct{}{}
        }
        r.logMu.Unlock()

        stale := make(map[string]struct{}, len(seen))
        for sess := range seen {
                if !isActive(sess) {
                        stale[sess] = struct{}{}
                }
        }
        if len(stale) == 0 {
                return
        }

        r.subMu.Lock()
        for uri, set := range r.subs {
                for sess := range stale {
                        delete(set, sess)
                }
                if len(set) == 0 {
                        delete(r.subs, uri)
                }
        }
        r.subMu.Unlock()

        r.logMu.Lock()
        for sess := range stale {
                delete(r.logLvl, sess)
        }
        r.logMu.Unlock()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package streamable

import (
        "log/slog"
        "sync"

        "golang.org/x/exp/jsonrpc2"
)

// serverStreamBufferSize is the buffer size for each registered stream's data
// channel.
const serverStreamBufferSize = 100

// serverStream is a single MCP session's standalone server->client stream.
// data delivers unsolicited notifications to the stream's consumer (the GET
// handler in streamable_proxy.go) and is NEVER closed -- eviction and
// shutdown are signaled exclusively via stop, so a concurrent dispatch can
// never race a close of data (see registry method docs for the panic this
// avoids).
type serverStream struct {
        // data is buffered (serverStreamBufferSize) and never closed.
        data chan jsonrpc2.Message
        // stop is closed exactly once, either by an evicting register call or by
        // closeAll, to signal the consumer to stop reading data and return.
        stop chan struct{}
}

// serverStreamRegistry is a transport-agnostic registry of live server->client
// streams, used to deliver unsolicited container notifications to the
// standalone stream for each MCP session. Per the MCP spec, a server MUST NOT
// deliver the same notification more than once per session: this registry
// enforces that by keying on session ID and allowing at most one stream per
// session. It has no knowledge of HTTP: callers (the GET handler in
// streamable_proxy.go) own translating a registered stream's messages into
// wire frames.
//
// Forward-compatibility note: the 2026-07-28 (Modern) revision's
// "subscriptions/listen" method (see mcp.MCPVersionModern in
// pkg/mcp/revision.go) is a reasonable candidate to fan out through this SAME
// decoupled registry rather than inventing a second mechanism, but it is NOT
// a drop-in reuse. Per the draft spec, a future Modern handler will need real
// adaptation on top of what is here today: "subscriptions/listen" is a
// long-lived POST (not a GET), Modern has no sessions so streams must be
// re-keyed per-subscription-id rather than per session ID, delivery requires
// per-notification-type AND per-URI opt-in filtering (not blanket fan-out),
// an initial notifications/subscriptions/acknowledged must be sent, and
// deliveries must be tagged with io.modelcontextprotocol/subscriptionId. Keep
// this type free of HTTP/GET-specific assumptions so that adaptation, when it
// happens, does not also require a rewrite of the fan-out primitives
// themselves.
//
// All fields are guarded by mu; do not add another synchronization primitive
// over streams (see the "one synchronization primitive per data structure"
// style rule). The only channel operation performed while mu is held is a
// non-blocking select-send in broadcast, dispatchTo, and deliverToMany, which
// cannot deadlock.
type serverStreamRegistry struct {
        mu      sync.Mutex
        streams map[string]*serverStream
}

// newServerStreamRegistry creates an empty serverStreamRegistry.
func newServerStreamRegistry() *serverStreamRegistry {
        return &serverStreamRegistry{streams: make(map[string]*serverStream)}
}

// register creates and stores a new stream for the given MCP session ID,
// returning it for the caller to consume from (data) and select on (stop).
// If a stream is already registered for sessionID (e.g. a second concurrent
// GET for the same session), the prior stream is evicted: its stop channel is
// closed so its consumer goroutine observes the eviction and returns, and the
// new stream replaces it in the map (last-writer-wins, one stream per
// session). The caller MUST call deregister(sessionID, s) when the stream
// ends (client disconnect, proxy shutdown) to avoid leaking the map entry.
func (r *serverStreamRegistry) register(sessionID string) *serverStream {
        r.mu.Lock()
        defer r.mu.Unlock()

        if old, ok := r.streams[sessionID]; ok {
                close(old.stop)
        }

        s := &serverStream{
                data: make(chan jsonrpc2.Message, serverStreamBufferSize),
                stop: make(chan struct{}),
        }
        r.streams[sessionID] = s
        return s
}

// deregister removes the stream for sessionID, but only if s is still the
// currently registered stream for that session (identity check). This makes
// a late deregister from an already-evicted handler (see register) a correct
// no-op instead of deleting the newer stream that replaced it. It never
// closes any channel: data is never closed (see serverStream's doc comment),
// and stop is closed exactly once by whichever of register (eviction) or
// closeAll observes it first.
func (r *serverStreamRegistry) deregister(sessionID string, s *serverStream) {
        r.mu.Lock()
        defer r.mu.Unlock()

        if r.streams[sessionID] == s {
                delete(r.streams, sessionID)
        }
}

// broadcast delivers msg once to EVERY currently registered session's stream.
// It is used only for GLOBAL notifications (see the */list_changed handling
// in dispatcher.go) that describe a server-wide capability change with no
// session-specific payload, so fanning out one copy to every session is
// correct and safe.
//
// Delivery is non-blocking and best-effort per the locked product decision: a
// stream whose data channel is full has a slow or stuck consumer, and
// dropping the message for that stream is preferable to blocking delivery to
// every other session's stream. Because data is never closed, this select can
// never panic on a send to a closed channel, even if a concurrent deregister
// or closeAll is racing this call.
func (r *serverStreamRegistry) broadcast(msg jsonrpc2.Message) {
        r.mu.Lock()
        defer r.mu.Unlock()

        for sessionID, s := range r.streams {
                select {
                case s.data <- msg:
                default:
                        slog.Warn("server stream channel full; dropping notification", "session_id", sessionID)
                }
        }
}

// dispatchTo delivers msg to the single stream registered for routeKey (a
// session ID), if one is currently connected. Delivery is non-blocking and
// best-effort, matching broadcast: an absent stream (no GET SSE connection for
// that session right now) or a full one both simply drop msg, logged at
// Debug/Warn respectively, rather than blocking or erroring. Callers use this
// for messages that are correlated to exactly one session (e.g. a
// resources/updated notification's subscribers) and MUST NOT ever reach a
// different session's stream.
func (r *serverStreamRegistry) dispatchTo(routeKey string, msg jsonrpc2.Message) {
        r.mu.Lock()
        defer r.mu.Unlock()

        s, ok := r.streams[routeKey]
        if !ok {
                slog.Debug("no server stream registered for route key; dropping notification", "route_key", routeKey)
                return
        }
        select {
        case s.data <- msg:
        default:
                slog.Warn("server stream channel full; dropping notification", "route_key", routeKey)
        }
}

// deliverToMany delivers msg to each of routeKeys' streams that is currently
// registered, taking mu once for the whole batch rather than once per key.
// Delivery to each present stream is non-blocking and best-effort, matching
// broadcast/dispatchTo. Callers use this for messages correlated to a known,
// bounded set of sessions (e.g. a resource's subscribers) -- routeKeys not
// present in the registry are silently skipped (no connected GET stream right
// now), not an error.
func (r *serverStreamRegistry) deliverToMany(routeKeys []string, msg jsonrpc2.Message) {
        r.mu.Lock()
        defer r.mu.Unlock()

        for _, routeKey := range routeKeys {
                s, ok := r.streams[routeKey]
                if !ok {
                        slog.Debug("no server stream registered for route key; dropping notification", "route_key", routeKey)
                        continue
                }
                select {
                case s.data <- msg:
                default:
                        slog.Warn("server stream channel full; dropping notification", "route_key", routeKey)
                }
        }
}

// streamCount reports how many sessions currently have a registered stream.
func (r *serverStreamRegistry) streamCount() int {
        r.mu.Lock()
        defer r.mu.Unlock()

        return len(r.streams)
}

// closeAll signals every currently registered stream to stop (by closing its
// stop channel) and resets the registry to empty. It is used on proxy Stop so
// that in-flight GET handlers observe a closed per-stream stop signal and
// return, in addition to (not instead of) selecting on shutdownCh.
func (r *serverStreamRegistry) closeAll() {
        r.mu.Lock()
        defer r.mu.Unlock()

        for _, s := range r.streams {
                close(s.stop)
        }
        r.streams = make(map[string]*serverStream)
}

// closeStream signals the single stream currently registered for routeKey (a
// session ID) to stop, mirroring closeAll's single-close discipline but
// scoped to one session: it closes stop (never data -- see serverStream's
// doc comment) exactly once and removes routeKey's entry, so the stream's
// consumer (handleGet) observes the closed stop channel and returns instead
// of surviving until its underlying socket happens to close on its own (see
// handleDelete and reapServerStreams, and FIX 2 in #5744's review). Unlike
// deregister, this does not take an identity parameter: callers here (a
// session teardown path) always want to evict WHATEVER stream is currently
// registered for routeKey, regardless of which register call produced it --
// there is no "stale handler" case to guard against as there is in
// register's eviction path. It reports whether a stream was found and
// closed, so callers can distinguish "nothing to tear down" from "torn down".
func (r *serverStreamRegistry) closeStream(routeKey string) bool {
        r.mu.Lock()
        defer r.mu.Unlock()

        s, ok := r.streams[routeKey]
        if !ok {
                return false
        }
        close(s.stop)
        delete(r.streams, routeKey)
        return true
}

// streamKeys returns a snapshot of every session ID that currently has a
// registered stream. Used by reapServerStreams to find candidate streams
// whose owning session may no longer be active, without holding mu while
// resolving each candidate's liveness (see reapServerStreams's doc comment
// for why that matters).
func (r *serverStreamRegistry) streamKeys() []string {
        r.mu.Lock()
        defer r.mu.Unlock()

        keys := make([]string, 0, len(r.streams))
        for k := range r.streams {
                keys = append(keys, k)
        }
        return keys
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package streamable

import (
        "context"
        "encoding/json"
        "errors"
        "sync/atomic"

        "golang.org/x/exp/jsonrpc2"
        "golang.org/x/sync/singleflight"

        sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp"
)

// errUnexpectedInitializeReply is returned when the backend answers initialize
// with something that is not a *jsonrpc2.Response, so it cannot be re-addressed
// to the calling client.
var errUnexpectedInitializeReply = errors.New("backend returned an unexpected reply to initialize")

// initializeCache holds the InitializeResult the backend produced for the
// first client handshake, so later handshakes are answered locally instead of
// reaching the backend again.
//
// A stdio MCP server is one process and therefore exactly one MCP session,
// while this proxy fans an arbitrary number of HTTP client sessions onto it.
// Forwarding every client's initialize to that single session means only the
// first can succeed: go-sdk v1.7+ rejects a second handshake on an
// already-initialized session with `duplicate "initialize" received`, which
// takes down every client after the first -- and, behind vMCP, every health
// check and tool call as well (see #5890, #1982).
//
// Answering from a cache restores the behaviour go-sdk v1.6.1 and earlier had,
// where a repeated initialize returned the original result, but places it in
// the component that actually owns the multiplexing rather than asking every
// MCP server to tolerate duplicate handshakes.
//
// The cached value is process-scoped rather than session-scoped: an
// InitializeResult carries the backend's protocol version, capabilities,
// instructions and server info, none of which vary by client.
type initializeCache struct {
        // flight collapses concurrent handshakes into a single upstream call whose
        // outcome every waiter shares. Serializing them behind a mutex instead
        // would queue each waiter's own round trip: with an unresponsive backend
        // the Nth client waits N*requestTimeout (60s each) before it may even
        // start.
        flight singleflight.Group

        // result is the raw InitializeResult of the first successful handshake,
        // nil until one completes.
        //
        // UNSYNCHRONISED BY DESIGN: only the flight callback in interceptInitialize
        // may read or write this field. singleflight runs at most one callback per
        // key at a time and a later callback starts only after the previous one has
        // completed through the group's own mutex, so callback-only access can
        // never overlap. Reading it anywhere else -- a status endpoint, a helper
        // accessor -- reintroduces a data race, which is why no accessor exists.
        // Once assigned it is never mutated, so the same bytes may be handed to
        // concurrent responses.
        result json.RawMessage

        // initializedForwarded records whether a notifications/initialized has
        // already been passed to the backend. This deliberately does NOT use the
        // flight: singleflight de-duplicates rather than serializes, so two
        // concurrent claims would share one callback's answer and both forward.
        // CompareAndSwap gives exactly one winner.
        initializedForwarded atomic.Bool
}

// handshake is the id-free half of an initialize outcome, shared by a flight
// with all of its waiters. Exactly one of the fields is set. Sharing this
// rather than a whole *jsonrpc2.Response means each caller builds its own
// reply, so a waiter cannot receive the forwarding caller's JSON-RPC id.
type handshake struct {
        result json.RawMessage
        rpcErr error
}

// initializeFlightKey is the sole singleflight key: the cached handshake is
// process-scoped, so all callers share one flight.
const initializeFlightKey = "initialize"

// newInitializeCache creates an empty initializeCache.
func newInitializeCache() *initializeCache {
        return &initializeCache{}
}

// interceptInitialize answers a client's initialize request: the first caller's
// request is forwarded to the backend and its result cached, and every later
// caller is served that cached result without an upstream round trip.
//
// Only a successful handshake is cached. A transport failure or a JSON-RPC
// error from the backend leaves the cache empty so the next client retries the
// handshake for real, rather than pinning every future client to one bad
// response.
//
// Concurrent first handshakes are collapsed by singleflight, so exactly one
// reaches the backend and the rest share its outcome. The upstream call runs on
// a context detached from the caller's request (context.WithoutCancel): the
// result belongs to every waiter, so the first caller disconnecting must not
// cancel the handshake for the others. forwardUpstream still applies
// p.requestTimeout to that detached context, so it cannot run unbounded.
//
// Every response is rebuilt with the calling request's own JSON-RPC id. A
// shared flight result carries the id of whichever caller performed the
// forward, and returning it verbatim would hand later waiters an id they never
// sent.
//
// One client-visible consequence: the protocol version negotiated by the first
// client is replayed to every later client. A client that asked for a
// different version receives the first client's instead, and per the MCP
// lifecycle spec should disconnect if it cannot speak it. This is inherent to
// collapsing N client sessions onto one backend session, and matches what
// go-sdk v1.6.1 and earlier did when replaying a cached result.
func (p *HTTPProxy) interceptInitialize(
        ctx context.Context, sessID string, req *jsonrpc2.Request,
) jsonrpc2.Message {
        c := p.initialize
        v, err, _ := c.flight.Do(initializeFlightKey, func() (any, error) {
                // A handshake already succeeded: replay it, no backend round trip.
                if c.result != nil {
                        return handshake{result: c.result}, nil
                }

                // Detached from the caller's context: this result belongs to every
                // waiter, so whichever caller performs the forward must not cancel the
                // handshake for the others by disconnecting. forwardUpstream still
                // bounds it with p.requestTimeout.
                msg, ferr := p.forwardUpstream(context.WithoutCancel(ctx), sessID, req)
                if ferr != nil {
                        return nil, ferr // nothing cached; the next client retries for real
                }

                resp, ok := msg.(*jsonrpc2.Response)
                if !ok {
                        return nil, errUnexpectedInitializeReply
                }
                if resp.Error != nil {
                        // Shared with this flight's waiters but deliberately not cached, so
                        // a transient failure cannot pin every future client to one
                        // response.
                        return handshake{rpcErr: resp.Error}, nil
                }
                if len(resp.Result) == 0 {
                        return nil, errUnexpectedInitializeReply // no result and no error
                }

                c.result = resp.Result
                return handshake{result: resp.Result}, nil
        })
        if err != nil {
                return upstreamErrorResponse(req.ID, err)
        }

        h := v.(handshake) // the callback returns a handshake or a non-nil error
        return &jsonrpc2.Response{ID: req.ID, Result: h.result, Error: h.rpcErr}
}

// claimInitializedForward reports whether a client's notifications/initialized
// should be passed to the backend, and records that it has been.
//
// The backend completes a single handshake and so expects a single initialized
// notification; the first client's is forwarded and every later one is
// swallowed. Swallowed notifications are still acknowledged with 202, which is
// all the sending client expects -- notifications carry no response.
func (p *HTTPProxy) claimInitializedForward() bool {
        return p.initialize.initializedForwarded.CompareAndSwap(false, true)
}

// isInitializedNotification reports whether msg is a client's
// notifications/initialized (a JSON-RPC notification, so no id).
func isInitializedNotification(msg jsonrpc2.Message) bool {
        req, ok := msg.(*jsonrpc2.Request)
        if !ok {
                return false
        }
        return req.ID.Raw() == nil && req.Method == string(sdkmcp.MethodNotificationInitialized)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package streamable

import "sync"

// keyedMutex provides per-key mutual exclusion: two callers locking the SAME
// key are fully serialized (the second blocks until the first releases),
// while callers locking DIFFERENT keys proceed concurrently, never blocking
// on each other. HTTPProxy uses it (as uriLocks) to order a
// resources/subscribe|unsubscribe ref-count decision (see
// sessionRouter.addSubscription/removeSubscription) atomically with the
// upstream forward it may trigger, per uri -- see handlePost and FIX 1 in
// #5744's review: without this, a concurrent last-unsubscribe and
// first-subscribe for the SAME uri from different sessions could reach the
// backend out of order, leaving the backend desynchronized from the routing
// table's ref count.
//
// Unlike sessionRouter's own mutexes (subMu et al.), which only ever guard a
// map lookup/mutation and are never held across a blocking call, a key's
// entry here IS deliberately held across a blocking upstream request/response
// round trip -- that is the whole point of this type. Do not reuse it to
// guard plain in-memory state that has no blocking call in its critical
// section; a plain mutex is simpler and sufficient there.
//
// Entries are ref-counted and removed once no caller holds or is waiting for
// them, bounding keyedMutex's memory footprint to the number of CURRENTLY
// in-flight keys rather than the number of keys ever seen (see the
// resource-leak style rule). This differs deliberately from, e.g.,
// pluginsvc's never-evicted pluginLock: that type's key cardinality is small
// and bounded by installed plugins, but uris here are arbitrary and
// client-supplied over the life of the proxy, so never evicting would leak
// memory.
type keyedMutex struct {
        mu    sync.Mutex
        locks map[string]*keyedMutexEntry
}

// keyedMutexEntry is one key's mutex plus a count of callers that currently
// hold or are waiting to acquire it. The count lets lock's returned unlock
// func safely delete the map entry once it reaches zero, without racing a
// concurrent lock(key) call that is (or is about to be) waiting on the same
// entry -- see lock's doc comment for why this two-phase acquire (register
// intent under mu, then block on the entry's own mutex) is race-free.
type keyedMutexEntry struct {
        mu       sync.Mutex
        refCount int
}

// newKeyedMutex creates an empty keyedMutex.
func newKeyedMutex() *keyedMutex {
        return &keyedMutex{locks: make(map[string]*keyedMutexEntry)}
}

// lock acquires exclusive ownership of key, blocking until any other caller
// currently holding (or already waiting for) the same key releases it. It
// returns an unlock func the caller MUST call exactly once (typically via
// defer, so it runs on every return/error path) to release key and, if no
// other caller is currently holding or waiting for it, remove its entry.
//
// Acquisition is two-phase: first, under k.mu, the entry is looked up (or
// created) and its refCount incremented -- this is what makes a concurrent
// unlock's decrement-and-maybe-delete race-free, since any new lock(key)
// caller either observes the entry before it is deleted (and increments it,
// preventing deletion) or observes it already deleted (and creates a fresh
// one), never a stale, about-to-be-deleted entry. Second, k.mu is released
// and the caller blocks on the entry's own mutex, so holding one key never
// blocks callers locking a different key.
//
// lock is reentrancy-free: calling lock(key) again for the same key from the
// same goroutine before releasing the first would deadlock, exactly like a
// plain sync.Mutex -- callers must never nest calls for the same key.
//
// The returned unlock is wrapped in sync.Once as defense-in-depth: it must
// still be called, but an accidental double-call becomes a no-op instead of
// unlocking a mutex a different caller now holds (releasing someone else's
// critical section) and driving refCount negative (a spurious later eviction).
func (k *keyedMutex) lock(key string) (unlock func()) {
        k.mu.Lock()
        e, ok := k.locks[key]
        if !ok {
                e = &keyedMutexEntry{}
                k.locks[key] = e
        }
        e.refCount++
        k.mu.Unlock()

        e.mu.Lock()

        var once sync.Once
        return func() {
                once.Do(func() {
                        e.mu.Unlock()

                        k.mu.Lock()
                        e.refCount--
                        if e.refCount == 0 {
                                delete(k.locks, key)
                        }
                        k.mu.Unlock()
                })
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package streamable provides a streamable HTTP proxy for MCP servers.
package streamable

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

        "github.com/google/uuid"
        "golang.org/x/exp/jsonrpc2"

        sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/sessionbinding"
        "github.com/stacklok/toolhive/pkg/bodylimit"
        "github.com/stacklok/toolhive/pkg/diagnostics"
        "github.com/stacklok/toolhive/pkg/healthcheck"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // StreamableHTTPEndpoint is the endpoint for streamable HTTP.
        StreamableHTTPEndpoint = "/mcp"

        // defaultRequestTimeout is the maximum time to wait for an MCP request to
        // complete. Override with TOOLHIVE_PROXY_REQUEST_TIMEOUT (e.g. "30s", "2m").
        defaultRequestTimeout = 60 * time.Second

        // proxyRequestTimeoutEnv is the environment variable that overrides the
        // default proxy request timeout.
        proxyRequestTimeoutEnv = "TOOLHIVE_PROXY_REQUEST_TIMEOUT"

        // defaultReadTimeout bounds reading the entire request (headers + body) on
        // the proxy http.Server, mitigating slow-upload connection exhaustion. It
        // does not affect responses, so SSE response streams are unaffected. Matches
        // the vMCP server default.
        defaultReadTimeout = 30 * time.Second

        // proxyChannelBufferSize is the buffer size for the messageCh and
        // responseCh channels connecting the HTTP proxy to the container runner.
        proxyChannelBufferSize = 100
)

// HTTPProxy implements a proxy for streamable HTTP transport.
type HTTPProxy struct {
        host              string
        port              int
        requestTimeout    time.Duration
        readTimeout       time.Duration
        shutdownCh        chan struct{}
        prometheusHandler http.Handler
        middlewares       []types.NamedMiddleware

        // Message channel for sending JSON-RPC to the container (from HTTP -> runner)
        messageCh chan jsonrpc2.Message
        // Response channel for receiving JSON-RPC from the container (runner -> HTTP)
        responseCh chan jsonrpc2.Message

        // Session manager for streamable HTTP sessions
        sessionManager *session.Manager

        // sessionTTL is the resolved inactivity timeout for the session manager.
        // Defaults to session.DefaultSessionTTL; overridable via WithSessionTTL.
        sessionTTL time.Duration

        // sessionStorage is the optional custom storage backend for the session manager.
        // When nil, in-memory LocalStorage is used. Set via WithSessionStorage.
        sessionStorage session.Storage

        // authInfoHandler is the optional RFC 9728 OAuth protected resource discovery handler.
        // When nil, /.well-known/ returns a clean JSON 404. Set via WithAuthInfoHandler.
        authInfoHandler http.Handler

        // prefixHandlers contains additional HTTP handlers mounted outside the middleware chain.
        // Keys are URL path prefixes. Set via WithPrefixHandlers.
        prefixHandlers map[string]http.Handler

        // Waiters keyed by compositeKey(sessID, per-request token) -> one-shot channel for response
        // delivery. The composite key MUST be unique per concurrent request; sharing it across
        // requests silently overwrites entries and crosses response payloads between requests.
        // createWaiter guarantees this by minting the token itself, because neither the session
        // ID (empty for sessionless requests) nor the client-chosen JSON-RPC ID is unique.
        waiters sync.Map // map[string]chan jsonrpc2.Message
        // Keyed by the same composite key; stores the original client JSON-RPC ID
        // to restore before replying. Same uniqueness requirement as `waiters`.
        idRestore sync.Map // map[string]jsonrpc2.ID

        // serverStreams delivers server->client notifications to connected
        // sessions' SSE streams (the standalone GET stream for global
        // notifications, or a specific session's stream for subscription-scoped
        // ones). See dispatcher_streams.go.
        serverStreams *serverStreamRegistry

        // routing owns the per-session routing state (progress tokens,
        // subscriptions, log levels) that dispatcher.go's routeNotification uses
        // to deliver server->client messages to the correct session(s) without
        // cross-session leakage. See dispatcher_routing.go.
        routing *sessionRouter

        // initialize caches the backend's InitializeResult so only the first
        // client handshake reaches the single stdio session behind this proxy and
        // every later one is answered locally. See initialize_cache.go.
        initialize *initializeCache

        // uriLocks serializes each uri's resources/subscribe|unsubscribe ref-count
        // decision with the upstream forward it may trigger, so concurrent
        // subscribe/unsubscribe calls for the SAME uri from different sessions can
        // never reach the backend out of order. See keyed_mutex.go and
        // handlePost/interceptSessionScopedRequest.
        uriLocks *keyedMutex

        // standaloneSSE controls whether GET /mcp opens a standalone SSE stream for
        // server->client notifications (the default) or returns 405. Set via
        // WithStandaloneSSE(false) to opt out.
        standaloneSSE bool

        // strictProtocolValidation controls whether handlePost rejects a request
        // whose MCP-Protocol-Version header names an unknown/unsupported MCP
        // revision (see isSupportedMCPVersion) with HTTP 400. Default false: any
        // version string is accepted, since this proxy is transport-level and does
        // not depend on a specific MCP revision. An absent header is always
        // accepted, in either mode, per the streamable HTTP spec's rule to assume
        // 2025-03-26 when the header is missing. Set via WithStrictProtocolValidation.
        strictProtocolValidation bool

        // Health checker
        healthChecker *healthcheck.HealthChecker

        server   *http.Server
        stopOnce sync.Once
}

// Option configures an HTTPProxy.
type Option func(*HTTPProxy)

// WithSessionStorage injects a custom storage backend into the session manager.
// When not provided, the proxy uses in-memory LocalStorage (single-replica default).
func WithSessionStorage(storage session.Storage) Option {
        return func(p *HTTPProxy) {
                if storage == nil {
                        return
                }
                p.sessionStorage = storage
        }
}

// WithSessionTTL overrides the session inactivity timeout used by this proxy.
// Zero or negative values are ignored so the constructor's default is preserved.
func WithSessionTTL(ttl time.Duration) Option {
        return func(p *HTTPProxy) {
                if ttl <= 0 {
                        return
                }
                p.sessionTTL = ttl
        }
}

// WithReadTimeout overrides http.Server.ReadTimeout for this proxy, which bounds
// reading the entire request (headers + body). Zero or negative values are
// ignored so the constructor's default (defaultReadTimeout) is preserved.
func WithReadTimeout(d time.Duration) Option {
        return func(p *HTTPProxy) {
                if d <= 0 {
                        return
                }
                p.readTimeout = d
        }
}

// WithAuthInfoHandler sets the handler for RFC 9728 OAuth protected resource discovery.
// When set, the handler is mounted at /.well-known/ outside the middleware chain.
// When nil, /.well-known/ returns a clean JSON 404 so OAuth clients parse it cleanly.
func WithAuthInfoHandler(h http.Handler) Option {
        return func(p *HTTPProxy) {
                p.authInfoHandler = h
        }
}

// WithPrefixHandlers registers additional HTTP handlers mounted before the MCP endpoint.
// These are mounted outside the middleware chain (RFC 9728, embedded auth server routes).
func WithPrefixHandlers(handlers map[string]http.Handler) Option {
        return func(p *HTTPProxy) {
                p.prefixHandlers = maps.Clone(handlers)
        }
}

// WithStandaloneSSE controls whether GET /mcp opens a standalone SSE stream
// for server->client notifications. It defaults to enabled; pass false to opt
// out and restore the prior behavior of returning 405 for GET requests.
func WithStandaloneSSE(enabled bool) Option {
        return func(p *HTTPProxy) {
                p.standaloneSSE = enabled
        }
}

// WithStrictProtocolValidation enables strict MCP-Protocol-Version checking.
// When enabled, a request whose MCP-Protocol-Version header names an
// unsupported/unknown MCP revision is rejected with HTTP 400. An absent
// header is still accepted (the streamable HTTP spec says to assume
// 2025-03-26). Default (false) preserves the version-agnostic behavior:
// any version string is accepted.
func WithStrictProtocolValidation(enabled bool) Option {
        return func(p *HTTPProxy) { p.strictProtocolValidation = enabled }
}

// NewHTTPProxy creates a new HTTPProxy for streamable HTTP transport.
func NewHTTPProxy(
        host string,
        port int,
        prometheusHandler http.Handler,
        middlewares []types.NamedMiddleware,
        opts ...Option,
) *HTTPProxy {
        // Use typed Streamable sessions
        sFactory := func(id string) session.Session { return session.NewStreamableSession(id) }

        proxy := &HTTPProxy{
                host:              host,
                port:              port,
                requestTimeout:    resolveRequestTimeout(),
                readTimeout:       defaultReadTimeout,
                shutdownCh:        make(chan struct{}),
                prometheusHandler: prometheusHandler,
                middlewares:       middlewares,
                messageCh:         make(chan jsonrpc2.Message, proxyChannelBufferSize),
                responseCh:        make(chan jsonrpc2.Message, proxyChannelBufferSize),
                sessionTTL:        session.DefaultSessionTTL,
                serverStreams:     newServerStreamRegistry(),
                routing:           newSessionRouter(),
                initialize:        newInitializeCache(),
                uriLocks:          newKeyedMutex(),
                standaloneSSE:     true,
        }

        for _, opt := range opts {
                opt(proxy)
        }

        // Construct the session manager once, after options have resolved sessionTTL and sessionStorage.
        if proxy.sessionStorage != nil {
                proxy.sessionManager = session.NewManagerWithStorage(proxy.sessionTTL, sFactory, proxy.sessionStorage)
        } else {
                proxy.sessionManager = session.NewManager(proxy.sessionTTL, sFactory)
        }

        // Create health checker without MCP pinger
        // Streamable transport doesn't support MCP ping, so health check only verifies proxy is running
        proxy.healthChecker = healthcheck.NewHealthChecker(string(types.TransportTypeStreamableHTTP), nil)

        return proxy
}

// Start starts the HTTPProxy server.
func (p *HTTPProxy) Start(_ context.Context) error {
        mux := http.NewServeMux()
        mux.Handle(StreamableHTTPEndpoint, p.applyMiddlewares(http.HandlerFunc(p.handleStreamableRequest)))

        // Add health check endpoint (no middlewares)
        if p.healthChecker != nil {
                mux.Handle("/health", p.healthChecker)
        }

        // Mount the metrics endpoint only if a handler is provided, otherwise return
        // 404 so /metrics is not proxied to the backend. The runner serves metrics on
        // a separate diagnostics listener instead (see pkg/diagnostics) so access can
        // be restricted by port rather than by HTTP path.
        if p.prometheusHandler != nil {
                mux.Handle("/metrics", p.prometheusHandler)
        } else {
                mux.Handle(diagnostics.MetricsPath, diagnostics.NotServedHereHandler())
        }

        // Mount prefix handlers (e.g. embedded auth server routes) outside the middleware chain.
        // RFC 9728 requires discovery endpoints to be reachable without authentication.
        for prefix, h := range p.prefixHandlers {
                mux.Handle(prefix, h)
                slog.Debug("mounted prefix handler", "prefix", prefix)
        }

        // Mount RFC 9728 OAuth protected resource discovery endpoint (no middlewares).
        // Always register so OAuth discovery gets a clean JSON 404 when auth is off.
        wellKnownHandler := auth.NewWellKnownHandler(p.authInfoHandler)
        mux.Handle("/.well-known/", wellKnownHandler)
        if p.authInfoHandler != nil {
                slog.Debug("rfc 9728 OAuth discovery endpoint enabled at /.well-known/oauth-protected-resource")
        }

        p.server = &http.Server{
                Addr:              fmt.Sprintf("%s:%d", p.host, p.port),
                Handler:           mux,
                ReadHeaderTimeout: 10 * time.Second,
                ReadTimeout:       p.readTimeout,
        }

        // Route container responses to matching waiter channels
        go p.dispatchResponses()

        // Periodically bound routing's subscription/log-level maps, and
        // serverStreams' standalone GET streams, to active sessions (see
        // reapRoutingState's doc comment for the isActive tradeoff).
        go p.reapRoutingState()

        go func() {
                slog.Debug("streamable HTTP proxy started", "port", p.port)
                //nolint:gosec // G706: logging configured host and port
                slog.Debug("streamable HTTP endpoint",
                        "url", fmt.Sprintf("http://%s:%d%s", p.host, p.port, StreamableHTTPEndpoint))
                if err := p.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
                        slog.Error("streamable HTTP server error", "error", err)
                }
        }()

        return nil
}

// Stop gracefully shuts down the HTTPProxy server.
func (p *HTTPProxy) Stop(ctx context.Context) error {
        var err error

        p.stopOnce.Do(func() {
                close(p.shutdownCh)

                // Signal every live GET SSE stream's stop channel so in-flight handlers
                // observe a closed per-stream stop signal and unblock/return promptly,
                // rather than relying solely on shutdownCh (which they also select on).
                p.serverStreams.closeAll()

                // Stop session manager cleanup; active sessions expire via TTL
                if p.sessionManager != nil {
                        if err := p.sessionManager.Stop(); err != nil {
                                slog.Error("failed to stop session manager", "error", err)
                        }
                }

                if p.server != nil {
                        if e := p.server.Shutdown(ctx); e != nil {
                                err = e
                        }
                }
        })

        return err
}

// IsRunning checks if the proxy is running.
func (p *HTTPProxy) IsRunning() (bool, error) {
        select {
        case <-p.shutdownCh:
                return false, nil
        default:
                return true, nil
        }
}

// GetMessageChannel returns the message channel for sending JSON-RPC to the container.
func (p *HTTPProxy) GetMessageChannel() chan jsonrpc2.Message {
        return p.messageCh
}

// GetResponseChannel returns the response channel for receiving JSON-RPC from the container.
func (p *HTTPProxy) GetResponseChannel() <-chan jsonrpc2.Message {
        return p.responseCh
}

// SendMessageToDestination sends a message to the container.
func (p *HTTPProxy) SendMessageToDestination(msg jsonrpc2.Message) error {
        select {
        case p.messageCh <- msg:
                return nil
        default:
                return fmt.Errorf("failed to send message to destination")
        }
}

// ForwardResponseToClients forwards a response from the container to the client.
func (p *HTTPProxy) ForwardResponseToClients(_ context.Context, msg jsonrpc2.Message) error {
        select {
        case p.responseCh <- msg:
                return nil
        default:
                return fmt.Errorf("failed to forward response to client")
        }
}

// SendResponseMessage is for compatibility with the Proxy interface.
func (p *HTTPProxy) SendResponseMessage(msg jsonrpc2.Message) error {
        return p.ForwardResponseToClients(context.Background(), msg)
}

// ------------------------- HTTP handlers -------------------------

// handleStreamableRequest handles HTTP POST requests to /mcp.
func (p *HTTPProxy) handleStreamableRequest(w http.ResponseWriter, r *http.Request) {
        switch r.Method {
        case http.MethodGet:
                p.handleGet(w, r)
        case http.MethodDelete:
                p.handleDelete(w, r)
        case http.MethodPost:
                p.handlePost(w, r)
        default:
                writeHTTPError(w, http.StatusMethodNotAllowed, "Method not allowed")
        }
}

// handleGet serves a standalone GET SSE stream for unsolicited server->client
// notifications (progress, resources/updated, */list_changed, logging/message).
// It is enabled by default; WithStandaloneSSE(false) restores the prior 405
// behavior. The caller MUST provide a known Mcp-Session-Id header: an empty
// header is rejected with 400, and an unknown one with 404 (mirroring
// handleDelete). A session's stream is one-per-session: a second GET for the
// same session evicts the first (see serverStreamRegistry.register). Per the
// locked product decision, a notification dispatched while no GET stream is
// connected for a session is dropped -- there is no replay or pending queue
// for this stream.
func (p *HTTPProxy) handleGet(w http.ResponseWriter, r *http.Request) {
        if !p.standaloneSSE {
                // SSE not offered here; explicit 405 is spec-compliant
                writeHTTPError(w, http.StatusMethodNotAllowed, "SSE not supported on this endpoint")
                return
        }

        flusher, ok := assertFlushable(w)
        if !ok {
                return
        }

        sessID, carrierErr := sessionbinding.RequestID(r, "")
        if carrierErr != nil {
                sessionbinding.WriteOwnershipError(w, nil, carrierErr)
                return
        }
        if sessID == "" {
                writeHTTPError(w, http.StatusBadRequest, "Mcp-Session-Id header required for standalone SSE")
                return
        }
        if err := sessionbinding.Check(r, sessID, p.sessionManager.LookupOwner); err != nil {
                sessionbinding.WriteOwnershipError(w, nil, err)
                return
        }

        setSSEHeaders(w)

        s := p.serverStreams.register(sessID)
        defer p.serverStreams.deregister(sessID, s)

        // Send headers immediately so the client knows the stream is open.
        flusher.Flush()

        keepAliveTicker := time.NewTicker(sseKeepAliveInterval)
        defer keepAliveTicker.Stop()

        for {
                select {
                case msg := <-s.data:
                        // No ok-check: data is never closed (see serverStream's doc comment).
                        data, err := jsonrpc2.EncodeMessage(msg)
                        if err != nil {
                                slog.Error("failed to encode server->client notification", "error", err)
                                continue
                        }
                        if err := writeSSEData(w, flusher, data); err != nil {
                                slog.Debug("failed to write notification to GET SSE stream", "error", err)
                                return
                        }
                case <-keepAliveTicker.C:
                        if err := writeSSEKeepAlive(w, flusher); err != nil {
                                slog.Debug("failed to write keep-alive", "error", err)
                                return
                        }
                case <-s.stop:
                        // Evicted by a second GET for this session, or proxy Stop's closeAll.
                        return
                case <-r.Context().Done():
                        return
                case <-p.shutdownCh:
                        return
                }
        }
}

func (p *HTTPProxy) handleDelete(w http.ResponseWriter, r *http.Request) {
        sessID, carrierErr := sessionbinding.RequestID(r, "")
        if carrierErr != nil {
                sessionbinding.WriteOwnershipError(w, nil, carrierErr)
                return
        }
        if sessID == "" {
                writeHTTPError(w, http.StatusBadRequest, "Mcp-Session-Id header required for DELETE")
                return
        }
        if err := sessionbinding.Check(r, sessID, p.sessionManager.LookupOwner); err != nil {
                sessionbinding.WriteOwnershipError(w, nil, err)
                return
        }
        identity, _ := auth.IdentityFromContext(r.Context())
        expectedOwner, err := sessionbinding.FromIdentity(identity)
        if err != nil {
                sessionbinding.WriteOwnershipError(w, nil, sessionbinding.ErrNotFound)
                return
        }
        if err := p.sessionManager.DeleteIfOwner(sessID, expectedOwner); err != nil {
                sessionbinding.WriteOwnershipError(w, nil, err)
                return
        }

        // Purge routing state (subscriptions, log level) for the deleted session so
        // it cannot outlive the session itself. maxChanged is intentionally not
        // reconciled upstream here: if the deleted session held the max-verbosity
        // log level, the shared backend simply stays MORE verbose than strictly
        // necessary for the remaining sessions until one of them changes the level
        // again -- extra log volume, not a correctness or security issue, so it is
        // not worth an extra upstream round-trip on every DELETE.
        p.routing.purgeSession(sessID)

        // Tear down the session's standalone GET SSE stream (if any), so its
        // handler goroutine (and the list_changed/subscription broadcasts it would
        // otherwise keep receiving) stops immediately instead of surviving until
        // the client eventually closes the underlying socket on its own (see FIX 2
        // in #5744's review). Ordering matches the durable-before-in-memory style
        // rule: the durable session delete and routing purge above have already
        // happened, so nothing durable or in routing state can reference this
        // session by the time its stream is closed.
        p.serverStreams.closeStream(sessID)

        w.WriteHeader(http.StatusNoContent)
}

func (p *HTTPProxy) handlePost(w http.ResponseWriter, r *http.Request) {

        // MCP-Protocol-Version validation is opt-in via strictProtocolValidation
        // (WithStrictProtocolValidation). Default (false) is version-agnostic:
        // any header value is accepted, since this proxy is transport-level and
        // does not depend on a specific MCP revision. When strict mode is
        // enabled, a present-but-unrecognized version is rejected with 400; an
        // absent header is always accepted in either mode, per the streamable
        // HTTP spec's rule to assume 2025-03-26 when the header is missing.
        protoVer := r.Header.Get("MCP-Protocol-Version")
        if p.strictProtocolValidation && protoVer != "" && !isSupportedMCPVersion(protoVer) {
                writeHTTPError(w, http.StatusBadRequest, "Unsupported MCP-Protocol-Version")
                return
        }

        // Read request body
        body, err := io.ReadAll(r.Body)
        if err != nil {
                // A body that exceeds the configured limit without a Content-Length
                // (e.g. chunked) trips http.MaxBytesReader here rather than at the
                // early Content-Length check. Surface it as 413, not 500.
                if bodylimit.IsRequestTooLarge(err) {
                        slog.Warn("rejected unreadable MCP request body", "reason", "body_too_large")
                        writeHTTPError(w, http.StatusRequestEntityTooLarge, "Request Entity Too Large")
                        return
                }
                slog.Warn("rejected unreadable MCP request body", "reason", "read_error")
                writeHTTPError(w, http.StatusBadRequest, "Invalid request body")
                return
        }

        msg, ok := decodeJSONRPCMessage(w, body)
        if !ok {
                return
        }

        p.handlePostMessage(w, r, msg)
}

func (p *HTTPProxy) handlePostMessage(w http.ResponseWriter, r *http.Request, msg jsonrpc2.Message) {
        ctx := r.Context()
        var legacySID string
        // Notifications and client responses are Legacy session-bearing traffic.
        // Id-bearing Modern requests are classified below and intentionally ignore SID.
        if req, ok := msg.(*jsonrpc2.Request); !ok || !req.ID.IsValid() {
                sid, err := sessionbinding.RequestID(r, "")
                if err != nil {
                        sessionbinding.WriteOwnershipError(w, nil, err)
                        return
                }
                if sid == "" {
                        writeHTTPError(w, http.StatusBadRequest, "Mcp-Session-Id header required")
                        return
                }
                if err := sessionbinding.Check(r, sid, p.sessionManager.LookupOwner); err != nil {
                        sessionbinding.WriteOwnershipError(w, nil, err)
                        return
                }
                legacySID = sid
        }
        // Notifications or client responses are accepted and forwarded (202)
        if p.handleNotificationOrClientResponse(w, legacySID, msg) {
                return
        }

        req, ok := msg.(*jsonrpc2.Request)
        if !ok || !req.ID.IsValid() {
                writeHTTPError(w, http.StatusBadRequest, "Invalid JSON-RPC request (missing id)")
                return
        }

        // Resolve session per spec (initialize vs ordinary)
        sessID, setSessionHeader, err := p.resolveSessionForRequest(w, r, req)
        if err != nil {
                return
        }

        // Intercept resources/subscribe, resources/unsubscribe, and
        // logging/setLevel for session-bearing (Legacy) requests, AFTER
        // applyMiddlewares/authz has admitted req (Start wraps handleStreamableRequest
        // in the middleware chain, and handlePost is only ever reached through it),
        // so an authz-denied request never gets recorded in routing.
        //
        // For resources/subscribe|unsubscribe, uriLocks serializes the ref-count
        // decision made inside interceptSessionScopedRequest WITH the upstream
        // forward it may perform itself (see interceptSubscribe/
        // interceptUnsubscribe), per uri, so a concurrent last-unsubscribe and
        // first-subscribe for the SAME uri from different sessions can never reach
        // the backend out of order (see uriLocks's doc comment and FIX 1 in
        // #5744's review). The lock is scoped to ONLY this decision+forward+record
        // -- acquired immediately before interceptSessionScopedRequest and released
        // immediately after it returns, BEFORE writeInterceptedResponse or any
        // client-facing streaming -- so a slow or SSE-streamed response never
        // head-of-line-blocks a different session's request for the SAME uri (see
        // FIX 3 in #5744's review). A different uri, or any non-subscription
        // request, is never blocked by this at all: see resourceSubscriptionURI.
        interceptedMsg, handled := func() (jsonrpc2.Message, bool) {
                uri, ok := resourceSubscriptionURI(req)
                if !ok {
                        return p.interceptSessionScopedRequest(ctx, sessID, req)
                }
                unlock := p.uriLocks.lock(uri)
                defer unlock()
                return p.interceptSessionScopedRequest(ctx, sessID, req)
        }()
        if handled {
                writeInterceptedResponse(w, r, interceptedMsg, sessID, setSessionHeader)
                return
        }

        // If client accepts SSE, stream the response on an SSE stream for this request
        if strings.Contains(r.Header.Get("Accept"), "text/event-stream") {
                p.handleSingleRequestSSE(ctx, w, sessID, req, setSessionHeader)
                return
        }

        // Request/response path with correlation (JSON response)
        p.handleSingleRequest(ctx, w, sessID, req, setSessionHeader)
}

// handleSingleRequest handles a single JSON-RPC request message end-to-end.
//
// If req carries a params._meta.progressToken, it is forwarded upstream
// UNREWRITTEN: a plain JSON POST response has no stream to carry interim
// progress on, so there is nothing to correlate a progress route to, and no
// progress routing is set up (see handleSingleRequestSSE for the SSE path,
// which does). Any notifications/progress the backend later sends for this
// request's original token will simply find no registered route in
// dispatcher.go's routeProgress and be dropped, which is correct: this
// non-streaming client could never have received them anyway.
func (p *HTTPProxy) handleSingleRequest(
        ctx context.Context,
        w http.ResponseWriter,
        sessID string,
        req *jsonrpc2.Request,
        setSessionHeader bool,
) {
        ctx, cancel := context.WithTimeout(ctx, p.requestTimeout)
        defer cancel()

        msg, err := p.doRequest(ctx, sessID, req)
        if err != nil {
                if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
                        //nolint:gosec // G706: method is from parsed JSON-RPC request
                        slog.Warn("timeout waiting for response", "method", req.Method)
                        writeHTTPError(w, http.StatusGatewayTimeout, "Timeout waiting for response from container")
                } else {
                        //nolint:gosec // G706: method is from parsed JSON-RPC request
                        slog.Error("failed to process request", "method", req.Method, "error", err)
                        writeHTTPError(w, http.StatusInternalServerError, "Failed to process request")
                }
                return
        }

        if setSessionHeader {
                w.Header().Set("Mcp-Session-Id", sessID)
        }
        if err := writeJSONRPC(w, msg); err != nil {
                slog.Error("failed to write JSON-RPC response", "error", err)
        }
}

// handleSingleRequestSSE handles a single JSON-RPC request whose client asked
// for an SSE response (Accept: text/event-stream). Unlike handleSingleRequest,
// this stream stays open for the lifetime of the request, so it is the
// correct destination for any interim notifications/progress the backend
// sends while it works -- per MCP, request-related messages belong on the
// originating request's stream, not the standalone GET stream (see
// dispatcher.go's routeProgress and docs/arch/03-transport-architecture.md).
//
// If req carries a params._meta.progressToken, setupProgressRouting mints a
// proxy-only token, rewrites the outgoing request to carry it, and registers a
// delivery route for it; the select loop below then interleaves any progress
// notifications delivered on that route with the final response, writing each
// as its own SSE data: frame, and only returns once the final response (or a
// context/shutdown signal) arrives.
func (p *HTTPProxy) handleSingleRequestSSE(
        ctx context.Context,
        w http.ResponseWriter,
        sessID string,
        req *jsonrpc2.Request,
        setSessionHeader bool,
) {
        ctx, cancel := context.WithTimeout(ctx, p.requestTimeout)
        defer cancel()

        // Prepare SSE response headers
        var setMcpSessionID func()
        if setSessionHeader {
                setMcpSessionID = func() { w.Header().Set("Mcp-Session-Id", sessID) }
        }
        flusher, ok := startSSEStream(w, setMcpSessionID)
        if !ok {
                return
        }

        outgoingReq, deliverCh, dropProgressRoute := p.setupProgressRouting(req)
        defer dropProgressRoute()

        waitCh, ck, cleanup, err := p.sendCorrelatedRequest(sessID, outgoingReq)
        if err != nil {
                //nolint:gosec // G706: method is from parsed JSON-RPC request
                slog.Error("failed to send request upstream", "method", req.Method, "error", err)
                writeSSEErrorEvent(w, flusher, req.ID, err)
                return
        }
        defer cleanup()

        for {
                select {
                case progressMsg := <-deliverCh:
                        // deliverCh is nil when req carried no progressToken (see
                        // setupProgressRouting); a nil channel's case is never selected,
                        // so this arm only fires when progress routing is active.
                        if !writeSSEProgressMessage(w, flusher, progressMsg) {
                                return
                        }
                        // Progress does not end the request; keep waiting for more
                        // progress or the final response.
                case msg := <-waitCh:
                        // dispatchResponses preserves backend message order. Therefore, any
                        // progress sent before this final response is already buffered in
                        // deliverCh and must be emitted before returning from this stream.
                        if !drainPendingSSEProgress(w, flusher, deliverCh) {
                                return
                        }
                        p.writeSingleRequestSSEFinalResponse(w, flusher, msg, ck)
                        return
                case <-ctx.Done():
                        writeSSEErrorEvent(w, flusher, req.ID, ctx.Err())
                        return
                case <-p.shutdownCh:
                        // Matches the prior (pre-progress) behavior: proxy shutdown while a
                        // request is in flight is reported to the client as a best-effort
                        // Timeout error event, the same as context.Canceled.
                        writeSSEErrorEvent(w, flusher, req.ID, context.Canceled)
                        return
                }
        }
}

// writeSSEProgressMessage encodes progressMsg and writes it as one SSE data
// frame. It reports whether the stream remains writable.
func writeSSEProgressMessage(w io.Writer, flusher http.Flusher, progressMsg jsonrpc2.Message) bool {
        data, err := jsonrpc2.EncodeMessage(progressMsg)
        if err != nil {
                slog.Error("failed to encode progress notification", "error", err)
                return true
        }
        if err := writeSSEData(w, flusher, data); err != nil {
                slog.Debug("failed to write progress notification to SSE stream", "error", err)
                return false
        }
        return true
}

// drainPendingSSEProgress writes progress messages that were already routed to
// deliverCh before a final response is emitted. It reports whether the stream
// remains writable.
func drainPendingSSEProgress(w io.Writer, flusher http.Flusher, deliverCh <-chan jsonrpc2.Message) bool {
        for {
                select {
                case progressMsg, ok := <-deliverCh:
                        if !ok {
                                return true
                        }
                        if !writeSSEProgressMessage(w, flusher, progressMsg) {
                                return false
                        }
                default:
                        return true
                }
        }
}

// writeSingleRequestSSEFinalResponse restores msg's original client ID (if it
// is a correlated *jsonrpc2.Response) and writes it as the final SSE data:
// frame for handleSingleRequestSSE's request. Errors encoding/restoring are
// logged; the client simply does not get a final frame in that case, matching
// the prior (pre-progress) behavior's error handling.
func (p *HTTPProxy) writeSingleRequestSSEFinalResponse(
        w http.ResponseWriter, flusher http.Flusher, msg jsonrpc2.Message, ck string,
) {
        finalMsg := msg
        if r, ok := msg.(*jsonrpc2.Response); ok && r.ID.IsValid() {
                restored, err := p.restoreResponseID(r, ck)
                if err != nil {
                        slog.Error("failed to restore response id", "error", err)
                        return
                }
                finalMsg = restored
        }

        data, err := jsonrpc2.EncodeMessage(finalMsg)
        if err != nil {
                slog.Error("failed to encode JSON-RPC response", "error", err)
                return
        }
        if err := writeSSEData(w, flusher, data); err != nil {
                slog.Debug("failed to write response", "error", err)
        }
}

// writeSSEErrorEvent writes a best-effort JSON-RPC error as a single SSE
// message event (via writeSSEData), for a request whose response headers
// (200 + text/event-stream) have already been sent -- an HTTP error status can
// no longer be set at this point, so the error must be communicated in-band as
// the response payload.
//
// The "id" key is included only if id.IsValid(); MCP narrows base JSON-RPC
// 2.0 here (schema/2025-11-25 types the error response id as optional,
// string|number, no null), so an invalid/absent id must be encoded by
// omitting the key, never as "id":null.
func writeSSEErrorEvent(w http.ResponseWriter, flusher http.Flusher, id jsonrpc2.ID, err error) {
        errMsg := "Internal error"
        code := -32603
        if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
                errMsg = "Timeout"
                code = -32000
        }
        errObj := map[string]any{
                "jsonrpc": "2.0",
                "error": map[string]any{
                        "code":    code,
                        "message": errMsg,
                },
        }
        if id.IsValid() {
                errObj["id"] = id.Raw()
        }
        data, mErr := json.Marshal(errObj)
        if mErr != nil {
                slog.Error("failed to encode SSE error event", "error", mErr)
                return
        }
        if err := writeSSEData(w, flusher, data); err != nil {
                slog.Debug("failed to write error message", "error", err)
        }
}

func encodeRequestWithID(req *jsonrpc2.Request, newID string) (jsonrpc2.Message, error) {
        data, err := jsonrpc2.EncodeMessage(req)
        if err != nil {
                return nil, err
        }
        var m map[string]any
        if err := json.Unmarshal(data, &m); err != nil {
                return nil, err
        }
        m["id"] = newID
        data2, err := json.Marshal(m)
        if err != nil {
                return nil, err
        }
        return jsonrpc2.DecodeMessage(data2)
}

// progressDeliverBufferSize bounds how many not-yet-written progress
// notifications may queue for a single in-flight POST-SSE request before
// further ones are dropped (see routeProgress's non-blocking send).
const progressDeliverBufferSize = 16

// rewriteMetaProgressToken returns a request identical to req except that, if
// params._meta.progressToken is present, it is replaced with ptGlobal. This
// lets dispatcher.go's routeProgress correlate the backend's later
// notifications/progress replies (which echo the token verbatim) back to the
// specific POST-SSE request that asked for them, without ever exposing
// ptGlobal -- or any other session's progress -- to more than the one client
// that requested it.
//
// It returns hadToken=false (with rewritten == req, unchanged) if req carried
// no progressToken, telling the caller not to set up progress routing at all.
// originalToken is the client's own progressToken value (string or number),
// to be restored before delivering progress back to the client.
//
// Per the "copy before mutating caller input" style rule, req's params map is
// decoded into fresh copies before any mutation, so req itself (and any
// caller-held reference to it) is never modified.
func rewriteMetaProgressToken(
        req *jsonrpc2.Request, ptGlobal string,
) (originalToken any, hadToken bool, rewritten *jsonrpc2.Request, err error) {
        if len(req.Params) == 0 {
                return nil, false, req, nil
        }
        var paramsMap map[string]any
        if err := json.Unmarshal(req.Params, &paramsMap); err != nil {
                return nil, false, nil, err
        }
        meta, ok := paramsMap["_meta"].(map[string]any)
        if !ok {
                return nil, false, req, nil
        }
        original, ok := meta["progressToken"]
        if !ok {
                return nil, false, req, nil
        }

        metaCopy := maps.Clone(meta)
        metaCopy["progressToken"] = ptGlobal
        paramsCopy := maps.Clone(paramsMap)
        paramsCopy["_meta"] = metaCopy

        data, err := json.Marshal(paramsCopy)
        if err != nil {
                return nil, false, nil, err
        }
        return original, true, &jsonrpc2.Request{ID: req.ID, Method: req.Method, Params: data}, nil
}

// setupProgressRouting mints a fresh proxy-only progress-routing token and
// rewrites req's params._meta.progressToken to it (see
// rewriteMetaProgressToken), registering a progressRoute so the backend's
// later notifications/progress messages that echo the token are delivered on
// the returned channel rather than leaking to the standalone GET stream or
// another session.
//
// If req carries no progressToken (or minting/rewriting fails, logged as a
// warning), the returned request is req itself, unmodified, and the returned
// channel is nil: a select on a nil channel never fires, so callers can select
// on it unconditionally without special-casing "no progress requested".
//
// The returned cleanup func drops the progress token and MUST be deferred by
// the caller; it is always safe to call, including when no route was
// registered.
func (p *HTTPProxy) setupProgressRouting(req *jsonrpc2.Request) (*jsonrpc2.Request, <-chan jsonrpc2.Message, func()) {
        noop := func() {}

        ptGlobal, err := uuid.NewRandom()
        if err != nil {
                slog.Warn("failed to mint progress routing token; interim progress will not be delivered for this request",
                        "error", err)
                return req, nil, noop
        }

        originalToken, hadToken, rewritten, err := rewriteMetaProgressToken(req, ptGlobal.String())
        if err != nil {
                slog.Warn("failed to rewrite progress token; interim progress will not be delivered for this request",
                        "error", err)
                return req, nil, noop
        }
        if !hadToken {
                return req, nil, noop
        }

        key := ptGlobal.String()
        deliver := make(chan jsonrpc2.Message, progressDeliverBufferSize)
        p.routing.recordProgressToken(key, progressRoute{deliver: deliver, originalToken: originalToken})
        return rewritten, deliver, func() { p.routing.dropProgressToken(key) }
}

func (p *HTTPProxy) restoreResponseID(resp *jsonrpc2.Response, ck string) (jsonrpc2.Message, error) {
        orig, ok := p.idRestore.Load(ck)
        if !ok {
                // No restore information; return as-is
                return resp, nil
        }
        origID, _ := orig.(jsonrpc2.ID)

        data, err := jsonrpc2.EncodeMessage(resp)
        if err != nil {
                return nil, err
        }
        var m map[string]any
        if err := json.Unmarshal(data, &m); err != nil {
                return nil, err
        }
        m["id"] = origID.Raw()
        data2, err := json.Marshal(m)
        if err != nil {
                return nil, err
        }
        return jsonrpc2.DecodeMessage(data2)
}

// sendCorrelatedRequest mints the composite-key waiter for req, rewrites its
// wire ID to that composite key, and sends it upstream. It returns the waiter
// channel and ck (needed to restore the original ID on the eventual response,
// see restoreResponseID), plus a cleanup func the caller MUST call (typically
// via defer) to release the waiter/idRestore entries once done waiting.
//
// This is split out from doRequest so callers that need to interleave OTHER
// events (e.g. handleSingleRequestSSE's progress notifications) with the wait
// for the final response can select over both, instead of being limited to
// doRequest's single blocking wait.
func (p *HTTPProxy) sendCorrelatedRequest(
        sessID string, req *jsonrpc2.Request,
) (waitCh chan jsonrpc2.Message, ck string, cleanup func(), err error) {
        waitCh, ck, cleanup, err = p.createWaiter(sessID, req.ID)
        if err != nil {
                return nil, "", nil, err
        }

        proxiedMsg, err := encodeRequestWithID(req, ck)
        if err != nil {
                cleanup()
                return nil, "", nil, fmt.Errorf("encode request: %w", err)
        }
        if err := p.SendMessageToDestination(proxiedMsg); err != nil {
                cleanup()
                return nil, "", nil, fmt.Errorf("send message: %w", err)
        }
        return waitCh, ck, cleanup, nil
}

func (p *HTTPProxy) doRequest(ctx context.Context, sessID string, req *jsonrpc2.Request) (jsonrpc2.Message, error) {
        waitCh, ck, cleanup, err := p.sendCorrelatedRequest(sessID, req)
        if err != nil {
                return nil, err
        }
        defer cleanup()

        select {
        case msg := <-waitCh:
                if r, ok := msg.(*jsonrpc2.Response); ok && r.ID.IsValid() {
                        restored, err := p.restoreResponseID(r, ck)
                        if err != nil {
                                return nil, fmt.Errorf("restore id: %w", err)
                        }
                        return restored, nil
                }
                return msg, nil
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-p.shutdownCh:
                return nil, context.Canceled
        }
}

// ------------------------- Helpers: middleware, parsing, correlation -------------------------

func (p *HTTPProxy) applyMiddlewares(handler http.Handler) http.Handler {
        for i := len(p.middlewares) - 1; i >= 0; i-- {
                handler = p.middlewares[i].Function(handler)
        }
        return handler
}

func (p *HTTPProxy) ensureSession(r *http.Request, id string) error {
        if _, err := uuid.Parse(id); err != nil {
                return sessionbinding.ErrNotFound
        }
        return sessionbinding.AddOwnedSession(r.Context(), p.sessionManager, session.NewProxySession(id))
}

// resolveSessionForRequest resolves session rules for a single JSON-RPC request.
//
// It first classifies the request as Modern or Legacy MCP (mcp.ClassifyRevision).
// A classification error is rejected outright with an HTTP 400 JSON-RPC error
// response. A Modern request is always sessionless: it gets a fresh per-request
// routing token and never touches sessionManager, regardless of any
// Mcp-Session-Id header it carries (see the confidentiality note below). A
// Legacy request falls through to the existing session rules: on initialize,
// assigns a new session ID if none is provided and returns setSessionHeader=true;
// a provided but unknown session ID returns 404.
//
// Sessionless requests (both the Modern branch and Legacy's sessionless
// non-initialize branch) receive a per-request UUID used solely as an
// in-process routing token (not registered with sessionManager). Sessionless
// routing tokens MUST be unique per request: sharing one (e.g. the empty string,
// or a stale/foreign Mcp-Session-Id) across concurrent sessionless requests with
// the same JSON-RPC id could otherwise share routing state and leak one
// client's response payload to another. This is a confidentiality bug, not a
// performance issue -- do not collapse the token, and do not let a Modern
// request fall through to the session-lookup path below. (The waiter key
// itself is additionally made unique per request by createWaiter, so two
// requests in one session reusing an id are isolated too.)
//
// Writes HTTP errors on failure and returns error to stop handling.
func (p *HTTPProxy) resolveSessionForRequest(
        w http.ResponseWriter,
        r *http.Request,
        req *jsonrpc2.Request,
) (string, bool, error) {
        // Classification/routing here applies only to the single id-bearing request
        // path; batches are rejected earlier in handlePost, and the
        // notification/client-response path is Legacy-only by construction.
        meta := mcp.ExtractMeta(req.Params)
        protoHeader := r.Header.Get("MCP-Protocol-Version")
        rev, err := mcp.ClassifyRevision(req.Method, meta, protoHeader)
        if err != nil {
                mcp.WriteClassificationError(w, req.ID.Raw(), err)
                return "", false, err
        }

        if rev == mcp.RevisionModern {
                // Modern is stateless: mint a fresh routing token unconditionally and
                // ignore any client-supplied Mcp-Session-Id. Never fall through to the
                // session-lookup path below with it -- see the confidentiality note
                // in this function's doc comment.
                //
                // This is safe only because, in this proxy, Mcp-Session-Id has no
                // downstream authority: it is purely an in-process response-correlation
                // token (the backend is stateless stdio, unlike the transparent proxy,
                // which maps it to a backend-assigned SID via session metadata). If this
                // proxy ever gains a shared session store or a stateful backend mapping,
                // "ignore the client SID on Modern" must stay true -- silently reusing or
                // looking up a known client SID here would reopen a session-validation
                // bypass. See TestModernNeverReusesClientSessionIDAsRoutingToken.
                token, err := uuid.NewRandom()
                if err != nil {
                        writeHTTPError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to generate routing token: %v", err))
                        return "", false, fmt.Errorf("generate routing token: %w", err)
                }
                return token.String(), false, nil
        }

        var setSessionHeader bool
        sessID, carrierErr := sessionbinding.RequestID(r, "")
        if carrierErr != nil {
                sessionbinding.WriteOwnershipError(w, req.ID.Raw(), carrierErr)
                return "", false, carrierErr
        }

        if req.Method == "initialize" {
                if sessID == "" {
                        newID, err := uuid.NewRandom()
                        if err != nil {
                                writeHTTPError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to generate session ID: %v", err))
                                return "", false, fmt.Errorf("generate session ID: %w", err)
                        }
                        sessID = newID.String()
                        setSessionHeader = true
                }
                if err := p.ensureSession(r, sessID); err != nil {
                        sessionbinding.WriteOwnershipError(w, req.ID.Raw(), err)
                        return "", false, err
                }
                return sessID, setSessionHeader, nil
        }

        // Sessionless non-initialize: generate a per-request routing token.
        // setSessionHeader stays false so the client never sees this UUID and the
        // next request remains sessionless.
        if sessID == "" {
                token, err := uuid.NewRandom()
                if err != nil {
                        writeHTTPError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to generate routing token: %v", err))
                        return "", false, fmt.Errorf("generate routing token: %w", err)
                }
                return token.String(), false, nil
        }

        if err := sessionbinding.Check(r, sessID, p.sessionManager.LookupOwner); err != nil {
                sessionbinding.WriteOwnershipError(w, req.ID.Raw(), err)
                return "", false, err
        }
        return sessID, false, nil
}

// resourceSubscriptionURI returns req's uri param and true if req is a
// resources/subscribe or resources/unsubscribe request carrying a non-empty
// uri -- the only requests whose ref-count decision and upstream forward (see
// interceptSubscribe/interceptUnsubscribe, called from
// interceptSessionScopedRequest) handlePost must serialize per-uri via
// uriLocks. Any other request -- including logging/setLevel (whose
// reconciliation is best-effort-ordered, see reconcileUpstreamLogLevel) and a
// subscribe/unsubscribe with no/empty uri (which interceptSessionScopedRequest
// itself no-ops on) -- returns ("", false), so handlePost skips locking
// entirely for it.
func resourceSubscriptionURI(req *jsonrpc2.Request) (string, bool) {
        switch req.Method {
        case methodResourcesSubscribe, methodResourcesUnsubscribe:
                uri, ok := extractStringParam(req.Params, "uri")
                if !ok || uri == "" {
                        return "", false
                }
                return uri, true
        default:
                return "", false
        }
}

// interceptSessionScopedRequest applies ref-counted resources/subscribe and
// resources/unsubscribe forwarding, and logging/setLevel max-verbosity
// reconciliation, for a request that belongs to a real, persisted Legacy
// session. It MUST be called only after applyMiddlewares/authz has admitted
// req (handlePost is only ever reached through Start's applyMiddlewares-
// wrapped mux entry), so an authz-denied subscribe is never recorded in
// routing -- see the "authz non-bypass" security test.
//
// For resources/subscribe|unsubscribe, the caller (handlePost) holds uriLocks
// for req's uri around this call, scoped to ONLY this call (see FIX 3 in
// #5744's review): interceptSubscribe/interceptUnsubscribe perform their own
// upstream round-trip (via p.doRequest, bounded by p.requestTimeout so the
// lock is never held indefinitely) synchronously, INSIDE this call, rather
// than reporting "forward it yourself" back to handlePost -- this is what
// lets the ref-count decision, the forward, and recording the subscription be
// atomic with respect to any other session's concurrent (un)subscribe for the
// SAME uri (see resourceSubscriptionURI and FIX 1 in #5744's review).
//
// sessID identifies "a real session" (as opposed to a per-request routing
// token minted for a Modern or Legacy-sessionless request, see
// resolveSessionForRequest) by whether sessionManager currently has it: only
// real sessions are ever stored there, so this single check both selects the
// right requests AND naturally excludes ones with nothing durable to track.
//
// It returns (msg, true) when the caller should write msg to the client and
// nothing further needs to happen: msg may be a synthesized success (a
// non-first subscribe, a non-last unsubscribe, or logging/setLevel), or the
// real upstream response (success OR error) for a first subscribe/last
// unsubscribe that this call forwarded itself. It returns (nil, false) only
// when the method isn't one of the three above, or the request isn't
// session-bearing -- in which case the caller must proceed with the normal
// upstream request/response flow.
func (p *HTTPProxy) interceptSessionScopedRequest(
        ctx context.Context, sessID string, req *jsonrpc2.Request,
) (jsonrpc2.Message, bool) {
        if _, isSession := p.sessionManager.Get(sessID); !isSession {
                return nil, false
        }

        switch req.Method {
        case string(sdkmcp.MethodInitialize):
                // Served from the shared InitializeResult cache after the first
                // handshake, so the single stdio session behind this proxy never
                // receives a second initialize. See initialize_cache.go.
                return p.interceptInitialize(ctx, sessID, req), true

        case methodResourcesSubscribe:
                uri, ok := extractStringParam(req.Params, "uri")
                if !ok || uri == "" {
                        return nil, false
                }
                return p.interceptSubscribe(ctx, sessID, uri, req), true

        case methodResourcesUnsubscribe:
                uri, ok := extractStringParam(req.Params, "uri")
                if !ok || uri == "" {
                        return nil, false
                }
                return p.interceptUnsubscribe(ctx, sessID, uri, req), true

        case string(sdkmcp.MethodSetLogLevel):
                // Unlike resources/subscribe|unsubscribe, this decision is NOT ordered
                // with its upstream forward via uriLocks (there is no per-uri key to
                // order on, and no uriLocks acquisition happens for this method -- see
                // resourceSubscriptionURI). Concurrent maxChanged transitions from
                // different sessions may therefore reach the backend out of order; see
                // reconcileUpstreamLogLevel's doc comment for why that is an accepted,
                // harmless tradeoff for now. The client's own synthesized success below
                // is likewise a fire-and-forget, optimistic ack -- like subscribe's
                // dedup path, it is served without waiting for reconcileUpstreamLogLevel
                // to complete, which is intentional and low-stakes: notifications/message
                // is unconditionally dropped (see routeNotification's SECURE-DROP), so
                // there is no client-visible effect of the level change to get wrong.
                level, _ := extractStringParam(req.Params, "level")
                if maxChanged, newMax := p.routing.setLogLevel(sessID, logLevelRank(level)); maxChanged {
                        p.reconcileUpstreamLogLevel(logLevelName(newMax))
                }
                return synthesizeSuccessResponse(req.ID), true

        default:
                return nil, false
        }
}

// interceptSubscribe handles a session-bearing resources/subscribe request
// for uri, called from interceptSessionScopedRequest while the caller
// (handlePost) holds uriLocks for uri. If uri already has at least one
// recorded subscriber, the backend has already admitted the FIRST subscribe
// and is sending updates for uri, so sess is simply recorded as an
// additional subscriber and a synthesized success is returned WITHOUT another
// upstream call -- dedup is safe here because the thing being deduped against
// is a call that is known to have succeeded.
//
// Otherwise sess is (candidate) first subscriber: the request is forwarded
// upstream via forwardUpstream, and sess is recorded as a subscriber ONLY if
// that forward succeeds with a non-error response. If the backend errors,
// times out, or the transport call itself fails, sess is NOT recorded and the
// real failure is returned to the client -- recording on failure would leave
// a phantom subscription entry that permanently starves every later
// session's subscribe for the SAME uri (deduping it against an upstream
// subscribe the backend never actually granted), and would violate the
// "write durable(upstream) storage before updating in-memory state" style
// rule. See FIX 1 in #5744's review.
func (p *HTTPProxy) interceptSubscribe(ctx context.Context, sessID, uri string, req *jsonrpc2.Request) jsonrpc2.Message {
        if len(p.routing.subscribersOf(uri)) > 0 {
                p.routing.addSubscription(uri, sessID)
                return synthesizeSuccessResponse(req.ID)
        }

        resp, err := p.forwardUpstream(ctx, sessID, req)
        if err != nil {
                return upstreamErrorResponse(req.ID, err)
        }
        if isErrorResponse(resp) {
                // The backend rejected the first subscribe (e.g. resource not found):
                // do not record. A later session's subscribe for the SAME uri must be
                // free to try again upstream, not be dedup-served a fake success.
                return resp
        }
        p.routing.addSubscription(uri, sessID)
        return resp
}

// interceptUnsubscribe handles a session-bearing resources/unsubscribe
// request for uri, called from interceptSessionScopedRequest while the
// caller (handlePost) holds uriLocks for uri. sess is removed from uri's
// recorded subscribers immediately, regardless of what happens next; the
// request is forwarded upstream only if sess was the LAST recorded
// subscriber (the backend should keep sending updates for uri as long as
// anyone still wants them).
//
// Unlike subscribe, a failed/erroring upstream unsubscribe is not rolled
// back: removeSubscription has already run, and there is no OTHER session to
// re-attribute the subscription to even if we wanted to undo it. The worst
// case is the backend keeps a subscription nobody is listening for anymore --
// wasted upstream notifications that are silently dropped by
// routeResourceUpdated finding no subscribers, not a correctness or security
// issue -- so it is not worth the complexity of reversing removeSubscription
// here.
func (p *HTTPProxy) interceptUnsubscribe(ctx context.Context, sessID, uri string, req *jsonrpc2.Request) jsonrpc2.Message {
        last := p.routing.removeSubscription(uri, sessID)
        if !last {
                return synthesizeSuccessResponse(req.ID)
        }

        resp, err := p.forwardUpstream(ctx, sessID, req)
        if err != nil {
                return upstreamErrorResponse(req.ID, err)
        }
        return resp
}

// forwardUpstream wraps ctx with p.requestTimeout and forwards req upstream
// via doRequest, for interceptSubscribe/interceptUnsubscribe's synchronous
// forward performed while still holding handlePost's per-uri uriLocks lock
// for req's uri: the timeout bounds how long that lock can ever be held, so a
// slow or unresponsive backend cannot indefinitely block other sessions'
// resources/subscribe|unsubscribe requests for the SAME uri (see FIX 1 in
// #5744's review).
func (p *HTTPProxy) forwardUpstream(ctx context.Context, sessID string, req *jsonrpc2.Request) (jsonrpc2.Message, error) {
        ctx, cancel := context.WithTimeout(ctx, p.requestTimeout)
        defer cancel()
        return p.doRequest(ctx, sessID, req)
}

// isErrorResponse reports whether msg is a *jsonrpc2.Response carrying a
// JSON-RPC error, used by interceptSubscribe to distinguish a successful
// upstream resources/subscribe from one the backend rejected.
func isErrorResponse(msg jsonrpc2.Message) bool {
        resp, ok := msg.(*jsonrpc2.Response)
        return ok && resp.Error != nil
}

// upstreamErrorResponse builds a JSON-RPC error response for id when the
// upstream round-trip inside interceptSubscribe/interceptUnsubscribe itself
// fails (timeout, proxy shutdown, or a transport-level send error, as opposed
// to a JSON-RPC error the backend returned normally), mirroring the
// timeout-vs-internal error code mapping writeSSEErrorEvent uses for the
// normal request path.
func upstreamErrorResponse(id jsonrpc2.ID, err error) jsonrpc2.Message {
        var code int64 = -32603
        msg := "Internal error"
        if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
                code = -32000
                msg = "Timeout"
        }
        return &jsonrpc2.Response{ID: id, Error: jsonrpc2.NewError(code, msg)}
}

// synthesizeSuccessResponse builds a locally-generated JSON-RPC success
// response for id with an empty result object, used by
// interceptSessionScopedRequest when a request is served without an upstream
// round-trip.
func synthesizeSuccessResponse(id jsonrpc2.ID) jsonrpc2.Message {
        resp, err := jsonrpc2.NewResponse(id, map[string]any{}, nil)
        if err != nil {
                // marshalToRaw(map[string]any{}) cannot fail; NewResponse's only
                // error path is result marshaling.
                slog.Error("failed to synthesize success response", "error", err)
                return &jsonrpc2.Response{ID: id, Error: jsonrpc2.NewError(-32603, "internal error")}
        }
        return resp
}

// writeInterceptedResponse writes msg (from interceptSessionScopedRequest) to
// the client, honoring the same Accept-header-driven JSON-vs-SSE choice and
// Mcp-Session-Id header behavior as the normal handleSingleRequest /
// handleSingleRequestSSE paths, but without any upstream correlation (msg is
// already the complete, final response).
func writeInterceptedResponse(
        w http.ResponseWriter, r *http.Request, msg jsonrpc2.Message, sessID string, setSessionHeader bool,
) {
        if setSessionHeader {
                w.Header().Set("Mcp-Session-Id", sessID)
        }

        if !strings.Contains(r.Header.Get("Accept"), "text/event-stream") {
                if err := writeJSONRPC(w, msg); err != nil {
                        slog.Error("failed to write JSON-RPC response", "error", err)
                }
                return
        }

        flusher, ok := startSSEStream(w, nil)
        if !ok {
                return
        }

        data, err := jsonrpc2.EncodeMessage(msg)
        if err != nil {
                slog.Error("failed to encode JSON-RPC response", "error", err)
                return
        }
        if err := writeSSEData(w, flusher, data); err != nil {
                slog.Debug("failed to write response", "error", err)
        }
}

// logLevelRanks assigns each MCP logging level a verbosity rank following the
// RFC 5424 syslog severity numbers MCP's logging capability is aligned with:
// higher number = more verbose = lower severity. "debug" (7) is the most
// verbose; "emergency" (0) is the least. This makes "the maximum rank across
// sessions" exactly the level that satisfies every session's request without
// under-serving the most verbose one (see sessionRouter.setLogLevel).
var logLevelRanks = map[string]int{
        string(sdkmcp.LoggingLevelEmergency): 0,
        string(sdkmcp.LoggingLevelAlert):     1,
        string(sdkmcp.LoggingLevelCritical):  2,
        string(sdkmcp.LoggingLevelError):     3,
        string(sdkmcp.LoggingLevelWarning):   4,
        string(sdkmcp.LoggingLevelNotice):    5,
        string(sdkmcp.LoggingLevelInfo):      6,
        string(sdkmcp.LoggingLevelDebug):     7,
}

// logLevelRank returns level's verbosity rank (see logLevelRanks). An unknown
// or empty level is treated as the MOST verbose (debug's rank): if we can't
// recognize what the client asked for, we would rather over-serve (extra log
// volume) than silently under-serve a session that asked for a level we
// failed to parse.
func logLevelRank(level string) int {
        if rank, ok := logLevelRanks[level]; ok {
                return rank
        }
        slog.Debug("unrecognized logging level; treating as most verbose", "level", level)
        return logLevelRanks[string(sdkmcp.LoggingLevelDebug)]
}

// logLevelName returns the level name for rank (the inverse of logLevelRank),
// used to build the reconciled upstream logging/setLevel request.
func logLevelName(rank int) string {
        for name, r := range logLevelRanks {
                if r == rank {
                        return name
                }
        }
        return string(sdkmcp.LoggingLevelDebug)
}

// reconcileUpstreamLogLevel sends a logging/setLevel request upstream with
// level, reconciling the shared backend's single process-wide log level to
// the maximum verbosity requested by any session (see
// interceptSessionScopedRequest). It is fire-and-forget from the calling
// client request's perspective: the reconciliation happens in the background
// via the normal composite-key request/response correlation (so the eventual
// response, if any, is routed and cleaned up correctly instead of leaking a
// waiter), but its result is not surfaced to any client -- the client that
// triggered reconciliation already received its own synthesized success (see
// interceptSessionScopedRequest).
//
// Ordering tradeoff: this reconciliation is best-effort-ordered, NOT
// serialized the way resources/subscribe|unsubscribe is via uriLocks (see FIX
// 1 in #5744's review). Each call runs in its own goroutine, so two
// concurrent maxChanged transitions from different sessions' logging/setLevel
// requests can reach the backend in either order, and the backend could
// briefly end up at a lower verbosity than the CURRENT true max until the
// next transition reconciles it again. This is harmless while
// notifications/message is secure-dropped (see routeNotification): the
// backend's log level is not observable by any client either way, only its
// log volume is affected. Proper ordering (e.g. via a keyed lock analogous to
// uriLocks, or a monotonic sequence check) is deferred to the per-session
// backend / log-delivery follow-up (see the vMCP architecture, #5744).
func (p *HTTPProxy) reconcileUpstreamLogLevel(level string) {
        id, err := uuid.NewRandom()
        if err != nil {
                slog.Warn("failed to mint id for upstream log level reconciliation", "error", err)
                return
        }
        req, err := jsonrpc2.NewCall(jsonrpc2.StringID(id.String()), string(sdkmcp.MethodSetLogLevel), map[string]any{"level": level})
        if err != nil {
                slog.Warn("failed to build upstream log level reconciliation request", "error", err)
                return
        }

        go func() {
                // internalReconciliationSessionID keys this request's waiter under a
                // session ID no real request ever uses: resolveSessionForRequest
                // always assigns a non-empty UUID (real session or per-request
                // routing token), so "" can never collide with one.
                const internalReconciliationSessionID = ""
                ctx, cancel := context.WithTimeout(context.Background(), p.requestTimeout)
                defer cancel()
                if _, err := p.doRequest(ctx, internalReconciliationSessionID, req); err != nil {
                        slog.Debug("upstream log level reconciliation did not complete", "error", err)
                }
        }()
}

// reapRoutingState periodically purges sessionRouter subscription and
// log-level entries, AND serverStreams' standalone GET streams, for sessions
// that are no longer active, bounding the otherwise-unbounded growth/leakage
// of that state for a session whose owning client disappears without ever
// sending DELETE (see handleDelete's purgeSession/closeStream calls for the
// explicit-teardown path). It ticks at the same cadence as the session
// manager's own cleanup routine (sessionTTL/2, see manager.go's
// cleanupRoutine), and exits when shutdownCh is closed.
//
// isActive uses LookupOwner's authoritative, non-touching metadata read. Reaping
// therefore observes liveness without extending the session TTL; only normal
// session activity through Manager.Get refreshes the sliding expiration window.
func (p *HTTPProxy) reapRoutingState() {
        ticker := time.NewTicker(p.sessionTTL / 2)
        defer ticker.Stop()

        isActive := func(sess string) bool {
                _, err := p.sessionManager.LookupOwner(context.Background(), sess)
                return err == nil
        }

        for {
                select {
                case <-ticker.C:
                        p.routing.reap(isActive)
                        p.reapServerStreams(isActive)
                case <-p.shutdownCh:
                        return
                }
        }
}

// reapServerStreams closes the standalone GET SSE stream (see
// serverStreamRegistry.closeStream) for any session serverStreams currently
// has registered that isActive no longer reports as live -- a session whose
// owning client disappeared (crashed, or otherwise never sent DELETE) without
// the proxy ever observing r.Context().Done() on its GET request. This mirrors
// handleDelete's explicit p.serverStreams.closeStream(sessID) call for the
// case where the client never disconnects and never calls DELETE either (see
// FIX 2 in #5744's review).
//
// Liveness is resolved for every candidate stream key in a first pass with no
// registry lock held, exactly like sessionRouter.reap's two-phase discipline:
// isActive may perform I/O (e.g. a Redis-backed session.Manager), and running
// it while holding serverStreams.mu would block every concurrent
// broadcast/dispatchTo/deliverToMany/register/deregister call for the
// (possibly slow) duration of every liveness check. closeStream itself only
// ever takes the lock for the in-memory work of removing one already-decided
// stale entry, never across isActive.
func (p *HTTPProxy) reapServerStreams(isActive func(sess string) bool) {
        for _, key := range p.serverStreams.streamKeys() {
                if isActive(key) {
                        continue
                }
                if p.serverStreams.closeStream(key) {
                        //nolint:gosec // G706: session ID is server-minted, not user-controlled free text
                        slog.Debug("reaped standalone GET stream for inactive session", "session_id", key)
                }
        }
}

// decodeJSONRPCMessage decodes a JSON-RPC message from the request body.
func decodeJSONRPCMessage(w http.ResponseWriter, body []byte) (jsonrpc2.Message, bool) {
        msg, err := mcp.DecodeMessage(body)
        if err != nil {
                mcp.WriteClassificationError(w, nil, err)
                return nil, false
        }
        return msg, true
}

// handleNotificationOrClientResponse handles notifications and client responses
// as Legacy today; Modern-aware handling of this path is deferred.
func (p *HTTPProxy) handleNotificationOrClientResponse(w http.ResponseWriter, sessID string, msg jsonrpc2.Message) bool {
        if isNotification(msg) || isClientResponse(msg) {
                // Refresh TTL so a client sending only notifications doesn't get evicted.
                if sessID != "" {
                        p.sessionManager.Get(sessID)
                }
                // The backend completed a single handshake and expects a single
                // notifications/initialized. Later clients' copies are acknowledged
                // here without being forwarded, mirroring the initialize interception
                // that answered their handshakes from cache.
                if isInitializedNotification(msg) && !p.claimInitializedForward() {
                        w.WriteHeader(http.StatusAccepted)
                        return true
                }
                if err := p.SendMessageToDestination(msg); err != nil {
                        slog.Error("failed to send message to destination", "error", err)
                }
                w.WriteHeader(http.StatusAccepted)
                return true
        }
        return false
}

// resolveRequestTimeout returns the proxy request timeout, reading from the
// TOOLHIVE_PROXY_REQUEST_TIMEOUT environment variable if set, otherwise
// returning defaultRequestTimeout.
func resolveRequestTimeout() time.Duration {
        v := os.Getenv(proxyRequestTimeoutEnv)
        if v == "" {
                return defaultRequestTimeout
        }
        d, _ := time.ParseDuration(v)
        if d > 0 {
                slog.Debug("using custom proxy request timeout", "timeout", d)
                return d
        }
        slog.Warn("invalid proxy request timeout, using default",
                "env_var", proxyRequestTimeoutEnv, "value", v, "default", defaultRequestTimeout)
        return defaultRequestTimeout
}

// createWaiter registers a waiter channel for a request and returns the channel,
// the composite key (ck) the request must carry upstream, and a cleanup fn.
//
// ck is minted here from a per-request random token, never from the
// client-chosen JSON-RPC id: a client may have several requests with the same id
// in flight in one session, and a key derived from the id would let them
// overwrite each other's waiter and receive each other's responses (including
// responses the authz response filter never saw). The client's original id is
// kept in idRestore and put back on the reply.
func (p *HTTPProxy) createWaiter(sessID string, id jsonrpc2.ID) (chan jsonrpc2.Message, string, func(), error) {
        token, err := uuid.NewRandom()
        if err != nil {
                return nil, "", nil, fmt.Errorf("generate request correlation token: %w", err)
        }
        ck := compositeKey(sessID, token.String())
        // store original client id to restore before replying
        p.idRestore.Store(ck, id)

        ch := make(chan jsonrpc2.Message, 1)
        p.waiters.Store(ck, ch)

        cleanup := func() {
                p.waiters.Delete(ck)
                p.idRestore.Delete(ck)
        }
        return ch, ck, cleanup, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package streamable

import (
        "encoding/json"
        "fmt"
        "io"
        "net/http"
        "time"

        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/mcp"
)

// sseKeepAliveInterval is the cadence at which an otherwise-idle SSE stream
// writes a ": keep-alive\n\n" comment (see writeSSEKeepAlive), so
// intermediary proxies/load balancers do not close the connection for
// inactivity.
const sseKeepAliveInterval = 30 * time.Second

// isNotification returns true if the JSON-RPC message is a notification (no ID).
func isNotification(msg jsonrpc2.Message) bool {
        if req, ok := msg.(*jsonrpc2.Request); ok {
                return req.ID.Raw() == nil
        }
        return false
}

// isClientResponse returns true if msg is a *jsonrpc2.Response, i.e. a
// client's reply to a server-initiated request (see
// rejectServerRequestToBackend) rather than a client-initiated request or
// notification.
func isClientResponse(msg jsonrpc2.Message) bool {
        _, ok := msg.(*jsonrpc2.Response)
        return ok
}

// writeHTTPError writes a plain HTTP error with status.
func writeHTTPError(w http.ResponseWriter, status int, msg string) {
        http.Error(w, msg, status)
}

// writeSSEData writes a single JSON-RPC SSE frame with an explicit
// "event: message" name plus a "data:" line, then flushes it. The event name
// matches MCP reference transports and ToolHive's own ssecommon serializer;
// clients that only dispatch named "message" events (for example @ai-sdk/mcp)
// drop data-only frames. It returns the underlying write error (if any) so
// the caller can decide whether to end the stream; it does not itself log or
// close anything.
func writeSSEData(w io.Writer, flusher http.Flusher, data []byte) error {
        //nolint:gosec // G705: SSE data from MCP protocol
        if _, err := fmt.Fprintf(w, "event: message\ndata: %s\n\n", data); err != nil {
                return err
        }
        flusher.Flush()
        return nil
}

// assertFlushable type-asserts w as http.Flusher, which every SSE response in
// this package requires so data can be pushed to the client as it is
// written, rather than buffered until the handler returns. On failure it
// writes the standard 500 "Streaming not supported" error to w and returns
// (nil, false); the caller must return immediately without writing anything
// else. On success it returns (w.(http.Flusher), true).
func assertFlushable(w http.ResponseWriter) (http.Flusher, bool) {
        flusher, ok := w.(http.Flusher)
        if !ok {
                writeHTTPError(w, http.StatusInternalServerError, "Streaming not supported")
                return nil, false
        }
        return flusher, true
}

// setSSEHeaders sets the standard response headers for an SSE stream:
// Content-Type: text/event-stream, Cache-Control: no-cache, and
// Connection: keep-alive.
func setSSEHeaders(w http.ResponseWriter) {
        w.Header().Set("Content-Type", "text/event-stream")
        w.Header().Set("Cache-Control", "no-cache")
        w.Header().Set("Connection", "keep-alive")
}

// startSSEStream asserts that w supports flushing (see assertFlushable), sets
// the standard SSE headers (see setSSEHeaders), and flushes them to the
// client immediately so it knows the stream is open without waiting for the
// first data: frame.
//
// If extra is non-nil, it is called after the SSE headers are set but before
// the flush, letting a caller add further response headers (e.g.
// Mcp-Session-Id) while they are still guaranteed to reach the client:
// net/http locks in headers at the first Flush/WriteHeader call, so any
// header set after this function returns would be silently dropped.
//
// It returns (flusher, true) on success. If w does not implement
// http.Flusher, it writes the "Streaming not supported" 500 itself (see
// assertFlushable) and returns (nil, false); the caller must return without
// writing anything else, and extra is never called.
func startSSEStream(w http.ResponseWriter, extra func()) (http.Flusher, bool) {
        flusher, ok := assertFlushable(w)
        if !ok {
                return nil, false
        }
        setSSEHeaders(w)
        if extra != nil {
                extra()
        }
        flusher.Flush()
        return flusher, true
}

// writeSSEKeepAlive writes a single SSE comment line (": keep-alive\n\n") to
// w and flushes it, so intermediary proxies/load balancers do not close an
// otherwise-idle SSE connection for inactivity. It returns the underlying
// write error (if any), mirroring writeSSEData, so the caller can decide
// whether to end the stream; it does not itself log or close anything.
func writeSSEKeepAlive(w io.Writer, flusher http.Flusher) error {
        if _, err := fmt.Fprint(w, ": keep-alive\n\n"); err != nil { //nolint:gosec // G705: fixed literal, not derived from any request
                return err
        }
        flusher.Flush()
        return nil
}

// writeJSONRPC writes a jsonrpc2.Message using the library's encoder to ensure proper serialization.
func writeJSONRPC(w http.ResponseWriter, msg jsonrpc2.Message) error {
        w.Header().Set("Content-Type", "application/json")
        data, err := jsonrpc2.EncodeMessage(msg)
        if err != nil {
                return err
        }
        _, err = w.Write(data) //nolint:gosec // G705: data is JSON-RPC from MCP protocol
        return err
}

// compositeKey builds the upstream correlation key from the session ID and a
// per-request token. The token must be unique per in-flight request (see
// createWaiter); it must never be derived from the client-chosen JSON-RPC id.
func compositeKey(sessID, requestToken string) string {
        return sessID + "|" + requestToken
}

// extractStringParam returns params[key] as a string, and whether it was
// present and string-typed. Absent params, params that don't decode as a JSON
// object, or a non-string value at key all yield ("", false) rather than an
// error -- callers treat a missing/malformed field as "nothing to route on".
func extractStringParam(params json.RawMessage, key string) (string, bool) {
        if len(params) == 0 {
                return "", false
        }
        var m map[string]any
        if err := json.Unmarshal(params, &m); err != nil {
                return "", false
        }
        s, ok := m[key].(string)
        return s, ok
}

// rewriteRequestParam returns a new *jsonrpc2.Request identical to req except
// that its top-level params[key] is set to value (params is created if req
// had none). It never mutates req: per the "copy before mutating caller
// input" style rule, the params map is freshly decoded from req.Params before
// the assignment, so no caller-held reference (including req itself) is
// modified.
func rewriteRequestParam(req *jsonrpc2.Request, key string, value any) (*jsonrpc2.Request, error) {
        m := map[string]any{}
        if len(req.Params) > 0 {
                if err := json.Unmarshal(req.Params, &m); err != nil {
                        return nil, err
                }
        }
        m[key] = value
        data, err := json.Marshal(m)
        if err != nil {
                return nil, err
        }
        return &jsonrpc2.Request{ID: req.ID, Method: req.Method, Params: data}, nil
}

// supportedMCPVersions is the set of MCP protocol revision dates this proxy
// recognizes when strict validation is enabled (WithStrictProtocolValidation).
// The upcoming stateless revision is sourced from mcp.MCPVersionModern rather
// than duplicated as a literal, so the strict gate and the ClassifyRevision
// routing path (which keys off the same constant) cannot drift if that draft
// date changes before it is finalized.
var supportedMCPVersions = map[string]struct{}{
        "2024-11-05":         {},
        "2025-03-26":         {},
        "2025-06-18":         {},
        "2025-11-25":         {},
        mcp.MCPVersionModern: {},
}

// isSupportedMCPVersion reports whether version is one of the known MCP
// protocol revision dates in supportedMCPVersions. It is only consulted when
// the proxy's strict protocol validation is enabled (see
// WithStrictProtocolValidation); in the default permissive mode, handlePost
// never calls it and any MCP-Protocol-Version header value is accepted, since
// this proxy is transport-level and does not depend on a specific MCP
// revision.
func isSupportedMCPVersion(version string) bool {
        _, ok := supportedMCPVersions[version]
        return ok
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package transparent provides MCP ping functionality for transparent proxies.
package transparent

import (
        "context"
        "fmt"
        "log/slog"
        "net/http"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/healthcheck"
        "github.com/stacklok/toolhive/pkg/networking"
)

// MCPPinger implements healthcheck.MCPPinger for transparent proxies
type MCPPinger struct {
        targetURL string
        client    *http.Client
}

const (
        // DefaultPingerTimeout is the default timeout for health check pings
        DefaultPingerTimeout = 5 * time.Second
)

// redirectPolicy is installed on both pinger clients. It refuses redirects to a
// different host (or an HTTPS-to-HTTP downgrade) rather than following Go's
// http.Client default of any host, any number of hosts, indefinitely. See
// networking.SameHostRedirectPolicy's doc comment for the full rationale; this
// package's data-plane forwarder (followRedirects, in transparent_proxy.go)
// enforces the same same-host-only policy and the two must stay in sync.
var redirectPolicy = networking.SameHostRedirectPolicy()

// NewMCPPingerWithTimeout creates a new MCP pinger with a custom timeout
//
// The configured MCP endpoint may be a remote server outside our control, so
// the client restricts redirects with SameHostRedirectPolicy: without it, a
// malicious or compromised remote could return a 3xx pointing this GET at a
// loopback, private, or otherwise host-reachable address, and the default
// http.Client would follow it (CWE-918, blind SSRF).
func NewMCPPingerWithTimeout(targetURL string, timeout time.Duration) healthcheck.MCPPinger {
        if timeout <= 0 {
                timeout = DefaultPingerTimeout
        }
        return &MCPPinger{
                targetURL: targetURL,
                client: &http.Client{
                        Timeout:       timeout,
                        CheckRedirect: redirectPolicy,
                },
        }
}

// Ping performs a simple HTTP health check for SSE transport servers
// For SSE transport, we don't send MCP ping requests because that would require
// establishing an SSE session first. Instead, we do a simple HTTP GET to check
// if the server is responding.
func (p *MCPPinger) Ping(ctx context.Context) (time.Duration, error) {
        start := time.Now()

        // Create a simple GET request to check if the server is responding
        req, err := http.NewRequestWithContext(ctx, http.MethodGet, p.targetURL, nil)
        if err != nil {
                return 0, fmt.Errorf("failed to create HTTP request: %w", err)
        }

        //nolint:gosec // G706: logging target URL from config
        slog.Debug("checking SSE server health", "target", p.targetURL)

        // Send the request
        resp, err := p.client.Do(req) // #nosec G704 -- targetURL is the local MCP server health endpoint
        if err != nil {
                return time.Since(start), fmt.Errorf("failed to connect to SSE server: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        duration := time.Since(start)

        // For Streamable HTTP servers, we expect various responses:
        // - 200 OK for root endpoints
        // - 404 for non-existent endpoints (but server is still alive)
        // - 401 for remote workloads (we may not be able to authenticate, but this response indicates the server is alive).
        // - Other 4xx/5xx may indicate server issues
        // For now, we accept any non 50x response for both local and remote.
        if resp.StatusCode >= 200 && resp.StatusCode < 500 {
                //nolint:gosec // G706: logging HTTP status code from health check response
                slog.Debug("sse server health check successful",
                        "duration", duration, "status", resp.StatusCode)
                return duration, nil
        }

        return duration, fmt.Errorf("SSE server health check failed with status %d", resp.StatusCode)
}

// StatelessMCPPinger health-checks stateless streamable-HTTP servers via POST ping.
// Stateless servers don't support GET, so we send a minimal JSON-RPC ping instead.
type StatelessMCPPinger struct {
        targetURL string
        client    *http.Client
}

// NewStatelessMCPPinger creates a pinger for stateless streamable-HTTP servers.
func NewStatelessMCPPinger(targetURL string) healthcheck.MCPPinger {
        return NewStatelessMCPPingerWithTimeout(targetURL, DefaultPingerTimeout)
}

// NewStatelessMCPPingerWithTimeout creates a stateless pinger with a custom timeout.
//
// Same redirect restriction as NewMCPPingerWithTimeout, and for the same
// reason: the target may be an untrusted remote endpoint. A 307/308 here would
// otherwise carry the POST body and method to whatever host the response
// names; 301/302/303 already lose the body under Go's default redirect
// handling, so refusing cross-host keeps that residual GET on the configured
// host too.
func NewStatelessMCPPingerWithTimeout(targetURL string, timeout time.Duration) healthcheck.MCPPinger {
        if timeout <= 0 {
                timeout = DefaultPingerTimeout
        }
        return &StatelessMCPPinger{
                targetURL: targetURL,
                client: &http.Client{
                        Timeout:       timeout,
                        CheckRedirect: redirectPolicy,
                },
        }
}

// Ping sends a JSON-RPC ping POST to check if the stateless server is reachable.
// Accepts any 2xx-4xx response as healthy; only network errors and 5xx indicate failure.
func (p *StatelessMCPPinger) Ping(ctx context.Context) (time.Duration, error) {
        start := time.Now()

        body := `{"jsonrpc":"2.0","id":0,"method":"ping","params":{}}`
        req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.targetURL, strings.NewReader(body))
        if err != nil {
                return 0, fmt.Errorf("failed to create stateless ping request: %w", err)
        }
        req.Header.Set("Content-Type", "application/json")
        req.Header.Set("Accept", "application/json, text/event-stream")

        //nolint:gosec // G706: logging target URL from config
        slog.Debug("checking stateless MCP server health via POST ping", "target", p.targetURL)

        resp, err := p.client.Do(req)
        if err != nil {
                return time.Since(start), fmt.Errorf("stateless ping failed to connect: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close ping response body", "error", err)
                }
        }()

        duration := time.Since(start)

        // Accept 2xx-4xx: even 401/403 means the server is reachable.
        // Only 5xx or network errors indicate the server is down.
        if resp.StatusCode >= 200 && resp.StatusCode < 500 {
                //nolint:gosec // G706: logging HTTP status code from health check response
                slog.Debug("stateless MCP server health check successful",
                        "duration", duration, "status", resp.StatusCode)
                return duration, nil
        }

        return duration, fmt.Errorf("stateless ping returned status %d", resp.StatusCode)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package transparent provides a transparent HTTP proxy implementation
// that forwards requests to a destination without modifying them.
package transparent

import (
        "net/http"

        "github.com/stacklok/toolhive/pkg/transport/types"
)

// ResponseProcessor defines the interface for processing and modifying HTTP responses
// based on transport-specific requirements.
type ResponseProcessor interface {
        // ProcessResponse modifies an HTTP response based on transport-specific logic.
        // Returns error if processing fails.
        ProcessResponse(resp *http.Response) error

        // ShouldProcess returns true if this processor should handle the given response.
        ShouldProcess(resp *http.Response) bool
}

// NoOpResponseProcessor is a processor that does nothing.
// Used for transports that don't require response processing (e.g., streamable-http).
type NoOpResponseProcessor struct{}

// ProcessResponse is a no-op implementation.
func (*NoOpResponseProcessor) ProcessResponse(_ *http.Response) error {
        return nil
}

// ShouldProcess always returns false for no-op processor.
func (*NoOpResponseProcessor) ShouldProcess(_ *http.Response) bool {
        return false
}

// createResponseProcessor is a factory function that creates the appropriate
// response processor based on transport type.
func createResponseProcessor(
        transportType string,
        proxy *TransparentProxy,
        endpointPrefix string,
        trustProxyHeaders bool,
) ResponseProcessor {
        switch transportType {
        case types.TransportTypeSSE.String():
                return NewSSEResponseProcessor(proxy, endpointPrefix, trustProxyHeaders)
        case types.TransportTypeStreamableHTTP.String():
                return &NoOpResponseProcessor{}
        default:
                // Default to no-op for unknown transport types
                return &NoOpResponseProcessor{}
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package transparent

import (
        "net/http"
        "net/url"
        "slices"
        "strings"

        "github.com/google/uuid"

        "github.com/stacklok/toolhive/pkg/auth/sessionbinding"
)

// sessionQueryKeys are the query parameter names an MCP SSE server may use to carry
// the session id in its endpoint event. The proxy does not get to choose the name,
// and the spellings in the wild genuinely differ:
//
//   - "sessionId"  the TypeScript SDK and mcp-go
//   - "sessionid"  mark3labs/mcp-go servers such as the yardstick test server,
//     whose endpoint event reads `data: /sse?sessionid=...`
//   - "session_id" ToolHive's own SSE proxy (pkg/transport/proxy/httpsse)
//
// Recognising only one spelling does not fail safe. extractSessionID treats an
// unrecognised carrier as an absent one and closes the stream, so the client never
// receives the endpoint event at all -- a total outage for that backend rather than
// a tightened check.
//
// Every site that reads, enforces or rewrites a session carrier must use this same
// list. A carrier bound under one spelling and checked under another is how a
// request gets authorized as one session and executed as another.
var sessionQueryKeys = []string{"sessionId", "sessionid", "session_id"}

// requestSessionID resolves a request's session id from the Mcp-Session-Id header or
// any recognised query spelling.
//
// Two spellings carrying different values are refused rather than resolved to one of
// them: this proxy and the backend could pick differently, and a disagreement there
// is a session-confusion bug. Returning an empty id is not an error -- it means the
// request names no session, which callers handle separately.
func requestSessionID(r *http.Request) (string, error) {
        var found string
        for _, key := range sessionQueryKeys {
                id, err := sessionbinding.RequestID(r, key)
                if err != nil {
                        return "", err
                }
                if id == "" {
                        continue
                }
                if found != "" && found != id {
                        return "", sessionbinding.ErrNotFound
                }
                found = id
        }
        return found, nil
}

// rewriteSessionQuery modifies only existing session query fields, preserving the
// byte representation and order of other query parameters. An empty ID strips them.
func rewriteSessionQuery(u *url.URL, id string) {
        parts := strings.Split(u.RawQuery, "&")
        kept := parts[:0]
        for _, part := range parts {
                key, _, _ := strings.Cut(part, "=")
                key, err := url.QueryUnescape(key)
                if err == nil && slices.Contains(sessionQueryKeys, key) {
                        if id != "" {
                                kept = append(kept, key+"="+url.QueryEscape(id))
                        }
                        continue
                }
                kept = append(kept, part)
        }
        u.RawQuery = strings.Join(kept, "&")
}

// mcpSessionNamespace is the UUID v5 namespace used when normalizing non-UUID
// Mcp-Session-Id values received from upstream MCP servers.
var mcpSessionNamespace = uuid.MustParse("6ba7b810-9dad-11d1-80b4-00c04fd430c8") // RFC 4122 URL namespace

// normalizeSessionID returns id unchanged if it is already a valid UUID.
// Otherwise it returns a deterministic UUID v5 derived from id, ensuring that
// the session manager (which requires UUID-format IDs) can store sessions whose
// Mcp-Session-Id was issued by an upstream server in a non-UUID format.
//
// The mapping is stable: the same external id always produces the same UUID,
// so the proxy can look up and delete sessions without maintaining a separate
// reverse-mapping table.
func normalizeSessionID(id string) string {
        if _, err := uuid.Parse(id); err == nil {
                return id
        }
        return uuid.NewSHA1(mcpSessionNamespace, []byte(id)).String()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package transparent provides a transparent HTTP proxy implementation
// that forwards requests to a destination without modifying them.
package transparent

import (
        "bufio"
        "context"
        "fmt"
        "io"
        "log/slog"
        "mime"
        "net/http"
        "net/url"
        "strings"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/sessionbinding"
        "github.com/stacklok/toolhive/pkg/transport/session"
)

// inboundRequestKey is the context key for storing the original inbound request.
// This is needed because httputil.ReverseProxy's ModifyResponse receives resp.Request
// as the *outbound* request, which has auto-injected X-Forwarded-* headers from
// SetXForwarded(). To read the client's original headers, we stash the inbound
// request in the outbound request's context during Rewrite.
type inboundRequestKey struct{}

// InboundRequestToContext returns a new context that carries the inbound request.
func InboundRequestToContext(ctx context.Context, req *http.Request) context.Context {
        return context.WithValue(ctx, inboundRequestKey{}, req)
}

// inboundRequestFromContext retrieves the inbound request from the context.
// Returns nil if not present.
func inboundRequestFromContext(ctx context.Context) *http.Request {
        req, _ := ctx.Value(inboundRequestKey{}).(*http.Request)
        return req
}

// sseRewriteConfig holds the configuration for rewriting SSE endpoint URLs.
// This is used to handle path-based ingress routing scenarios where the ingress
// strips a path prefix before forwarding to the backend MCP server.
type sseRewriteConfig struct {
        // prefix is the path prefix to prepend to endpoint URLs (e.g., "/playwright")
        prefix string
        // scheme is the URL scheme to use (e.g., "https"), empty means preserve original
        scheme string
        // host is the host to use (e.g., "public.example.com"), empty means preserve original
        host string
}

// hasRewriteConfig returns true if any rewriting is configured.
func (c sseRewriteConfig) hasRewriteConfig() bool {
        return c.prefix != "" || c.scheme != "" || c.host != ""
}

// SSEResponseProcessor handles SSE-specific response processing including:
// - Session ID extraction from SSE streams
// - Endpoint URL rewriting for path-based routing
type SSEResponseProcessor struct {
        proxy             *TransparentProxy
        endpointPrefix    string
        trustProxyHeaders bool
}

// NewSSEResponseProcessor creates a new SSE response processor.
func NewSSEResponseProcessor(
        proxy *TransparentProxy,
        endpointPrefix string,
        trustProxyHeaders bool,
) *SSEResponseProcessor {
        return &SSEResponseProcessor{
                proxy:             proxy,
                endpointPrefix:    endpointPrefix,
                trustProxyHeaders: trustProxyHeaders,
        }
}

// ShouldProcess returns true if the response is an SSE stream.
func (*SSEResponseProcessor) ShouldProcess(resp *http.Response) bool {
        mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
        return mediaType == "text/event-stream"
}

// ProcessResponse modifies SSE responses to:
// 1. Extract session IDs from endpoint events for session tracking
// 2. Rewrite endpoint URLs when X-Forwarded-Prefix or endpointPrefix is configured
//
// SSE Event Format:
//
//        event: endpoint
//        data: /sse?sessionId=abc
//
// Only "event: endpoint" events have their data field rewritten.
// Other events (e.g., "event: message") are passed through unchanged.
func (s *SSEResponseProcessor) ProcessResponse(resp *http.Response) error {
        if !s.ShouldProcess(resp) {
                return nil
        }

        // Get rewrite config from the request headers
        var rewriteConfig sseRewriteConfig
        if resp.Request != nil {
                rewriteConfig = s.getSSERewriteConfig(resp.Request)
        }

        ctx := context.Background()
        if resp.Request != nil {
                identity, _ := auth.IdentityFromContext(resp.Request.Context())
                ctx = auth.WithIdentity(ctx, identity)
        }
        identity, _ := auth.IdentityFromContext(ctx)
        if _, err := sessionbinding.FromIdentity(identity); err != nil {
                return err
        }
        pr, pw := io.Pipe()
        originalBody := resp.Body
        resp.Body = pr

        // NOTE: it would be better to have a proper function instead of a goroutine, as this
        // makes it harder to debug and test.
        go func() {
                defer func() {
                        if err := pw.Close(); err != nil {
                                slog.Debug("failed to close pipe writer", "error", err)
                        }
                }()
                defer func() {
                        if err := originalBody.Close(); err != nil {
                                slog.Debug("failed to close upstream SSE body", "error", err)
                        }
                }()
                s.processSSEStream(ctx, originalBody, pw, rewriteConfig)
        }()

        return nil
}

// getSSERewriteConfig determines the SSE endpoint URL rewrite configuration based on priority:
// 1. Explicit endpointPrefix configuration (highest priority)
// 2. X-Forwarded-Prefix header (only when trustProxyHeaders is true)
// 3. No rewriting (default)
//
// IMPORTANT: req is the outbound request from httputil.ReverseProxy, which has
// auto-injected X-Forwarded-* headers via SetXForwarded(). To read the client's
// original headers we use the inbound request stashed in the context during Rewrite.
func (s *SSEResponseProcessor) getSSERewriteConfig(req *http.Request) sseRewriteConfig {
        config := sseRewriteConfig{}

        // Use the inbound (client-facing) request for reading forwarded headers,
        // because the outbound request has auto-injected X-Forwarded-* values
        // from httputil.ReverseProxy.SetXForwarded().
        inbound := inboundRequestFromContext(req.Context())
        if inbound == nil {
                // Fallback: if no inbound request in context (e.g. in tests), use
                // the outbound request directly.
                inbound = req
        }

        // Priority 1: Explicit endpointPrefix configuration
        if s.endpointPrefix != "" {
                config.prefix = s.endpointPrefix
        } else if s.trustProxyHeaders {
                // Priority 2: X-Forwarded-Prefix header from the original client request
                if prefix := inbound.Header.Get("X-Forwarded-Prefix"); prefix != "" {
                        config.prefix = prefix
                }
        }

        // Also check for X-Forwarded-Proto and X-Forwarded-Host if trustProxyHeaders is enabled
        if s.trustProxyHeaders {
                if scheme := inbound.Header.Get("X-Forwarded-Proto"); scheme != "" {
                        config.scheme = scheme
                }
                if host := inbound.Header.Get("X-Forwarded-Host"); host != "" {
                        config.host = host
                }
        }

        return config
}

// rewriteEndpointURL rewrites an SSE endpoint URL with the given configuration.
// It handles both relative URLs (e.g., "/sse?sessionId=abc") and absolute URLs
// (e.g., "http://backend:8080/sse?sessionId=abc").
func rewriteEndpointURL(originalURL string, config sseRewriteConfig) (string, error) {
        if !config.hasRewriteConfig() {
                return originalURL, nil
        }

        parsed, err := url.Parse(originalURL)
        if err != nil {
                return originalURL, fmt.Errorf("failed to parse URL: %w", err)
        }

        // Prepend prefix to path
        if config.prefix != "" {
                // Ensure prefix starts with "/" and doesn't end with "/"
                prefix := config.prefix
                if !strings.HasPrefix(prefix, "/") {
                        prefix = "/" + prefix
                }
                prefix = strings.TrimSuffix(prefix, "/")
                parsed.Path = prefix + parsed.Path
        }

        // Override host if configured
        if config.host != "" {
                parsed.Host = config.host

                // Override scheme if configured
                if config.scheme != "" {
                        parsed.Scheme = config.scheme
                }
        }

        return parsed.String(), nil
}

// sseLineProcessor handles line-by-line processing of SSE streams.
// It tracks event types and processes data lines for session extraction and URL rewriting.
type sseLineProcessor struct {
        proxy            *TransparentProxy
        rewriteConfig    sseRewriteConfig
        currentEventType string
        ctx              context.Context
        err              error
}

// processLine processes a single SSE line and returns the potentially modified line.
func (s *sseLineProcessor) processLine(line string) string {
        // Parse SSE event type
        if strings.HasPrefix(line, "event:") {
                s.currentEventType = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
                return line
        }

        // Empty line marks the end of an SSE event, reset event type
        if line == "" {
                s.currentEventType = ""
                return line
        }

        // Process data lines
        if strings.HasPrefix(line, "data:") {
                return s.processDataLine(line)
        }

        return line
}

// processDataLine handles SSE data lines for session extraction and URL rewriting.
func (s *sseLineProcessor) processDataLine(line string) string {
        dataContent := strings.TrimSpace(strings.TrimPrefix(line, "data:"))

        // Bind every endpoint before exposing its ID, including later endpoint events.
        if s.currentEventType == "endpoint" {
                s.err = s.extractSessionID(dataContent)
        }

        // Rewrite endpoint URLs only for "endpoint" events
        if s.currentEventType == "endpoint" && s.rewriteConfig.hasRewriteConfig() {
                return s.rewriteDataLine(line, dataContent)
        }

        return line
}

// extractSessionID extracts and stores the session ID from a data line.
func (s *sseLineProcessor) extractSessionID(endpoint string) error {
        parsed, err := url.Parse(endpoint)
        if err != nil {
                return sessionbinding.ErrNotFound
        }
        sid, err := requestSessionID(&http.Request{URL: parsed})
        if err != nil || sid == "" {
                return sessionbinding.ErrNotFound
        }
        sess := session.NewProxySession(normalizeSessionID(sid))
        if err := sessionbinding.AddOwnedSession(s.ctx, s.proxy.sessionManager, sess); err != nil {
                return err
        }
        s.proxy.setServerInitialized()
        return nil
}

// rewriteDataLine rewrites the URL in an endpoint event's data line.
func (s *sseLineProcessor) rewriteDataLine(line, dataContent string) string {
        rewrittenURL, err := rewriteEndpointURL(dataContent, s.rewriteConfig)
        if err != nil {
                //nolint:gosec // G706: logging endpoint URL from SSE stream
                slog.Warn("failed to rewrite endpoint URL",
                        "url", dataContent, "error", err)
                return line
        }
        if rewrittenURL != dataContent {
                //nolint:gosec // G706: logging endpoint URLs from SSE stream
                slog.Debug("rewrote SSE endpoint URL",
                        "from", dataContent, "to", rewrittenURL)
                return "data: " + rewrittenURL
        }
        return line
}

// processSSEStream processes an SSE stream, extracting session IDs and rewriting URLs.
func (s *SSEResponseProcessor) processSSEStream(
        ctx context.Context, originalBody io.Reader, pw *io.PipeWriter, rewriteConfig sseRewriteConfig,
) {
        scanner := bufio.NewScanner(originalBody)
        // NOTE: The following line mitigates the issue of the response body being too large.
        // By default, the maximum token size of the scanner is 64KB, which is too small in
        // the case of e.g. images. This raises the limit to 1MB. This is a workaround, and
        // not a proper fix.
        scanner.Buffer(make([]byte, 0, 1024), 1024*1024*1)

        processor := &sseLineProcessor{
                ctx:           ctx,
                proxy:         s.proxy,
                rewriteConfig: rewriteConfig,
        }

        for scanner.Scan() {
                line := processor.processLine(scanner.Text())
                if processor.err != nil {
                        _ = pw.CloseWithError(processor.err)
                        return
                }
                if _, err := pw.Write([]byte(line + "\n")); err != nil {
                        return
                }
        }

        if err := scanner.Err(); err != nil {
                slog.Error("failed to scan response body", "error", err)
        }

        if readCloser, ok := originalBody.(io.ReadCloser); ok {
                if _, err := io.Copy(pw, readCloser); err != nil && err != io.EOF {
                        slog.Error("failed to copy response body", "error", err)
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package transparent provides a transparent HTTP proxy implementation
// that forwards requests to a destination without modifying them.
package transparent

import (
        "bytes"
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net"
        "net/http"
        "net/http/httptrace"
        "net/http/httputil"
        "net/url"
        "os"
        "strconv"
        "strings"
        "sync"
        "sync/atomic"
        "time"

        "go.opentelemetry.io/otel"
        "go.opentelemetry.io/otel/propagation"
        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/auth/sessionbinding"
        "github.com/stacklok/toolhive/pkg/bodylimit"
        "github.com/stacklok/toolhive/pkg/diagnostics"
        "github.com/stacklok/toolhive/pkg/healthcheck"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/proxy/socket"
        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// TransparentProxy implements the Proxy interface as a transparent HTTP proxy
// that forwards requests to a destination.
// It's used by the SSE transport to forward requests to the container's HTTP server.
//
//nolint:revive // Intentionally named TransparentProxy despite package name
type TransparentProxy struct {
        // Basic configuration
        host      string
        port      int
        targetURI string

        // HTTP server
        server *http.Server

        // Middleware chain
        middlewares []types.NamedMiddleware

        // Mutex for protecting shared state
        mutex sync.Mutex

        // Track if Stop() has been called
        stopped bool

        // Shutdown channel
        shutdownCh chan struct{}

        // Health checker
        healthChecker *healthcheck.HealthChecker

        // Optional Prometheus metrics handler
        prometheusHandler http.Handler

        // Optional auth info handler
        authInfoHandler http.Handler

        // prefixHandlers is a map of path prefixes to HTTP handlers
        // mounted before the catch-all proxy handler
        prefixHandlers map[string]http.Handler

        // Sessions for tracking state
        sessionManager *session.Manager

        // sessionTTL is the resolved inactivity timeout for the session manager.
        // Defaults to session.DefaultSessionTTL; overridable via WithSessionTTL.
        sessionTTL time.Duration

        // readTimeout is the resolved http.Server.ReadTimeout for this proxy.
        // Defaults to defaultReadTimeout; overridable via WithReadTimeout.
        readTimeout time.Duration

        // sessionStorage is the optional custom storage backend for the session manager.
        // When nil, in-memory LocalStorage is used. Set via WithSessionStorage.
        sessionStorage session.Storage

        // If mcp server has been initialized (atomic access)
        isServerInitialized atomic.Bool

        // Listener for the HTTP server
        listener net.Listener

        // Whether the target URI is remote
        isRemote bool

        // Transport type (sse, streamable-http)
        transportType string

        // stateless indicates the server is POST-only (no SSE/GET support)
        stateless bool

        // stripConsentCookie prevents the embedded AS session from reaching the backend.
        stripConsentCookie bool

        // Callback when health check fails (for remote servers)
        onHealthCheckFailed types.HealthCheckFailedCallback

        // Callback when 401 Unauthorized response is received (for bearer token authentication)
        onUnauthorizedResponse types.UnauthorizedResponseCallback

        // Response processor for transport-specific logic
        responseProcessor ResponseProcessor

        // Deprecated: SSE endpoint URL rewriting configuration (moved to SSEResponseProcessor)
        // endpointPrefix is an explicit prefix to prepend to SSE endpoint URLs
        endpointPrefix string

        // remoteBasePath is the path prefix from the remote URL that must be prepended
        // to incoming request paths before forwarding. For example, if the remote URL is
        // https://mcp.asana.com/v2/mcp and a client sends to /mcp, the proxy must
        // forward to /v2/mcp. Without this, the path prefix is lost because the target
        // URI only contains the scheme and host.
        remoteBasePath string

        // remoteRawQuery holds the raw query string from the remote URL (e.g.,
        // "toolsets=core,alerting" from "https://mcp.example.com/mcp?toolsets=core,alerting").
        // When set, it is merged into every outbound request so query parameters
        // from the original registration URL are never silently dropped.
        remoteRawQuery string

        // Deprecated: trustProxyHeaders indicates whether to trust X-Forwarded-* headers (moved to SSEResponseProcessor)
        trustProxyHeaders bool

        // Health check interval (default: 10 seconds)
        healthCheckInterval time.Duration

        // Health check retry delay (default: 5 seconds)
        healthCheckRetryDelay time.Duration

        // Health check ping timeout (default: 5 seconds)
        healthCheckPingTimeout time.Duration

        // Health check failure threshold: consecutive failures before shutdown (default: 5)
        healthCheckFailureThreshold int

        // Shutdown timeout for graceful HTTP server shutdown (default: 30 seconds)
        shutdownTimeout time.Duration
}

const (
        // DefaultHealthCheckInterval is the default interval for health checks
        DefaultHealthCheckInterval = 10 * time.Second

        // DefaultHealthCheckRetryDelay is the default delay between retry attempts
        DefaultHealthCheckRetryDelay = 5 * time.Second

        // defaultShutdownTimeout is the maximum time to wait for graceful HTTP server
        // shutdown before force-closing connections.
        defaultShutdownTimeout = 30 * time.Second

        // defaultIdleTimeout is the maximum time to wait for the next request on a
        // keep-alive connection. Matches the value used by the vMCP server.
        defaultIdleTimeout = 120 * time.Second

        // defaultReadTimeout bounds reading the entire request (headers + body) on
        // the proxy http.Server, mitigating slow-upload connection exhaustion. It
        // does not affect responses, so streamed backend responses are unaffected.
        // Matches the value used by the vMCP server.
        defaultReadTimeout = 30 * time.Second

        // HealthCheckIntervalEnvVar is the environment variable name for configuring health check interval.
        HealthCheckIntervalEnvVar = "TOOLHIVE_HEALTH_CHECK_INTERVAL"

        // sessionMetadataBackendURL is the session metadata key that stores the backend pod URL.
        // It is written on initialize and read in the Rewrite closure to route follow-up requests
        // to the same backend pod that handled the session's initialize request.
        sessionMetadataBackendURL = "backend_url"

        // sessionMetadataInitBody stores the raw JSON-RPC initialize request body.
        // It is used to transparently re-initialize a backend session when the pod that
        // originally handled initialize has been replaced (new IP or lost in-memory state).
        sessionMetadataInitBody = "init_body"

        // sessionMetadataBackendSID stores the backend's assigned Mcp-Session-Id when it
        // diverges from the client-facing session ID after a transparent re-initialization.
        // tracingTransport.RoundTrip rewrites the outbound Mcp-Session-Id header to this
        // value so the backend sees its own session ID while the client keeps its original one.
        sessionMetadataBackendSID = "backend_sid"

        // HealthCheckPingTimeoutEnvVar is the environment variable name for configuring health check ping timeout.
        HealthCheckPingTimeoutEnvVar = "TOOLHIVE_HEALTH_CHECK_PING_TIMEOUT"

        // HealthCheckRetryDelayEnvVar is the environment variable name for configuring health check retry delay.
        HealthCheckRetryDelayEnvVar = "TOOLHIVE_HEALTH_CHECK_RETRY_DELAY"

        // HealthCheckFailureThresholdEnvVar is the environment variable name for configuring
        // the number of consecutive health check failures before shutdown.
        HealthCheckFailureThresholdEnvVar = "TOOLHIVE_HEALTH_CHECK_FAILURE_THRESHOLD"

        // DefaultHealthCheckFailureThreshold is the default number of consecutive health check
        // failures before the proxy initiates shutdown.
        DefaultHealthCheckFailureThreshold = 5

        // maxRedirects is the maximum number of HTTP redirects to follow when
        // forwarding requests to a remote MCP server. Uses the same limit as
        // http.Client (10), but unlike http.Client the HTTP method is always
        // preserved (POST never becomes GET) because MCP uses JSON-RPC over POST.
        maxRedirects = 10
)

// Option is a functional option for configuring TransparentProxy
type Option func(*TransparentProxy)

// withHealthCheckInterval sets the health check interval.
// This is primarily useful for testing with shorter intervals.
// Ignores non-positive intervals; default will be used.
func withHealthCheckInterval(interval time.Duration) Option {
        return func(p *TransparentProxy) {
                if interval > 0 {
                        p.healthCheckInterval = interval
                }
        }
}

// withHealthCheckRetryDelay sets the health check retry delay.
// This is primarily useful for testing with shorter delays.
// Ignores non-positive delays; default will be used.
func withHealthCheckRetryDelay(delay time.Duration) Option {
        return func(p *TransparentProxy) {
                if delay > 0 {
                        p.healthCheckRetryDelay = delay
                }
        }
}

// WithRemoteBasePath sets the base path prefix from the remote URL.
// When set, incoming request paths are rewritten to include this prefix
// before forwarding to the remote server.
func WithRemoteBasePath(basePath string) Option {
        return func(p *TransparentProxy) {
                p.remoteBasePath = basePath
        }
}

// WithRemoteRawQuery sets the raw query string from the remote URL.
// When set, these query parameters are merged into every outbound request,
// ensuring query parameters from the original registration URL are always forwarded.
// Ignores empty strings; default (no query forwarding) will be used.
func WithRemoteRawQuery(rawQuery string) Option {
        return func(p *TransparentProxy) {
                if rawQuery != "" {
                        p.remoteRawQuery = rawQuery
                }
        }
}

// WithStripConsentCookie removes the entire embedded auth server browser session
// cookie family from outbound backend requests. Enable only when that server is mounted.
func WithStripConsentCookie() Option {
        return func(p *TransparentProxy) {
                p.stripConsentCookie = true
        }
}

// WithStateless configures the proxy for stateless streamable-HTTP servers.
// In stateless mode, incoming GET, HEAD, and DELETE requests receive 405 Method Not Allowed
// instead of being forwarded, and health checks use POST ping instead of GET.
func WithStateless() Option {
        return func(p *TransparentProxy) {
                p.stateless = true
        }
}

// withHealthCheckPingTimeout sets the health check ping timeout.
// This is primarily useful for testing with shorter timeouts.
// Ignores non-positive timeouts; default will be used.
func withHealthCheckPingTimeout(timeout time.Duration) Option {
        return func(p *TransparentProxy) {
                if timeout > 0 {
                        p.healthCheckPingTimeout = timeout
                }
        }
}

// withHealthCheckFailureThreshold sets the consecutive failure count before shutdown.
// This is primarily useful for testing with lower thresholds.
// Ignores non-positive values; default will be used.
func withHealthCheckFailureThreshold(threshold int) Option {
        return func(p *TransparentProxy) {
                if threshold > 0 {
                        p.healthCheckFailureThreshold = threshold
                }
        }
}

// withShutdownTimeout sets the graceful shutdown timeout for the HTTP server.
// This is primarily useful for testing with shorter timeouts.
// Ignores non-positive timeouts; default will be used.
func withShutdownTimeout(timeout time.Duration) Option {
        return func(p *TransparentProxy) {
                if timeout > 0 {
                        p.shutdownTimeout = timeout
                }
        }
}

// WithSessionStorage injects a custom storage backend into the session manager.
// When not provided, the proxy uses in-memory LocalStorage (single-replica default).
// Provide a Redis-backed storage for multi-replica deployments so all replicas
// share the same session store.
func WithSessionStorage(storage session.Storage) Option {
        return func(p *TransparentProxy) {
                if storage == nil {
                        return
                }
                p.sessionStorage = storage
        }
}

// WithSessionTTL overrides the session inactivity timeout used by this proxy.
// Zero or negative values are ignored so the constructor's default is preserved.
func WithSessionTTL(ttl time.Duration) Option {
        return func(p *TransparentProxy) {
                if ttl <= 0 {
                        return
                }
                p.sessionTTL = ttl
        }
}

// WithReadTimeout overrides http.Server.ReadTimeout for this proxy, which bounds
// reading the entire request (headers + body). Zero or negative values are
// ignored so the constructor's default (defaultReadTimeout) is preserved.
func WithReadTimeout(d time.Duration) Option {
        return func(p *TransparentProxy) {
                if d <= 0 {
                        return
                }
                p.readTimeout = d
        }
}

// NewTransparentProxy creates a new transparent proxy with optional middlewares.
// The endpointPrefix parameter specifies an explicit prefix to prepend to SSE endpoint URLs.
// The trustProxyHeaders parameter indicates whether to trust X-Forwarded-* headers from reverse proxies.
// The prefixHandlers parameter is a map of path prefixes to HTTP handlers mounted before the catch-all proxy handler.
func NewTransparentProxy(
        host string,
        port int,
        targetURI string,
        prometheusHandler http.Handler,
        authInfoHandler http.Handler,
        prefixHandlers map[string]http.Handler,
        enableHealthCheck bool,
        isRemote bool,
        transportType string,
        onHealthCheckFailed types.HealthCheckFailedCallback,
        onUnauthorizedResponse types.UnauthorizedResponseCallback,
        endpointPrefix string,
        trustProxyHeaders bool,
        middlewares ...types.NamedMiddleware,
) *TransparentProxy {
        return NewTransparentProxyWithOptions(
                host,
                port,
                targetURI,
                prometheusHandler,
                authInfoHandler,
                prefixHandlers,
                enableHealthCheck,
                isRemote,
                transportType,
                onHealthCheckFailed,
                onUnauthorizedResponse,
                endpointPrefix,
                trustProxyHeaders,
                middlewares,
        )
}

// getHealthCheckInterval returns the health check interval to use.
// Uses TOOLHIVE_HEALTH_CHECK_INTERVAL environment variable if set and valid,
// otherwise returns the default interval.
func getHealthCheckInterval() time.Duration {
        if val := os.Getenv(HealthCheckIntervalEnvVar); val != "" {
                if d, err := time.ParseDuration(val); err == nil && d > 0 {
                        slog.Debug("using custom health check interval", "interval", d)
                        return d
                }
                slog.Warn("invalid health check interval, using default",
                        "env_var", HealthCheckIntervalEnvVar, "value", val, "default", DefaultHealthCheckInterval)
        }
        return DefaultHealthCheckInterval
}

// getHealthCheckPingTimeout returns the health check ping timeout to use.
// Uses TOOLHIVE_HEALTH_CHECK_PING_TIMEOUT environment variable if set and valid,
// otherwise returns the default timeout.
func getHealthCheckPingTimeout() time.Duration {
        if val := os.Getenv(HealthCheckPingTimeoutEnvVar); val != "" {
                if d, err := time.ParseDuration(val); err == nil && d > 0 {
                        slog.Debug("using custom health check ping timeout", "timeout", d)
                        return d
                }
                slog.Warn("invalid health check ping timeout, using default",
                        "env_var", HealthCheckPingTimeoutEnvVar, "value", val, "default", DefaultPingerTimeout)
        }
        return DefaultPingerTimeout
}

// getHealthCheckRetryDelay returns the health check retry delay to use.
// Uses TOOLHIVE_HEALTH_CHECK_RETRY_DELAY environment variable if set and valid,
// otherwise returns the default delay.
func getHealthCheckRetryDelay() time.Duration {
        if val := os.Getenv(HealthCheckRetryDelayEnvVar); val != "" {
                if d, err := time.ParseDuration(val); err == nil && d > 0 {
                        slog.Debug("using custom health check retry delay", "delay", d)
                        return d
                }
                slog.Warn("invalid health check retry delay, using default",
                        "env_var", HealthCheckRetryDelayEnvVar, "value", val, "default", DefaultHealthCheckRetryDelay)
        }
        return DefaultHealthCheckRetryDelay
}

// getHealthCheckFailureThreshold returns the consecutive failure threshold.
// Uses TOOLHIVE_HEALTH_CHECK_FAILURE_THRESHOLD environment variable if set and valid,
// otherwise returns the default threshold.
func getHealthCheckFailureThreshold() int {
        if val := os.Getenv(HealthCheckFailureThresholdEnvVar); val != "" {
                if n, err := strconv.Atoi(val); err == nil && n > 0 {
                        slog.Debug("using custom health check failure threshold", "threshold", n)
                        return n
                }
                slog.Warn("invalid health check failure threshold, using default",
                        "env_var", HealthCheckFailureThresholdEnvVar, "value", val, "default", DefaultHealthCheckFailureThreshold)
        }
        return DefaultHealthCheckFailureThreshold
}

// NewTransparentProxyWithOptions creates a new transparent proxy with optional configuration.
func NewTransparentProxyWithOptions(
        host string,
        port int,
        targetURI string,
        prometheusHandler http.Handler,
        authInfoHandler http.Handler,
        prefixHandlers map[string]http.Handler,
        enableHealthCheck bool,
        isRemote bool,
        transportType string,
        onHealthCheckFailed types.HealthCheckFailedCallback,
        onUnauthorizedResponse types.UnauthorizedResponseCallback,
        endpointPrefix string,
        trustProxyHeaders bool,
        middlewares []types.NamedMiddleware,
        options ...Option,
) *TransparentProxy {
        proxy := &TransparentProxy{
                host:                        host,
                port:                        port,
                targetURI:                   targetURI,
                middlewares:                 middlewares,
                shutdownCh:                  make(chan struct{}),
                prometheusHandler:           prometheusHandler,
                authInfoHandler:             authInfoHandler,
                prefixHandlers:              prefixHandlers,
                sessionTTL:                  session.DefaultSessionTTL,
                readTimeout:                 defaultReadTimeout,
                isRemote:                    isRemote,
                transportType:               transportType,
                onHealthCheckFailed:         onHealthCheckFailed,
                onUnauthorizedResponse:      onUnauthorizedResponse,
                endpointPrefix:              endpointPrefix,
                trustProxyHeaders:           trustProxyHeaders,
                healthCheckInterval:         getHealthCheckInterval(),
                healthCheckRetryDelay:       getHealthCheckRetryDelay(),
                healthCheckPingTimeout:      getHealthCheckPingTimeout(),
                healthCheckFailureThreshold: getHealthCheckFailureThreshold(),
                shutdownTimeout:             defaultShutdownTimeout,
        }

        // Apply options
        for _, opt := range options {
                opt(proxy)
        }

        // Construct the session manager once, after options have resolved sessionTTL and sessionStorage.
        proxyFactory := func(id string) session.Session { return session.NewProxySession(id) }
        if proxy.sessionStorage != nil {
                proxy.sessionManager = session.NewManagerWithStorage(proxy.sessionTTL, proxyFactory, proxy.sessionStorage)
        } else {
                proxy.sessionManager = session.NewManager(proxy.sessionTTL, proxyFactory)
        }

        // Create appropriate response processor based on transport type
        proxy.responseProcessor = createResponseProcessor(
                transportType,
                proxy,
                endpointPrefix,
                trustProxyHeaders,
        )

        // Create health checker always for Kubernetes probes
        var mcpPinger healthcheck.MCPPinger
        if enableHealthCheck {
                if proxy.stateless {
                        mcpPinger = NewStatelessMCPPingerWithTimeout(targetURI, proxy.healthCheckPingTimeout)
                } else {
                        mcpPinger = NewMCPPingerWithTimeout(targetURI, proxy.healthCheckPingTimeout)
                }
        }
        proxy.healthChecker = healthcheck.NewHealthChecker(transportType, mcpPinger)

        return proxy
}

// recoverySessionStore is the subset of session.Manager that backendRecovery needs.
type recoverySessionStore interface {
        sessionbinding.Store
        LoadIfOwner(id, expectedOwner string) (session.Session, error)
        UpsertSessionIfOwner(sess session.Session, expectedOwner string) error
        DeleteIfOwner(id, expectedOwner string) error
}

// backendRecovery handles transparent re-initialization of backend sessions when the
// target pod is unreachable (dial error) or has lost its in-memory session state (404).
// It depends only on a narrow session interface and a forward function, so it can be
// tested without standing up a full proxy.
type backendRecovery struct {
        targetURI string
        forward   func(*http.Request) (*http.Response, error)
        sessions  recoverySessionStore
}

type tracingTransport struct {
        p        *TransparentProxy
        recovery *backendRecovery
}

func newTracingTransport(base http.RoundTripper, p *TransparentProxy) *tracingTransport {
        return &tracingTransport{
                p: p,
                recovery: &backendRecovery{
                        targetURI: p.targetURI,
                        forward:   base.RoundTrip,
                        sessions:  p.sessionManager,
                },
        }
}

func (p *TransparentProxy) setServerInitialized() {
        if p.isServerInitialized.CompareAndSwap(false, true) {
                //nolint:gosec // G706: logging target URI from config
                slog.Debug("server was initialized successfully", "target", p.targetURI)
        }
}

// serverInitialized returns whether the server has been initialized (thread-safe)
func (p *TransparentProxy) serverInitialized() bool {
        return p.isServerInitialized.Load()
}

// nolint:gocyclo // This function handles multiple request types and is complex by design
func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
        // Always rewrite Host header to match the target URL to avoid "Invalid Host" errors
        // from backends with strict host validation (e.g., Django ALLOWED_HOSTS, FastAPI TrustedHostMiddleware).
        // Without this, the backend receives the proxy's Host header (e.g., Kubernetes service DNS name)
        // instead of its own hostname, causing validation failures.
        // See: https://github.com/stacklok/stacklok-epics/issues/231
        if req.URL.Host != req.Host {
                req.Host = req.URL.Host
        }

        slog.Debug("outbound request to upstream",
                "method", req.Method,
                "url", req.URL.String(),
                "host", req.Host,
                "accept", req.Header.Get("Accept"),
                "content_type", req.Header.Get("Content-Type"),
        )

        reqBody, err := readRequestBody(req)
        if err != nil {
                // Never forward a partial body, even on non-size read errors.
                if bodylimit.IsRequestTooLarge(err) {
                        slog.Warn("rejected unreadable MCP request body", "reason", "body_too_large")
                        return plainResponse(req, http.StatusRequestEntityTooLarge, "Request Entity Too Large"), nil
                }
                slog.Warn("rejected unreadable MCP request body", "reason", "read_error")
                return plainResponse(req, http.StatusBadRequest, "Invalid request body"), nil
        }

        sawInitialize := false
        revision := mcp.RevisionLegacy
        // Only ID-bearing requests are classified and supply an ID for local errors.
        // Notifications, client responses and bodiless GET/DELETE leave it absent.
        var requestID any

        if req.Method == http.MethodPost {
                msg, err := mcp.DecodeMessage(reqBody)
                if err != nil {
                        return mcp.ClassificationErrorResponse(req, nil, err), nil
                }
                if rpcRequest, ok := msg.(*jsonrpc2.Request); ok {
                        sawInitialize = rpcRequest.Method == "initialize"
                        if rpcRequest.ID.IsValid() {
                                requestID = rpcRequest.ID.Raw()
                                meta := mcp.ExtractMeta(rpcRequest.Params)
                                rev, cerr := mcp.ClassifyRevision(rpcRequest.Method, meta, req.Header.Get("MCP-Protocol-Version"))
                                if cerr != nil {
                                        // Malformed Modern request: reject before the backend is ever contacted.
                                        return mcp.ClassificationErrorResponse(req, requestID, cerr), nil
                                }
                                revision = rev
                        }
                }
        }
        //nolint:gosec // G706: logging target URI from config
        slog.Debug("classified request revision", "modern", revision == mcp.RevisionModern, "target", t.p.targetURI)

        // Transparent requests forward backend session authority even for Modern MCP.
        clientSID, carrierErr := requestSessionID(req)
        if carrierErr != nil {
                return session.NotFoundResponse(req, requestID), nil
        }
        // Legacy SSE establishes its session through an endpoint event before initialize.
        // An unspecified transport (thv proxy) uses the endpoint's query carrier; it
        // must pass the same ownership check below, not start a streamable session.
        legacySSE := t.p.transportType == types.TransportTypeSSE.String()
        if t.p.transportType == "" {
                for _, key := range sessionQueryKeys {
                        legacySSE = legacySSE || req.URL.Query().Has(key)
                }
        }
        startsSession := sawInitialize && !legacySSE
        var (
                routedSession session.Session
                expectedOwner string
        )
        if startsSession {
                clientSID = ""
                req.Header.Del("Mcp-Session-Id")
                rewriteSessionQuery(req.URL, "")
        } else if clientSID != "" {
                identity, _ := auth.IdentityFromContext(req.Context())
                expectedOwner, err = sessionbinding.FromIdentity(identity)
                if err != nil {
                        return session.NotFoundResponse(req, requestID), nil
                }
                routedSession, err = t.p.sessionManager.LoadIfOwner(normalizeSessionID(clientSID), expectedOwner)
                if err != nil {
                        return ownershipErrorResponse(req, requestID, err), nil
                }
        }

        if clientSID != "" {
                // Route only from the same snapshot whose ownership was atomically validated.
                if backendURL, exists := routedSession.GetMetadataValue(sessionMetadataBackendURL); exists {
                        if parsed, err := url.Parse(backendURL); err == nil && parsed.Scheme != "" && parsed.Host != "" {
                                req.URL.Scheme, req.URL.Host, req.Host = parsed.Scheme, parsed.Host, parsed.Host
                        }
                }
                if backendSID, exists := routedSession.GetMetadataValue(sessionMetadataBackendSID); exists && backendSID != "" {
                        req.Header.Set("Mcp-Session-Id", backendSID)
                        rewriteSessionQuery(req.URL, backendSID)
                }
        }

        // Attach an httptrace to capture the actual backend pod IP after kube-proxy
        // DNAT resolves the ClusterIP to a specific pod. The captured address is stored
        // as backend_url so follow-up requests always reach the same pod, even after a
        // proxy runner restart that would otherwise lose the in-memory routing state.
        var capturedPodAddr string
        if sawInitialize {
                trace := &httptrace.ClientTrace{
                        GotConn: func(info httptrace.GotConnInfo) {
                                capturedPodAddr = info.Conn.RemoteAddr().String()
                        },
                }
                req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
        }

        resp, err := followRedirects(t.recovery.forward, req, reqBody)
        if err != nil {
                if errors.Is(err, context.Canceled) {
                        // Expected during shutdown or client disconnect—silently ignore
                        return nil, err
                }
                // Never recover during initialize: streamable HTTP has discarded the
                // supplied SID, and legacy SSE is still completing its first handshake.
                if isDialError(err) && !sawInitialize {
                        req.Header.Set("Mcp-Session-Id", clientSID)
                        if reInitResp, reInitErr := t.recovery.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil {
                                return reInitResp, reInitErr
                        }
                }
                slog.Error("failed to forward request", "error", err)
                return nil, err
        }

        slog.Debug("upstream response received",
                "status", resp.StatusCode,
                "url", req.URL.String(),
                "content_type", resp.Header.Get("Content-Type"),
                "mcp_session_id", resp.Header.Get("Mcp-Session-Id"),
        )

        // Check for 401 Unauthorized response (bearer token authentication failure)
        if resp.StatusCode == http.StatusUnauthorized {
                //nolint:gosec // G706: logging target URI from config
                slog.Debug("received 401 Unauthorized response, bearer token may be invalid",
                        "target", t.p.targetURI)
                if t.p.onUnauthorizedResponse != nil {
                        t.p.onUnauthorizedResponse()
                }
        }

        // Clean up session on DELETE so the transparent proxy's session manager
        // doesn't hold references until TTL expiry (#4062).
        // Remove on 2xx (successful termination) and 404 (upstream already
        // considers the session gone), since in both cases keeping a local
        // reference would only waste memory.
        if req.Method == http.MethodDelete &&
                (resp.StatusCode >= 200 && resp.StatusCode < 300 || resp.StatusCode == http.StatusNotFound) {
                if clientSID != "" {
                        if err := t.p.sessionManager.DeleteIfOwner(normalizeSessionID(clientSID), expectedOwner); err != nil {
                                _, _ = io.Copy(io.Discard, resp.Body)
                                _ = resp.Body.Close()
                                return ownershipErrorResponse(req, requestID, err), nil
                        }
                }
        }

        // Backend returned 404 for a non-initialize, non-DELETE request whose session IS
        // known to the proxy. This means the backend pod lost its in-memory session state
        // (e.g. it was restarted but got the same IP). Attempt transparent re-initialization
        // so the client sees no error. DELETE is excluded because the session has already
        // been cleaned up above and the 404 is the expected terminal response.
        if resp.StatusCode == http.StatusNotFound && !sawInitialize && req.Method != http.MethodDelete {
                if clientSID != "" {
                        req.Header.Set("Mcp-Session-Id", clientSID)
                        if reInitResp, reInitErr := t.recovery.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil {
                                _, _ = io.Copy(io.Discard, resp.Body)
                                _ = resp.Body.Close()
                                return reInitResp, reInitErr
                        }
                }
        }

        if resp.StatusCode == http.StatusOK {
                // check if we saw a valid mcp header
                ct := resp.Header.Get("Mcp-Session-Id")
                if ct != "" {
                        //nolint:gosec // G706: logging session ID from HTTP response header
                        slog.Debug("detected Mcp-Session-Id header", "session_id", ct)
                        internalID := normalizeSessionID(ct)
                        sess := session.NewProxySession(internalID)
                        sess.SetMetadata(sessionMetadataBackendURL, t.recovery.podBackendURL(capturedPodAddr))
                        if sawInitialize && len(reqBody) > 0 {
                                sess.SetMetadata(sessionMetadataInitBody, string(reqBody))
                        }
                        if err := sessionbinding.AddOwnedSession(req.Context(), t.p.sessionManager, sess); err != nil {
                                _, _ = io.Copy(io.Discard, resp.Body)
                                _ = resp.Body.Close()
                                return ownershipErrorResponse(req, requestID, err), nil
                        }
                        t.p.setServerInitialized()
                        return resp, nil
                }
                // status was ok and we saw an initialize call
                if sawInitialize && !t.p.serverInitialized() {
                        t.p.setServerInitialized()
                        return resp, nil
                }

                // Modern (2026-07-28, stateless) has no initialize handshake, so use the
                // first successful (HTTP 200) Modern request as the readiness signal: it
                // lets the background health monitor begin probing this backend, which
                // otherwise never starts for a pure-Modern server. See #5831.
                if revision == mcp.RevisionModern && !t.p.serverInitialized() {
                        t.p.setServerInitialized()
                }
        }

        return resp, nil
}

func ownershipErrorResponse(req *http.Request, requestID any, err error) *http.Response {
        if errors.Is(err, sessionbinding.ErrNotFound) {
                return session.NotFoundResponse(req, requestID)
        }
        return plainResponse(req, http.StatusServiceUnavailable, "Session store unavailable")
}

// plainResponse builds a minimal text/plain *http.Response for the given status,
// used to synthesize error responses from within RoundTrip.
func plainResponse(req *http.Request, status int, body string) *http.Response {
        hdr := make(http.Header)
        hdr.Set("Content-Type", "text/plain; charset=utf-8")
        return &http.Response{
                StatusCode: status,
                Status:     fmt.Sprintf("%d %s", status, http.StatusText(status)),
                Proto:      "HTTP/1.1",
                ProtoMajor: 1,
                ProtoMinor: 1,
                Header:     hdr,
                Body:       io.NopCloser(strings.NewReader(body + "\n")),
                Request:    req,
        }
}

func readRequestBody(req *http.Request) ([]byte, error) {
        reqBody := []byte{}
        if req.Body != nil {
                buf, err := io.ReadAll(req.Body)
                if err != nil {
                        return nil, err
                }
                reqBody = buf
                req.Body = io.NopCloser(bytes.NewReader(reqBody))
        }
        return reqBody, nil
}

// podBackendURL constructs a backend URL that targets the specific pod IP captured
// via httptrace.GotConn, using the scheme from targetURI. Falls back to targetURI
// when no address was captured (e.g. single-replica, connection reuse without a new conn),
// or when targetURI uses HTTPS — IP-literal HTTPS URLs fail TLS verification because
// server certificates are issued for hostnames, not pod IPs.
func (r *backendRecovery) podBackendURL(capturedAddr string) string {
        if capturedAddr == "" {
                return r.targetURI
        }
        parsed, err := url.Parse(r.targetURI)
        if err != nil {
                return r.targetURI
        }
        if parsed.Scheme == "https" { //nolint:goconst // protocol name, not a magic string
                return r.targetURI
        }
        parsed.Host = capturedAddr
        return parsed.String()
}

// isDialError reports whether err is a TCP dial failure, indicating that the
// target host is unreachable (pod has been terminated or rescheduled).
func isDialError(err error) bool {
        var opErr *net.OpError
        return errors.As(err, &opErr) && opErr.Op == "dial"
}

// followRedirects wraps a forward call with same-host HTTP redirect following.
// MCP clients expect JSON-RPC responses and cannot handle 3xx redirects, so the
// proxy must resolve them before returning the response. Only same-host redirects
// are followed to prevent SSRF. The HTTP method and request body are always
// preserved (POST never becomes GET), which is correct for JSON-RPC semantics.
func followRedirects(
        forward func(*http.Request) (*http.Response, error),
        req *http.Request,
        body []byte,
) (*http.Response, error) {
        resp, err := forward(req)
        if err != nil {
                return nil, err
        }

        originalHost := req.URL.Host
        for redirectsFollowed := 0; redirectsFollowed < maxRedirects &&
                isRedirectStatus(resp.StatusCode); redirectsFollowed++ {
                location := resp.Header.Get("Location")
                if location == "" {
                        break
                }

                redirectURL, parseErr := req.URL.Parse(location)
                if parseErr != nil {
                        slog.Warn("failed to parse redirect Location header",
                                "location", location, "error", parseErr)
                        break
                }

                // Block cross-host redirects to prevent SSRF and credential leakage.
                if redirectURL.Host != originalHost {
                        slog.Warn("refusing cross-host redirect from remote MCP server; update the configured target URL",
                                "from_host", originalHost, "to_host", redirectURL.Host)
                        break
                }

                // Block HTTPS-to-HTTP downgrades to prevent silent loss of transport security.
                //nolint:goconst // "https" is a protocol name, not a magic string worth extracting
                if req.URL.Scheme == "https" && redirectURL.Scheme == "http" {
                        slog.Warn("refusing redirect that downgrades from HTTPS to HTTP",
                                "from", req.URL.String(), "to", redirectURL.String())
                        break
                }

                slog.Info("following HTTP redirect from remote MCP server; consider updating the server URL",
                        "status", resp.StatusCode,
                        "from", req.URL.String(),
                        "to", redirectURL.String(),
                        "redirect_number", redirectsFollowed+1)

                // Drain and close the redirect response body to release the
                // underlying connection back to the transport's connection pool.
                _, _ = io.Copy(io.Discard, resp.Body)
                _ = resp.Body.Close()

                // Clone preserves Method and all headers. We intentionally do not
                // change the method to GET for 301/302 (as browsers do) because
                // MCP JSON-RPC requires POST with a body on every request.
                req = req.Clone(req.Context())
                req.URL = redirectURL
                req.Host = redirectURL.Host
                if len(body) > 0 {
                        req.Body = io.NopCloser(bytes.NewReader(body))
                        req.ContentLength = int64(len(body))
                }

                resp, err = forward(req)
                if err != nil {
                        return nil, err
                }
        }

        return resp, nil
}

// isRedirectStatus reports whether the HTTP status code is a redirect
// that should be followed. This excludes 300 (Multiple Choices), 303
// (See Other), and 304 (Not Modified) which are not standard redirects
// or would require changing the request method.
func isRedirectStatus(code int) bool {
        switch code {
        case http.StatusMovedPermanently, // 301
                http.StatusFound,             // 302
                http.StatusTemporaryRedirect, // 307
                http.StatusPermanentRedirect: // 308
                return true
        default:
                return false
        }
}

func (r *backendRecovery) loadOwnedSession(req *http.Request, id string) (session.Session, string, error) {
        identity, _ := auth.IdentityFromContext(req.Context())
        expectedOwner, err := sessionbinding.FromIdentity(identity)
        if err != nil {
                return nil, "", sessionbinding.ErrNotFound
        }
        stored, err := r.sessions.LoadIfOwner(id, expectedOwner)
        if err != nil {
                return nil, "", err
        }
        return stored, expectedOwner, nil
}

// reinitializeAndReplay is called when the proxy detects that the backend pod
// that owned a session is no longer reachable (dial error) or has lost its
// in-memory session state (backend returned 404). It transparently:
//  1. Re-sends the stored initialize body to the ClusterIP service so kube-proxy
//     selects a healthy pod and the backend creates a new session.
//  2. Captures the new pod IP via httptrace.GotConn and stores it as backend_url.
//  3. Maps the client's original session ID to the new backend session ID.
//  4. Replays the original client request so the client sees no error.
//
// Returns (nil, nil) when there is no SID or no stored initialize body. Missing,
// unowned, or foreign sessions fail closed before any backend work.
func (r *backendRecovery) reinitializeAndReplay(req *http.Request, origBody []byte) (*http.Response, error) {
        sid := req.Header.Get("Mcp-Session-Id")
        if sid == "" {
                return nil, nil
        }
        internalSID := normalizeSessionID(sid)
        stored, expectedOwner, err := r.loadOwnedSession(req, internalSID)
        if err != nil {
                return ownershipErrorResponse(req, nil, err), nil
        }
        // Do not mutate a value returned by the authoritative store before the
        // replacement has been persisted successfully.
        sess := session.NewProxySession(internalSID)
        for key, value := range stored.GetMetadata() {
                sess.SetMetadata(key, value)
        }

        initBody, hasInit := sess.GetMetadataValue(sessionMetadataInitBody)
        if !hasInit || initBody == "" {
                // No stored init body — cannot re-initialize transparently.
                // Reset backend_url to ClusterIP so the next request goes through
                // kube-proxy and lets the client receive a clean 404 to re-initialize.
                sess.SetMetadata(sessionMetadataBackendURL, r.targetURI)
                if err := r.sessions.UpsertSessionIfOwner(sess, expectedOwner); err != nil {
                        return ownershipErrorResponse(req, nil, err), nil
                }
                return nil, nil
        }

        slog.Debug("backend session lost; transparently re-initializing",
                "session_id", sid, "target", r.targetURI)

        // Capture the new pod IP via GotConn on the re-initialize connection.
        var capturedPodAddr string
        trace := &httptrace.ClientTrace{
                GotConn: func(info httptrace.GotConnInfo) {
                        capturedPodAddr = info.Conn.RemoteAddr().String()
                },
        }
        initCtx := httptrace.WithClientTrace(req.Context(), trace)

        // Build a fresh initialize request to the ClusterIP (no Mcp-Session-Id —
        // the backend assigns a new session ID in the response).
        parsedTarget, err := url.Parse(r.targetURI)
        if err != nil {
                return nil, nil
        }
        initURL := *req.URL
        initURL.Scheme = parsedTarget.Scheme
        initURL.Host = parsedTarget.Host
        rewriteSessionQuery(&initURL, "")

        initReq, err := http.NewRequestWithContext(initCtx, http.MethodPost, initURL.String(), bytes.NewReader([]byte(initBody)))
        if err != nil {
                return nil, nil
        }
        // Propagate headers from the original request (Authorization, tenant headers, etc.)
        // so the backend accepts the re-initialize. Mcp-Session-Id must not be forwarded —
        // the backend assigns a new session ID in the response. Content-Length and
        // Transfer-Encoding are deleted because http.NewRequestWithContext already set
        // ContentLength from the body; leaving stale header values would be misleading
        // (Go's transport ignores them in favour of the struct field, but clarity matters).
        initReq.Header = req.Header.Clone()
        initReq.Header.Del("Mcp-Session-Id")
        initReq.Header.Del("Content-Length")
        initReq.Header.Del("Transfer-Encoding")
        initReq.Header.Set("Content-Type", "application/json")

        initResp, err := r.forward(initReq)
        if err != nil {
                slog.Error("transparent re-initialize failed", "error", err)
                return nil, err
        }
        _, _ = io.Copy(io.Discard, initResp.Body)
        _ = initResp.Body.Close()

        newBackendSID := initResp.Header.Get("Mcp-Session-Id")
        if newBackendSID == "" {
                slog.Debug("re-initialize response contained no Mcp-Session-Id; falling back to ClusterIP")
                sess.SetMetadata(sessionMetadataBackendURL, r.targetURI)
                if err := r.sessions.UpsertSessionIfOwner(sess, expectedOwner); err != nil {
                        return ownershipErrorResponse(req, nil, err), nil
                }
                return nil, nil
        }

        // Reserve the backend ID before routing to it, so recovery cannot adopt an
        // ID already owned by another principal. Preserve the original owner record.
        backendInternalSID := normalizeSessionID(newBackendSID)
        reservedBackend := backendInternalSID != internalSID
        if reservedBackend {
                backendSession := session.NewProxySession(backendInternalSID)
                if err := sessionbinding.AddOwnedSession(req.Context(), r.sessions, backendSession); err != nil {
                        return ownershipErrorResponse(req, nil, err), nil
                }
        }

        // Update session: point backend_url at the newly-discovered pod and record
        // the backend session ID so tracingTransport.RoundTrip rewrites Mcp-Session-Id on outbound requests.
        newPodURL := r.podBackendURL(capturedPodAddr)
        sess.SetMetadata(sessionMetadataBackendURL, newPodURL)
        // Store the raw backend session ID (not normalized) because the Rewrite closure
        // uses this value verbatim as the outbound Mcp-Session-Id header. Normalizing
        // would change non-UUID IDs to a UUID v5 hash the backend never issued.
        sess.SetMetadata(sessionMetadataBackendSID, newBackendSID)
        if upsertErr := r.sessions.UpsertSessionIfOwner(sess, expectedOwner); upsertErr != nil {
                return r.failedRecoveryResponse(req, upsertErr, reservedBackend, backendInternalSID, expectedOwner)
        }

        return r.replay(req, origBody, newBackendSID, newPodURL)
}

func (r *backendRecovery) failedRecoveryResponse(
        req *http.Request,
        upsertErr error,
        reservedBackend bool,
        backendInternalSID, expectedOwner string,
) (*http.Response, error) {
        if reservedBackend {
                if cleanupErr := r.sessions.DeleteIfOwner(backendInternalSID, expectedOwner); cleanupErr != nil {
                        slog.Warn("failed to remove backend session reservation after recovery failure", "error", cleanupErr)
                }
                // The backend-created session is left for backend expiry. Terminating it
                // cannot be made atomic with releasing and potentially reusing its ID.
        }
        return ownershipErrorResponse(req, nil, upsertErr), nil
}

// replay uses only the newly reserved backend SID. podBackendURL preserves the
// original hostname for HTTPS, where an IP literal would fail TLS verification.
func (r *backendRecovery) replay(req *http.Request, origBody []byte, backendSID, backendURL string) (*http.Response, error) {
        target, err := url.Parse(backendURL)
        if err != nil {
                return nil, err
        }
        replayReq := req.Clone(req.Context())
        replayReq.URL.Scheme = target.Scheme
        replayReq.URL.Host = target.Host
        replayReq.Host = target.Host
        replayReq.Header.Set("Mcp-Session-Id", backendSID)
        rewriteSessionQuery(replayReq.URL, backendSID)
        replayReq.Body = io.NopCloser(bytes.NewReader(origBody))
        replayReq.ContentLength = int64(len(origBody))
        // origBody is fully buffered, so chunked encoding is unnecessary and would
        // suppress the Content-Length header. Clear any TransferEncoding copied from
        // the original request so net/http sends Content-Length instead.
        replayReq.TransferEncoding = nil

        slog.Debug("replaying original request after transparent re-initialization")
        resp, err := r.forward(replayReq)
        if err != nil {
                return nil, err
        }
        if returnedSID := resp.Header.Get("Mcp-Session-Id"); returnedSID != "" {
                returnedSession := session.NewProxySession(normalizeSessionID(returnedSID))
                if err := sessionbinding.AddOwnedSession(req.Context(), r.sessions, returnedSession); err != nil {
                        _, _ = io.Copy(io.Discard, resp.Body)
                        _ = resp.Body.Close()
                        return ownershipErrorResponse(req, nil, err), nil
                }
        }
        return resp, nil
}

// modifyResponse modifies HTTP responses based on transport-specific requirements.
// Delegates to the appropriate ResponseProcessor based on transport type.
func (p *TransparentProxy) modifyResponse(resp *http.Response) error {
        return p.responseProcessor.ProcessResponse(resp)
}

// setXForwardedHeaders populates the standard X-Forwarded-* headers on the
// outbound request. For remote upstreams X-Forwarded-Host is removed: a
// third-party server may use it to construct redirect URLs pointing back at
// the proxy, producing 307 redirect loops. X-Forwarded-For is kept so remote
// backends can still log the client IP. Client-supplied X-Forwarded-* values
// never pass through either way — httputil strips them from the outbound
// request before Rewrite runs.
//
// For remote upstreams X-Forwarded-Proto is also rewritten to the scheme of
// the actual upstream connection. SetXForwarded derives it from the inbound
// connection, which behind a TLS-terminating load balancer is plain HTTP even
// though the upstream is reached over HTTPS. Upstreams that redirect when
// X-Forwarded-Proto != https would otherwise 301-loop forever. The upstream
// scheme is supplied by the caller (parsed once from targetURI in Start) to
// avoid re-parsing on every request.
func (p *TransparentProxy) setXForwardedHeaders(pr *httputil.ProxyRequest, upstreamScheme string) {
        pr.SetXForwarded()
        if p.isRemote {
                pr.Out.Header.Del("X-Forwarded-Host")
                if upstreamScheme != "" {
                        pr.Out.Header.Set("X-Forwarded-Proto", upstreamScheme)
                        slog.Debug("set X-Forwarded-Proto for remote upstream",
                                "scheme", upstreamScheme, "target", p.targetURI)
                }
        }
}

// stripConsentCookie removes the embedded AS cookie family, retaining other cookie
// pairs across all Cookie header lines. The caller supplies the outbound header.
func stripConsentCookie(header http.Header) {
        var kept []string
        removed := false
        for key, values := range header {
                if !strings.EqualFold(key, "Cookie") {
                        continue
                }
                for _, value := range values {
                        var pairs []string
                        for _, pair := range strings.Split(value, ";") {
                                name, _, ok := strings.Cut(strings.TrimSpace(pair), "=")
                                if ok && strings.HasPrefix(strings.TrimSpace(name), "__Host-thv_consent_") {
                                        removed = true
                                        continue
                                }
                                pairs = append(pairs, pair)
                        }
                        if len(pairs) > 0 {
                                kept = append(kept, strings.Join(pairs, ";"))
                        }
                }
        }
        if !removed {
                return
        }
        for key := range header {
                if strings.EqualFold(key, "Cookie") {
                        delete(header, key)
                }
        }
        if len(kept) > 0 {
                header["Cookie"] = kept
        }
}

// Start starts the transparent proxy.
// nolint:gocyclo // This function handles multiple startup scenarios and is complex by design
func (p *TransparentProxy) Start(ctx context.Context) error {
        p.mutex.Lock()
        defer p.mutex.Unlock()

        // Guard against calling Start() after Stop()
        if p.stopped {
                return fmt.Errorf("proxy has been stopped")
        }

        // Parse the target URI
        targetURL, err := url.Parse(p.targetURI)
        if err != nil {
                return fmt.Errorf("failed to parse target URI: %w", err)
        }

        // Create a reverse proxy
        proxy := &httputil.ReverseProxy{
                FlushInterval: -1,
                Rewrite: func(pr *httputil.ProxyRequest) {
                        pr.SetURL(targetURL)
                        if p.stripConsentCookie {
                                stripConsentCookie(pr.Out.Header)
                        }
                        p.setXForwardedHeaders(pr, targetURL.Scheme)

                        // Stash the original inbound request in the outbound request's
                        // context so that ModifyResponse (SSE response processor) can
                        // read the client's real headers instead of the auto-injected
                        // X-Forwarded-* values that SetXForwarded() wrote to pr.Out.
                        ctx := InboundRequestToContext(pr.Out.Context(), pr.In)
                        pr.Out = pr.Out.WithContext(ctx)

                        // Rewrite path to the remote server's path when configured.
                        // When the remote URL has a path (e.g., /v2/mcp), the target URI only
                        // contains the scheme+host. The client sends to /mcp (default MCP
                        // endpoint) but the remote server expects /v2/mcp. We replace the
                        // request path with the remote server's configured path.
                        if p.remoteBasePath != "" {
                                pr.Out.URL.Path = p.remoteBasePath
                                pr.Out.URL.RawPath = ""
                        }

                        // Merge query parameters from the remote URL into the outbound request.
                        // Remote params are prepended so they appear first; most HTTP servers
                        // adopt first-value-wins semantics for duplicate keys, ensuring operator
                        // configured values (e.g., toolsets=core,alerting) take precedence over
                        // any client-supplied params with the same key.
                        // Raw string concatenation is intentional: url.Values.Encode() would
                        // percent-encode characters like commas that some APIs expect as literals.
                        if p.remoteRawQuery != "" {
                                merged := p.remoteRawQuery
                                if pr.Out.URL.RawQuery != "" {
                                        merged += "&" + pr.Out.URL.RawQuery
                                }
                                pr.Out.URL.RawQuery = merged
                        }

                        // Inject OpenTelemetry trace propagation headers for downstream tracing
                        if pr.Out.Context() != nil {
                                otel.GetTextMapPropagator().Inject(pr.Out.Context(), propagation.HeaderCarrier(pr.Out.Header))
                        }
                },
        }

        proxy.Transport = newTracingTransport(http.DefaultTransport, p)
        proxy.ModifyResponse = func(resp *http.Response) error {
                return p.modifyResponse(resp)
        }

        handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                proxy.ServeHTTP(w, r) // #nosec G704 -- target URL is the configured backend MCP server
        })

        // Create a mux to handle both proxy and health endpoints
        mux := http.NewServeMux()

        // Apply middleware chain in reverse order (last middleware is applied first)
        var finalHandler http.Handler = handler
        for i := len(p.middlewares) - 1; i >= 0; i-- {
                finalHandler = p.middlewares[i].Function(finalHandler)
                slog.Debug("applied middleware", "name", p.middlewares[i].Name)
        }

        // 1. Mount prefix handlers first (user-specified, most specific paths)
        // These are registered first but Go's ServeMux longest-match routing ensures
        // more specific paths take precedence regardless of registration order.
        for prefix, prefixHandler := range p.prefixHandlers {
                mux.Handle(prefix, prefixHandler)
                slog.Debug("mounted prefix handler", "prefix", prefix)
        }

        // 2. Mount health check endpoint if enabled, otherwise return 404
        // (prevents /health from being proxied to the backend)
        if p.healthChecker != nil {
                mux.Handle("/health", p.healthChecker)
        } else {
                mux.HandleFunc("/health", http.NotFound)
        }

        // 3. Mount Prometheus metrics endpoint if handler is provided (no middlewares),
        // otherwise return 404 (prevents /metrics from being proxied to the backend).
        //
        // The runner no longer supplies a handler here: metrics are served on a
        // separate diagnostics listener (see pkg/diagnostics) so access can be
        // restricted by port, which NetworkPolicy can express and a shared port
        // cannot. Registering "/metrics" on this mux would also outrank the "/"
        // catch-all below and so sit outside the middleware chain.
        if p.prometheusHandler != nil {
                mux.Handle("/metrics", p.prometheusHandler)
                slog.Debug("prometheus metrics endpoint enabled at /metrics")
        } else {
                mux.Handle(diagnostics.MetricsPath, diagnostics.NotServedHereHandler())
        }

        // 4. Mount RFC 9728 OAuth Protected Resource discovery endpoint (no middlewares)
        // Note: This is DIFFERENT from the auth server's /.well-known/oauth-authorization-server
        // Always register so OAuth discovery gets a clean 404 JSON when auth is off,
        // instead of falling through to the proxy catch-all.
        wellKnownHandler := auth.NewWellKnownHandler(p.authInfoHandler)
        mux.Handle("/.well-known/", wellKnownHandler)
        if p.authInfoHandler != nil {
                slog.Debug("rfc 9728 OAuth discovery endpoint enabled at /.well-known/oauth-protected-resource")
        }

        // 5. Catch-all proxy handler (least specific - ServeMux routing handles precedence)
        // Note: No manual path checking needed - ServeMux longest-match routing ensures
        // more specific paths registered above take precedence over this catch-all.
        finalHandler = p.methodGate(finalHandler)
        mux.Handle("/", finalHandler)

        // Use ListenConfig with SO_REUSEADDR to allow port reuse after unclean shutdown
        // (e.g., after laptop sleep where zombie processes may hold ports)
        lc := socket.ListenConfig()
        ln, err := lc.Listen(context.Background(), "tcp", fmt.Sprintf("%s:%d", p.host, p.port))
        if err != nil {
                return fmt.Errorf("failed to listen: %w", err)
        }
        p.listener = ln

        // Create the server
        p.server = &http.Server{
                Addr:              fmt.Sprintf("%s:%d", p.host, p.port),
                Handler:           mux,
                ReadHeaderTimeout: 10 * time.Second,   // Prevent Slowloris attacks
                ReadTimeout:       p.readTimeout,      // Bound slow body uploads
                IdleTimeout:       defaultIdleTimeout, // Prevent idle keep-alive connections from blocking Shutdown()
        }

        // Capture server in local variable to avoid race with Stop()
        server := p.server
        go func() {
                err := server.Serve(ln)
                if err != nil && !errors.Is(err, http.ErrServerClosed) {
                        var opErr *net.OpError
                        if errors.As(err, &opErr) && opErr.Op == "accept" {
                                // Expected when listener is closed—silently return
                                return
                        }
                        slog.Error("transparent proxy error", "error", err)
                }
        }()
        // Start health-check monitoring only if health checker is enabled
        if p.healthChecker != nil {
                go p.monitorHealth(ctx)
        }

        return nil
}

// ListenerAddr returns the network address the proxy is listening on.
// Returns an empty string if the proxy has not been started.
func (p *TransparentProxy) ListenerAddr() string {
        if p.listener == nil {
                return ""
        }
        return p.listener.Addr().String()
}

// CloseListener closes the listener for the transparent proxy.
func (p *TransparentProxy) CloseListener() error {
        if p.listener != nil {
                return p.listener.Close()
        }
        return nil
}

// performHealthCheckRetry performs a retry health check after a delay
// Returns true if the retry was successful (health check recovered), false otherwise
func (p *TransparentProxy) performHealthCheckRetry(ctx context.Context) bool {
        retryTimer := time.NewTimer(p.healthCheckRetryDelay)
        defer retryTimer.Stop()

        select {
        case <-ctx.Done():
                return false
        case <-p.shutdownCh:
                return false
        case <-retryTimer.C:
                retryAlive := p.healthChecker.CheckHealth(ctx)
                if retryAlive.Status == healthcheck.StatusHealthy {
                        //nolint:gosec // G706: logging target URI from config
                        slog.Debug("health check recovered after retry", "target", p.targetURI)
                        return true
                }
                return false
        }
}

// handleHealthCheckFailure handles a failed health check, including retry logic and shutdown.
// Returns (updatedFailureCount, shouldContinue) - true if monitoring should continue, false if it should stop.
func (p *TransparentProxy) handleHealthCheckFailure(
        ctx context.Context,
        consecutiveFailures int,
        status healthcheck.HealthStatus,
) (int, bool) {
        consecutiveFailures++
        //nolint:gosec // G706: logging target URI from config and health status
        slog.Warn("health check failed",
                "target", p.targetURI,
                "attempt", consecutiveFailures,
                "max_attempts", p.healthCheckFailureThreshold,
                "status", status)

        if consecutiveFailures < p.healthCheckFailureThreshold {
                if p.performHealthCheckRetry(ctx) {
                        consecutiveFailures = 0
                }
                return consecutiveFailures, true
        }

        // All retries exhausted, initiate shutdown
        //nolint:gosec // G706: logging target URI from config
        slog.Error("health check failed after consecutive attempts; initiating proxy shutdown",
                "target", p.targetURI, "attempts", p.healthCheckFailureThreshold)
        if p.onHealthCheckFailed != nil {
                p.onHealthCheckFailed()
        }
        if err := p.Stop(ctx); err != nil {
                slog.Error("failed to stop proxy",
                        "target", p.targetURI, "error", err)
        }
        return consecutiveFailures, false
}

func (p *TransparentProxy) monitorHealth(parentCtx context.Context) {
        ticker := time.NewTicker(p.healthCheckInterval)
        defer ticker.Stop()

        consecutiveFailures := 0

        for {
                select {
                case <-parentCtx.Done():
                        //nolint:gosec // G706: logging target URI from config
                        slog.Debug("context cancelled, stopping health monitor", "target", p.targetURI)
                        return
                case <-p.shutdownCh:
                        //nolint:gosec // G706: logging target URI from config
                        slog.Debug("shutdown initiated, stopping health monitor", "target", p.targetURI)
                        return
                case <-ticker.C:
                        if !p.serverInitialized() {
                                //nolint:gosec // G706: logging target URI from config
                                slog.Debug("mcp server not initialized yet, skipping health check",
                                        "target", p.targetURI)
                                continue
                        }

                        alive := p.healthChecker.CheckHealth(parentCtx)
                        if alive.Status != healthcheck.StatusHealthy {
                                var shouldContinue bool
                                consecutiveFailures, shouldContinue = p.handleHealthCheckFailure(parentCtx, consecutiveFailures, alive.Status)
                                if !shouldContinue {
                                        return
                                }
                                continue
                        }

                        // Reset failure count on successful health check
                        if consecutiveFailures > 0 {
                                //nolint:gosec // G706: logging target URI from config
                                slog.Debug("health check recovered",
                                        "target", p.targetURI, "previous_failures", consecutiveFailures)
                        }
                        consecutiveFailures = 0
                }
        }
}

// Stop stops the transparent proxy.
func (p *TransparentProxy) Stop(ctx context.Context) error {
        p.mutex.Lock()

        // Check if already stopped
        if p.stopped {
                p.mutex.Unlock()
                //nolint:gosec // G706: logging target URI from config
                slog.Debug("proxy is already stopped, skipping", "target", p.targetURI)
                return nil
        }

        // Mark as stopped and signal shutdown under the lock
        p.stopped = true
        close(p.shutdownCh)

        // Capture server reference and nil it out under the lock so no other
        // goroutine can race on p.server after we release the mutex.
        server := p.server
        p.server = nil

        // Release the lock before server.Shutdown() so IsRunning() is not blocked
        // while long-lived connections drain.
        p.mutex.Unlock()

        if server != nil {
                // Use the caller's context if still valid; fall back to a fresh one
                // when the caller's context is already cancelled (e.g. the health
                // monitor calls Stop() after its parent context is done).
                base := ctx
                if base.Err() != nil {
                        base = context.Background()
                }
                shutdownCtx, cancel := context.WithTimeout(base, p.shutdownTimeout)
                defer cancel()

                err := server.Shutdown(shutdownCtx)
                if err != nil {
                        if errors.Is(err, context.DeadlineExceeded) {
                                // Graceful shutdown timed out — force-close remaining connections
                                slog.Warn("graceful shutdown timed out, force-closing connections",
                                        "target", p.targetURI, "timeout", p.shutdownTimeout)
                                if closeErr := server.Close(); closeErr != nil {
                                        slog.Warn("error during forced server close", "error", closeErr)
                                }
                        } else if !errors.Is(err, http.ErrServerClosed) {
                                slog.Warn("error during proxy shutdown", "error", err)
                                return err
                        }
                }
                //nolint:gosec // G706: logging target URI from config
                slog.Debug("server stopped successfully", "target", p.targetURI)
        }

        // Stop the session manager to terminate its cleanup goroutine and close any
        // underlying storage connections (e.g. Redis client) opened via WithSessionStorage.
        if p.sessionManager != nil {
                if err := p.sessionManager.Stop(); err != nil {
                        slog.Warn("error stopping session manager", "error", err)
                }
        }

        return nil
}

// IsRunning checks if the proxy is running.
func (p *TransparentProxy) IsRunning() (bool, error) {
        // No mutex needed: shutdownCh is closed under the lock in Stop(),
        // and a select on a closed channel is goroutine-safe by design.
        select {
        case <-p.shutdownCh:
                return false, nil
        default:
                return true, nil
        }
}

// GetMessageChannel returns the channel for messages to/from the destination.
// This is not used in the TransparentProxy implementation as it forwards HTTP requests directly.
func (*TransparentProxy) GetMessageChannel() chan jsonrpc2.Message {
        return nil
}

// SendMessageToDestination sends a message to the destination.
// This is not used in the TransparentProxy implementation as it forwards HTTP requests directly.
func (*TransparentProxy) SendMessageToDestination(_ jsonrpc2.Message) error {
        return fmt.Errorf("SendMessageToDestination not implemented for TransparentProxy")
}

// ForwardResponseToClients forwards a response from the destination to clients.
// This is not used in the TransparentProxy implementation as it forwards HTTP requests directly.
func (*TransparentProxy) ForwardResponseToClients(_ context.Context, _ jsonrpc2.Message) error {
        return fmt.Errorf("ForwardResponseToClients not implemented for TransparentProxy")
}

// methodGate wraps a handler to reject GET, HEAD, and DELETE requests with 405
// when the proxy is running stateless (--stateless) or the request is tagged
// with the Modern (stateless-only) MCP-Protocol-Version: neither has any
// protocol use for GET/HEAD/DELETE, since Modern is unary request/response
// over POST only.
//
// This check runs outermost, before auth and any other middleware: a method
// rejection needs no identity and reveals nothing auth-gated. It is
// deliberately header-only and never reads or otherwise touches r.Body, so it
// must not call mcp.ClassifyRevision (which requires a body) or consume the
// body in any way — doing so would break the body for downstream handlers.
//
// HEAD is blocked alongside GET because HEAD is semantically a GET without a
// response body; a server that cannot handle GET will not handle HEAD either.
func (p *TransparentProxy) methodGate(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                isGateableMethod := r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodDelete
                isModern := r.Header.Get("MCP-Protocol-Version") == mcp.MCPVersionModern
                if isGateableMethod && (p.stateless || isModern) {
                        w.Header().Set("Allow", "POST, OPTIONS")
                        http.Error(w, "method not allowed: server is stateless / stateless MCP revision (POST only)", http.StatusMethodNotAllowed)
                        return
                }
                next.ServeHTTP(w, r)
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "bytes"
        "encoding/json"
        "fmt"
        "io"
        "net/http"
)

const (
        // CodeSessionNotFound is the JSON-RPC error code for expired or unknown sessions.
        // This matches the MCP TypeScript SDK reference server convention and falls within
        // the JSON-RPC 2.0 implementation-defined server-errors range (-32000 to -32099).
        // MCP clients use this code to trigger automatic session recovery.
        CodeSessionNotFound int64 = -32001

        // MessageSessionNotFound is the JSON-RPC error message for session-not-found.
        MessageSessionNotFound = "Session not found"
)

// NotFoundBody returns the JSON-encoded body for a session-not-found
// JSON-RPC error response. The requestID is the "id" from the incoming
// JSON-RPC request, echoed so a client correlating by id can match this error
// to the request that caused it.
//
// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error
// response as `id?: RequestId` where RequestId = string | number, so an
// absent id is encoded by omitting the "id" key entirely, never as null. The
// reference TypeScript SDK enforces this with a .strict() schema and a
// throwing parse(), so an "id":null response crashes a conformant client's
// transport. See HasJSONRPCID.
//
// Pass nil only when the request genuinely carries no id: a bodiless GET
// (standalone SSE) or DELETE, a notification, or a batch. Do NOT pass nil
// merely because threading the id to the call site is inconvenient — that
// was the asymmetry fixed in #5945.
func NotFoundBody(requestID any) []byte {
        resp := map[string]any{
                "jsonrpc": "2.0",
                "error": map[string]any{
                        "code":    CodeSessionNotFound,
                        "message": MessageSessionNotFound,
                },
        }
        if HasJSONRPCID(requestID) {
                resp["id"] = requestID
        }
        data, err := json.Marshal(resp)
        if err != nil {
                // This should never happen with simple map types, but return a
                // hand-crafted fallback to guarantee a valid JSON-RPC error.
                return []byte(`{"jsonrpc":"2.0","error":{"code":-32001,"message":"Session not found"}}`)
        }
        return data
}

// HasJSONRPCID reports whether requestID is a usable JSON-RPC id, i.e. whether
// callers should emit an "id" key at all.
//
// MCP narrows base JSON-RPC here: schema/2025-11-25 types the error response as
// `id?: RequestId` where RequestId = string | number, so an absent id is encoded
// by omitting the key, never as null. The reference TypeScript SDK enforces this
// with a .strict() schema and a throwing parse(), so "id":null crashes a
// conformant client's transport.
//
// A json.RawMessage holding literal null -- or nothing -- is an absent id
// even though the interface value is non-nil.
//
// Do not pass a typed jsonrpc2.ID here: its zero value is a non-nil interface
// holding an empty struct, so this predicate would (wrongly) report it as
// present and callers would emit "id":{}. Use id.IsValid() instead.
func HasJSONRPCID(requestID any) bool {
        if requestID == nil {
                return false
        }
        if raw, ok := requestID.(json.RawMessage); ok {
                return len(raw) > 0 && string(raw) != "null"
        }
        return true
}

// WriteNotFound writes an HTTP 404 response with a JSON-RPC error body
// for session-not-found. Use this with http.ResponseWriter in the streamable
// HTTP and SSE proxies.
func WriteNotFound(w http.ResponseWriter, requestID any) {
        w.Header().Set("Content-Type", "application/json")
        w.WriteHeader(http.StatusNotFound)
        //nolint:gosec // G104: writing a static JSON error response to an HTTP client
        _, _ = w.Write(NotFoundBody(requestID))
}

// NotFoundResponse constructs an *http.Response with HTTP 404 and a
// JSON-RPC error body. Use this in httputil.ReverseProxy.ModifyResponse
// (transparent proxy) where no http.ResponseWriter is available.
//
// requestID follows NotFoundBody: pass the incoming request's JSON-RPC id so the
// error echoes it, or nil when the request has none. Callers that reject a
// request before parsing a body (GET/DELETE) pass nil.
func NotFoundResponse(req *http.Request, requestID any) *http.Response {
        body := NotFoundBody(requestID)
        hdr := make(http.Header)
        hdr.Set("Content-Type", "application/json")
        return &http.Response{
                StatusCode:    http.StatusNotFound,
                Status:        fmt.Sprintf("%d %s", http.StatusNotFound, http.StatusText(http.StatusNotFound)),
                Proto:         "HTTP/1.1",
                ProtoMajor:    1,
                ProtoMinor:    1,
                Header:        hdr,
                ContentLength: int64(len(body)),
                Body:          io.NopCloser(bytes.NewReader(body)),
                Request:       req,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package session provides a session manager with TTL cleanup.
package session

import (
        "context"
        "fmt"
        "log/slog"
        "time"

        "github.com/google/uuid"

        "github.com/stacklok/toolhive-core/redisconn"
)

const (
        // defaultOperationTimeout is the timeout for standard storage operations
        defaultOperationTimeout = 5 * time.Second
        // cleanupOperationTimeout is the timeout for cleanup operations which may take longer
        cleanupOperationTimeout = 30 * time.Second
)

// Session interface defines the contract for all session types
type Session interface {
        ID() string
        Type() SessionType
        CreatedAt() time.Time
        UpdatedAt() time.Time

        // Data and metadata methods
        GetData() interface{}
        SetData(data interface{})
        GetMetadata() map[string]string
        GetMetadataValue(key string) (string, bool)
        SetMetadata(key, value string)
}

// Manager holds sessions with TTL cleanup.
type Manager struct {
        storage Storage
        ttl     time.Duration
        stopCh  chan struct{}
        factory Factory
}

// Factory defines a function type for creating new sessions.
// It now returns the Session interface to support different session types.
type Factory func(id string) Session

// LegacyFactory is the old factory type for backward compatibility
type LegacyFactory func(id string) *ProxySession

// NewManager creates a session manager with TTL and starts cleanup worker.
// It accepts either the new Factory or the legacy factory for backward compatibility.
func NewManager(ttl time.Duration, factory interface{}) *Manager {
        var f Factory

        switch factoryFunc := factory.(type) {
        case Factory:
                f = factoryFunc
        case LegacyFactory:
                // Wrap legacy factory to return Session interface
                f = func(id string) Session {
                        return factoryFunc(id)
                }
        case func(id string) *ProxySession:
                // Also support direct function for backward compatibility
                f = func(id string) Session {
                        return factoryFunc(id)
                }
        default:
                // Default to creating basic ProxySession
                f = func(id string) Session {
                        return NewProxySession(id)
                }
        }

        m := &Manager{
                storage: NewLocalStorage(),
                ttl:     ttl,
                stopCh:  make(chan struct{}),
                factory: f,
        }
        go m.cleanupRoutine()
        return m
}

// NewTypedManager creates a session manager for a specific session type.
func NewTypedManager(ttl time.Duration, sessionType SessionType) *Manager {
        factory := func(id string) Session {
                switch sessionType {
                case SessionTypeSSE:
                        return NewSSESession(id)
                case SessionTypeMCP:
                        return NewProxySession(id)
                case SessionTypeStreamable:
                        return NewTypedProxySession(id, sessionType)
                default:
                        return NewTypedProxySession(id, sessionType)
                }
        }

        return NewManager(ttl, factory)
}

// NewManagerWithStorage creates a session manager with a custom storage backend.
func NewManagerWithStorage(ttl time.Duration, factory Factory, storage Storage) *Manager {
        m := &Manager{
                storage: storage,
                ttl:     ttl,
                stopCh:  make(chan struct{}),
                factory: factory,
        }
        go m.cleanupRoutine()
        return m
}

// NewManagerWithRedis creates a session manager backed by Redis.
// ctx is used for the initial Ping during construction and should carry any
// deadline appropriate for the connection attempt (e.g. a startup timeout).
// cfg supplies the Redis connection configuration (delegated to the shared
// toolhive-core redis package); keyPrefix namespaces the keys and must end
// with ':'. ttl is applied as both the manager's cleanup interval and the
// Redis key TTL. Returns an error if the Redis client cannot be constructed
// (e.g. invalid config or TLS error). Callers that do not require Redis
// should use NewManager or NewTypedManager instead.
func NewManagerWithRedis(
        ctx context.Context,
        ttl time.Duration,
        factory Factory,
        cfg redisconn.Config,
        keyPrefix string,
) (*Manager, error) {
        storage, err := NewRedisStorage(ctx, cfg, keyPrefix, ttl)
        if err != nil {
                return nil, fmt.Errorf("creating redis storage: %w", err)
        }
        return NewManagerWithStorage(ttl, factory, storage), nil
}

func (m *Manager) cleanupRoutine() {
        ticker := time.NewTicker(m.ttl / 2)
        defer ticker.Stop()
        for {
                select {
                case <-ticker.C:
                        cutoff := time.Now().Add(-m.ttl)
                        ctx, cancel := context.WithTimeout(context.Background(), cleanupOperationTimeout)
                        if err := m.storage.DeleteExpired(ctx, cutoff); err != nil {
                                slog.Error("failed to delete expired sessions", "error", err)
                        }
                        cancel()
                case <-m.stopCh:
                        return
                }
        }
}

// validateSessionID returns an error if id is empty or not a valid UUID.
// UUID format is enforced across all storage backends to keep ID semantics consistent.
func validateSessionID(id string) error {
        if id == "" {
                return fmt.Errorf("session ID cannot be empty")
        }
        if _, err := uuid.Parse(id); err != nil {
                return fmt.Errorf("invalid session ID format: %w", err)
        }
        return nil
}

// AddWithID creates (and adds) a new session with the provided ID.
// Returns error if ID is empty, not a valid UUID, or already exists.
func (m *Manager) AddWithID(id string) error {
        if err := validateSessionID(id); err != nil {
                return err
        }
        return m.AddSession(m.factory(id))
}

// AddSession adds an existing session to the manager.
// This is useful when you need to create a session with specific properties.
func (m *Manager) AddSession(session Session) error {
        if session == nil {
                return fmt.Errorf("session cannot be nil")
        }
        if err := validateSessionID(session.ID()); err != nil {
                return err
        }

        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()
        return m.storage.Create(ctx, session)
}

// Get retrieves a session by ID. Returns (session, true) if found.
// For LocalStorage, the storage backend updates the session's last-access timestamp on Load.
func (m *Manager) Get(id string) (Session, bool) {
        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()

        sess, err := m.storage.Load(ctx, id)
        if err != nil {
                return nil, false
        }
        return sess, true
}

// LoadIfOwner returns the session only if its current ownership binding matches expectedOwner.
func (m *Manager) LoadIfOwner(id, expectedOwner string) (Session, error) {
        if err := validateSessionID(id); err != nil {
                return nil, err
        }
        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()
        return m.storage.LoadIfOwner(ctx, id, expectedOwner)
}

// GetWithError retrieves a session by ID and returns the underlying storage
// error when the lookup fails. Callers that need to distinguish between
// ErrSessionNotFound (session absent or expired) and other storage errors
// (e.g. Redis connectivity failures) should use this method so they can
// return an appropriate HTTP status code (400 vs 503) to the client.
func (m *Manager) GetWithError(id string) (Session, error) {
        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()
        return m.storage.Load(ctx, id)
}

// UpsertSession inserts or updates a session in storage, replacing any existing
// session with the same ID. Used by SessionManager to replace placeholder sessions
// with fully-formed MultiSession objects after phase-2 construction.
func (m *Manager) UpsertSession(session Session) error {
        if session == nil {
                return fmt.Errorf("session cannot be nil")
        }
        if err := validateSessionID(session.ID()); err != nil {
                return err
        }
        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()
        return m.storage.Store(ctx, session)
}

// UpsertSessionIfOwner replaces a session only if the current record still has expectedOwner.
func (m *Manager) UpsertSessionIfOwner(session Session, expectedOwner string) error {
        if session == nil {
                return fmt.Errorf("session cannot be nil")
        }
        if err := validateSessionID(session.ID()); err != nil {
                return err
        }
        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()
        updated, err := m.storage.StoreIfOwner(ctx, session, expectedOwner)
        if err != nil {
                return err
        }
        if !updated {
                return ErrSessionNotFound
        }
        return nil
}

// Delete removes a session by ID.
// Returns an error if the ID is invalid or the deletion fails.
func (m *Manager) Delete(id string) error {
        if err := validateSessionID(id); err != nil {
                return err
        }
        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()
        return m.storage.Delete(ctx, id)
}

// DeleteIfOwner removes a session only if the current record still has expectedOwner.
func (m *Manager) DeleteIfOwner(id, expectedOwner string) error {
        if err := validateSessionID(id); err != nil {
                return err
        }
        ctx, cancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
        defer cancel()
        deleted, err := m.storage.DeleteIfOwner(ctx, id, expectedOwner)
        if err != nil {
                return err
        }
        if !deleted {
                return ErrSessionNotFound
        }
        return nil
}

// Stop stops the cleanup worker and closes the storage backend.
// Returns an error if closing the storage backend fails.
func (m *Manager) Stop() error {
        close(m.stopCh)
        if m.storage != nil {
                return m.storage.Close()
        }
        return nil
}

// Count returns the number of active sessions.
//
// Note: This method only works with LocalStorage backend and returns 0 for
// other storage backends. Count is not part of the Storage interface because
// it's not feasible for distributed storage backends like Redis where counting
// all keys can be prohibitively expensive.
//
// For distributed storage, consider maintaining a counter or using approximate
// count mechanisms provided by the storage backend.
func (m *Manager) Count() int {
        if localStorage, ok := m.storage.(*LocalStorage); ok {
                return localStorage.Count()
        }
        return 0
}

func (m *Manager) cleanupExpiredOnce() error {
        cutoff := time.Now().Add(-m.ttl)
        ctx, cancel := context.WithTimeout(context.Background(), cleanupOperationTimeout)
        defer cancel()
        return m.storage.DeleteExpired(ctx, cutoff)
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import "context"

// MetadataKeyIdentityBinding is the canonical persisted ownership key.
const MetadataKeyIdentityBinding = "vmcp.identity.binding"

// LookupOwner loads authoritative metadata, not a live stream or backend session.
// Redis shares ownership but does not provide cross-replica stream delivery.
func (m *Manager) LookupOwner(ctx context.Context, id string) (string, error) {
        ctx, cancel := context.WithTimeout(ctx, defaultOperationTimeout)
        defer cancel()
        metadata, err := m.storage.LoadMetadata(ctx, id)
        if err != nil {
                return "", err
        }
        return metadata[MetadataKeyIdentityBinding], nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "sync"
        "time"
)

// SessionType represents the type of session
//
//revive:disable-next-line:exported
type SessionType string

const (
        // SessionTypeMCP represents a standard MCP session
        SessionTypeMCP SessionType = "mcp"
        // SessionTypeSSE represents an SSE (Server-Sent Events) session
        SessionTypeSSE SessionType = "sse"
        // SessionTypeStreamable represents a streamable HTTP session
        SessionTypeStreamable SessionType = "streamable"
)

const (
        // DefaultSessionTTL is the default time-to-live for sessions (2 hours)
        DefaultSessionTTL = 2 * time.Hour
)

// ProxySession implements the Session interface for proxy sessions.
// It now includes support for session types, metadata, and custom data.
type ProxySession struct {
        id       string
        created  time.Time
        updated  time.Time
        sessType SessionType
        data     interface{}
        metadata map[string]string
        mu       sync.RWMutex // Protect concurrent access to metadata and data
}

// NewProxySession creates a new ProxySession with the given ID.
// It defaults to SessionTypeMCP for backward compatibility.
func NewProxySession(id string) *ProxySession {
        now := time.Now()
        return &ProxySession{
                id:       id,
                created:  now,
                updated:  now,
                sessType: SessionTypeMCP,
                metadata: make(map[string]string),
        }
}

// NewTypedProxySession creates a new ProxySession with the given ID and type.
func NewTypedProxySession(id string, sessType SessionType) *ProxySession {
        now := time.Now()
        return &ProxySession{
                id:       id,
                created:  now,
                updated:  now,
                sessType: sessType,
                metadata: make(map[string]string),
        }
}

// ID returns the session ID.
func (s *ProxySession) ID() string { return s.id }

// CreatedAt returns the creation time of the session.
func (s *ProxySession) CreatedAt() time.Time { return s.created }

// UpdatedAt returns the last updated time of the session.
func (s *ProxySession) UpdatedAt() time.Time {
        s.mu.RLock()
        defer s.mu.RUnlock()
        return s.updated
}

// Touch updates the session's last updated time to the current time.
func (s *ProxySession) Touch() {
        s.mu.Lock()
        defer s.mu.Unlock()
        s.updated = time.Now()
}

// Type returns the session type.
func (s *ProxySession) Type() SessionType { return s.sessType }

// GetData returns the session-specific data.
func (s *ProxySession) GetData() interface{} {
        s.mu.RLock()
        defer s.mu.RUnlock()
        return s.data
}

// SetData sets the session-specific data.
func (s *ProxySession) SetData(data interface{}) {
        s.mu.Lock()
        defer s.mu.Unlock()
        s.data = data
}

// GetMetadata returns all metadata as a map.
func (s *ProxySession) GetMetadata() map[string]string {
        s.mu.RLock()
        defer s.mu.RUnlock()
        // Return a copy to prevent external modification
        result := make(map[string]string, len(s.metadata))
        for k, v := range s.metadata {
                result[k] = v
        }
        return result
}

// SetMetadata sets a metadata key-value pair.
func (s *ProxySession) SetMetadata(key, value string) {
        s.mu.Lock()
        defer s.mu.Unlock()
        if s.metadata == nil {
                s.metadata = make(map[string]string)
        }
        s.metadata[key] = value
}

// GetMetadataValue gets a specific metadata value.
func (s *ProxySession) GetMetadataValue(key string) (string, bool) {
        s.mu.RLock()
        defer s.mu.RUnlock()
        value, ok := s.metadata[key]
        return value, ok
}

// DeleteMetadata removes a metadata key.
func (s *ProxySession) DeleteMetadata(key string) {
        s.mu.Lock()
        defer s.mu.Unlock()
        delete(s.metadata, key)
}

// setTimestamps updates the created and updated timestamps.
// This is used internally for deserialization to restore session state.
func (s *ProxySession) setTimestamps(created, updated time.Time) {
        s.mu.Lock()
        defer s.mu.Unlock()
        s.created = created
        s.updated = updated
}

// setMetadataMap replaces the entire metadata map.
// This is used internally for deserialization to restore session state.
func (s *ProxySession) setMetadataMap(metadata map[string]string) {
        s.mu.Lock()
        defer s.mu.Unlock()
        if metadata == nil {
                s.metadata = make(map[string]string)
        } else {
                s.metadata = metadata
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "log/slog"

        "github.com/stacklok/toolhive-core/redisconn"
)

// RedisPasswordEnvVar is the environment variable name for the Redis session storage password.
// The operator injects this as a SecretKeyRef when sessionStorage.provider is "redis"
// and passwordRef is set.
// #nosec G101 -- This is an environment variable name, not a hardcoded credential
const RedisPasswordEnvVar = "THV_SESSION_REDIS_PASSWORD"

// warnInsecureRedisTransport emits one startup WARN when the session store
// connection is not protected by verified TLS: either TLS is not configured
// (session data, and the password if one is set, cross the network in
// plaintext) or server certificate verification is disabled. Both remain
// allowed so existing deployments keep starting (for example, ones relying on
// a service mesh for encryption), but the configuration is made visible in
// logs. The password itself is never logged. It is called by the Redis session
// store constructors so every caller gets the warning.
func warnInsecureRedisTransport(cfg *redisconn.Config) {
        if cfg == nil {
                return
        }
        switch {
        case cfg.TLS == nil && cfg.Password != "":
                slog.Warn("Redis session storage sends its password and session data without TLS; "+
                        "configure TLS for the session storage connection to encrypt it",
                        "store", cfg.Addr)
        case cfg.TLS == nil:
                slog.Warn("Redis session storage sends session data without TLS; "+
                        "configure TLS for the session storage connection to encrypt it",
                        "store", cfg.Addr)
        case cfg.TLS.InsecureSkipVerify:
                slog.Warn("Redis session storage uses TLS without server certificate verification; "+
                        "the connection is encrypted but the server is not authenticated",
                        "store", cfg.Addr)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "encoding/json"
        "fmt"
        "time"
)

// sessionData is the JSON representation of a session.
// This structure is used for serializing sessions to/from storage backends.
type sessionData struct {
        ID        string            `json:"id"`
        Type      SessionType       `json:"type"`
        CreatedAt time.Time         `json:"created_at"`
        UpdatedAt time.Time         `json:"updated_at"`
        Data      json.RawMessage   `json:"data,omitempty"`
        Metadata  map[string]string `json:"metadata,omitempty"`
}

// serializeSession converts a Session to its JSON representation.
func serializeSession(s Session) ([]byte, error) {
        if s == nil {
                return nil, fmt.Errorf("cannot serialize nil session")
        }

        data := sessionData{
                ID:        s.ID(),
                Type:      s.Type(),
                CreatedAt: s.CreatedAt(),
                UpdatedAt: s.UpdatedAt(),
                Metadata:  s.GetMetadata(),
        }

        // Handle session-specific data
        if sessionData := s.GetData(); sessionData != nil {
                jsonData, err := json.Marshal(sessionData)
                if err != nil {
                        return nil, fmt.Errorf("failed to marshal session data: %w", err)
                }
                data.Data = jsonData
        }

        return json.Marshal(data)
}

// deserializeSession reconstructs a Session from its JSON representation.
// It creates the appropriate session type based on the Type field.
func deserializeSession(data []byte) (Session, error) {
        if len(data) == 0 {
                return nil, fmt.Errorf("cannot deserialize empty data")
        }

        var sd sessionData
        if err := json.Unmarshal(data, &sd); err != nil {
                return nil, fmt.Errorf("failed to unmarshal session data: %w", err)
        }

        // Create appropriate session type using existing constructors
        var session Session
        switch sd.Type {
        case SessionTypeSSE:
                // Use existing NewSSESession constructor
                sseSession := NewSSESession(sd.ID)
                // Update timestamps to match stored values
                sseSession.setTimestamps(sd.CreatedAt, sd.UpdatedAt)
                // Restore metadata
                sseSession.setMetadataMap(sd.Metadata)
                // Note: SSE channels and client info will be recreated when reconnected
                session = sseSession

        case SessionTypeStreamable:
                // Use existing NewStreamableSession constructor
                sess := NewStreamableSession(sd.ID)
                streamSession, ok := sess.(*StreamableSession)
                if !ok {
                        return nil, fmt.Errorf("failed to create StreamableSession")
                }
                // Update timestamps to match stored values
                streamSession.setTimestamps(sd.CreatedAt, sd.UpdatedAt)
                // Restore metadata
                streamSession.setMetadataMap(sd.Metadata)
                session = streamSession

        case SessionTypeMCP:
                fallthrough
        default:
                // Use existing NewTypedProxySession constructor
                proxySession := NewTypedProxySession(sd.ID, sd.Type)
                // Update timestamps to match stored values
                proxySession.setTimestamps(sd.CreatedAt, sd.UpdatedAt)
                // Restore metadata
                proxySession.setMetadataMap(sd.Metadata)
                session = proxySession
        }

        // Restore session-specific data if present
        if len(sd.Data) > 0 {
                // For now, we store the raw JSON. Session-specific implementations
                // can unmarshal this as needed.
                session.SetData(sd.Data)
        }

        return session, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "context"
        "fmt"
        "time"
)

// DataStorage stores session metadata as plain key-value pairs.
//
// Unlike the Session-based Storage interface, DataStorage never attempts
// to round-trip live session objects (MultiSession, StreamableSession, etc.).
// It stores only serializable metadata, keeping data storage and live-object
// lifecycle as separate concerns.
//
// This separation avoids the type-assertion bug where a Redis round-trip
// deserializes a MultiSession as a plain *StreamableSession, losing all
// backend connections and routing state.
//
// # Contract
//
//   - Create atomically creates metadata for id only if it does not already exist.
//     Returns (true, nil) if created, (false, nil) if the key already existed.
//   - Update overwrites metadata only if the key already exists (SET XX semantics).
//     Returns (true, nil) if updated, (false, nil) if the session was not found.
//     Use this instead of Load→mutate→unconditional-write to avoid TOCTOU
//     resurrection races (a concurrent Delete between Load and write would be
//     silently undone by an unconditional write).
//   - Load retrieves metadata and refreshes the TTL (sliding-window expiry).
//     Returns ErrSessionNotFound if the session does not exist.
//   - Delete removes the session. It is not an error if the session is absent.
//   - Close releases any resources held by the backend (connections, goroutines).
//
// # Implementations
//
// Two concrete implementations are provided:
//   - LocalSessionDataStorage (in-memory, single-process)
//   - RedisSessionDataStorage (Redis/Valkey, multi-process)
type DataStorage interface {
        // Create atomically creates session metadata only if the session ID
        // does not already exist. Returns (true, nil) if the entry was created,
        // (false, nil) if it already existed, or (false, err) on storage errors.
        Create(ctx context.Context, id string, metadata map[string]string) (bool, error)

        // Update overwrites session metadata only if the session ID already exists
        // (conditional write, equivalent to Redis SET XX). Returns (true, nil) if
        // the entry was updated, (false, nil) if it was not found, or (false, err)
        // on storage errors. Use this (SET XX) for writes that must not recreate a
        // key that was concurrently deleted — an unconditional write would silently
        // resurrect a session the caller intentionally terminated.
        Update(ctx context.Context, id string, metadata map[string]string) (bool, error)

        // Load retrieves session metadata and refreshes its TTL.
        // Returns ErrSessionNotFound if the session does not exist.
        Load(ctx context.Context, id string) (map[string]string, error)

        // LoadMetadata retrieves session metadata without refreshing its TTL, for
        // checks (such as ownership) that must not keep a session alive on behalf
        // of a caller that may be rejected. Returns ErrSessionNotFound if the
        // session does not exist.
        LoadMetadata(ctx context.Context, id string) (map[string]string, error)

        // Delete removes session metadata. Not an error if absent.
        Delete(ctx context.Context, id string) error

        // Close releases resources (connections, background goroutines).
        Close() error
}

// NewLocalSessionDataStorage creates a LocalSessionDataStorage with the given TTL.
// ttl must be positive; a zero or negative value returns an error.
// A background cleanup goroutine is started and runs until Close is called.
func NewLocalSessionDataStorage(ttl time.Duration) (*LocalSessionDataStorage, error) {
        if ttl <= 0 {
                return nil, fmt.Errorf("ttl must be a positive duration")
        }
        s := &LocalSessionDataStorage{
                sessions: make(map[string]*localDataEntry),
                ttl:      ttl,
                stopCh:   make(chan struct{}),
        }
        go s.cleanupRoutine()
        return s, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "context"
        "fmt"
        "maps"
        "sync"
        "sync/atomic"
        "time"
)

// localDataEntry wraps session metadata with a storage-owned last-access
// timestamp used for TTL-based eviction.
type localDataEntry struct {
        metadata       map[string]string
        lastAccessNano atomic.Int64
}

func newLocalDataEntry(metadata map[string]string) *localDataEntry {
        e := &localDataEntry{metadata: metadata}
        e.lastAccessNano.Store(time.Now().UnixNano())
        return e
}

func (e *localDataEntry) lastAccess() time.Time {
        return time.Unix(0, e.lastAccessNano.Load())
}

// LocalSessionDataStorage implements DataStorage using an in-memory
// map with TTL-based eviction.
//
// Sessions are evicted if they have not been accessed within the configured TTL.
// A background goroutine runs until Close is called.
type LocalSessionDataStorage struct {
        sessions map[string]*localDataEntry // guarded by mu
        mu       sync.Mutex
        ttl      time.Duration
        stopCh   chan struct{}
        stopOnce sync.Once
}

// Load retrieves session metadata and refreshes its last-access timestamp.
// Returns ErrSessionNotFound if the session does not exist.
func (s *LocalSessionDataStorage) Load(_ context.Context, id string) (map[string]string, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session data with empty ID")
        }
        s.mu.Lock()
        entry, ok := s.sessions[id]
        if ok {
                entry.lastAccessNano.Store(time.Now().UnixNano())
        }
        s.mu.Unlock()
        if !ok {
                return nil, ErrSessionNotFound
        }
        return maps.Clone(entry.metadata), nil
}

// LoadMetadata retrieves session metadata without refreshing its last-access
// timestamp. Returns ErrSessionNotFound if the session does not exist.
func (s *LocalSessionDataStorage) LoadMetadata(_ context.Context, id string) (map[string]string, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session data with empty ID")
        }
        s.mu.Lock()
        entry, ok := s.sessions[id]
        s.mu.Unlock()
        if !ok {
                return nil, ErrSessionNotFound
        }
        return maps.Clone(entry.metadata), nil
}

// Create creates session metadata only if the session ID does not already exist.
// Returns (true, nil) if created, (false, nil) if the key already existed.
func (s *LocalSessionDataStorage) Create(_ context.Context, id string, metadata map[string]string) (bool, error) {
        if id == "" {
                return false, fmt.Errorf("cannot write session data with empty ID")
        }
        if metadata == nil {
                metadata = make(map[string]string)
        }
        s.mu.Lock()
        defer s.mu.Unlock()
        if _, exists := s.sessions[id]; exists {
                return false, nil
        }
        s.sessions[id] = newLocalDataEntry(maps.Clone(metadata))
        return true, nil
}

// Update overwrites session metadata only if the session ID already exists.
// Returns (true, nil) if updated, (false, nil) if not found.
func (s *LocalSessionDataStorage) Update(_ context.Context, id string, metadata map[string]string) (bool, error) {
        if id == "" {
                return false, fmt.Errorf("cannot write session data with empty ID")
        }
        if metadata == nil {
                metadata = make(map[string]string)
        }
        s.mu.Lock()
        defer s.mu.Unlock()
        if _, ok := s.sessions[id]; !ok {
                return false, nil
        }
        s.sessions[id] = newLocalDataEntry(maps.Clone(metadata))
        return true, nil
}

// Delete removes session metadata. Not an error if absent.
func (s *LocalSessionDataStorage) Delete(_ context.Context, id string) error {
        if id == "" {
                return fmt.Errorf("cannot delete session data with empty ID")
        }
        s.mu.Lock()
        delete(s.sessions, id)
        s.mu.Unlock()
        return nil
}

// Close stops the background cleanup goroutine and clears all stored metadata.
func (s *LocalSessionDataStorage) Close() error {
        s.stopOnce.Do(func() { close(s.stopCh) })
        s.mu.Lock()
        s.sessions = make(map[string]*localDataEntry)
        s.mu.Unlock()
        return nil
}

// minCleanupInterval is the floor applied to the cleanup ticker interval.
// time.NewTicker panics when given a duration ≤ 0, so any TTL smaller than
// 2ns would produce ttl/2 == 0 without this guard. 1ms is a practical floor
// that avoids the panic without restricting legitimate short TTLs in tests.
const minCleanupInterval = time.Millisecond

func (s *LocalSessionDataStorage) cleanupRoutine() {
        if s.ttl <= 0 {
                return
        }
        interval := s.ttl / 2
        if interval < minCleanupInterval {
                interval = minCleanupInterval
        }
        ticker := time.NewTicker(interval)
        defer ticker.Stop()
        for {
                select {
                case <-ticker.C:
                        s.deleteExpired()
                case <-s.stopCh:
                        return
                }
        }
}

func (s *LocalSessionDataStorage) deleteExpired() {
        cutoff := time.Now().Add(-s.ttl)
        s.mu.Lock()
        defer s.mu.Unlock()
        for id, entry := range s.sessions {
                if entry.lastAccess().Before(cutoff) {
                        delete(s.sessions, id)
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "time"

        "github.com/redis/go-redis/v9"

        "github.com/stacklok/toolhive-core/redisconn"
)

// RedisSessionDataStorage implements DataStorage backed by Redis/Valkey.
//
// Metadata is serialized as a JSON object and stored with a sliding-window TTL:
// each Create/Update/Load call refreshes the key's expiry so that active sessions
// never expire while they are in use.
//
// Because only plain map[string]string is stored, there are no type assertions
// or deserialization tricks — Redis round-trips the map cleanly.
type RedisSessionDataStorage struct {
        client    redis.UniversalClient
        keyPrefix string
        ttl       time.Duration
}

// NewRedisSessionDataStorage constructs a RedisSessionDataStorage. Connection-mode
// topology, timeouts, TLS, and credentials are configured through cfg; keyPrefix
// is the per-tenant key prefix (e.g. "thv:vmcp:session:") and must end with ':'
// to avoid key collisions. ttl is the sliding-window expiry applied on every
// Create/Update and Load. The caller must call Close when done.
//
// Connection-mode validation, timeout defaults, client construction (standalone,
// cluster, or sentinel), TLS plumbing, and connectivity verification are
// delegated to the shared toolhive-core redis package.
// A connection without verified TLS is allowed but logged at WARN.
func NewRedisSessionDataStorage(
        ctx context.Context,
        cfg redisconn.Config,
        keyPrefix string,
        ttl time.Duration,
) (*RedisSessionDataStorage, error) {
        if err := validateSessionInvariants(keyPrefix, ttl); err != nil {
                return nil, err
        }
        warnInsecureRedisTransport(&cfg)
        client, err := redisconn.NewClient(ctx, &cfg)
        if err != nil {
                return nil, err
        }
        return &RedisSessionDataStorage{
                client:    client,
                keyPrefix: keyPrefix,
                ttl:       ttl,
        }, nil
}

func (s *RedisSessionDataStorage) key(id string) string {
        return s.keyPrefix + id
}

// Load retrieves metadata from Redis and refreshes the key's TTL via GETEX.
// Returns ErrSessionNotFound if the key does not exist.
func (s *RedisSessionDataStorage) Load(ctx context.Context, id string) (map[string]string, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session data with empty ID")
        }
        data, err := s.client.GetEx(ctx, s.key(id), s.ttl).Bytes()
        return decodeRedisSessionData(data, err)
}

// LoadMetadata retrieves metadata from Redis with a plain GET, leaving the
// key's TTL unchanged. Returns ErrSessionNotFound if the key does not exist.
func (s *RedisSessionDataStorage) LoadMetadata(ctx context.Context, id string) (map[string]string, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session data with empty ID")
        }
        data, err := s.client.Get(ctx, s.key(id)).Bytes()
        return decodeRedisSessionData(data, err)
}

// decodeRedisSessionData maps a Redis read result to session metadata.
func decodeRedisSessionData(data []byte, err error) (map[string]string, error) {
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, ErrSessionNotFound
                }
                return nil, fmt.Errorf("failed to load session metadata: %w", err)
        }
        var metadata map[string]string
        if err := json.Unmarshal(data, &metadata); err != nil {
                return nil, fmt.Errorf("failed to deserialize session metadata: %w", err)
        }
        return metadata, nil
}

// Update overwrites session metadata only if the key already exists.
// Uses Redis SET XX (set-if-exists) to prevent resurrecting a session that
// was deleted by a concurrent Delete call (e.g. from another pod).
// Returns (true, nil) if updated, (false, nil) if the key was not found.
func (s *RedisSessionDataStorage) Update(ctx context.Context, id string, metadata map[string]string) (bool, error) {
        if id == "" {
                return false, fmt.Errorf("cannot write session data with empty ID")
        }
        if metadata == nil {
                metadata = make(map[string]string)
        }
        data, err := json.Marshal(metadata)
        if err != nil {
                return false, fmt.Errorf("failed to serialize session metadata: %w", err)
        }
        // Mode "XX" means "only set if the key already exists".
        res, err := s.client.SetArgs(ctx, s.key(id), data, redis.SetArgs{
                Mode: "XX",
                TTL:  s.ttl,
        }).Result()
        if err != nil {
                // go-redis surfaces the "key does not exist" nil bulk reply as redis.Nil.
                if errors.Is(err, redis.Nil) {
                        return false, nil
                }
                return false, fmt.Errorf("failed to conditionally update session metadata: %w", err)
        }
        // SetArgs with Mode "XX" returns "" when the key does not exist and "OK"
        // when the write succeeded.
        return res == "OK", nil
}

// Create atomically creates session metadata only if the key does not
// already exist. Uses Redis SET NX (set-if-not-exists) to avoid the
// read-then-write TOCTOU race that a Load-then-unconditional-write pattern
// would introduce in multi-pod deployments.
func (s *RedisSessionDataStorage) Create(ctx context.Context, id string, metadata map[string]string) (bool, error) {
        if id == "" {
                return false, fmt.Errorf("cannot write session data with empty ID")
        }
        if metadata == nil {
                metadata = make(map[string]string)
        }
        data, err := json.Marshal(metadata)
        if err != nil {
                return false, fmt.Errorf("failed to serialize session metadata: %w", err)
        }
        ok, err := s.client.SetNX(ctx, s.key(id), data, s.ttl).Result()
        if err != nil {
                return false, fmt.Errorf("failed to atomically store session metadata: %w", err)
        }
        return ok, nil
}

// Delete removes the Redis key. A missing key is not an error.
func (s *RedisSessionDataStorage) Delete(ctx context.Context, id string) error {
        if id == "" {
                return fmt.Errorf("cannot delete session data with empty ID")
        }
        if err := s.client.Del(ctx, s.key(id)).Err(); err != nil {
                return fmt.Errorf("failed to delete session metadata: %w", err)
        }
        return nil
}

// Close closes the underlying Redis client.
func (s *RedisSessionDataStorage) Close() error {
        return s.client.Close()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "time"

        "github.com/stacklok/toolhive/pkg/transport/ssecommon"
)

// SSESession represents an SSE (Server-Sent Events) session.
// It embeds ProxySession and adds SSE-specific functionality.
type SSESession struct {
        *ProxySession

        // SSE-specific fields
        MessageCh   chan string
        ClientInfo  *ssecommon.SSEClient
        IsConnected bool
}

// NewSSESession creates a new SSE session with the given ID.
func NewSSESession(id string) *SSESession {
        return &SSESession{
                ProxySession: NewTypedProxySession(id, SessionTypeSSE),
                MessageCh:    make(chan string, 100),
                IsConnected:  true,
        }
}

// NewSSESessionWithClient creates a new SSE session with the given ID and client info.
func NewSSESessionWithClient(id string, client *ssecommon.SSEClient) *SSESession {
        sess := NewSSESession(id)
        sess.ClientInfo = client
        if client != nil {
                sess.MessageCh = client.MessageCh
        }
        return sess
}

// Disconnect marks the session as disconnected and closes the message channel.
func (s *SSESession) Disconnect() {
        s.mu.Lock()
        defer s.mu.Unlock()

        if s.IsConnected {
                s.IsConnected = false
                if s.MessageCh != nil {
                        close(s.MessageCh)
                }
        }
}

// SendMessage sends a message to the SSE client if connected.
func (s *SSESession) SendMessage(msg string) error {
        s.mu.RLock()
        defer s.mu.RUnlock()

        if !s.IsConnected {
                return ErrSessionDisconnected
        }

        select {
        case s.MessageCh <- msg:
                return nil
        default:
                return ErrMessageChannelFull
        }
}

// GetClientInfo returns the SSE client information.
func (s *SSESession) GetClientInfo() *ssecommon.SSEClient {
        s.mu.RLock()
        defer s.mu.RUnlock()
        return s.ClientInfo
}

// SetClientInfo sets the SSE client information.
func (s *SSESession) SetClientInfo(client *ssecommon.SSEClient) {
        s.mu.Lock()
        defer s.mu.Unlock()
        s.ClientInfo = client
        if client != nil && client.MessageCh != nil {
                s.MessageCh = client.MessageCh
        }
}

// GetConnectionStatus returns whether the SSE session is connected.
func (s *SSESession) GetConnectionStatus() bool {
        s.mu.RLock()
        defer s.mu.RUnlock()
        return s.IsConnected
}

// GetCreatedAt returns when the SSE session was created.
// This is useful for tracking connection duration.
func (s *SSESession) GetCreatedAt() time.Time {
        if s.ClientInfo != nil {
                return s.ClientInfo.CreatedAt
        }
        return s.CreatedAt()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "context"
        "fmt"
        "io"
        "log/slog"
        "sync"
        "sync/atomic"
        "time"
)

// localEntry wraps a session with a storage-owned last-access timestamp.
// All eviction decisions in LocalStorage are based on this timestamp, not on
// any field carried by the Session itself. This ensures every session type gets
// correct TTL behaviour regardless of its own implementation details.
type localEntry struct {
        session        Session
        lastAccessNano atomic.Int64
}

func newLocalEntry(session Session) *localEntry {
        e := &localEntry{session: session}
        e.lastAccessNano.Store(time.Now().UnixNano())
        return e
}

func (e *localEntry) lastAccess() time.Time {
        return time.Unix(0, e.lastAccessNano.Load())
}

// LocalStorage implements the Storage interface using an in-memory sync.Map.
// This is the default storage backend for single-instance deployments.
type LocalStorage struct {
        sessions sync.Map // map[string]*localEntry
}

// NewLocalStorage creates a new local in-memory storage backend.
func NewLocalStorage() *LocalStorage {
        return &LocalStorage{}
}

// Create inserts a session without replacing an existing owner.
func (s *LocalStorage) Create(_ context.Context, session Session) error {
        if session == nil || session.ID() == "" {
                return fmt.Errorf("cannot create nil session or empty ID")
        }
        if _, loaded := s.sessions.LoadOrStore(session.ID(), newLocalEntry(session)); loaded {
                return ErrSessionAlreadyExists
        }
        return nil
}

// Store saves a session to the local storage.
// For local storage, we store the session object directly without serialization.
func (s *LocalStorage) Store(_ context.Context, session Session) error {
        if session == nil {
                return fmt.Errorf("cannot store nil session")
        }
        if session.ID() == "" {
                return fmt.Errorf("cannot store session with empty ID")
        }

        s.sessions.Store(session.ID(), newLocalEntry(session))
        return nil
}

// Load retrieves a session from local storage and refreshes its last-access timestamp.
// The timestamp update happens inside LocalStorage so that eviction is correct for
// all session types, not just those that implement a Touch() method.
func (s *LocalStorage) Load(_ context.Context, id string) (Session, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session with empty ID")
        }

        val, ok := s.sessions.Load(id)
        if !ok {
                return nil, ErrSessionNotFound
        }

        entry, ok := val.(*localEntry)
        if !ok {
                return nil, fmt.Errorf("invalid session type in storage")
        }

        // Refresh last-access time by swapping in a new entry pointer. This is
        // intentional: if we mutated lastAccessNano in-place, DeleteExpired could
        // still evict the session via CompareAndDelete (it holds the same pointer).
        // Swapping the pointer makes CompareAndDelete fail for any DeleteExpired
        // goroutine that snapshotted the old pointer, preventing eviction of active
        // sessions under concurrent load.
        newEntry := newLocalEntry(entry.session)
        s.sessions.CompareAndSwap(id, entry, newEntry)
        // If CAS fails, another goroutine already replaced this entry (e.g. a
        // concurrent Store or Load). Either way the map holds a fresh pointer, so
        // DeleteExpired will not evict it incorrectly.

        return entry.session, nil
}

// LoadIfOwner atomically validates ownership and refreshes last access for the validated entry.
func (s *LocalStorage) LoadIfOwner(_ context.Context, id, expectedOwner string) (Session, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session with empty ID")
        }
        for {
                val, ok := s.sessions.Load(id)
                if !ok {
                        return nil, ErrSessionNotFound
                }
                entry, ok := val.(*localEntry)
                if !ok {
                        return nil, fmt.Errorf("invalid session type in storage")
                }
                owner, _ := entry.session.GetMetadataValue(MetadataKeyIdentityBinding)
                if owner != expectedOwner {
                        return nil, ErrSessionNotFound
                }
                if s.sessions.CompareAndSwap(id, entry, newLocalEntry(entry.session)) {
                        return entry.session, nil
                }
        }
}

// LoadMetadata returns a copy of authoritative metadata without refreshing last access.
func (s *LocalStorage) LoadMetadata(_ context.Context, id string) (map[string]string, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session metadata with empty ID")
        }
        val, ok := s.sessions.Load(id)
        if !ok {
                return nil, ErrSessionNotFound
        }
        entry, ok := val.(*localEntry)
        if !ok {
                return nil, fmt.Errorf("invalid session type in storage")
        }
        return entry.session.GetMetadata(), nil
}

// StoreIfOwner atomically replaces a session only while its ownership binding matches.
func (s *LocalStorage) StoreIfOwner(_ context.Context, session Session, expectedOwner string) (bool, error) {
        if session == nil || session.ID() == "" {
                return false, fmt.Errorf("cannot store nil session or empty ID")
        }
        if owner, _ := session.GetMetadataValue(MetadataKeyIdentityBinding); owner != expectedOwner {
                return false, nil
        }
        for {
                val, ok := s.sessions.Load(session.ID())
                if !ok {
                        return false, nil
                }
                entry, ok := val.(*localEntry)
                if !ok {
                        return false, fmt.Errorf("invalid session type in storage")
                }
                if owner, _ := entry.session.GetMetadataValue(MetadataKeyIdentityBinding); owner != expectedOwner {
                        return false, nil
                }
                if s.sessions.CompareAndSwap(session.ID(), entry, newLocalEntry(session)) {
                        return true, nil
                }
        }
}

// Delete removes a session from local storage.
func (s *LocalStorage) Delete(_ context.Context, id string) error {
        if id == "" {
                return fmt.Errorf("cannot delete session with empty ID")
        }

        s.sessions.Delete(id)
        return nil
}

// DeleteIfOwner atomically removes a session only while its ownership binding matches.
func (s *LocalStorage) DeleteIfOwner(_ context.Context, id, expectedOwner string) (bool, error) {
        if id == "" {
                return false, fmt.Errorf("cannot delete session with empty ID")
        }
        for {
                val, ok := s.sessions.Load(id)
                if !ok {
                        return false, nil
                }
                entry, ok := val.(*localEntry)
                if !ok {
                        return false, fmt.Errorf("invalid session type in storage")
                }
                if owner, _ := entry.session.GetMetadataValue(MetadataKeyIdentityBinding); owner != expectedOwner {
                        return false, nil
                }
                if s.sessions.CompareAndDelete(id, entry) {
                        return true, nil
                }
        }
}

// DeleteExpired removes all sessions whose last-access time is before the given cutoff.
func (s *LocalStorage) DeleteExpired(ctx context.Context, before time.Time) error {
        var toDelete []struct {
                id    string
                entry *localEntry
        }

        // First pass: collect expired entries
        s.sessions.Range(func(key, val any) bool {
                // Check for context cancellation
                select {
                case <-ctx.Done():
                        return false
                default:
                }

                if entry, ok := val.(*localEntry); ok {
                        if entry.lastAccess().Before(before) {
                                if id, ok := key.(string); ok {
                                        toDelete = append(toDelete, struct {
                                                id    string
                                                entry *localEntry
                                        }{id, entry})
                                }
                        }
                }
                return true
        })

        // Second pass: close and delete expired entries
        for _, item := range toDelete {
                // Check for context cancellation before processing each session
                select {
                case <-ctx.Done():
                        return ctx.Err()
                default:
                }

                // Re-check expiration and use CompareAndDelete to handle race conditions:
                // - Entry may have been touched via LocalStorage.Load and is no longer expired
                // - Entry may have been replaced via Store/UpsertSession with a new object
                // Only proceed if the stored value is still the same entry and still expired.
                if item.entry.lastAccess().Before(before) {
                        // CompareAndDelete ensures we only delete if the entry hasn't been replaced
                        if deleted := s.sessions.CompareAndDelete(item.id, item.entry); deleted {
                                // Successfully deleted - now close if implements io.Closer
                                if closer, ok := item.entry.session.(io.Closer); ok {
                                        if err := closer.Close(); err != nil {
                                                slog.Warn("failed to close session during cleanup",
                                                        "session_id", item.id,
                                                        "error", err)
                                        }
                                }
                        }
                        // If CompareAndDelete returned false, the entry was already replaced/deleted - skip it
                }
                // If re-check shows entry is no longer expired (was touched via Load), skip it
        }

        return nil
}

// Close clears all sessions from local storage.
func (s *LocalStorage) Close() error {
        // Collect keys first to avoid modifying map during iteration
        var toDelete []any
        s.sessions.Range(func(key, _ any) bool {
                toDelete = append(toDelete, key)
                return true
        })
        // Clear all sessions
        for _, key := range toDelete {
                s.sessions.Delete(key)
        }
        return nil
}

// Count returns the number of sessions in storage.
// This is a helper method not part of the Storage interface.
func (s *LocalStorage) Count() int {
        count := 0
        s.sessions.Range(func(_, _ interface{}) bool {
                count++
                return true
        })
        return count
}

// Range iterates over all sessions in storage, passing the session (not the
// internal wrapper) to f. This is a helper method not part of the Storage interface.
func (s *LocalStorage) Range(f func(key, value interface{}) bool) {
        s.sessions.Range(func(key, val interface{}) bool {
                if entry, ok := val.(*localEntry); ok {
                        return f(key, entry.session)
                }
                return f(key, val)
        })
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "time"

        "github.com/redis/go-redis/v9"

        "github.com/stacklok/toolhive-core/redisconn"
)

var (
        loadIfOwnerScript = redis.NewScript(`
local current = redis.call("GET", KEYS[1])
if not current then return false end
local decoded = cjson.decode(current)
if not decoded.metadata or decoded.metadata[ARGV[1]] ~= ARGV[2] then return false end
redis.call("PEXPIRE", KEYS[1], ARGV[3])
return current`)
        storeIfOwnerScript = redis.NewScript(`
local current = redis.call("GET", KEYS[1])
if not current then return 0 end
local decoded = cjson.decode(current)
if not decoded.metadata or decoded.metadata[ARGV[1]] ~= ARGV[2] then return 0 end
local replacement = cjson.decode(ARGV[3])
if not replacement.metadata or replacement.metadata[ARGV[1]] ~= ARGV[2] then return 0 end
redis.call("SET", KEYS[1], ARGV[3], "PX", ARGV[4])
return 1`)
        deleteIfOwnerScript = redis.NewScript(`
local current = redis.call("GET", KEYS[1])
if not current then return 0 end
local decoded = cjson.decode(current)
if not decoded.metadata or decoded.metadata[ARGV[1]] ~= ARGV[2] then return 0 end
return redis.call("DEL", KEYS[1])`)
)

// RedisStorage implements the Storage interface backed by Redis.
type RedisStorage struct {
        client    redis.UniversalClient
        keyPrefix string
        ttl       time.Duration
}

// NewRedisStorage constructs a RedisStorage. Connection-mode topology, timeouts,
// TLS, and credentials are configured through cfg; keyPrefix is the per-tenant
// key prefix (e.g. "thv:proxy:session:") and must end with ':' to avoid key
// collisions. ttl is the sliding-window expiry applied to every key on Store
// and refreshed on every Load (sliding window): sessions remain valid
// indefinitely while actively used; revocation requires an explicit Delete call.
//
// Connection-mode validation, timeout defaults, client construction (standalone,
// cluster, or sentinel), TLS plumbing, and connectivity verification are
// delegated to the shared toolhive-core redisconn package.
// A connection without verified TLS is allowed but logged at WARN.
func NewRedisStorage(ctx context.Context, cfg redisconn.Config, keyPrefix string, ttl time.Duration) (*RedisStorage, error) {
        if err := validateSessionInvariants(keyPrefix, ttl); err != nil {
                return nil, err
        }
        warnInsecureRedisTransport(&cfg)
        client, err := redisconn.NewClient(ctx, &cfg)
        if err != nil {
                return nil, err
        }
        return &RedisStorage{
                client:    client,
                keyPrefix: keyPrefix,
                ttl:       ttl,
        }, nil
}

// newRedisStorageWithClient creates a RedisStorage with a pre-configured client.
// Intended for tests only (bypasses Ping); production callers must use NewRedisStorage.
func newRedisStorageWithClient(client redis.UniversalClient, keyPrefix string, ttl time.Duration) *RedisStorage {
        return &RedisStorage{
                client:    client,
                keyPrefix: keyPrefix,
                ttl:       ttl,
        }
}

// validateSessionInvariants checks the per-storage invariants that the shared
// toolhive-core redis package does not enforce (key-prefix conventions, TTL
// bounds).
func validateSessionInvariants(keyPrefix string, ttl time.Duration) error {
        if keyPrefix == "" {
                return errors.New("invalid redis configuration: key prefix is required")
        }
        if keyPrefix[len(keyPrefix)-1] != ':' {
                return errors.New(
                        `invalid redis configuration: key prefix must end with ':' to avoid key collisions (e.g. "thv:vmcp:session:")`,
                )
        }
        if ttl <= 0 {
                return fmt.Errorf("ttl must be a positive duration")
        }
        return nil
}

func (s *RedisStorage) key(id string) string {
        return s.keyPrefix + id
}

// Create persists the complete session atomically with SET NX and an expiry.
func (s *RedisStorage) Create(ctx context.Context, session Session) error {
        if session == nil || session.ID() == "" {
                return fmt.Errorf("cannot create nil session or empty ID")
        }
        data, err := serializeSession(session)
        if err != nil {
                return fmt.Errorf("failed to serialize session: %w", err)
        }
        created, err := s.client.SetNX(ctx, s.key(session.ID()), data, s.ttl).Result()
        if err != nil {
                return err
        }
        if !created {
                return ErrSessionAlreadyExists
        }
        return nil
}

// Store serializes the session and persists it with SET … EX, refreshing the TTL on every call.
func (s *RedisStorage) Store(ctx context.Context, session Session) error {
        if session == nil {
                return fmt.Errorf("cannot store nil session")
        }
        if session.ID() == "" {
                return fmt.Errorf("cannot store session with empty ID")
        }

        data, err := serializeSession(session)
        if err != nil {
                return fmt.Errorf("failed to serialize session: %w", err)
        }

        return s.client.Set(ctx, s.key(session.ID()), data, s.ttl).Err()
}

// Load retrieves a session by ID. Returns ErrSessionNotFound when the key does not exist.
// The Redis eviction TTL is refreshed atomically via GETEX on every read so that active
// sessions are not evicted between accesses. The session's UpdatedAt timestamp is not modified.
//
// Lifetime note: this implements a sliding-window TTL. A session accessed at least once per
// TTL window will never expire and can live indefinitely while the client keeps making requests.
// Session revocation therefore depends entirely on explicit Delete calls (e.g. on logout or
// token invalidation); there is no absolute maximum session lifetime enforced here. If a hard
// cap is required in future, a MaxLifetime field checked against the session's CreatedAt would
// be the path forward.
func (s *RedisStorage) Load(ctx context.Context, id string) (Session, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session with empty ID")
        }

        data, err := s.client.GetEx(ctx, s.key(id), s.ttl).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, ErrSessionNotFound
                }
                return nil, fmt.Errorf("failed to load session: %w", err)
        }

        session, err := deserializeSession(data)
        if err != nil {
                return nil, fmt.Errorf("failed to deserialize session: %w", err)
        }

        return session, nil
}

// LoadIfOwner atomically validates ownership, refreshes the TTL, and returns the validated record.
func (s *RedisStorage) LoadIfOwner(ctx context.Context, id, expectedOwner string) (Session, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session with empty ID")
        }
        result, err := loadIfOwnerScript.Run(ctx, s.client, []string{s.key(id)},
                MetadataKeyIdentityBinding, expectedOwner, max(s.ttl.Milliseconds(), 1)).Result()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, ErrSessionNotFound
                }
                return nil, fmt.Errorf("failed to conditionally load session: %w", err)
        }
        data, ok := result.(string)
        if !ok {
                return nil, fmt.Errorf("failed to conditionally load session: unexpected response type %T", result)
        }
        sess, err := deserializeSession([]byte(data))
        if err != nil {
                return nil, fmt.Errorf("failed to deserialize session: %w", err)
        }
        return sess, nil
}

// LoadMetadata retrieves authoritative metadata with GET, without refreshing the key TTL.
func (s *RedisStorage) LoadMetadata(ctx context.Context, id string) (map[string]string, error) {
        if id == "" {
                return nil, fmt.Errorf("cannot load session metadata with empty ID")
        }
        data, err := s.client.Get(ctx, s.key(id)).Bytes()
        if err != nil {
                if errors.Is(err, redis.Nil) {
                        return nil, ErrSessionNotFound
                }
                return nil, fmt.Errorf("failed to load session metadata: %w", err)
        }
        var stored sessionData
        if err := json.Unmarshal(data, &stored); err != nil {
                return nil, fmt.Errorf("failed to deserialize session metadata: %w", err)
        }
        return stored.Metadata, nil
}

// StoreIfOwner atomically replaces a session only while its ownership binding matches.
func (s *RedisStorage) StoreIfOwner(ctx context.Context, session Session, expectedOwner string) (bool, error) {
        if session == nil || session.ID() == "" {
                return false, fmt.Errorf("cannot store nil session or empty ID")
        }
        data, err := serializeSession(session)
        if err != nil {
                return false, fmt.Errorf("failed to serialize session: %w", err)
        }
        ttlMilliseconds := max(s.ttl.Milliseconds(), 1)
        result, err := storeIfOwnerScript.Run(ctx, s.client, []string{s.key(session.ID())},
                MetadataKeyIdentityBinding, expectedOwner, data, ttlMilliseconds).Int()
        if err != nil {
                return false, fmt.Errorf("failed to conditionally store session: %w", err)
        }
        return result == 1, nil
}

// Delete removes the Redis key. A missing key is not an error.
func (s *RedisStorage) Delete(ctx context.Context, id string) error {
        if id == "" {
                return fmt.Errorf("cannot delete session with empty ID")
        }

        if err := s.client.Del(ctx, s.key(id)).Err(); err != nil {
                return fmt.Errorf("failed to delete session: %w", err)
        }

        return nil
}

// DeleteIfOwner atomically removes a session only while its ownership binding matches.
func (s *RedisStorage) DeleteIfOwner(ctx context.Context, id, expectedOwner string) (bool, error) {
        if id == "" {
                return false, fmt.Errorf("cannot delete session with empty ID")
        }
        result, err := deleteIfOwnerScript.Run(ctx, s.client, []string{s.key(id)},
                MetadataKeyIdentityBinding, expectedOwner).Int()
        if err != nil {
                return false, fmt.Errorf("failed to conditionally delete session: %w", err)
        }
        return result == 1, nil
}

// DeleteExpired is a no-op. Redis TTL handles key expiry natively.
func (*RedisStorage) DeleteExpired(_ context.Context, _ time.Time) error {
        return nil
}

// Close closes the underlying Redis client connection.
func (s *RedisStorage) Close() error {
        return s.client.Close()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package session

// StreamableSession represents a Streamable HTTP session
type StreamableSession struct {
        *ProxySession
        disconnected bool
}

// NewStreamableSession constructs a new streamable session.
func NewStreamableSession(id string) Session {
        return &StreamableSession{
                ProxySession: NewTypedProxySession(id, SessionTypeStreamable),
        }
}

// Type identifies this as a streamable session
func (*StreamableSession) Type() SessionType {
        return SessionTypeStreamable
}

// GetData returns nil for StreamableSession.
func (*StreamableSession) GetData() interface{} {
        return nil
}

// SetData is a no-op for StreamableSession.
func (*StreamableSession) SetData(interface{}) {}

// Disconnect marks session as disconnected.
func (s *StreamableSession) Disconnect() {
        s.disconnected = true
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package ssecommon provides common types and utilities for Server-Sent Events (SSE)
// used in communication between the client and MCP server.
package ssecommon

import (
        "fmt"
        "strings"
        "time"
)

const (
        // HTTPSSEEndpoint is the endpoint for SSE connections
        HTTPSSEEndpoint = "/sse"
        // HTTPMessagesEndpoint is the endpoint for JSON-RPC messages
        HTTPMessagesEndpoint = "/messages"
)

// SSEMessage represents a Server-Sent Event message
type SSEMessage struct {
        // EventType is the type of event (e.g., "message", "endpoint")
        EventType string
        // Data is the event data
        Data string
        // TargetClientID is the ID of the target client (if any)
        TargetClientID string
        // CreatedAt is the time the message was created
        CreatedAt time.Time
}

// NewSSEMessage creates a new SSE message
func NewSSEMessage(eventType, data string) *SSEMessage {
        return &SSEMessage{
                EventType: eventType,
                Data:      data,
                CreatedAt: time.Now(),
        }
}

// WithTargetClientID sets the target client ID for the message
func (m *SSEMessage) WithTargetClientID(clientID string) *SSEMessage {
        m.TargetClientID = clientID
        return m
}

// ToSSEString converts the message to an SSE-formatted string
func (m *SSEMessage) ToSSEString() string {
        var sb strings.Builder

        // Add event type
        fmt.Fprintf(&sb, "event: %s\n", m.EventType)

        // Add data (split by newlines to ensure proper formatting)
        for _, line := range strings.Split(m.Data, "\n") {
                fmt.Fprintf(&sb, "data: %s\n", line)
        }

        // End the message with a blank line
        sb.WriteString("\n")

        return sb.String()
}

// PendingSSEMessage represents an SSE message that is pending delivery
type PendingSSEMessage struct {
        // Message is the SSE message
        Message *SSEMessage
        // CreatedAt is the time the message was created
        CreatedAt time.Time
}

// NewPendingSSEMessage creates a new pending SSE message
func NewPendingSSEMessage(message *SSEMessage) *PendingSSEMessage {
        return &PendingSSEMessage{
                Message:   message,
                CreatedAt: time.Now(),
        }
}

// SSEClient represents a connected SSE client
type SSEClient struct {
        // MessageCh is the channel for sending messages to the client
        MessageCh chan string
        // CreatedAt is the time the client connected
        CreatedAt time.Time
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package transport provides utilities for handling different transport modes
// for communication between the client and MCP server, including stdio transport
// with automatic re-attachment on Docker/container restarts.
package transport

import (
        "bytes"
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "maps"
        "net"
        "net/http"
        "strings"
        "sync"
        "time"
        "unicode"

        "github.com/cenkalti/backoff/v7"
        "golang.org/x/exp/jsonrpc2"
        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/container"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        transporterrors "github.com/stacklok/toolhive/pkg/transport/errors"
        "github.com/stacklok/toolhive/pkg/transport/proxy/httpsse"
        "github.com/stacklok/toolhive/pkg/transport/proxy/streamable"
        "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // Retry configuration constants
        // defaultMaxRetries is the maximum number of re-attachment attempts after a connection loss.
        // Set to 10 to allow sufficient time for Docker/Rancher Desktop to restart (~5 minutes with backoff).
        defaultMaxRetries = 10

        // defaultInitialRetryDelay is the starting delay for exponential backoff.
        // Starts at 2 seconds to quickly recover from transient issues without overwhelming the system.
        defaultInitialRetryDelay = 2 * time.Second

        // defaultMaxRetryDelay caps the maximum delay between retry attempts.
        // Set to 30 seconds to balance between responsiveness and resource usage during extended outages.
        defaultMaxRetryDelay = 30 * time.Second

        // shutdownTimeout is the maximum time to wait for graceful shutdown operations.
        shutdownTimeout = 30 * time.Second
)

// StdioTransport implements the Transport interface using standard input/output.
// It acts as a proxy between the MCP client and the container's stdin/stdout.
type StdioTransport struct {
        host              string
        proxyPort         int
        containerName     string
        deployer          rt.Deployer
        debug             bool
        middlewares       []types.NamedMiddleware
        prometheusHandler http.Handler
        trustProxyHeaders bool
        sessionStorage    session.Storage
        sessionTTL        time.Duration
        readTimeout       time.Duration
        authInfoHandler   http.Handler
        prefixHandlers    map[string]http.Handler

        // strictProtocolValidation controls whether the streamable HTTP proxy
        // rejects an unknown/unsupported MCP-Protocol-Version header with 400.
        // See streamable.WithStrictProtocolValidation.
        strictProtocolValidation bool

        // Mutex for protecting shared state
        mutex sync.Mutex

        // Channels for communication
        shutdownCh chan struct{}
        errorCh    <-chan error

        // Proxy (SSE or Streamable HTTP)
        httpProxy types.Proxy
        proxyMode types.ProxyMode

        // Container I/O
        stdin  io.WriteCloser
        stdout io.ReadCloser

        // Container monitor
        monitor rt.Monitor

        // Container exit error (for determining if restart is needed)
        containerExitErr error
        exitErrMutex     sync.Mutex

        // Retry configuration (for testing)
        retryConfig *retryConfig
}

// retryConfig holds configuration for retry behavior
type retryConfig struct {
        maxRetries   int
        initialDelay time.Duration
        maxDelay     time.Duration
}

// defaultRetryConfig returns the default retry configuration
func defaultRetryConfig() *retryConfig {
        return &retryConfig{
                maxRetries:   defaultMaxRetries,
                initialDelay: defaultInitialRetryDelay,
                maxDelay:     defaultMaxRetryDelay,
        }
}

// NewStdioTransport creates a new stdio transport.
func NewStdioTransport(
        host string,
        proxyPort int,
        deployer rt.Deployer,
        debug bool,
        trustProxyHeaders bool,
        prometheusHandler http.Handler,
        middlewares ...types.NamedMiddleware,
) *StdioTransport {
        return &StdioTransport{
                host:              host,
                proxyPort:         proxyPort,
                deployer:          deployer,
                debug:             debug,
                trustProxyHeaders: trustProxyHeaders,
                middlewares:       middlewares,
                prometheusHandler: prometheusHandler,
                shutdownCh:        make(chan struct{}),
                proxyMode:         types.ProxyModeStreamableHTTP, // default to streamable-http
                retryConfig:       defaultRetryConfig(),
        }
}

// SetProxyMode allows configuring the proxy mode (SSE or Streamable HTTP)
func (t *StdioTransport) SetProxyMode(mode types.ProxyMode) {
        t.proxyMode = mode
}

// SetStrictProtocolValidation configures whether the streamable HTTP proxy
// rejects requests carrying an unknown/unsupported MCP-Protocol-Version
// header with HTTP 400. Default false preserves the version-agnostic
// behavior of accepting any header value.
func (t *StdioTransport) SetStrictProtocolValidation(strict bool) {
        t.strictProtocolValidation = strict
}

// SetSessionStorage configures a custom session storage backend.
// When set, the underlying proxy will use this storage instead of the default
// in-memory store, enabling session sharing across replicas (e.g. Redis-backed).
func (t *StdioTransport) SetSessionStorage(storage session.Storage) {
        t.sessionStorage = storage
}

// SetSessionTTL configures the inactivity timeout for sessions managed by the
// underlying proxy. Zero is valid and means "use the proxy's default".
func (t *StdioTransport) SetSessionTTL(ttl time.Duration) {
        t.sessionTTL = ttl
}

// SetReadTimeout configures http.Server.ReadTimeout on the underlying proxy.
// Zero is valid and means "use the proxy's default".
func (t *StdioTransport) SetReadTimeout(d time.Duration) {
        t.readTimeout = d
}

// SetAuthInfoHandler sets the RFC 9728 OAuth protected resource discovery handler.
func (t *StdioTransport) SetAuthInfoHandler(h http.Handler) {
        t.authInfoHandler = h
}

// SetPrefixHandlers sets additional HTTP handlers mounted outside the middleware chain.
func (t *StdioTransport) SetPrefixHandlers(m map[string]http.Handler) {
        t.prefixHandlers = maps.Clone(m)
}

// Mode returns the transport mode.
func (*StdioTransport) Mode() types.TransportType {
        return types.TransportTypeStdio
}

// ProxyPort returns the proxy port used by the transport.
func (t *StdioTransport) ProxyPort() int {
        return t.proxyPort
}

// setContainerName configures the transport with the container name.
// This is an unexported method used by the option pattern.
func (t *StdioTransport) setContainerName(containerName string) {
        t.mutex.Lock()
        defer t.mutex.Unlock()
        t.containerName = containerName
}

// setTargetURI configures the transport with the target URI for proxying.
// For stdio transport, this is a no-op as stdio doesn't use a target URI.
// This is an unexported method used by the option pattern.
func (*StdioTransport) setTargetURI(_ string) {
        // No-op for stdio transport
}

// Start initializes the transport and begins processing messages.
// The transport is responsible for attaching to the container.
func (t *StdioTransport) Start(ctx context.Context) error {
        t.mutex.Lock()
        defer t.mutex.Unlock()

        if t.containerName == "" {
                return transporterrors.ErrContainerNameNotSet
        }

        if t.deployer == nil {
                return fmt.Errorf("container deployer not set")
        }

        // Attach to the container
        var err error
        t.stdin, t.stdout, err = t.deployer.AttachToWorkload(ctx, t.containerName)
        if err != nil {
                return fmt.Errorf("failed to attach to container: %w", err)
        }

        // Create and start the correct proxy with middlewares
        switch t.proxyMode {
        case types.ProxyModeStreamableHTTP:
                t.httpProxy = streamable.NewHTTPProxy(
                        t.host, t.proxyPort, t.prometheusHandler, t.middlewares, t.streamableProxyOptions()...)
                if err := t.httpProxy.Start(ctx); err != nil {
                        return err
                }
                slog.Debug("streamable HTTP proxy started, processing messages")
        case types.ProxyModeSSE:
                t.httpProxy = httpsse.NewHTTPSSEProxy(
                        t.host, t.proxyPort, t.trustProxyHeaders, t.prometheusHandler, t.middlewares, t.sseProxyOptions()...)
                if err := t.httpProxy.Start(ctx); err != nil {
                        return err
                }
                slog.Debug("http SSE proxy started, processing messages")
        default:
                return fmt.Errorf("unsupported proxy mode: %v", t.proxyMode)
        }

        // Start processing messages in a goroutine
        go t.processMessages(ctx, t.stdin, t.stdout)

        // Create a container monitor
        monitorRuntime, err := container.NewFactory().Create(ctx)
        if err != nil {
                return fmt.Errorf("failed to create container monitor: %w", err)
        }
        t.monitor = container.NewMonitor(monitorRuntime, t.containerName)

        // Start monitoring the container
        t.errorCh, err = t.monitor.StartMonitoring(ctx)
        if err != nil {
                return fmt.Errorf("failed to start container monitoring: %w", err)
        }

        // Start a goroutine to handle container exit
        go t.handleContainerExit(ctx) //nolint:gosec // G118 - background goroutine manages container lifecycle, outlives request

        return nil
}

// streamableProxyOptions assembles the options for the streamable HTTP proxy
// from the transport's configured settings. Zero-valued timeouts/TTL are omitted
// so the proxy keeps its own defaults.
func (t *StdioTransport) streamableProxyOptions() []streamable.Option {
        var opts []streamable.Option
        if t.sessionTTL > 0 {
                opts = append(opts, streamable.WithSessionTTL(t.sessionTTL))
        }
        if t.readTimeout > 0 {
                opts = append(opts, streamable.WithReadTimeout(t.readTimeout))
        }
        if t.sessionStorage != nil {
                opts = append(opts, streamable.WithSessionStorage(t.sessionStorage))
        }
        return append(opts,
                streamable.WithAuthInfoHandler(t.authInfoHandler),
                streamable.WithPrefixHandlers(t.prefixHandlers),
                streamable.WithStrictProtocolValidation(t.strictProtocolValidation),
        )
}

// sseProxyOptions assembles the options for the SSE proxy from the transport's
// configured settings. Zero-valued TTL is omitted so the proxy keeps its default.
func (t *StdioTransport) sseProxyOptions() []httpsse.Option {
        var opts []httpsse.Option
        if t.sessionTTL > 0 {
                opts = append(opts, httpsse.WithSessionTTL(t.sessionTTL))
        }
        if t.readTimeout > 0 {
                opts = append(opts, httpsse.WithReadTimeout(t.readTimeout))
        }
        if t.sessionStorage != nil {
                opts = append(opts, httpsse.WithSessionStorage(t.sessionStorage))
        }
        return append(opts,
                httpsse.WithAuthInfoHandler(t.authInfoHandler),
                httpsse.WithPrefixHandlers(t.prefixHandlers),
        )
}

// Stop gracefully shuts down the transport and the container.
func (t *StdioTransport) Stop(ctx context.Context) error {
        // First check if the transport is already stopped without locking
        // to avoid deadlocks if Stop is called from multiple goroutines
        select {
        case <-t.shutdownCh:
                // Channel is already closed, transport is already stopping or stopped
                // Just return without doing anything else
                return nil
        default:
                // Channel is still open, proceed with stopping
        }

        // Now lock the mutex for the actual stopping process
        t.mutex.Lock()
        defer t.mutex.Unlock()

        // Check again after locking to handle race conditions
        select {
        case <-t.shutdownCh:
                // Channel was closed between our first check and acquiring the lock
                return nil
        default:
                // Channel is still open, close it to signal shutdown
                close(t.shutdownCh)
        }

        // Stop the monitor if it's running and we haven't already stopped it
        if t.monitor != nil {
                t.monitor.StopMonitoring()
                t.monitor = nil
        }

        // Stop the HTTP proxy
        if t.httpProxy != nil {
                if err := t.httpProxy.Stop(ctx); err != nil {
                        slog.Warn("failed to stop HTTP proxy", "error", err)
                }
        }

        // Close stdin and stdout if they're open
        if t.stdin != nil {
                if err := t.stdin.Close(); err != nil {
                        slog.Warn("failed to close stdin", "error", err)
                }
                t.stdin = nil
        }

        // Stop the container if deployer is available and we haven't already stopped it
        if t.deployer != nil && t.containerName != "" {
                // Check if the workload is still running before trying to stop it
                running, err := t.deployer.IsWorkloadRunning(ctx, t.containerName)
                if err != nil {
                        // If there's an error checking the workload status, it might be gone already
                        slog.Warn("failed to check workload status", "error", err)
                } else if running {
                        // Only try to stop the workload if it's still running
                        if err := t.deployer.StopWorkload(ctx, t.containerName); err != nil {
                                slog.Warn("failed to stop workload", "error", err)
                        }
                }
        }

        return nil
}

// IsRunning checks if the transport is currently running.
func (t *StdioTransport) IsRunning() (bool, error) {
        t.mutex.Lock()
        defer t.mutex.Unlock()

        // Check if the shutdown channel is closed
        select {
        case <-t.shutdownCh:
                return false, nil
        default:
                return true, nil
        }
}

// SetRemoteURL sets the remote URL for the MCP server.
// This is a no-op for stdio transport as it doesn't support remote servers.
func (*StdioTransport) SetRemoteURL(_ string) {
        // No-op: stdio transport doesn't support remote servers
}

// SetTokenSource sets the OAuth token source for remote authentication.
// This is a no-op for stdio transport as it doesn't support remote authentication.
func (*StdioTransport) SetTokenSource(_ oauth2.TokenSource) {
        // No-op: stdio transport doesn't support remote authentication
}

// SetOnHealthCheckFailed sets the callback for health check failures.
// This is a no-op for stdio transport as it doesn't support health checks.
func (*StdioTransport) SetOnHealthCheckFailed(_ types.HealthCheckFailedCallback) {
        // No-op: stdio transport doesn't support health checks
}

// SetOnUnauthorizedResponse sets the callback for 401 Unauthorized responses.
// This is a no-op for stdio transport as it doesn't handle HTTP responses.
func (*StdioTransport) SetOnUnauthorizedResponse(_ types.UnauthorizedResponseCallback) {
        // No-op: stdio transport doesn't handle HTTP responses
}

// isDockerSocketError checks if an error indicates Docker socket unavailability using typed error detection
func isDockerSocketError(err error) bool {
        if err == nil {
                return false
        }

        // Check for EOF errors
        if errors.Is(err, io.EOF) {
                return true
        }

        // Check for network-related errors
        var netErr *net.OpError
        if errors.As(err, &netErr) {
                // Connection refused typically indicates Docker daemon is not running
                return true
        }

        // Fallback to string matching for errors that don't implement standard interfaces
        // This handles Docker SDK errors that may not wrap standard error types
        errStr := err.Error()
        return strings.Contains(errStr, "EOF") ||
                strings.Contains(errStr, "connection refused") ||
                strings.Contains(errStr, "Cannot connect to the Docker daemon")
}

// processMessages handles the message exchange between the client and container.
func (t *StdioTransport) processMessages(ctx context.Context, _ io.WriteCloser, stdout io.ReadCloser) {
        // Create a context that will be canceled when shutdown is signaled
        ctx, cancel := context.WithCancel(ctx)
        defer cancel()

        // Monitor for shutdown signal
        go func() {
                select {
                case <-t.shutdownCh:
                        cancel()
                case <-ctx.Done():
                        // Context was canceled elsewhere
                }
        }()

        // Start a goroutine to read from stdout
        go t.processStdout(ctx, stdout)
        // Process incoming messages and send them to the container
        messageCh := t.httpProxy.GetMessageChannel()

        for {
                select {
                case <-ctx.Done():
                        return
                case msg := <-messageCh:
                        slog.Debug("processing incoming message and sending to container")
                        // Use t.stdin instead of parameter so it uses the current stdin after re-attachment
                        t.mutex.Lock()
                        currentStdin := t.stdin
                        t.mutex.Unlock()
                        if err := t.sendMessageToContainer(ctx, currentStdin, msg); err != nil {
                                slog.Error("error sending message to container", "error", err)
                        }
                        slog.Debug("message processed")
                }
        }
}

// attemptReattachment tries to re-attach to a container that has lost its stdout connection.
// Returns true if re-attachment was successful, false otherwise.
func (t *StdioTransport) attemptReattachment(ctx context.Context, stdout io.ReadCloser) bool {
        if t.deployer == nil || t.containerName == "" {
                return false
        }

        // Create an exponential backoff with the configured parameters
        expBackoff := backoff.NewExponentialBackOff()
        expBackoff.InitialInterval = t.retryConfig.initialDelay
        expBackoff.MaxInterval = t.retryConfig.maxDelay
        // Reset to allow unlimited elapsed time - we control retries via MaxTries
        expBackoff.Reset()

        var attemptCount int
        maxRetries := t.retryConfig.maxRetries

        operation := func() (any, error) {
                attemptCount++

                // Check if context is cancelled
                select {
                case <-ctx.Done():
                        return nil, backoff.Permanent(ctx.Err())
                default:
                }

                running, checkErr := t.deployer.IsWorkloadRunning(ctx, t.containerName)
                if checkErr != nil {
                        // Check if error is due to Docker being unavailable
                        if isDockerSocketError(checkErr) {
                                slog.Warn("docker socket unavailable, will retry",
                                        "attempt", attemptCount, "max_retries", maxRetries, "error", checkErr)
                                return nil, checkErr // Retry
                        }
                        slog.Warn("error checking if container is running",
                                "attempt", attemptCount, "max_retries", maxRetries, "error", checkErr)
                        return nil, checkErr // Retry
                }

                if !running {
                        slog.Info("container not running",
                                "attempt", attemptCount, "max_retries", maxRetries)
                        return nil, backoff.Permanent(fmt.Errorf("container not running"))
                }

                slog.Warn("container is still running after stdout EOF, attempting to re-attach")

                // Try to re-attach to the container
                newStdin, newStdout, attachErr := t.deployer.AttachToWorkload(ctx, t.containerName)
                if attachErr != nil {
                        slog.Error("failed to re-attach to container",
                                "attempt", attemptCount, "max_retries", maxRetries, "error", attachErr)
                        return nil, attachErr // Retry
                }

                slog.Debug("successfully re-attached to container, restarting message processing")

                // Close old stdout and log any errors
                if closeErr := stdout.Close(); closeErr != nil {
                        slog.Warn("error closing old stdout during re-attachment", "error", closeErr)
                }

                // Update stdio references with proper synchronization
                t.mutex.Lock()
                t.stdin = newStdin
                t.stdout = newStdout
                t.mutex.Unlock()

                // Start ONLY the stdout reader, not the full processMessages
                // The existing processMessages goroutine is still running and handling stdin
                go t.processStdout(ctx, newStdout)
                slog.Debug("restarted stdout processing with new pipe")
                return nil, nil // Success
        }

        // Execute the operation with retry
        // Safe conversion: maxRetries is constrained by defaultMaxRetries constant (10)
        _, err := backoff.Retry(ctx, operation,
                backoff.WithBackOff(expBackoff),
                backoff.WithMaxTries(uint(maxRetries)), // #nosec G115
                backoff.WithNotify(func(_ error, duration time.Duration) {
                        slog.Info("retry attempt",
                                "attempt", attemptCount+1, "max_retries", maxRetries, "after", duration)
                }),
        )

        if err != nil {
                if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
                        slog.Warn("re-attachment cancelled or timed out", "error", err)
                } else {
                        slog.Warn("failed to re-attach after all retry attempts")
                }
                return false
        }

        return true
}

// processStdout reads from the container's stdout and processes JSON-RPC messages.
func (t *StdioTransport) processStdout(ctx context.Context, stdout io.ReadCloser) {
        // Create a buffer for accumulating data
        var buffer bytes.Buffer

        // Create a buffer for reading
        readBuffer := make([]byte, 4096)

        for {
                select {
                case <-ctx.Done():
                        return
                default:
                        // Read data from stdout
                        n, err := stdout.Read(readBuffer)
                        if err != nil {
                                if err == io.EOF {
                                        slog.Warn("container stdout closed, checking if container is still running")

                                        // Try to re-attach to the container
                                        if t.attemptReattachment(ctx, stdout) {
                                                return
                                        }

                                        slog.Debug("container stdout closed, exiting read loop")
                                } else {
                                        slog.Error("error reading from container stdout", "error", err)
                                }
                                return
                        }

                        if n > 0 {
                                // Write the data to the buffer
                                buffer.Write(readBuffer[:n])

                                // Process the buffer
                                t.processBuffer(ctx, &buffer)
                        }
                }
        }
}

// processBuffer processes the accumulated data in the buffer.
func (t *StdioTransport) processBuffer(ctx context.Context, buffer *bytes.Buffer) {
        // Process complete lines
        for {
                line, err := buffer.ReadString('\n')
                if err == io.EOF {
                        // No complete line found, put the data back in the buffer
                        buffer.WriteString(line)
                        break
                }

                // Verify if new line character is present as last character
                // If so, remove it
                if len(line) > 0 && line[len(line)-1] == '\n' {
                        // Remove the trailing newline
                        line = line[:len(line)-1]
                }
                t.parseAndForwardJSONRPC(ctx, line)
        }
}

// sanitizeJSONString extracts the first valid JSON object from a string
func sanitizeJSONString(input string) string {
        return sanitizeBinaryString(input)
}

// sanitizeBinaryString removes all non-JSON characters and whitespace from a string
func sanitizeBinaryString(input string) string {
        // Find the first opening brace
        startIdx := strings.Index(input, "{")
        if startIdx == -1 {
                return "" // No JSON object found
        }

        // Find the last closing brace
        endIdx := strings.LastIndex(input, "}")
        if endIdx == -1 || endIdx < startIdx {
                return "" // No valid JSON object found
        }

        // Extract just the JSON object, discarding everything else
        jsonObj := input[startIdx : endIdx+1]

        // Remove all whitespace, control characters, and replacement characters
        var buffer bytes.Buffer

        for _, r := range jsonObj {
                // Skip replacement character (U+FFFD) and non-printable characters
                if r != '\uFFFD' && (unicode.IsPrint(r) || isSpace(r)) {
                        buffer.WriteRune(r)
                }
        }

        return buffer.String()
}

// isSpace reports whether r is a space character as defined by JSON.
// These are the valid space characters in JSON:
//   - ' ' (U+0020, SPACE)
//   - '\t' (U+0009, HORIZONTAL TAB)
//   - '\n' (U+000A, LINE FEED)
//   - '\r' (U+000D, CARRIAGE RETURN)
func isSpace(r rune) bool {
        return r == ' ' || r == '\t' || r == '\n' || r == '\r'
}

// parseAndForwardJSONRPC parses a JSON-RPC message and forwards it.
func (t *StdioTransport) parseAndForwardJSONRPC(ctx context.Context, line string) {
        //nolint:gosec // G706: logging raw JSON-RPC data from container stdout
        slog.Debug("JSON-RPC raw", "line", line)
        jsonData := sanitizeJSONString(line)
        //nolint:gosec // G706: logging sanitized JSON data from container stdout
        slog.Debug("Sanitized JSON", "data", jsonData)

        if jsonData == "" || jsonData == "[]" {
                return
        }

        // Try to parse the JSON
        msg, err := jsonrpc2.DecodeMessage([]byte(jsonData))
        if err != nil {
                slog.Error("error parsing JSON-RPC message", "error", err)
                return
        }

        slog.Debug("received JSON-RPC message", "type", fmt.Sprintf("%T", msg))

        if err := t.httpProxy.ForwardResponseToClients(ctx, msg); err != nil {
                if t.proxyMode == types.ProxyModeStreamableHTTP {
                        slog.Error("error forwarding to streamable-http client", "error", err)
                } else {
                        slog.Error("error dispatching backend message to SSE proxy", "error", err)
                }
        }
}

// sendMessageToContainer sends a JSON-RPC message to the container.
func (*StdioTransport) sendMessageToContainer(_ context.Context, stdin io.Writer, msg jsonrpc2.Message) error {
        // Serialize the message
        data, err := jsonrpc2.EncodeMessage(msg)
        if err != nil {
                return fmt.Errorf("failed to encode JSON-RPC message: %w", err)
        }

        // Add newline
        data = append(data, '\n')

        // Write to stdin
        slog.Debug("writing to container stdin")
        if _, err := stdin.Write(data); err != nil {
                return fmt.Errorf("failed to write to container stdin: %w", err)
        }
        slog.Debug("wrote to container stdin")

        return nil
}

// handleContainerExit handles container exit events.
func (t *StdioTransport) handleContainerExit(ctx context.Context) {
        select {
        case <-ctx.Done():
                return
        case err, ok := <-t.errorCh:
                // Check if the channel is closed
                if !ok {
                        slog.Debug("container monitor channel closed",
                                "container", t.containerName)
                        return
                }

                // Store the exit error so runner can check if restart is needed
                t.exitErrMutex.Lock()
                t.containerExitErr = err
                t.exitErrMutex.Unlock()

                //nolint:gosec // G706: logging container name from config
                slog.Warn("container exited", "container", t.containerName, "error", err)

                // Check if container was removed (not just exited) using typed error
                if errors.Is(err, rt.ErrContainerRemoved) {
                        //nolint:gosec // G706: logging container name from config
                        slog.Debug("container was removed, stopping proxy and cleaning up",
                                "container", t.containerName)
                } else {
                        //nolint:gosec // G706: logging container name from config
                        slog.Debug("container exited, will attempt automatic restart",
                                "container", t.containerName)
                }

                // Check if the transport is already stopped before trying to stop it
                select {
                case <-t.shutdownCh:
                        // Transport is already stopping or stopped
                        slog.Debug("transport is already stopping or stopped",
                                "container", t.containerName)
                        return
                default:
                        // Transport is still running, stop it
                        // Create a context with timeout for stopping the transport
                        stopCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
                        defer cancel()

                        if stopErr := t.Stop(stopCtx); stopErr != nil {
                                slog.Error("error stopping transport after container exit", "error", stopErr)
                        }
                }
        }
}

// ShouldRestart returns true if the container exited and should be restarted.
// Returns false if the container was removed (intentionally deleted) or
// restarted by Docker (already running, no ToolHive restart needed).
func (t *StdioTransport) ShouldRestart() bool {
        t.exitErrMutex.Lock()
        defer t.exitErrMutex.Unlock()

        if t.containerExitErr == nil {
                return false // No exit error, normal shutdown
        }

        // Don't restart if container was removed or restarted by Docker (use typed error check)
        return !errors.Is(t.containerExitErr, rt.ErrContainerRemoved) &&
                !errors.Is(t.containerExitErr, rt.ErrContainerRestarted)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package ngrok provides an implementation of the TunnelProvider interface using ngrok.
package ngrok

import (
        "context"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "strings"

        "go.yaml.in/yaml/v3"
        "golang.ngrok.com/ngrok/v2"
)

// TunnelProvider implements the TunnelProvider interface for ngrok.
type TunnelProvider struct {
        config TunnelConfig
}

// TunnelConfig holds configuration options for the ngrok tunnel provider.
type TunnelConfig struct {
        AuthToken      string //nolint:gosec // G117: field legitimately holds sensitive data
        URL            string // Optional: specify custom URL
        TrafficPolicy  string // Optional: specify traffic policy
        PoolingEnabled bool   // Optional: enable pooling
        DryRun         bool
}

// loadTrafficPolicyFile reads a YAML file, ensures it's .yml/.yaml,
// validates its contents, and returns its text.
func loadTrafficPolicyFile(path string) (string, error) {
        ext := strings.ToLower(filepath.Ext(path))
        if ext != ".yml" && ext != ".yaml" {
                return "", fmt.Errorf("traffic policy file must be .yml or .yaml, got %q", ext)
        }

        cleanPath := filepath.Clean(path)
        b, err := os.ReadFile(cleanPath)
        if err != nil {
                return "", fmt.Errorf("reading traffic policy file: %w", err)
        }

        var tmp any
        if err := yaml.Unmarshal(b, &tmp); err != nil {
                return "", fmt.Errorf("invalid YAML in traffic policy file: %w", err)
        }

        return string(b), nil
}

// ParseConfig parses the configuration for the ngrok tunnel provider from a map.
func (p *TunnelProvider) ParseConfig(raw map[string]any) error {
        token, ok := raw["auth-token"].(string)
        if !ok || token == "" {
                return fmt.Errorf("auth-token is required")
        }

        cfg := TunnelConfig{
                AuthToken: token,
        }

        // optional settings: url, traffic policy, pooling
        if url, ok := raw["url"].(string); ok {
                cfg.URL = url
        }
        if path, ok := raw["traffic-policy-file"].(string); ok && path != "" {
                policyText, err := loadTrafficPolicyFile(path)
                if err != nil {
                        return err
                }
                cfg.TrafficPolicy = policyText
        }
        if pooling, ok := raw["pooling"].(bool); ok {
                cfg.PoolingEnabled = pooling
        }

        p.config = cfg

        if dr, ok := raw["dry-run"].(bool); ok {
                p.config.DryRun = dr
        }

        return nil
}

// StartTunnel starts a tunnel using ngrok to the specified target URI.
func (p *TunnelProvider) StartTunnel(ctx context.Context, name, targetURI string) error {
        if p.config.DryRun {
                // behave like an active tunnel that exits on ctx cancel
                <-ctx.Done()
                return nil
        }
        //nolint:gosec // G706: logging tunnel name and target URI from config
        slog.Info("starting ngrok tunnel", "name", name, "target", targetURI)

        agent, err := ngrok.NewAgent(
                ngrok.WithAuthtoken(p.config.AuthToken),
                ngrok.WithEventHandler(func(e ngrok.Event) {
                        //nolint:gosec // G706: logging ngrok event details
                        slog.Info("ngrok event",
                                "type", e.EventType(),
                                "timestamp", e.Timestamp(),
                        )
                }),
        )

        if err != nil {
                return fmt.Errorf("failed to create ngrok agent: %w", err)
        }

        // Set up only the necessary endpoint options
        endpointOpts := []ngrok.EndpointOption{
                ngrok.WithDescription("tunnel proxy for " + name),
        }
        if p.config.URL != "" {
                endpointOpts = append(endpointOpts, ngrok.WithURL(p.config.URL))
        }
        if p.config.TrafficPolicy != "" {
                endpointOpts = append(endpointOpts, ngrok.WithTrafficPolicy(p.config.TrafficPolicy))
        }
        if p.config.PoolingEnabled {
                endpointOpts = append(endpointOpts, ngrok.WithPoolingEnabled(true))
        }

        forwarder, err := agent.Forward(ctx,
                ngrok.WithUpstream(targetURI),
                endpointOpts...,
        )
        if err != nil {
                return fmt.Errorf("ngrok.Forward error: %w", err)
        }

        //nolint:gosec // G706: logging ngrok forwarder URL from runtime
        slog.Info("ngrok forwarding live", "url", forwarder.URL())

        // Run in background, non-blocking on `.Done()`
        go func() {
                <-forwarder.Done()
                //nolint:gosec // G706: logging ngrok forwarder URL from runtime
                slog.Info("ngrok forwarding stopped", "url", forwarder.URL())
        }()

        // Return immediately
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package types provides common types and interfaces for the transport package
// used in communication between the client and MCP server.
package types

//go:generate go run go.uber.org/mock/mockgen -package mocks -destination=mocks/mock_transport.go -source=transport.go MiddlewareRunner,RunnerConfig

import (
        "context"
        "encoding/json"
        "net/http"
        "time"

        "golang.org/x/exp/jsonrpc2"
        "golang.org/x/oauth2"

        "github.com/stacklok/toolhive/pkg/auth/upstreamtoken"
        "github.com/stacklok/toolhive/pkg/authserver/server/keys"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/transport/errors"
        "github.com/stacklok/toolhive/pkg/transport/session"
)

// MiddlewareFunction is a function that wraps an http.Handler with additional functionality.
type MiddlewareFunction func(http.Handler) http.Handler

// NamedMiddleware pairs a middleware function with its name for logging purposes.
type NamedMiddleware struct {
        Name     string
        Function MiddlewareFunction
}

// Middleware defines a middleware interceptor and a close method.
type Middleware interface {
        // Handler returns the middleware function used by the proxy.
        Handler() MiddlewareFunction
        // Close cleans up any resources used by the middleware.
        Close() error
}

// MiddlewareConfig represents the configuration for a middleware.
// This is stored in the run config, and is used to instantiate the middleware.
type MiddlewareConfig struct {
        // Type is a string representing the middleware type.
        Type string `json:"type"`
        // Parameters is a JSON object containing the middleware parameters.
        // It is stored as a raw message to allow flexible parameter types.
        Parameters json.RawMessage `json:"parameters" swaggertype:"object"`
}

// NewMiddlewareConfig creates a new MiddlewareConfig with the given type and parameters.
// The parameters are marshaled to JSON and stored in the Parameters field.
func NewMiddlewareConfig(middlewareType string, parameters any) (*MiddlewareConfig, error) {
        // Marshal the parameters to JSON
        paramsJSON, err := json.Marshal(parameters)
        if err != nil {
                return nil, err
        }

        return &MiddlewareConfig{
                Type:       middlewareType,
                Parameters: paramsJSON,
        }, nil
}

// MiddlewareRunner defines the interface that middleware can use to interact with the runner.
// This unified interface replaces the ad-hoc interfaces defined in each middleware package.
type MiddlewareRunner interface {
        // AddMiddleware adds a middleware instance to the runner's middleware chain with a name for logging
        AddMiddleware(name string, middleware Middleware)

        // SetAuthInfoHandler sets the authentication info handler (used by auth middleware)
        SetAuthInfoHandler(handler http.Handler)

        // SetPrometheusHandler sets the Prometheus metrics handler (used by telemetry middleware)
        SetPrometheusHandler(handler http.Handler)

        // GetConfig returns a config interface for middleware to access runner configuration
        GetConfig() RunnerConfig

        // GetUpstreamTokenReader returns a TokenReader for identity enrichment.
        // Returns nil if the embedded auth server is not configured.
        GetUpstreamTokenReader() upstreamtoken.TokenReader

        // GetKeyProvider returns the embedded auth server's public key provider
        // for in-process JWKS key lookups. Returns nil if no embedded auth server
        // is configured.
        GetKeyProvider() keys.PublicKeyProvider
}

// RunnerConfig defines the config interface needed by middleware to access runner configuration
type RunnerConfig interface {
        GetName() string
        GetPort() int
}

// MiddlewareFactory is a function that creates a Middleware instance based on the provided configuration
// and configures the runner with the middleware and any additional handlers it provides.
type MiddlewareFactory func(config *MiddlewareConfig, runner MiddlewareRunner) error

// Transport defines the interface for MCP transport implementations.
// It provides methods for handling communication between the client and server.
type Transport interface {
        // Mode returns the transport mode.
        Mode() TransportType

        // ProxyPort returns the port used by the transport.
        ProxyPort() int

        // Start initializes the transport and begins processing messages.
        // The transport is responsible for container operations like attaching to stdin/stdout if needed.
        Start(ctx context.Context) error

        // Stop gracefully shuts down the transport.
        Stop(ctx context.Context) error

        // IsRunning checks if the transport is currently running.
        IsRunning() (bool, error)

        // SetRemoteURL sets the remote URL for the MCP server.
        // For transports that don't support remote servers (e.g., stdio), this is a no-op.
        SetRemoteURL(remoteURL string)

        // SetTokenSource sets the OAuth token source for remote authentication.
        // For transports that don't support remote authentication (e.g., stdio), this is a no-op.
        SetTokenSource(tokenSource oauth2.TokenSource)

        // SetOnHealthCheckFailed sets the callback for health check failures.
        // For transports that don't support health checks (e.g., stdio), this is a no-op.
        SetOnHealthCheckFailed(callback HealthCheckFailedCallback)

        // SetOnUnauthorizedResponse sets the callback for 401 Unauthorized responses.
        // For transports that don't support this (e.g., stdio), this is a no-op.
        SetOnUnauthorizedResponse(callback UnauthorizedResponseCallback)
}

// TransportType represents the type of transport to use.
//
//nolint:revive // Intentionally named TransportType despite package name
type TransportType string

const (
        // TransportTypeStdio represents the stdio transport.
        TransportTypeStdio TransportType = "stdio"

        // TransportTypeSSE represents the SSE transport.
        TransportTypeSSE TransportType = "sse"

        // TransportTypeStreamableHTTP represents the streamable HTTP transport.
        TransportTypeStreamableHTTP TransportType = "streamable-http"

        // TransportTypeInspector represents the transport mode for MCP Inspector.
        TransportTypeInspector TransportType = "inspector"
)

// String returns the string representation of the transport type.
func (t TransportType) String() string {
        return string(t)
}

// ParseTransportType parses a string into a transport type.
func ParseTransportType(s string) (TransportType, error) {
        switch s {
        case "stdio", "STDIO":
                return TransportTypeStdio, nil
        case "sse", "SSE":
                return TransportTypeSSE, nil
        case "streamable-http", "STREAMABLE-HTTP":
                return TransportTypeStreamableHTTP, nil
        case "inspector", "INSPECTOR":
                return TransportTypeInspector, nil
        default:
                return "", errors.ErrUnsupportedTransport
        }
}

// Proxy defines the interface for proxying messages between clients and destinations.
type Proxy interface {
        // Start starts the proxy.
        Start(ctx context.Context) error

        // Stop stops the proxy.
        Stop(ctx context.Context) error

        // IsRunning checks if the proxy is currently running.
        // This is used by HTTPTransport to detect when the proxy has stopped
        // (e.g., due to health check failure) even if the transport itself hasn't been stopped.
        IsRunning() (bool, error)

        // GetMessageChannel returns the channel for messages to/from the destination.
        GetMessageChannel() chan jsonrpc2.Message

        // SendMessageToDestination sends a message to the destination.
        SendMessageToDestination(msg jsonrpc2.Message) error

        // ForwardResponseToClients delivers a message from the destination to the
        // client session(s) it belongs to. Implementations must never deliver a
        // message to a session that did not originate it.
        ForwardResponseToClients(ctx context.Context, msg jsonrpc2.Message) error
}

// HealthCheckFailedCallback is a function that is called when a health check fails.
// This allows the transport to notify the runner/status manager when remote servers become unhealthy.
type HealthCheckFailedCallback func()

// UnauthorizedResponseCallback is a function that is called when a 401 Unauthorized response is received.
// This allows the transport to notify the runner/status manager when bearer tokens become invalid.
type UnauthorizedResponseCallback func()

// Config contains configuration options for a transport.
type Config struct {
        // Type is the type of transport to use.
        Type TransportType

        // ProxyPort is the port to use for network transports (host port).
        ProxyPort int

        // TargetPort is the port that the container will expose (container port).
        // This is only applicable to SSE transport.
        TargetPort int

        // TargetHost is the host to forward traffic to.
        // This is only applicable to SSE transport.
        TargetHost string

        // Host is the host to use for network transports.
        Host string

        // Deployer is the container runtime to use.
        // This is used for container operations like creating, starting, and attaching.
        Deployer rt.Deployer

        // Debug indicates whether debug mode is enabled.
        // If debug mode is enabled, containers will not be removed when stopped.
        Debug bool

        // Middlewares is a list of named middleware to apply to the transport.
        // These are applied in order, with the first middleware being the outermost wrapper.
        Middlewares []NamedMiddleware

        // PrometheusHandler is an optional HTTP handler for Prometheus metrics endpoint.
        // If provided, it will be exposed at /metrics on the transport's HTTP server.
        PrometheusHandler http.Handler

        // AuthInfoHandler is an optional HTTP handler for authentication information endpoint.
        // If provided, it will be exposed at /.well-known/oauth-protected-resource on the transport's HTTP server.
        AuthInfoHandler http.Handler

        // TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies
        TrustProxyHeaders bool

        // StrictProtocolValidation enables strict MCP-Protocol-Version validation
        // on the streamable HTTP proxy: a present-but-unsupported header value is
        // rejected with HTTP 400. Default false accepts any version string. Only
        // consulted for the stdio/streamable-HTTP proxy path (see factory.go).
        StrictProtocolValidation bool

        // ProxyMode is the proxy mode for stdio transport ("sse" or "streamable-http")
        ProxyMode ProxyMode

        // EndpointPrefix is an explicit prefix to prepend to SSE endpoint URLs.
        // This is used to handle path-based ingress routing scenarios.
        EndpointPrefix string

        // PrefixHandlers is a map of path prefixes to HTTP handlers.
        // These handlers are mounted on the transport's HTTP server before
        // the catch-all proxy handler. Go's ServeMux longest-match routing
        // ensures more specific paths take precedence.
        //
        // Note: When integrating the auth server, the Handler() method returns
        // a single combined handler that internally routes all OAuth/OIDC endpoints.
        // Mount it at specific paths or use http.StripPrefix as needed.
        //
        // Example:
        //
        //        {
        //          "/oauth/": authServerHandler,
        //          "/.well-known/oauth-authorization-server": authServerHandler,
        //        }
        PrefixHandlers map[string]http.Handler

        // StripConsentCookie prevents the embedded auth server's browser session
        // cookie from reaching the MCP backend. Only set when the embedded AS is mounted.
        StripConsentCookie bool

        // SessionStorage overrides the default in-memory session store when set.
        // Used for Redis-backed session sharing across replicas.
        // When nil, transports use their default in-memory LocalStorage.
        SessionStorage session.Storage

        // SessionTTL is the inactivity timeout for sessions managed by this proxy.
        // Sessions idle for longer than this duration are cleaned up by the session
        // manager's background worker. Zero uses session.DefaultSessionTTL.
        SessionTTL time.Duration

        // ReadTimeout bounds reading the entire request (headers + body) on the proxy
        // http.Server. Zero uses the proxy package default; negative values are
        // rejected. Applies to all HTTP transports; it never affects SSE responses,
        // which stream on the response side.
        ReadTimeout time.Duration
}

// ProxyMode represents the proxy mode for stdio transport.
type ProxyMode string

const (
        // ProxyModeSSE is the proxy mode for SSE.
        // Deprecated: SSE proxy mode is deprecated and will be removed in a future release.
        // Use ProxyModeStreamableHTTP instead.
        ProxyModeSSE ProxyMode = "sse"
        // ProxyModeStreamableHTTP is the proxy mode for streamable HTTP.
        ProxyModeStreamableHTTP ProxyMode = "streamable-http"
)

// IsValidProxyMode returns true if the given mode is a valid ProxyMode.
func IsValidProxyMode(mode string) bool {
        return mode == ProxyModeSSE.String() || mode == ProxyModeStreamableHTTP.String()
}

func (p ProxyMode) String() string {
        return string(p)
}

// EffectiveProxyMode determines the actual HTTP protocol the proxy is using.
// For stdio transports, this returns the proxy mode (sse or streamable-http).
// For direct transports (sse/streamable-http), this returns the transport type
// since the transport itself is the protocol.
func EffectiveProxyMode(transportType TransportType, proxyMode ProxyMode) ProxyMode {
        if transportType == TransportTypeStdio {
                if proxyMode == "" {
                        return ProxyModeStreamableHTTP
                }
                return proxyMode
        }
        return ProxyMode(transportType.String())
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package types

import (
        "context"

        "github.com/stacklok/toolhive/pkg/transport/tunnel/ngrok"
)

//go:generate mockgen -destination=mocks/mock_tunnel_provider.go -package=mocks -source=tunnel.go

// SupportedTunnelProviders maps provider names to their implementations.
var SupportedTunnelProviders = map[string]TunnelProvider{
        "ngrok": &ngrok.TunnelProvider{},
}

// TunnelProvider defines the interface for tunnel providers.
type TunnelProvider interface {
        ParseConfig(config map[string]any) error
        StartTunnel(ctx context.Context, name string, targetURI string) error
}

// GetSupportedProviderNames returns a list of supported tunnel provider names.
func GetSupportedProviderNames() []string {
        names := make([]string, 0, len(SupportedTunnelProviders))
        for name := range SupportedTunnelProviders {
                names = append(names, name)
        }
        return names
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package transport provides utilities for MCP transport operations
package transport

import (
        "fmt"
        "log/slog"
        "net/url"

        "github.com/stacklok/toolhive/pkg/transport/ssecommon"
        "github.com/stacklok/toolhive/pkg/transport/streamable"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// GenerateMCPServerURL generates the URL for an MCP server.
// If remoteURL is provided, the remote server's path will be used as the path of the proxy.
// For SSE/STDIO transports, a "#<containerName>" fragment is appended.
// For StreamableHTTP, no fragment is appended.
func GenerateMCPServerURL(transportType string, proxyMode string, host string, port int, containerName, remoteURL string) string {
        base := fmt.Sprintf("http://%s:%d", host, port)

        var isSSE, isStreamable bool

        if transportType == types.TransportTypeStdio.String() {
                // For stdio, the proxy mode determines the HTTP endpoint
                // Default to streamable-http if proxyMode is empty (matches CRD default)
                effectiveProxyMode := proxyMode
                if effectiveProxyMode == "" {
                        effectiveProxyMode = types.ProxyModeStreamableHTTP.String()
                }

                // Map proxy mode to endpoint type
                if effectiveProxyMode == types.ProxyModeSSE.String() {
                        isSSE = true
                } else {
                        // streamable-http or any other value
                        isStreamable = true
                }
        } else if transportType == types.TransportTypeSSE.String() {
                // Native SSE transport
                isSSE = true
        } else if transportType == types.TransportTypeStreamableHTTP.String() {
                // Native streamable-http transport
                isStreamable = true
        }

        // ---- Remote path case ----
        if remoteURL != "" {
                return generateRemoteMCPServerURL(base, containerName, remoteURL, isSSE, isStreamable)
        }

        // ---- Local path case (use constants as-is) ----
        if isSSE {
                // ssecommon.HTTPSSEEndpoint already includes "/sse"
                return fmt.Sprintf("%s%s#%s", base, ssecommon.HTTPSSEEndpoint, url.PathEscape(containerName))
        }

        if isStreamable {
                // streamable.HTTPStreamableHTTPEndpoint is "mcp"
                return fmt.Sprintf("%s/%s", base, streamable.HTTPStreamableHTTPEndpoint)
        }

        return ""
}

// generateRemoteMCPServerURL builds the proxy URL for a remote MCP server,
// using only the path from the remote URL.
//
// Query parameters are intentionally excluded from the generated client URL.
// The transparent proxy forwards them on every outbound request via
// WithRemoteRawQuery, so including them here would cause duplication —
// the upstream would receive the same parameter twice (e.g.
// "toolsets=core&toolsets=core"). Clients connect to the clean proxy
// URL; the proxy transparently appends the configured query string.
func generateRemoteMCPServerURL(base, containerName, remoteURL string, isSSE, isStreamable bool) string {
        targetURL, err := url.Parse(remoteURL)
        if err != nil {
                slog.Error("failed to parse target URI", "error", err)
                return ""
        }

        // Use remote path as-is; treat "/" as empty
        path := targetURL.EscapedPath()
        if path == "/" {
                path = ""
        }

        if isSSE {
                if path == "" {
                        path = ssecommon.HTTPSSEEndpoint
                }
                return fmt.Sprintf("%s%s#%s", base, path, url.PathEscape(containerName))
        }
        if isStreamable {
                return fmt.Sprintf("%s%s", base, path)
        }
        return ""
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package updates contains logic for checking if an update is available for ToolHive.
package updates

import (
        "encoding/json"
        "fmt"
        "os"
        "regexp"
        "strings"
        "time"

        "github.com/adrg/xdg"
        "github.com/google/uuid"
        "golang.org/x/mod/semver"

        "github.com/stacklok/toolhive/pkg/desktop"
        "github.com/stacklok/toolhive/pkg/lockfile"
        "github.com/stacklok/toolhive/pkg/versions"
)

// UpdateChecker is an interface for checking if a new version of ToolHive is available.
type UpdateChecker interface {
        // CheckLatestVersion checks if a new version of ToolHive is available
        // and prints the result to the console.
        CheckLatestVersion() error
}

// NewUpdateChecker creates a new instance of UpdateChecker.
func NewUpdateChecker(versionClient VersionClient) (UpdateChecker, error) {
        path, err := xdg.DataFile(updateFilePathSuffix)
        if err != nil {
                return nil, fmt.Errorf("unable to access update file path %w", err)
        }

        // Get component name for component-specific data
        component := getComponentFromVersionClient(versionClient)

        // Check to see if the file already exists. Read the instance ID and component-specific data from the
        // file if it does. If it doesn't exist, create a new instance ID.
        var instanceID, previousVersion string
        // #nosec G304: File path is not configurable at this time.
        rawContents, err := os.ReadFile(path)
        if err != nil {
                if os.IsNotExist(err) {
                        instanceID = uuid.NewString()
                } else {
                        return nil, fmt.Errorf("failed to read update file: %w", err)
                }
        } else {
                var contents updateFile
                err = json.Unmarshal(rawContents, &contents)
                if err != nil {
                        // If the file is corrupted, attempt to recover
                        if recoveredFile, recoverErr := recoverCorruptedJSON(rawContents); recoverErr == nil {
                                contents = recoveredFile
                                // Note: Update file is corrupted, attempting to preserve instance ID
                        } else {
                                return nil, fmt.Errorf("failed to deserialize update file: %w", err)
                        }
                }
                instanceID = contents.InstanceID
                previousVersion = contents.LatestVersion
        }

        return &defaultUpdateChecker{
                currentVersion:      versions.GetVersionInfo().Version,
                instanceID:          instanceID,
                updateFilePath:      path,
                versionClient:       versionClient,
                previousAPIResponse: previousVersion,
                component:           component,
        }, nil
}

const (
        updateFilePathSuffix = "toolhive/updates.json"
        updateInterval       = 30 * time.Minute
)

// TryGetAnonymousID returns the instance ID from the updates file if it exists.
// This is a read-only operation - it never generates a new ID.
// Returns empty string if the file doesn't exist or doesn't contain an instance ID.
// Use this for optional features like metrics that shouldn't trigger ID generation.
// TODO this should probably be extracted into its own package to handle instance ID generation.
func TryGetAnonymousID() (string, error) {
        path, err := xdg.DataFile(updateFilePathSuffix)
        if err != nil {
                return "", fmt.Errorf("unable to access update file path: %w", err)
        }

        // #nosec G304: File path is not configurable at this time.
        rawContents, err := os.ReadFile(path)
        if err != nil {
                if os.IsNotExist(err) {
                        // File doesn't exist yet - return empty (don't generate)
                        return "", nil
                }
                return "", fmt.Errorf("failed to read update file: %w", err)
        }

        var contents updateFile
        if err := json.Unmarshal(rawContents, &contents); err != nil {
                // If corrupted, try to recover the instance ID
                if recoveredFile, recoverErr := recoverCorruptedJSON(rawContents); recoverErr == nil {
                        return recoveredFile.InstanceID, nil
                }
                return "", fmt.Errorf("failed to deserialize update file: %w", err)
        }

        // Return whatever is in the file, even if empty
        return contents.InstanceID, nil
}

// componentInfo represents component-specific update timing information.
type componentInfo struct {
        LastCheck time.Time `json:"last_check"`
}

// updateFile represents the structure of the update file.
type updateFile struct {
        InstanceID    string                   `json:"instance_id"`
        LatestVersion string                   `json:"latest_version"`
        Components    map[string]componentInfo `json:"components"`
}

type defaultUpdateChecker struct {
        instanceID          string
        currentVersion      string
        previousAPIResponse string
        updateFilePath      string
        versionClient       VersionClient
        component           string
}

func (d *defaultUpdateChecker) CheckLatestVersion() error {
        // Read the current update file to get component-specific data
        var currentFile updateFile
        // #nosec G304: File path is not configurable at this time.
        rawContents, err := os.ReadFile(d.updateFilePath)
        if err != nil && !os.IsNotExist(err) {
                return fmt.Errorf("failed to read update file: %w", err)
        }

        // Initialize file structure if it doesn't exist or is empty
        if os.IsNotExist(err) || len(rawContents) == 0 {
                currentFile = updateFile{
                        InstanceID: d.instanceID,
                        Components: make(map[string]componentInfo),
                }
        } else {
                if err := json.Unmarshal(rawContents, &currentFile); err != nil {
                        // If the file is corrupted, attempt to recover
                        if recoveredFile, recoverErr := recoverCorruptedJSON(rawContents); recoverErr == nil {
                                currentFile = recoveredFile
                                // Note: Recovered corrupted update file, preserving instance ID
                        } else {
                                return fmt.Errorf("failed to deserialize update file: %w", err)
                        }
                }

                // Initialize components map if it doesn't exist (for backward compatibility)
                if currentFile.Components == nil {
                        currentFile.Components = make(map[string]componentInfo)
                }

                // Use the instance ID from file, but fallback to the one we generated
                if currentFile.InstanceID == "" {
                        currentFile.InstanceID = d.instanceID
                }
        }

        // Check component-specific timing
        if componentData, exists := currentFile.Components[d.component]; exists {
                if time.Since(componentData.LastCheck) < updateInterval {
                        // If it is too soon - notify the user if we already know there is
                        // an update, then exit.
                        notifyIfUpdateAvailable(d.currentVersion, currentFile.LatestVersion)
                        return nil
                }
        }

        // If the component data is stale or does not exist - get the latest version
        // from the API.
        latestVersion, err := d.versionClient.GetLatestVersion(currentFile.InstanceID, d.currentVersion)
        if err != nil {
                return fmt.Errorf("failed to check for updates: %w", err)
        }

        notifyIfUpdateAvailable(d.currentVersion, latestVersion)

        // Update shared latest version and component-specific timing
        currentFile.LatestVersion = latestVersion
        currentFile.Components[d.component] = componentInfo{
                LastCheck: time.Now().UTC(),
        }

        // Write the updated file
        updatedData, err := json.Marshal(currentFile)
        if err != nil {
                return fmt.Errorf("failed to marshal updated data: %w", err)
        }

        // Acquire lock just before writing to minimize lock time
        lockPath := d.updateFilePath + ".lock"
        lockFile := lockfile.NewTrackedLock(lockPath)
        if err := lockFile.Lock(); err != nil {
                return fmt.Errorf("failed to acquire lock on update file: %w", err)
        }
        defer lockfile.ReleaseTrackedLock(lockPath, lockFile)

        //nolint:gosec // G703 - path from trusted app config directory
        if err := os.WriteFile(d.updateFilePath, updatedData, 0600); err != nil {
                return fmt.Errorf("failed to write updated file: %w", err)
        }

        return nil
}

// getComponentFromVersionClient extracts the component name from a VersionClient.
func getComponentFromVersionClient(versionClient VersionClient) string {
        return versionClient.GetComponent()
}

func notifyIfUpdateAvailable(current, latest string) {
        // Desktop app manages its own updates, suppress CLI update message
        if desktop.IsDesktopManagedCLI() {
                return
        }
        // Print a meaningful message for people running local builds.
        if strings.HasPrefix(current, "build-") {
                // No need to compare versions, user is already aware they are not on the latest release.
                return
        }
        // Ensure both versions have the 'v' prefix for proper semantic version comparison
        if !semver.IsValid(current) {
                current = fmt.Sprintf("v%s", current)
        }
        if !semver.IsValid(latest) {
                latest = fmt.Sprintf("v%s", latest)
        }
        // Compare the versions ensuring their canonical forms
        if semver.Compare(semver.Canonical(current), semver.Canonical(latest)) < 0 {
                fmt.Fprintf(os.Stderr, "A new version of ToolHive is available: %s\nCurrently running: %s\n", latest, current)
        }
}

// recoverCorruptedJSON attempts to recover from common JSON corruption issues
// while preserving the instance_id to avoid regenerating it.
func recoverCorruptedJSON(rawContents []byte) (updateFile, error) {
        content := string(rawContents)

        // Extract the instance_id from the corrupted JSON and regenerate the file
        instanceIDRegex := regexp.MustCompile(`"instance_id":"([^"]+)"`)
        if matches := instanceIDRegex.FindStringSubmatch(content); len(matches) > 1 {
                return updateFile{
                        InstanceID: matches[1],
                        Components: make(map[string]componentInfo),
                }, nil
        }

        return updateFile{}, fmt.Errorf("unable to recover corrupted JSON")
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package updates

import (
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "os"
        "runtime"
        "strings"
        "time"

        "github.com/stacklok/toolhive/pkg/versions"
)

// VersionClient is an interface for calling the update service API.
type VersionClient interface {
        GetLatestVersion(instanceID string, currentVersion string) (string, error)
        GetComponent() string
}

// NewVersionClient creates a new instance of VersionClient.
func NewVersionClient() VersionClient {
        return NewVersionClientForComponent("CLI", "", false)
}

// NewVersionClientForComponent creates a new instance of VersionClient for a specific component.
func NewVersionClientForComponent(component, version string, uiReleaseBuild bool) VersionClient {
        return &defaultVersionClient{
                versionEndpoint: defaultVersionAPI,
                component:       component,
                customVersion:   version,
                uiReleaseBuild:  uiReleaseBuild,
        }
}

type defaultVersionClient struct {
        versionEndpoint string
        component       string
        customVersion   string
        uiReleaseBuild  bool // For the UI component, tracks if the UI is calling from a release build, false otherwise
}

const (
        instanceIDHeader  = "X-Instance-ID"
        userAgentHeader   = "User-Agent"
        defaultVersionAPI = "https://updates.stacklok.com/api/v1/version"
        defaultTimeout    = 3 * time.Second

        buildTypeRelease    = "release"
        buildTypeLocalBuild = "local_build"
)

// EnvVarSkipUpdateCheck disables update checks when set to "true"
// (case-insensitive). Disabling update checks also disables update-derived
// usage metrics, because usagemetrics.ShouldEnableMetrics gates on
// ShouldSkipUpdateChecks.
const EnvVarSkipUpdateCheck = "TOOLHIVE_SKIP_UPDATE_CHECK"

// ciEnvVars contains environment variables that indicate CI environments
var ciEnvVars = []string{
        "GITHUB_ACTIONS",
        "CI",
        "GITLAB_CI",
        "CIRCLECI",
        "TRAVIS",
        "BUILDKITE",
        "DRONE",
        "CONTINUOUS_INTEGRATION",
}

// GetLatestVersion sends a GET request to the update API endpoint and returns the version from the response.
// It returns an error if the request fails or if the response status code is not 200.
func (d *defaultVersionClient) GetLatestVersion(instanceID string, currentVersion string) (string, error) {
        // Create a new request
        req, err := http.NewRequest(http.MethodGet, d.versionEndpoint, nil)
        if err != nil {
                return "", fmt.Errorf("failed to create request: %w", err)
        }

        // Generate user agent in format: toolhive/[component] [vXX] [release/local_build] ([operating_system])

        // Use custom version if set, otherwise use the passed currentVersion
        version := currentVersion
        if d.customVersion != "" {
                version = d.customVersion
        }

        // Format version with 'v' prefix if it doesn't start with 'v'
        if !strings.HasPrefix(version, "v") {
                version = "v" + version
        }

        buildType := buildTypeLocalBuild
        if d.component == "UI" {
                // For UI: only use "release" if both server and UI are release builds
                if versions.BuildType == buildTypeRelease && d.uiReleaseBuild {
                        buildType = buildTypeRelease
                }
        } else {
                // For other components: use server build type
                if versions.BuildType == buildTypeRelease {
                        buildType = buildTypeRelease
                }
        }

        // Get platform info as OperatingSystem/Architecture
        platform := fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH)

        // Format: toolhive/[component] [vXX] [release/local_build] ([operating_system])
        userAgent := fmt.Sprintf("toolhive/%s %s %s (%s)", d.component, version, buildType, platform)

        req.Header.Set(instanceIDHeader, instanceID)
        req.Header.Set(userAgentHeader, userAgent)

        // Send the request with a reasonable timeout
        client := &http.Client{
                Timeout: defaultTimeout,
        }
        resp, err := client.Do(req) // #nosec G704 -- URL is constructed from hardcoded update API base URL
        if err != nil {
                return "", fmt.Errorf("failed to send request to update API: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("Failed to close response body", "error", err)
                }
        }()

        // Check if status code is 200
        if resp.StatusCode != http.StatusOK {
                return "", fmt.Errorf("update API returned non-200 status code: %d", resp.StatusCode)
        }

        // Read and parse the response body
        body, err := io.ReadAll(resp.Body)
        if err != nil {
                return "", fmt.Errorf("failed to read response body: %w", err)
        }

        // Parse JSON response
        var response struct {
                Version string `json:"version"`
        }
        if err := json.Unmarshal(body, &response); err != nil {
                return "", fmt.Errorf("failed to parse JSON response: %w", err)
        }

        return response.Version, nil
}

// GetComponent returns the component name for this version client.
func (d *defaultVersionClient) GetComponent() string {
        return d.component
}

// ShouldSkipUpdateChecks returns true if update checks should be skipped.
// This includes CI environments, an explicit opt-out via TOOLHIVE_SKIP_UPDATE_CHECK,
// and other scenarios where automated update checking is undesirable.
func ShouldSkipUpdateChecks() bool {
        if strings.EqualFold(os.Getenv(EnvVarSkipUpdateCheck), "true") {
                return true
        }
        // Check if running in any known CI environment
        for _, envVar := range ciEnvVars {
                if os.Getenv(envVar) != "" {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package usagemetrics

import (
        "bytes"
        "encoding/json"
        "fmt"
        "log/slog"
        "net/http"
        "runtime"
        "time"

        "github.com/stacklok/toolhive-core/env"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/updates"
        "github.com/stacklok/toolhive/pkg/versions"
)

const (
        defaultEndpoint   = "https://updates.stacklok.com/api/v1/toolcount"
        defaultTimeout    = 5 * time.Second
        instanceIDHeader  = "X-Instance-ID"
        anonymousIDHeader = "X-Anonymous-Id"
        userAgentHeader   = "User-Agent"
)

// Client sends usage metrics to the API
type Client struct {
        endpoint    string
        client      *http.Client
        anonymousID string
}

// NewClient creates a new metrics client
func NewClient(endpoint string) *Client {
        if endpoint == "" {
                endpoint = defaultEndpoint
        }

        // Get anonymous ID once at client creation and cache for process duration
        anonymousID, err := updates.TryGetAnonymousID()
        if err != nil {
                slog.Debug("failed to get anonymous ID during client creation",
                        "error", err)
        }

        return &Client{
                endpoint:    endpoint,
                anonymousID: anonymousID,
                client: &http.Client{
                        Timeout: defaultTimeout,
                },
        }
}

// SendMetrics sends the metrics record to the API
func (c *Client) SendMetrics(instanceID string, record MetricRecord) error {
        // Use cached anonymous ID (set at client creation)
        // For operator-deployed proxies without filesystem access, anonymous_id will be empty,
        // but we still send metrics with a default value.
        anonymousID := c.anonymousID
        if anonymousID == "" {
                // Only use default for operator-deployed proxies (detected via K8s env vars)
                if rt.IsKubernetesRuntimeWithEnv(&env.OSReader{}) {
                        anonymousID = "operator-proxy"
                } else {
                        // For local deployments, empty anonymous_id means file doesn't exist yet - skip sending
                        return nil
                }
        }

        data, err := json.Marshal(record)
        if err != nil {
                return fmt.Errorf("failed to marshal metrics record: %w", err)
        }

        req, err := http.NewRequest(http.MethodPost, c.endpoint, bytes.NewBuffer(data))
        if err != nil {
                return fmt.Errorf("failed to create request: %w", err)
        }

        req.Header.Set("Content-Type", "application/json")
        req.Header.Set(instanceIDHeader, instanceID)   // Proxy instance ID
        req.Header.Set(anonymousIDHeader, anonymousID) // User anonymous ID (or default for operator)
        req.Header.Set(userAgentHeader, generateUserAgent())

        resp, err := c.client.Do(req) // #nosec G704 -- URL is the hardcoded usage metrics endpoint
        if err != nil {
                return fmt.Errorf("failed to send request: %w", err)
        }
        defer func() {
                if err := resp.Body.Close(); err != nil {
                        slog.Debug("failed to close response body", "error", err)
                }
        }()

        if resp.StatusCode < 200 || resp.StatusCode >= 300 {
                return fmt.Errorf("API returned non-2xx status code: %d", resp.StatusCode)
        }

        return nil
}

// generateUserAgent creates the user agent string using the OS environment
// Format: toolhive/[local|operator] [version] [build_type] (os/arch)
func generateUserAgent() string {
        return generateUserAgentWithEnv(&env.OSReader{})
}

// generateUserAgentWithEnv creates the user agent string using the provided environment reader
// Format: toolhive/[local|operator] [version] [build_type] (os/arch)
func generateUserAgentWithEnv(envReader env.Reader) string {
        // Determine component type
        envType := "local"
        if rt.IsKubernetesRuntimeWithEnv(envReader) {
                envType = "operator"
        }

        version := versions.GetVersionInfo().Version
        if version != "" && version[0] != 'v' {
                version = "v" + version
        }

        // Get build type, buildType is set at building time
        buildType := versions.BuildType
        if buildType == "" {
                buildType = "local_build"
        }

        platform := fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH)

        return fmt.Sprintf("toolhive/%s %s %s (%s)", envType, version, buildType, platform)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package usagemetrics

import (
        "context"
        "log/slog"
        "os"
        "time"

        "github.com/google/uuid"

        "github.com/stacklok/toolhive/pkg/updates"
)

const (
        flushInterval = 15 * time.Minute

        // EnvVarUsageMetricsEnabled is the environment variable to disable usage metrics
        EnvVarUsageMetricsEnabled = "TOOLHIVE_USAGE_METRICS_ENABLED"
)

// ShouldEnableMetrics returns true if metrics collection should be enabled
// Checks: CI detection > Config disable > Environment variable > Default (true)
// If either config or env var explicitly disables metrics, they stay disabled
func ShouldEnableMetrics(configDisabled bool) bool {
        return shouldEnableMetrics(configDisabled, updates.ShouldSkipUpdateChecks, os.Getenv)
}

// shouldEnableMetrics is an internal testable version that accepts dependencies
func shouldEnableMetrics(configDisabled bool, isCI func() bool, getEnv func(string) string) bool {
        // 1. Skip metrics in CI environments (automatic)
        if isCI() {
                return false
        }

        // 2. Check config for explicit disable
        if configDisabled {
                return false
        }

        // 3. Check environment variable for explicit opt-out
        if envValue := getEnv(EnvVarUsageMetricsEnabled); envValue == "false" {
                return false
        }

        // 4. Default: enabled
        return true
}

// NewCollector creates a new metrics collector
func NewCollector() (*Collector, error) {
        client := NewClient("")
        collector := &Collector{
                instanceID:  uuid.NewString(), // Generate new instance ID for this process
                currentDate: getCurrentDateUTC(),
                client:      client,
                stopCh:      make(chan struct{}),
                doneCh:      make(chan struct{}),
                flushCh:     make(chan struct{}, 1),
        }

        // Counter starts at 0
        collector.counter.Store(0)

        return collector, nil
}

// Start begins the background goroutine for periodic flush
// If already started, this is a no-op to prevent goroutine leaks
func (c *Collector) Start() {
        if c.started.Swap(true) {
                return // Already started
        }
        go c.flushLoop()
}

// IncrementToolCall increments the tool call counter atomically
func (c *Collector) IncrementToolCall() {
        c.counter.Add(1)
}

// Flush sends the current metrics to the API
// Checks for midnight boundary crossing and handles daily reset
func (c *Collector) Flush() error {
        currentDate := getCurrentDateUTC()
        count := c.counter.Load()

        //nolint:gosec // G706: instance ID and date from internal state
        slog.Debug("usage metrics flush triggered",
                "count", count, "date", currentDate, "instance_id", c.instanceID)

        // Check if we crossed midnight UTC
        if c.currentDate != currentDate {
                // Midnight crossed! We need to flush the previous day's count
                // IMPORTANT: We send a fake timestamp (previous day at 23:59:00Z) to ensure
                // the backend attributes these calls to the correct day. The count includes
                // calls from the entire previous day plus ~0-15 minutes from the current day
                // (depending on when the flush runs), but we accept this small misattribution
                // to avoid backend complexity.

                //nolint:gosec // G706: date values from internal state
                slog.Debug("date changed, flushing previous day's count",
                        "previous_date", c.currentDate, "current_date", currentDate)

                if count > 0 {
                        // Create fake timestamp for end of previous day
                        previousDayTimestamp := c.currentDate + "T23:59:00Z"

                        record := MetricRecord{
                                Count:     count,
                                Timestamp: previousDayTimestamp,
                        }

                        //nolint:gosec // G706: date from internal state
                        slog.Debug("flushing tool calls for previous day at midnight boundary",
                                "count", count, "date", c.currentDate)

                        if err := c.client.SendMetrics(c.instanceID, record); err != nil {
                                slog.Debug("failed to send metrics for previous day",
                                        "error", err)
                                // Continue anyway - we'll reset for the new day
                        }
                }

                // Reset for the new day
                c.currentDate = currentDate
                c.counter.Store(0)

                return nil
        }

        // Normal flush - not a midnight boundary
        // Send even if count is 0 to provide presence signal

        // Send with real current timestamp
        now := time.Now().UTC()
        timestamp := now.Format(time.RFC3339) // ISO 8601 format

        record := MetricRecord{
                Count:     count,
                Timestamp: timestamp,
        }

        //nolint:gosec // G706: timestamp from time.Now()
        slog.Debug("flushing tool calls", "count", count, "timestamp", timestamp)

        if err := c.client.SendMetrics(c.instanceID, record); err != nil {
                slog.Debug("failed to send metrics", "error", err)
                // Don't return error - we'll retry on next flush
                return nil
        }

        return nil
}

// Shutdown performs final flush and stops background goroutines
func (c *Collector) Shutdown(ctx context.Context) {
        if c.shutdown.Load() {
                return
        }

        slog.Debug("shutting down usage metrics collector")
        c.shutdown.Store(true)

        // Signal stop to background goroutines (if started)
        if c.started.Load() {
                close(c.stopCh)
        }

        // Perform final flush
        if err := c.Flush(); err != nil {
                slog.Warn("failed to flush metrics on shutdown", "error", err)
        }

        // Wait for goroutines to finish with timeout (only if started)
        if c.started.Load() {
                select {
                case <-c.doneCh:
                        slog.Debug("collector stopped cleanly")
                case <-ctx.Done():
                        slog.Warn("collector shutdown timed out", "error", ctx.Err())
                }
        }
}

// GetCurrentCount returns the current count (for testing/debugging)
func (c *Collector) GetCurrentCount() int64 {
        return c.counter.Load()
}

// flushLoop periodically flushes metrics
func (c *Collector) flushLoop() {
        ticker := time.NewTicker(flushInterval)
        defer ticker.Stop()
        defer close(c.doneCh)

        for {
                select {
                case <-ticker.C:
                        if err := c.Flush(); err != nil {
                                slog.Warn("periodic flush failed", "error", err)
                        }
                case <-c.flushCh:
                        if err := c.Flush(); err != nil {
                                slog.Warn("manual flush failed", "error", err)
                        }
                case <-c.stopCh:
                        return
                }
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package usagemetrics

import (
        "context"
        "log/slog"
        "net/http"
        "time"

        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

const (
        // MiddlewareType is the type identifier for usage metrics middleware
        MiddlewareType = "usagemetrics"

        // shutdownTimeout is the maximum time to wait for collector shutdown
        shutdownTimeout = 10 * time.Second
)

// MiddlewareParams represents the parameters for usage metrics middleware
type MiddlewareParams struct {
        // No parameters needed
}

// Middleware implements the types.Middleware interface
type Middleware struct {
        collector *Collector
}

// Handler returns the middleware function
func (m *Middleware) Handler() types.MiddlewareFunction {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // Check if this is a tool call by examining the parsed MCP request
                        if parsed := mcp.GetParsedMCPRequest(r.Context()); parsed != nil {
                                if parsed.Method == "tools/call" {
                                        // Increment the tool call counter
                                        if m.collector != nil {
                                                m.collector.IncrementToolCall()
                                        }
                                }
                        }

                        // Pass to next handler
                        next.ServeHTTP(w, r)
                })
        }
}

// Close cleans up any resources
func (m *Middleware) Close() error {
        if m.collector != nil {
                ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
                defer cancel()
                m.collector.Shutdown(ctx)
        }
        return nil
}

// CreateMiddleware is the factory function for creating usage metrics middleware
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        // Create a new collector instance for this middleware
        collector, err := NewCollector()
        if err != nil {
                slog.Warn("failed to initialize usage metrics", "error", err)
                // Continue - metrics are non-critical, register no-op middleware
                mw := &Middleware{}
                runner.AddMiddleware(config.Type, mw)
                return nil
        }

        // Start the collector's background flush loop
        collector.Start()

        mw := &Middleware{
                collector: collector,
        }
        runner.AddMiddleware(config.Type, mw)
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package usagemetrics provides internal usage metrics for tracking usage and adoption.
package usagemetrics

import (
        "sync/atomic"
        "time"
)

// MetricRecord is the payload sent to the metrics API
type MetricRecord struct {
        Count     int64  `json:"count"`
        Timestamp string `json:"timestamp"` // ISO 8601 format in UTC (e.g., "2025-01-01T23:50:00Z")
}

// Collector manages tool call counting and reporting
type Collector struct {
        // Unique identifier for this proxy instance (UUID)
        instanceID string
        // Atomic counter for thread-safe increments
        counter atomic.Int64
        // Current date in YYYY-MM-DD format (UTC)
        currentDate string
        // HTTP client
        client *Client
        // Lifecycle management
        stopCh   chan struct{}
        doneCh   chan struct{}
        flushCh  chan struct{}
        started  atomic.Bool
        shutdown atomic.Bool
}

// getCurrentDateUTC returns the current date in YYYY-MM-DD format (UTC)
func getCurrentDateUTC() string {
        return time.Now().UTC().Format("2006-01-02")
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package versions provides version information for the ToolHive application.
package versions

import (
        "fmt"
        "runtime"
        "runtime/debug"
        "strings"
        "time"
)

const (
        unknownStr = "unknown"
)

// Version information set by build using -ldflags
var (
        // Version is the current version of ToolHive
        Version = "dev"
        // Commit is the git commit hash of the build
        //nolint:goconst // This is a placeholder for the commit hash
        Commit = unknownStr
        // BuildDate is the date when the binary was built
        // nolint:goconst // This is a placeholder for the build date
        BuildDate = unknownStr
        // BuildType indicates if this is a release build.
        // Set to "release" only in official release builds, everything else is considered "development".
        BuildType = "development"
)

// VersionInfo represents the version information
type VersionInfo struct {
        Version   string `json:"version"`
        Commit    string `json:"commit"`
        BuildDate string `json:"build_date"`
        GoVersion string `json:"go_version"`
        Platform  string `json:"platform"`
}

// GetVersionInfo returns the version information
func GetVersionInfo() VersionInfo {
        return getVersionInfoWithValues(Version, Commit, BuildDate)
}

// GetUserAgent returns a User-Agent string for HTTP clients
// Format: ToolHive/version (platform) go/version
func GetUserAgent() string {
        info := GetVersionInfo()
        return fmt.Sprintf("ToolHive/%s (%s) go/%s", info.Version, info.Platform, info.GoVersion)
}

// getVersionInfoWithValues returns version info with provided values (for testing)
func getVersionInfoWithValues(version, commit, buildDate string) VersionInfo {
        ver := version
        commitVal := commit
        buildDateVal := buildDate

        if strings.HasPrefix(ver, "dev") {
                if info, ok := debug.ReadBuildInfo(); ok {
                        // Try to get version from build info
                        for _, setting := range info.Settings {
                                switch setting.Key {
                                case "vcs.revision":
                                        if commitVal == unknownStr {
                                                commitVal = setting.Value
                                        }
                                case "vcs.time":
                                        if buildDateVal == unknownStr {
                                                buildDateVal = setting.Value
                                        }
                                }
                        }
                }
        }

        // Format the build date if it's a timestamp
        if buildDateVal != unknownStr {
                if t, err := time.Parse(time.RFC3339, buildDateVal); err == nil {
                        buildDateVal = t.Format("2006-01-02 15:04:05 MST")
                }
        }

        // If the version is just "dev" - manufacture a version string using the commit.
        // NOTE: Ignore any IDE warnings about this condition always being true - it is
        // overridden by the build flags.
        if ver == "dev" {
                // Truncate commit to 8 characters for brevity.
                ver = fmt.Sprintf("build-%s", fmt.Sprintf("%.*s", 8, commitVal))
        }

        return VersionInfo{
                Version:   ver,
                Commit:    commitVal,
                BuildDate: buildDateVal,
                GoVersion: runtime.Version(),
                Platform:  fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH),
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package types

import "errors"

// ErrAmbiguousSubjectProvider is returned when an xaa strategy's
// SubjectProviderName is empty and more than one upstream is configured, so
// the first-upstream default would be ambiguous.
var ErrAmbiguousSubjectProvider = errors.New(
        "SubjectProviderName must be set explicitly for xaa when more than one upstream is configured")

// DefaultSubjectProviderName defaults strategy's SubjectProviderName to
// providerName for token_exchange, aws_sts, and xaa strategies whose field is
// empty, returning a deep copy with the field defaulted; the caller's
// strategy is never mutated.
//
// All three strategies share the same wrong-subject risk: picking the wrong
// upstream token via SubjectProviderName is equally security-relevant for
// each. Only xaa hard-errors on ambiguity (when hasMultipleUpstreams is true
// and SubjectProviderName is empty, it returns ErrAmbiguousSubjectProvider
// instead of defaulting), because xaa has no existing deployments to break,
// so there is no compatibility cost. token_exchange and aws_sts keep
// defaulting silently, since hard-erroring those would be a breaking change
// requiring a deprecation path (see issue #5697).
//
// Strategies that are nil, of a different type, missing their per-type
// sub-config, or that already have SubjectProviderName set are returned
// unchanged via the original pointer.
func DefaultSubjectProviderName(
        strategy *BackendAuthStrategy,
        providerName string,
        hasMultipleUpstreams bool,
) (*BackendAuthStrategy, error) {
        if strategy == nil {
                return nil, nil
        }

        switch strategy.Type {
        case StrategyTypeTokenExchange:
                if strategy.TokenExchange == nil || strategy.TokenExchange.SubjectProviderName != "" {
                        return strategy, nil
                }
                copied := strategy.DeepCopy()
                copied.TokenExchange.SubjectProviderName = providerName
                return copied, nil

        case StrategyTypeAwsSts:
                if strategy.AwsSts == nil || strategy.AwsSts.SubjectProviderName != "" {
                        return strategy, nil
                }
                copied := strategy.DeepCopy()
                copied.AwsSts.SubjectProviderName = providerName
                return copied, nil

        case StrategyTypeXAA:
                if strategy.XAA == nil || strategy.XAA.SubjectProviderName != "" {
                        return strategy, nil
                }
                if hasMultipleUpstreams {
                        return nil, ErrAmbiguousSubjectProvider
                }
                copied := strategy.DeepCopy()
                copied.XAA.SubjectProviderName = providerName
                return copied, nil

        default:
                return strategy, nil
        }
}

//go:build !ignore_autogenerated

/*
Copyright 2025 Stacklok

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Code generated by controller-gen. DO NOT EDIT.

package types

import ()

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AwsStsConfig) DeepCopyInto(out *AwsStsConfig) {
        *out = *in
        if in.RoleMappings != nil {
                in, out := &in.RoleMappings, &out.RoleMappings
                *out = make([]RoleMapping, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.SessionDuration != nil {
                in, out := &in.SessionDuration, &out.SessionDuration
                *out = new(int32)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AwsStsConfig.
func (in *AwsStsConfig) DeepCopy() *AwsStsConfig {
        if in == nil {
                return nil
        }
        out := new(AwsStsConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *BackendAuthStrategy) DeepCopyInto(out *BackendAuthStrategy) {
        *out = *in
        if in.HeaderInjection != nil {
                in, out := &in.HeaderInjection, &out.HeaderInjection
                *out = new(HeaderInjectionConfig)
                **out = **in
        }
        if in.TokenExchange != nil {
                in, out := &in.TokenExchange, &out.TokenExchange
                *out = new(TokenExchangeConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.UpstreamInject != nil {
                in, out := &in.UpstreamInject, &out.UpstreamInject
                *out = new(UpstreamInjectConfig)
                **out = **in
        }
        if in.AwsSts != nil {
                in, out := &in.AwsSts, &out.AwsSts
                *out = new(AwsStsConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.OBO != nil {
                in, out := &in.OBO, &out.OBO
                *out = new(OBOConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.XAA != nil {
                in, out := &in.XAA, &out.XAA
                *out = new(XAAConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new BackendAuthStrategy.
func (in *BackendAuthStrategy) DeepCopy() *BackendAuthStrategy {
        if in == nil {
                return nil
        }
        out := new(BackendAuthStrategy)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *HeaderInjectionConfig) DeepCopyInto(out *HeaderInjectionConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new HeaderInjectionConfig.
func (in *HeaderInjectionConfig) DeepCopy() *HeaderInjectionConfig {
        if in == nil {
                return nil
        }
        out := new(HeaderInjectionConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OBOConfig) DeepCopyInto(out *OBOConfig) {
        *out = *in
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.CacheSkewSeconds != nil {
                in, out := &in.CacheSkewSeconds, &out.CacheSkewSeconds
                *out = new(int32)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OBOConfig.
func (in *OBOConfig) DeepCopy() *OBOConfig {
        if in == nil {
                return nil
        }
        out := new(OBOConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *RoleMapping) DeepCopyInto(out *RoleMapping) {
        *out = *in
        if in.Priority != nil {
                in, out := &in.Priority, &out.Priority
                *out = new(int)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RoleMapping.
func (in *RoleMapping) DeepCopy() *RoleMapping {
        if in == nil {
                return nil
        }
        out := new(RoleMapping)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TokenExchangeConfig) DeepCopyInto(out *TokenExchangeConfig) {
        *out = *in
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TokenExchangeConfig.
func (in *TokenExchangeConfig) DeepCopy() *TokenExchangeConfig {
        if in == nil {
                return nil
        }
        out := new(TokenExchangeConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *UpstreamInjectConfig) DeepCopyInto(out *UpstreamInjectConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new UpstreamInjectConfig.
func (in *UpstreamInjectConfig) DeepCopy() *UpstreamInjectConfig {
        if in == nil {
                return nil
        }
        out := new(UpstreamInjectConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *XAAConfig) DeepCopyInto(out *XAAConfig) {
        *out = *in
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new XAAConfig.
func (in *XAAConfig) DeepCopy() *XAAConfig {
        if in == nil {
                return nil
        }
        out := new(XAAConfig)
        in.DeepCopyInto(out)
        return out
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import "github.com/stacklok/toolhive/pkg/vmcp"

// ToAnnotations converts the wire-side override into the vmcp-domain
// ToolAnnotations. It returns nil when o is nil. Pointer fields are shared
// (callers that need a private copy should deep-copy the result).
//
// This lives in the config package (not in the compositetools derivation module)
// so the pure derivation logic does not depend on the wire format — config owns
// the mapping from its override type to the domain type it already imports.
func (o *ToolAnnotationsOverride) ToAnnotations() *vmcp.ToolAnnotations {
        if o == nil {
                return nil
        }
        ann := &vmcp.ToolAnnotations{
                ReadOnlyHint:    o.ReadOnlyHint,
                DestructiveHint: o.DestructiveHint,
                IdempotentHint:  o.IdempotentHint,
                OpenWorldHint:   o.OpenWorldHint,
        }
        if o.Title != nil {
                ann.Title = *o.Title
        }
        return ann
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package config provides configuration types and validation for VirtualMCP.
package config

import (
        "encoding/json"
        "fmt"
        "regexp"
        "strings"
        "text/template"

        "github.com/xeipuuv/gojsonschema"

        thvjson "github.com/stacklok/toolhive/pkg/json"
        "github.com/stacklok/toolhive/pkg/templates"
)

// Constants for workflow step types
const (
        WorkflowStepTypeToolCall    = "tool"
        WorkflowStepTypeElicitation = "elicitation"
        WorkflowStepTypeForEach     = "forEach"
)

// Constants for error actions
const (
        ErrorActionAbort    = "abort"
        ErrorActionContinue = "continue"
        ErrorActionRetry    = "retry"
)

// Constants for elicitation response actions
const (
        ElicitationResponseActionAbort         = "abort"
        ElicitationResponseActionContinue      = "continue"
        ElicitationResponseActionSkipRemaining = "skip_remaining"
)

// ValidateCompositeToolConfig validates a CompositeToolConfig.
// This is the primary entry point for composite tool validation, used by both
// webhooks (VirtualMCPServer, VirtualMCPCompositeToolDefinition) and runtime validation.
//
// Note: the tool.Annotations field requires no structural check here — all of
// its fields are optional pointers with no invalid value. The annotations
// safety-floor guardrail (rejecting explicit hints that would make the
// composite tool look safer than its step tools allow) runs at runtime
// (advertise time), not config-load, because step-tool annotations are only
// known once backends are aggregated.
func ValidateCompositeToolConfig(pathPrefix string, tool *CompositeToolConfig) error {
        var errors []string

        // Validate required fields
        if tool.Name == "" {
                errors = append(errors, fmt.Sprintf("%s.name is required", pathPrefix))
        }
        if tool.Description == "" {
                errors = append(errors, fmt.Sprintf("%s.description is required", pathPrefix))
        }
        if len(tool.Steps) == 0 {
                errors = append(errors, fmt.Sprintf("%s.steps must have at least one step", pathPrefix))
        }

        // Timeout validation: Duration handles parsing, but check for negative
        if tool.Timeout < 0 {
                errors = append(errors, fmt.Sprintf("%s.timeout cannot be negative", pathPrefix))
        }

        // Validate parameters if present
        if err := ValidateParameters(pathPrefix, tool.Parameters); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate steps
        if len(tool.Steps) > 0 {
                if err := ValidateWorkflowSteps(pathPrefix+".steps", tool.Steps); err != nil {
                        errors = append(errors, err.Error())
                }

                // Validate defaultResults for skippable steps
                if err := ValidateDefaultResultsForSteps(pathPrefix+".steps", tool.Steps, tool.Output); err != nil {
                        errors = append(errors, err.Error())
                }
        }

        if len(errors) > 0 {
                return fmt.Errorf("validation failed: %s", strings.Join(errors, "; "))
        }

        return nil
}

// ValidateParameters validates the parameter schema (JSON Schema format).
func ValidateParameters(pathPrefix string, params thvjson.Map) error {
        if params.IsEmpty() {
                return nil
        }

        paramsMap, err := params.ToMap()
        if err != nil {
                return fmt.Errorf("%s.parameters: invalid JSON: %w", pathPrefix, err)
        }

        // Validate type field
        typeVal, hasType := paramsMap["type"]
        if !hasType {
                return fmt.Errorf("%s.parameters: must have 'type' field (should be 'object' for JSON Schema)", pathPrefix)
        }

        typeStr, ok := typeVal.(string)
        if !ok {
                return fmt.Errorf("%s.parameters: 'type' field must be a string", pathPrefix)
        }

        if typeStr != "object" {
                return fmt.Errorf("%s.parameters: 'type' must be 'object' (got '%s')", pathPrefix, typeStr)
        }

        // Validate using JSON Schema validator
        schemaBytes, err := params.MarshalJSON()
        if err != nil {
                return fmt.Errorf("%s.parameters: failed to marshal: %w", pathPrefix, err)
        }
        if err := ValidateJSONSchema(schemaBytes); err != nil {
                return fmt.Errorf("%s.parameters: invalid JSON Schema: %w", pathPrefix, err)
        }

        return nil
}

// ValidateWorkflowSteps validates all workflow steps.
func ValidateWorkflowSteps(pathPrefix string, steps []WorkflowStepConfig) error {
        stepIDs := make(map[string]bool)
        stepIndices := make(map[string]int)

        // First pass: collect step IDs
        for i, step := range steps {
                if step.ID == "" {
                        return fmt.Errorf("%s[%d].id is required", pathPrefix, i)
                }
                if stepIDs[step.ID] {
                        return fmt.Errorf("%s[%d].id %q is duplicated", pathPrefix, i, step.ID)
                }
                stepIDs[step.ID] = true
                stepIndices[step.ID] = i
        }

        // Second pass: validate each step
        for i := range steps {
                if err := ValidateWorkflowStep(pathPrefix, i, &steps[i], stepIDs); err != nil {
                        return err
                }
        }

        // Third pass: validate no dependency cycles
        return ValidateDependencyCycles(pathPrefix, steps)
}

// ValidateWorkflowStep validates a single workflow step.
func ValidateWorkflowStep(pathPrefix string, index int, step *WorkflowStepConfig, stepIDs map[string]bool) error {
        // Validate step type
        if err := ValidateStepType(pathPrefix, index, step); err != nil {
                return err
        }

        // Validate templates
        if err := ValidateStepTemplates(pathPrefix, index, step); err != nil {
                return err
        }

        // Validate dependencies
        if err := ValidateStepDependencies(pathPrefix, index, step, stepIDs); err != nil {
                return err
        }

        // Validate error handling
        if step.OnError != nil {
                if err := ValidateStepErrorHandling(pathPrefix, index, step.OnError); err != nil {
                        return err
                }
        }

        // Validate elicitation response handlers
        stepType := step.Type
        if stepType == "" {
                stepType = WorkflowStepTypeToolCall
        }
        if stepType == WorkflowStepTypeElicitation {
                if step.OnDecline != nil {
                        if err := ValidateElicitationResponseHandler(pathPrefix, index, "onDecline", step.OnDecline); err != nil {
                                return err
                        }
                }
                if step.OnCancel != nil {
                        if err := ValidateElicitationResponseHandler(pathPrefix, index, "onCancel", step.OnCancel); err != nil {
                                return err
                        }
                }
        }

        return nil
}

// ValidateStepType validates step type and type-specific required fields.
func ValidateStepType(pathPrefix string, index int, step *WorkflowStepConfig) error {
        // Check for ambiguous configuration: both tool and message fields present without explicit type
        if step.Type == "" && step.Tool != "" && step.Message != "" {
                return fmt.Errorf(
                        "%s[%d] cannot have both tool and message fields - use explicit type to clarify intent",
                        pathPrefix, index)
        }

        stepType := step.Type
        if stepType == "" {
                stepType = WorkflowStepTypeToolCall // default
        }

        validTypes := map[string]bool{
                WorkflowStepTypeToolCall:    true,
                WorkflowStepTypeElicitation: true,
                WorkflowStepTypeForEach:     true,
        }
        if !validTypes[stepType] {
                return fmt.Errorf("%s[%d].type must be one of: tool, elicitation, forEach", pathPrefix, index)
        }

        if stepType == WorkflowStepTypeToolCall {
                if step.Tool == "" {
                        return fmt.Errorf("%s[%d].tool is required when type is tool", pathPrefix, index)
                }
                if !IsValidToolReference(step.Tool) {
                        return fmt.Errorf("%s[%d].tool must be a valid tool name", pathPrefix, index)
                }
        }

        if stepType == WorkflowStepTypeElicitation && step.Message == "" {
                return fmt.Errorf("%s[%d].message is required when type is elicitation", pathPrefix, index)
        }

        if stepType == WorkflowStepTypeForEach {
                if err := ValidateForEachStep(pathPrefix, index, step); err != nil {
                        return err
                }
        }

        return nil
}

// MaxForEachIterations is the hard cap on forEach iterations.
const MaxForEachIterations = 1000

// ValidateForEachStep validates forEach-specific configuration.
func ValidateForEachStep(pathPrefix string, index int, step *WorkflowStepConfig) error {
        // forEach must not have tool or message fields
        if step.Tool != "" {
                return fmt.Errorf("%s[%d]: forEach step must not have 'tool' field", pathPrefix, index)
        }
        if step.Message != "" {
                return fmt.Errorf("%s[%d]: forEach step must not have 'message' field", pathPrefix, index)
        }

        // collection is required and must be a valid template
        if step.Collection == "" {
                return fmt.Errorf("%s[%d].collection is required for forEach steps", pathPrefix, index)
        }
        if err := ValidateTemplate(step.Collection); err != nil {
                return fmt.Errorf("%s[%d].collection: invalid template: %w", pathPrefix, index, err)
        }

        // inner step is required
        if step.InnerStep == nil {
                return fmt.Errorf("%s[%d].step is required for forEach steps", pathPrefix, index)
        }

        if err := validateForEachInnerStep(pathPrefix, index, step.InnerStep); err != nil {
                return err
        }

        return validateForEachLimits(pathPrefix, index, step)
}

// validateForEachInnerStep validates the inner step of a forEach.
func validateForEachInnerStep(pathPrefix string, index int, inner *WorkflowStepConfig) error {
        innerType := inner.Type
        if innerType == "" {
                innerType = WorkflowStepTypeToolCall
        }
        if innerType != WorkflowStepTypeToolCall {
                return fmt.Errorf(
                        "%s[%d].step.type must be 'tool' (got %q); only tool inner steps are supported",
                        pathPrefix, index, innerType)
        }

        if inner.Tool == "" {
                return fmt.Errorf("%s[%d].step.tool is required for tool inner steps", pathPrefix, index)
        }
        if !IsValidToolReference(inner.Tool) {
                return fmt.Errorf("%s[%d].step.tool must be a valid tool name", pathPrefix, index)
        }

        if !inner.Arguments.IsEmpty() {
                args, err := inner.Arguments.ToMap()
                if err != nil {
                        return fmt.Errorf("%s[%d].step.arguments: invalid JSON: %w", pathPrefix, index, err)
                }
                for argName, argValue := range args {
                        if strValue, ok := argValue.(string); ok {
                                if err := ValidateTemplate(strValue); err != nil {
                                        return fmt.Errorf("%s[%d].step.arguments[%s]: invalid template: %w", pathPrefix, index, argName, err)
                                }
                        }
                }
        }

        return nil
}

// maxForEachParallel is the hard cap on forEach parallelism.
const maxForEachParallel = 50

// validateForEachLimits validates itemVar, maxParallel, and maxIterations.
func validateForEachLimits(pathPrefix string, index int, step *WorkflowStepConfig) error {
        if step.ItemVar != "" {
                if !isValidGoIdentifier(step.ItemVar) {
                        return fmt.Errorf("%s[%d].itemVar must be a valid Go identifier (got %q)", pathPrefix, index, step.ItemVar)
                }
                if step.ItemVar == "index" {
                        return fmt.Errorf("%s[%d].itemVar cannot be 'index' (reserved)", pathPrefix, index)
                }
        }
        if step.MaxParallel < 0 {
                return fmt.Errorf("%s[%d].maxParallel must be non-negative", pathPrefix, index)
        }
        if step.MaxParallel > maxForEachParallel {
                return fmt.Errorf("%s[%d].maxParallel must be <= %d (got %d)",
                        pathPrefix, index, maxForEachParallel, step.MaxParallel)
        }
        if step.MaxIterations < 0 {
                return fmt.Errorf("%s[%d].maxIterations must be non-negative", pathPrefix, index)
        }
        if step.MaxIterations > MaxForEachIterations {
                return fmt.Errorf("%s[%d].maxIterations must be <= %d (got %d)",
                        pathPrefix, index, MaxForEachIterations, step.MaxIterations)
        }
        return nil
}

// goIdentifierRegex matches valid Go identifiers.
var goIdentifierRegex = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)

// isValidGoIdentifier checks if s is a valid Go identifier.
func isValidGoIdentifier(s string) bool {
        return s != "" && goIdentifierRegex.MatchString(s)
}

// ValidateStepTemplates validates all template fields in a step.
func ValidateStepTemplates(pathPrefix string, index int, step *WorkflowStepConfig) error {
        // Validate arguments
        if !step.Arguments.IsEmpty() {
                args, err := step.Arguments.ToMap()
                if err != nil {
                        return fmt.Errorf("%s[%d].arguments: invalid JSON: %w", pathPrefix, index, err)
                }
                for argName, argValue := range args {
                        if strValue, ok := argValue.(string); ok {
                                if err := ValidateTemplate(strValue); err != nil {
                                        return fmt.Errorf("%s[%d].arguments[%s]: invalid template: %w", pathPrefix, index, argName, err)
                                }
                        }
                }
        }

        // Validate condition
        if step.Condition != "" {
                if err := ValidateTemplate(step.Condition); err != nil {
                        return fmt.Errorf("%s[%d].condition: invalid template: %w", pathPrefix, index, err)
                }
        }

        // Validate message
        if step.Message != "" {
                if err := ValidateTemplate(step.Message); err != nil {
                        return fmt.Errorf("%s[%d].message: invalid template: %w", pathPrefix, index, err)
                }
        }

        // Validate JSON Schema for elicitation steps
        if !step.Schema.IsEmpty() {
                schemaBytes, err := step.Schema.MarshalJSON()
                if err != nil {
                        return fmt.Errorf("%s[%d].schema: failed to marshal: %w", pathPrefix, index, err)
                }
                if err := ValidateJSONSchema(schemaBytes); err != nil {
                        return fmt.Errorf("%s[%d].schema: invalid JSON Schema: %w", pathPrefix, index, err)
                }
        }

        return nil
}

// ValidateStepDependencies validates step dependencies reference existing steps.
func ValidateStepDependencies(pathPrefix string, index int, step *WorkflowStepConfig, stepIDs map[string]bool) error {
        for _, depID := range step.DependsOn {
                if !stepIDs[depID] {
                        return fmt.Errorf("%s[%d].dependsOn references unknown step %q", pathPrefix, index, depID)
                }
        }
        return nil
}

// ValidateStepErrorHandling validates error handling configuration.
func ValidateStepErrorHandling(pathPrefix string, index int, onError *StepErrorHandling) error {
        if onError.Action == "" {
                return nil // Action is optional, defaults to abort
        }

        validActions := map[string]bool{
                ErrorActionAbort:    true,
                ErrorActionContinue: true,
                ErrorActionRetry:    true,
        }
        if !validActions[onError.Action] {
                return fmt.Errorf("%s[%d].onError.action must be one of: abort, continue, retry", pathPrefix, index)
        }

        if onError.Action == ErrorActionRetry && onError.RetryCount < 1 {
                return fmt.Errorf("%s[%d].onError.retryCount must be at least 1 when action is retry", pathPrefix, index)
        }

        return nil
}

// ValidateElicitationResponseHandler validates elicitation response handlers.
func ValidateElicitationResponseHandler(
        pathPrefix string, index int, handlerName string, handler *ElicitationResponseConfig,
) error {
        if handler.Action == "" {
                return fmt.Errorf("%s[%d].%s.action is required", pathPrefix, index, handlerName)
        }

        validActions := map[string]bool{
                ElicitationResponseActionAbort:         true,
                ElicitationResponseActionContinue:      true,
                ElicitationResponseActionSkipRemaining: true,
        }
        if !validActions[handler.Action] {
                return fmt.Errorf(
                        "%s[%d].%s.action must be one of: abort, continue, skip_remaining",
                        pathPrefix, index, handlerName)
        }

        return nil
}

// ValidateDependencyCycles validates that step dependencies don't create cycles.
func ValidateDependencyCycles(pathPrefix string, steps []WorkflowStepConfig) error {
        // Build adjacency list
        graph := make(map[string][]string)
        for _, step := range steps {
                graph[step.ID] = step.DependsOn
        }

        // DFS cycle detection
        visited := make(map[string]bool)
        recStack := make(map[string]bool)

        var hasCycle func(string) bool
        hasCycle = func(stepID string) bool {
                visited[stepID] = true
                recStack[stepID] = true

                for _, depID := range graph[stepID] {
                        if !visited[depID] {
                                if hasCycle(depID) {
                                        return true
                                }
                        } else if recStack[depID] {
                                return true
                        }
                }

                recStack[stepID] = false
                return false
        }

        for stepID := range graph {
                if !visited[stepID] {
                        if hasCycle(stepID) {
                                return fmt.Errorf("%s: dependency cycle detected involving step %q", pathPrefix, stepID)
                        }
                }
        }

        return nil
}

// stepFieldRef represents a reference to a specific field on a step's output.
type stepFieldRef struct {
        stepID string
        field  string
}

// ValidateDefaultResultsForSteps validates that defaultResults is specified for steps that:
// 1. May be skipped (have a condition or onError.action == "continue")
// 2. Are referenced by downstream steps
//
// nolint:gocyclo // multiple passes of the workflow are required to validate references are safe.
func ValidateDefaultResultsForSteps(pathPrefix string, steps []WorkflowStepConfig, output *OutputConfig) error {
        // 1. Compute all skippable step IDs
        skippableStepIDs := make(map[string]struct{})
        for _, step := range steps {
                if stepMayBeSkipped(step) {
                        skippableStepIDs[step.ID] = struct{}{}
                }
        }

        if len(skippableStepIDs) == 0 {
                return nil
        }

        // 2. Compute map from skippable step ID to set of fields with default values
        skippableStepDefaults := make(map[string]map[string]struct{})
        for _, step := range steps {
                if _, ok := skippableStepIDs[step.ID]; ok {
                        skippableStepDefaults[step.ID] = make(map[string]struct{})
                        if !step.DefaultResults.IsEmpty() {
                                defaultsMap, err := step.DefaultResults.ToMap()
                                if err == nil {
                                        for key := range defaultsMap {
                                                skippableStepDefaults[step.ID][key] = struct{}{}
                                        }
                                }
                        }
                }
        }

        // 3. Check references in steps
        for _, step := range steps {
                refs, err := extractStepFieldRefsFromStep(step)
                if err != nil {
                        return fmt.Errorf("failed to extract step references from step %s: %w", step.ID, err)
                }

                for _, ref := range refs {
                        defaultFields, isSkippable := skippableStepDefaults[ref.stepID]
                        if !isSkippable {
                                continue
                        }
                        if _, hasDefault := defaultFields[ref.field]; !hasDefault {
                                return fmt.Errorf(
                                        "%s[%s].defaultResults[%s] is required: step %q may be skipped and field %q is referenced by step %s",
                                        pathPrefix, ref.stepID, ref.field, ref.stepID, ref.field, step.ID)
                        }
                }
        }

        // 4. Check references in output
        if output != nil {
                outputRefs, err := extractStepFieldRefsFromOutput(output)
                if err != nil {
                        return fmt.Errorf("failed to extract step references from output: %w", err)
                }

                for _, ref := range outputRefs {
                        defaultFields, isSkippable := skippableStepDefaults[ref.stepID]
                        if !isSkippable {
                                continue
                        }
                        if _, hasDefault := defaultFields[ref.field]; !hasDefault {
                                return fmt.Errorf(
                                        "%s[%s].defaultResults[%s] is required: step %q may be skipped and field %q is referenced by output",
                                        pathPrefix, ref.stepID, ref.field, ref.stepID, ref.field)
                        }
                }
        }

        return nil
}

// stepMayBeSkipped returns true if a step may be skipped during execution.
func stepMayBeSkipped(step WorkflowStepConfig) bool {
        if step.Condition != "" {
                return true
        }
        if step.OnError != nil && step.OnError.Action == ErrorActionContinue {
                return true
        }
        return false
}

// extractStepFieldRefsFromStep extracts step field references from a step's templates.
func extractStepFieldRefsFromStep(step WorkflowStepConfig) ([]stepFieldRef, error) {
        var allRefs []stepFieldRef

        if step.Condition != "" {
                refs, err := extractStepFieldRefsFromTemplate(step.Condition)
                if err != nil {
                        return nil, err
                }
                allRefs = append(allRefs, refs...)
        }

        if !step.Arguments.IsEmpty() {
                args, err := step.Arguments.ToMap()
                if err == nil {
                        for _, argValue := range args {
                                if strValue, ok := argValue.(string); ok {
                                        refs, err := extractStepFieldRefsFromTemplate(strValue)
                                        if err != nil {
                                                return nil, err
                                        }
                                        allRefs = append(allRefs, refs...)
                                }
                        }
                }
        }

        if step.Message != "" {
                refs, err := extractStepFieldRefsFromTemplate(step.Message)
                if err != nil {
                        return nil, err
                }
                allRefs = append(allRefs, refs...)
        }

        return uniqueStepFieldRefs(allRefs), nil
}

// extractStepFieldRefsFromOutput extracts step field references from output templates.
func extractStepFieldRefsFromOutput(output *OutputConfig) ([]stepFieldRef, error) {
        if output == nil {
                return nil, nil
        }

        var allRefs []stepFieldRef

        for _, prop := range output.Properties {
                if prop.Value != "" {
                        refs, err := extractStepFieldRefsFromTemplate(prop.Value)
                        if err != nil {
                                return nil, err
                        }
                        allRefs = append(allRefs, refs...)
                }

                if len(prop.Properties) > 0 {
                        nestedOutput := &OutputConfig{Properties: prop.Properties}
                        nestedRefs, err := extractStepFieldRefsFromOutput(nestedOutput)
                        if err != nil {
                                return nil, err
                        }
                        allRefs = append(allRefs, nestedRefs...)
                }
        }

        return uniqueStepFieldRefs(allRefs), nil
}

// extractStepFieldRefsFromTemplate extracts step output field references from a template string.
func extractStepFieldRefsFromTemplate(tmplStr string) ([]stepFieldRef, error) {
        refs, err := templates.ExtractReferences(tmplStr)
        if err != nil {
                return nil, err
        }

        var stepRefs []stepFieldRef
        for _, ref := range refs {
                if strings.HasPrefix(ref, ".steps.") {
                        parts := strings.SplitN(ref, ".", 6)
                        if len(parts) >= 5 && parts[3] == "output" {
                                stepRefs = append(stepRefs, stepFieldRef{
                                        stepID: parts[2],
                                        field:  parts[4],
                                })
                        }
                }
        }

        return uniqueStepFieldRefs(stepRefs), nil
}

// uniqueStepFieldRefs returns a deduplicated slice of stepFieldRefs.
func uniqueStepFieldRefs(refs []stepFieldRef) []stepFieldRef {
        seen := make(map[stepFieldRef]struct{})
        result := make([]stepFieldRef, 0, len(refs))
        for _, r := range refs {
                if _, ok := seen[r]; !ok {
                        seen[r] = struct{}{}
                        result = append(result, r)
                }
        }
        return result
}

// ValidateTemplate validates Go template syntax including custom functions.
// It uses the same FuncMap as the runtime template expander to ensure
// templates using json, quote, or fromJson are validated correctly.
func ValidateTemplate(tmpl string) error {
        _, err := template.New("validation").Funcs(templates.FuncMap()).Parse(tmpl)
        if err != nil {
                return fmt.Errorf("invalid template syntax: %w", err)
        }
        return nil
}

// ValidateJSONSchema validates that bytes contain a valid JSON Schema.
func ValidateJSONSchema(schemaBytes []byte) error {
        if len(schemaBytes) == 0 {
                return nil
        }

        var schemaDoc interface{}
        if err := json.Unmarshal(schemaBytes, &schemaDoc); err != nil {
                return fmt.Errorf("failed to parse JSON: %w", err)
        }

        schemaLoader := gojsonschema.NewBytesLoader(schemaBytes)
        documentLoader := gojsonschema.NewStringLoader("{}")

        _, err := gojsonschema.Validate(schemaLoader, documentLoader)
        if err != nil {
                return fmt.Errorf("invalid JSON Schema: %w", err)
        }

        return nil
}

// IsValidToolReference validates tool reference format.
// Accepts multiple formats:
//   - "workload.tool_name" (semantic format specifying which backend's tool)
//   - "workload_toolname" (aggregated format used with prefix conflict resolution)
//   - "toolname" (simple format when there's no ambiguity)
func IsValidToolReference(tool string) bool {
        if tool == "" {
                return false
        }
        // Accept any reasonable tool name format: alphanumeric with dots, underscores, and hyphens
        pattern := `^[a-zA-Z0-9][a-zA-Z0-9._-]*$`
        matched, _ := regexp.MatchString(pattern, tool)
        return matched
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package config provides the configuration model for Virtual MCP Server.
//
// This package defines a platform-agnostic configuration model that works
// for both CLI (YAML) and Kubernetes (CRD) deployments. Platform-specific
// adapters transform their native formats into this unified model.
package config

import (
        "encoding/json"
        "fmt"
        "log/slog"
        "time"

        "github.com/stacklok/toolhive/pkg/audit"
        thvjson "github.com/stacklok/toolhive/pkg/json"
        ratelimittypes "github.com/stacklok/toolhive/pkg/ratelimit/types"
        "github.com/stacklok/toolhive/pkg/redisconfig"
        "github.com/stacklok/toolhive/pkg/telemetry"
        "github.com/stacklok/toolhive/pkg/vmcp"
        authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
)

// RedisPasswordEnvVar is the environment variable name for the Redis session storage password.
// The operator injects this as a SecretKeyRef when sessionStorage.provider is "redis"
// and passwordRef is set. The vMCP process reads this at startup to authenticate to Redis.
// #nosec G101 -- This is an environment variable name, not a hardcoded credential
const RedisPasswordEnvVar = "THV_SESSION_REDIS_PASSWORD"

// Transport type constants for static backend configuration.
// These define the allowed network transport protocols for vMCP backends in static mode.
const (
        // TransportSSE is the Server-Sent Events transport protocol.
        TransportSSE = "sse"
        // TransportStreamableHTTP is the streamable HTTP transport protocol.
        TransportStreamableHTTP = "streamable-http"
)

// StaticModeAllowedTransports lists all transport types allowed for static backend configuration.
// This must be kept in sync with the CRD enum validation in StaticBackendConfig.Transport.
var StaticModeAllowedTransports = []string{TransportSSE, TransportStreamableHTTP}

// Duration is a wrapper around time.Duration that marshals/unmarshals as a duration string.
// This ensures duration values are serialized as "30s", "1m", etc. instead of nanosecond integers.
// +gendoc
// +kubebuilder:validation:Type=string
// +kubebuilder:validation:Pattern=`^([0-9]+(\.[0-9]+)?(ns|us|µs|ms|s|m|h))+$`
type Duration time.Duration

// MarshalJSON implements json.Marshaler.
func (d Duration) MarshalJSON() ([]byte, error) {
        return json.Marshal(time.Duration(d).String())
}

// UnmarshalJSON implements json.Unmarshaler.
func (d *Duration) UnmarshalJSON(data []byte) error {
        var s string
        if err := json.Unmarshal(data, &s); err != nil {
                return err
        }
        dur, err := time.ParseDuration(s)
        if err != nil {
                return fmt.Errorf("invalid duration: %w", err)
        }
        *d = Duration(dur)
        return nil
}

// MarshalYAML implements yaml.Marshaler.
func (d Duration) MarshalYAML() (interface{}, error) {
        return time.Duration(d).String(), nil
}

// UnmarshalYAML implements yaml.Unmarshaler.
func (d *Duration) UnmarshalYAML(unmarshal func(interface{}) error) error {
        var s string
        if err := unmarshal(&s); err != nil {
                return err
        }
        dur, err := time.ParseDuration(s)
        if err != nil {
                return fmt.Errorf("invalid duration: %w", err)
        }
        *d = Duration(dur)
        return nil
}

// Config is the unified configuration model for Virtual MCP Server.
// This is platform-agnostic and used by both CLI and Kubernetes deployments.
//
// Platform-specific adapters (CLI YAML loader, Kubernetes CRD converter)
// transform their native formats into this model.
// +kubebuilder:object:generate=true
// +kubebuilder:pruning:PreserveUnknownFields
// +kubebuilder:validation:Type=object
// +gendoc
type Config struct {
        // Name is the virtual MCP server name.
        // +optional
        Name string `json:"name,omitempty" yaml:"name,omitempty"`

        // Group references an existing MCPGroup that defines backend workloads.
        // In standalone CLI mode, this is set from the YAML config file.
        // In Kubernetes, the operator populates this from spec.groupRef during conversion.
        // +optional
        Group string `json:"groupRef,omitempty" yaml:"groupRef,omitempty"`

        // Backends defines pre-configured backend servers for static mode.
        // When OutgoingAuth.Source is "inline", this field contains the full list of backend
        // servers with their URLs and transport types, eliminating the need for K8s API access.
        // When OutgoingAuth.Source is "discovered", this field is empty and backends are
        // discovered at runtime via Kubernetes API.
        // +optional
        Backends []StaticBackendConfig `json:"backends,omitempty" yaml:"backends,omitempty"`

        // IncomingAuth configures how clients authenticate to the virtual MCP server.
        // When using the Kubernetes operator, this is populated by the converter from
        // VirtualMCPServerSpec.IncomingAuth and any values set here will be superseded.
        // +optional
        IncomingAuth *IncomingAuthConfig `json:"incomingAuth,omitempty" yaml:"incomingAuth,omitempty"`

        // OutgoingAuth configures how the virtual MCP server authenticates to backends.
        // When using the Kubernetes operator, this is populated by the converter from
        // VirtualMCPServerSpec.OutgoingAuth and any values set here will be superseded.
        // +optional
        OutgoingAuth *OutgoingAuthConfig `json:"outgoingAuth,omitempty" yaml:"outgoingAuth,omitempty"`

        // Aggregation defines tool aggregation and conflict resolution strategies.
        // Supports ToolConfigRef for Kubernetes-native MCPToolConfig resource references.
        // +optional
        Aggregation *AggregationConfig `json:"aggregation,omitempty" yaml:"aggregation,omitempty"`

        // CompositeTools defines inline composite tool workflows.
        // Full workflow definitions are embedded in the configuration.
        // For Kubernetes, complex workflows can also reference VirtualMCPCompositeToolDefinition CRDs.
        // +optional
        CompositeTools []CompositeToolConfig `json:"compositeTools,omitempty" yaml:"compositeTools,omitempty"`

        // CompositeToolRefs references VirtualMCPCompositeToolDefinition resources
        // for complex, reusable workflows. Only applicable when running in Kubernetes.
        // Referenced resources must be in the same namespace as the VirtualMCPServer.
        // +optional
        CompositeToolRefs []CompositeToolRef `json:"compositeToolRefs,omitempty" yaml:"compositeToolRefs,omitempty"`

        // Operational configures operational settings.
        Operational *OperationalConfig `json:"operational,omitempty" yaml:"operational,omitempty"`

        // Metadata stores additional configuration metadata.
        Metadata map[string]string `json:"metadata,omitempty" yaml:"metadata,omitempty"`

        // Telemetry configures OpenTelemetry-based observability for the Virtual MCP server
        // including distributed tracing, OTLP metrics export, and Prometheus metrics endpoint.
        // Deprecated (Kubernetes operator only): When deploying via the operator, use
        // VirtualMCPServer.spec.telemetryConfigRef to reference a shared MCPTelemetryConfig
        // resource instead. This field remains valid for standalone (non-operator) deployments.
        // +optional
        Telemetry *telemetry.Config `json:"telemetry,omitempty" yaml:"telemetry,omitempty"`

        // Audit configures audit logging for the Virtual MCP server.
        // When present, audit logs include MCP protocol operations.
        // See audit.Config for available configuration options.
        // +optional
        Audit *audit.Config `json:"audit,omitempty" yaml:"audit,omitempty"`

        // Optimizer configures the MCP optimizer for context optimization on large toolsets.
        // When enabled, vMCP exposes only find_tool and call_tool operations to clients
        // instead of all backend tools directly. This reduces token usage by allowing
        // LLMs to discover relevant tools on demand rather than receiving all tool definitions.
        // Enabling the optimizer currently pins this instance to MCP 2025-11-25: the
        // find_tool/call_tool meta-tools are session-scoped, so Modern-capable (2026-07-28)
        // clients are negotiated down to the Legacy revision (see
        // pkg/vmcp/server/modern_gate.go; Modern parity is tracked in stacklok/toolhive#6089).
        // +optional
        Optimizer *OptimizerConfig `json:"optimizer,omitempty" yaml:"optimizer,omitempty"`

        // CodeMode configures vMCP code mode: server-side execution of Starlark scripts that
        // orchestrate multiple backend tool calls in a single request via the execute_tool_script
        // virtual tool. When enabled, execute_tool_script is advertised alongside the backend
        // tools; a script's inner tool calls are authorized individually, so a script can only
        // reach tools the caller is already permitted to use. Disabled by default.
        // +optional
        CodeMode *CodeModeConfig `json:"codeMode,omitempty" yaml:"codeMode,omitempty"`

        // SessionStorage configures session storage for stateful horizontal scaling.
        // When provider is "redis", the operator injects Redis connection parameters
        // (address, db, keyPrefix) here. The Redis password is provided separately via
        // the THV_SESSION_REDIS_PASSWORD environment variable.
        // +optional
        SessionStorage *SessionStorageConfig `json:"sessionStorage,omitempty" yaml:"sessionStorage,omitempty"`

        // RateLimiting defines rate limiting configuration for the Virtual MCP server.
        // Requires Redis session storage to be configured for distributed rate limiting.
        // +optional
        RateLimiting *ratelimittypes.RateLimitConfig `json:"rateLimiting,omitempty" yaml:"rateLimiting,omitempty"`

        // PassthroughHeaders is an allowlist of incoming client request header names
        // forwarded verbatim to all backends. Captured per request at the vMCP
        // incoming edge by headerforward.CaptureMiddleware: tool calls, resource
        // reads and prompt gets carry the values from the request that triggered
        // them, and a header absent on that request is not forwarded. Session-scoped
        // work (the session's backend connections, its advertised capability list,
        // and background capability resyncs) uses the values from the request that
        // created the session. Names must not be in the restricted set (Host,
        // hop-by-hop, X-Forwarded-*, etc.).
        // +optional
        // +listType=atomic
        PassthroughHeaders []string `json:"passthroughHeaders,omitempty" yaml:"passthroughHeaders,omitempty"`
}

// IncomingAuthConfig configures client authentication to the virtual MCP server.
//
// Note: When using the Kubernetes operator (VirtualMCPServer CRD), the
// VirtualMCPServerSpec.IncomingAuth field is the authoritative source for
// authentication configuration. The operator's converter will resolve the CRD's
// IncomingAuth (which supports Kubernetes-native references like SecretKeyRef,
// ConfigMapRef, etc.) and populate this IncomingAuthConfig with the resolved values.
// Any values set here directly will be superseded by the CRD configuration.
//
// +kubebuilder:object:generate=true
// +gendoc
type IncomingAuthConfig struct {
        // Type is the auth type: "oidc", "local", "anonymous"
        Type string `json:"type" yaml:"type"`

        // OIDC contains OIDC configuration (when Type = "oidc").
        OIDC *OIDCConfig `json:"oidc,omitempty" yaml:"oidc,omitempty"`

        // Authz contains authorization configuration (optional).
        Authz *AuthzConfig `json:"authz,omitempty" yaml:"authz,omitempty"`
}

// OIDCConfig configures OpenID Connect authentication.
// +kubebuilder:object:generate=true
// +gendoc
type OIDCConfig struct {
        // Issuer is the OIDC issuer URL.
        // +kubebuilder:validation:Pattern=`^https?://`
        Issuer string `json:"issuer" yaml:"issuer"`

        // ClientID is the OAuth client ID.
        ClientID string `json:"clientId" yaml:"clientId"`

        // ClientSecretEnv is the name of the environment variable containing the client secret.
        // This is the secure way to reference secrets - the actual secret value is never stored
        // in configuration files, only the environment variable name.
        // The secret value will be resolved from this environment variable at runtime.
        ClientSecretEnv string `json:"clientSecretEnv,omitempty" yaml:"clientSecretEnv,omitempty"`

        // Audience is the required token audience.
        Audience string `json:"audience" yaml:"audience"`

        // Resource is the OAuth 2.0 resource indicator (RFC 8707).
        // Used in WWW-Authenticate header and OAuth discovery metadata (RFC 9728).
        // If not specified, defaults to Audience.
        Resource string `json:"resource,omitempty" yaml:"resource,omitempty"`

        // JWKSURL is the explicit JWKS endpoint URL.
        // When set, skips OIDC discovery and fetches the JWKS directly from this URL.
        // This is useful when the OIDC issuer does not serve a /.well-known/openid-configuration.
        // +optional
        JWKSURL string `json:"jwksUrl,omitempty" yaml:"jwksUrl,omitempty"`

        // IntrospectionURL is the token introspection endpoint URL (RFC 7662).
        // When set, enables token introspection for opaque (non-JWT) tokens.
        // +optional
        IntrospectionURL string `json:"introspectionUrl,omitempty" yaml:"introspectionUrl,omitempty"`

        // Scopes are the required OAuth scopes.
        Scopes []string `json:"scopes,omitempty" yaml:"scopes,omitempty"`

        // ProtectedResourceAllowPrivateIP allows protected resource endpoint on private IP addresses
        // Use with caution - only enable for trusted internal IDPs or testing
        ProtectedResourceAllowPrivateIP bool `json:"protectedResourceAllowPrivateIp,omitempty" yaml:"protectedResourceAllowPrivateIp,omitempty"` //nolint:lll

        // JwksAllowPrivateIP allows OIDC discovery and JWKS fetches to private IP addresses.
        // Enable when the embedded auth server runs on a loopback address and
        // the OIDC middleware needs to fetch its JWKS from that address.
        // Use with caution - only enable for trusted internal IDPs or testing.
        JwksAllowPrivateIP bool `json:"jwksAllowPrivateIp,omitempty" yaml:"jwksAllowPrivateIp,omitempty"`

        // InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing
        // WARNING: This is insecure and should NEVER be used in production
        InsecureAllowHTTP bool `json:"insecureAllowHttp,omitempty" yaml:"insecureAllowHttp,omitempty"`

        // CABundlePath is the absolute file path to a PEM-encoded CA certificate bundle
        // used when the OIDC middleware performs HTTPS requests to the issuer
        // (OIDC discovery, JWKS fetch, token introspection). When set, the CA bundle
        // at this path is added to the trust store used for verifying the issuer's
        // TLS certificate. Typically populated by the Kubernetes operator from
        // MCPOIDCConfig.spec.inline.caBundleRef (ConfigMap) or from the in-cluster
        // service-account CA when using Kubernetes service-account auth.
        // +optional
        CABundlePath string `json:"caBundlePath,omitempty" yaml:"caBundlePath,omitempty"`
}

// AuthzConfig configures authorization.
// +kubebuilder:object:generate=true
// +gendoc
type AuthzConfig struct {
        // Type is the authz type: "cedar", "none"
        Type string `json:"type" yaml:"type"`

        // Policies contains Cedar policy definitions (when Type = "cedar").
        Policies []string `json:"policies,omitempty" yaml:"policies,omitempty"`

        // EntitiesJSON is a JSON string representing Cedar entities. Required for
        // enterprise policies that rely on transitive relationships (e.g.
        // `ClaimGroup → PlatformRole`) — without it the Cedar authorizer is
        // constructed with an empty entity store and `in` checks against absent
        // entities silently evaluate to false. Defaults to "[]" when empty.
        // +optional
        EntitiesJSON string `json:"entitiesJson,omitempty" yaml:"entitiesJson,omitempty"`

        // PrimaryUpstreamProvider names the upstream IDP provider whose access
        // token should be used as the source of JWT claims for Cedar evaluation.
        // When empty, claims from the ToolHive-issued token are used.
        // Must match an upstream provider name configured in the embedded auth server
        // (e.g. "default", "github"). Only relevant when the embedded auth server is active.
        // +optional
        PrimaryUpstreamProvider string `json:"primaryUpstreamProvider,omitempty" yaml:"primaryUpstreamProvider,omitempty"`

        // GroupClaimName is the JWT claim key that contains group membership for
        // the principal. When set, takes priority over the well-known defaults
        // ("groups", "roles", "cognito:groups"). Use this for IDPs that place
        // groups under a URI-style claim (e.g. "https://example.com/groups").
        // When empty, only the well-known claim names are checked.
        // +optional
        GroupClaimName string `json:"groupClaimName,omitempty" yaml:"groupClaimName,omitempty"`

        // RoleClaimName is the JWT claim key that contains role membership for the
        // principal. When set, the claim is extracted separately from GroupClaimName
        // and both are mapped to the configured group entity type. When empty, no
        // role extraction is performed.
        // +optional
        RoleClaimName string `json:"roleClaimName,omitempty" yaml:"roleClaimName,omitempty"`

        // GroupEntityType is the Cedar entity type name used for principal parent
        // UIDs synthesised from JWT group/role claims. Defaults to "THVGroup" when
        // empty. Must match the entity type used in EntitiesJSON for transitive
        // `in` checks to resolve. Namespaced names (`Foo::Bar`) are not yet supported.
        // +optional
        GroupEntityType string `json:"groupEntityType,omitempty" yaml:"groupEntityType,omitempty"`
}

// StaticBackendConfig defines a pre-configured backend server for static mode.
// This allows vMCP to operate without Kubernetes API access by embedding all backend
// information directly in the configuration.
// +gendoc
// +kubebuilder:object:generate=true
type StaticBackendConfig struct {
        // Name is the backend identifier.
        // Must match the backend name from the MCPGroup for auth config resolution.
        // +kubebuilder:validation:Required
        Name string `json:"name" yaml:"name"`

        // URL is the backend's MCP server base URL.
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Pattern=`^https?://`
        URL string `json:"url" yaml:"url"`

        // Transport is the MCP transport protocol: "sse" or "streamable-http"
        // Only network transports supported by vMCP client are allowed.
        // +kubebuilder:validation:Enum=sse;streamable-http
        // +kubebuilder:validation:Required
        Transport string `json:"transport" yaml:"transport"`

        // Type is the backend workload type: "entry" for MCPServerEntry backends, or empty
        // for container/proxy backends. Entry backends connect directly to remote MCP servers.
        // +kubebuilder:validation:Enum=entry;""
        // +optional
        Type string `json:"type,omitempty" yaml:"type,omitempty"`

        // CABundlePath is the file path to a custom CA certificate bundle for TLS verification.
        // Only valid when Type is "entry". The operator mounts CA bundles at
        // /etc/toolhive/ca-bundles/<name>/ca.crt.
        // +optional
        CABundlePath string `json:"caBundlePath,omitempty" yaml:"caBundlePath,omitempty"`

        // Metadata is a custom key-value map for storing additional backend information
        // such as labels, tags, or other arbitrary data (e.g., "env": "prod", "region": "us-east-1").
        // This is NOT Kubernetes ObjectMeta - it's a simple string map for user-defined metadata.
        // Reserved keys: "group" is automatically set by vMCP and any user-provided value will be overridden.
        // +optional
        Metadata map[string]string `json:"metadata,omitempty" yaml:"metadata,omitempty"`
}

// OutgoingAuthConfig configures backend authentication.
//
// Note: When using the Kubernetes operator (VirtualMCPServer CRD), the
// VirtualMCPServerSpec.OutgoingAuth field is the authoritative source for
// backend authentication configuration. The operator's converter will resolve
// the CRD's OutgoingAuth (which supports Kubernetes-native references like
// SecretKeyRef, ConfigMapRef, etc.) and populate this OutgoingAuthConfig with
// the resolved values. Any values set here directly will be superseded by the
// CRD configuration.
//
// +kubebuilder:object:generate=true
// +gendoc
type OutgoingAuthConfig struct {
        // Source defines how to discover backend auth: "inline", "discovered"
        // - inline: Explicit configuration in OutgoingAuth
        // - discovered: Auto-discover from backend MCPServer.externalAuthConfigRef (Kubernetes only)
        Source string `json:"source" yaml:"source"`

        // Default is the default auth strategy for backends without explicit config.
        Default *authtypes.BackendAuthStrategy `json:"default,omitempty" yaml:"default,omitempty"`

        // Backends contains per-backend auth configuration.
        Backends map[string]*authtypes.BackendAuthStrategy `json:"backends,omitempty" yaml:"backends,omitempty"`
}

// ResolveForBackend returns the auth strategy for a given backend ID.
// It checks for backend-specific config first, then falls back to default.
// Returns nil if no authentication is configured.
func (c *OutgoingAuthConfig) ResolveForBackend(backendID string) *authtypes.BackendAuthStrategy {
        if c == nil {
                return nil
        }

        // Check for backend-specific configuration
        if strategy, exists := c.Backends[backendID]; exists && strategy != nil {
                return strategy
        }

        // Fall back to default configuration
        if c.Default != nil {
                return c.Default
        }

        // No authentication configured
        return nil
}

// ApplyToBackend applies this outgoing auth configuration to a backend, choosing
// between the backend's own discovered auth and config-based auth. It is the single
// implementation shared by startup discovery (pkg/vmcp/aggregator) and the Kubernetes
// backend watcher's reconciler (pkg/vmcp/k8s), so both paths resolve auth identically.
//
// Auth resolution logic:
//   - "discovered" mode: keep discovered auth if the backend has any, otherwise fall
//     back to config-based auth via ResolveForBackend (backend-specific entry, then
//     Default, then no auth)
//   - "inline" mode (or ""): always use config-based auth, ignore discovered auth
//   - unknown mode: default to config-based auth for safety
//
// A nil receiver is a no-op: the backend keeps whatever auth it already carries.
func (c *OutgoingAuthConfig) ApplyToBackend(backend *vmcp.Backend, backendName string) {
        if c == nil {
                return
        }

        // Determine if we should use discovered auth or config-based auth
        var useDiscoveredAuth bool
        switch c.Source {
        case "discovered":
                // In discovered mode, use auth discovered from MCPServer (if any exists)
                // If no auth is discovered, fall back to config-based auth via ResolveForBackend
                // which will use backend-specific config, then Default, then no auth
                useDiscoveredAuth = backend.AuthConfig != nil
        case "inline", "":
                // For inline mode or empty source, always use config-based auth
                // Ignore any discovered auth from backends
                useDiscoveredAuth = false
        default:
                // Unknown source mode - default to config-based auth for safety
                slog.Warn("unknown auth source mode, defaulting to config-based auth", "source", c.Source)
                useDiscoveredAuth = false
        }

        if useDiscoveredAuth {
                // Keep the auth discovered from MCPServer (already populated in backend)
                slog.Debug("backend using discovered auth strategy", "backend", backendName, "strategy", backend.AuthConfig.Type)
        } else {
                // Use auth from config (inline mode)
                if authConfig := c.ResolveForBackend(backendName); authConfig != nil {
                        backend.AuthConfig = authConfig
                        slog.Debug("backend configured with auth strategy from config", "backend", backendName, "strategy", authConfig.Type)
                }
        }
}

// AggregationConfig defines tool aggregation, filtering, and conflict resolution strategies.
//
// Tool Visibility vs Routing:
//   - ExcludeAllTools, DefaultToolVisibility, per-workload ExcludeAll, and Filter control
//     which tools are advertised to MCP clients (visible in tools/list responses).
//   - ALL backend tools remain available in the internal routing table, allowing
//     composite tools to call hidden backend tools.
//   - This enables curated experiences where raw backend tools are hidden from
//     MCP clients but accessible through composite tool workflows.
//
// +kubebuilder:object:generate=true
// +gendoc
type AggregationConfig struct {
        // ConflictResolution defines the strategy for resolving tool name conflicts.
        // - prefix: Automatically prefix tool names with workload identifier
        // - priority: Highest-priority workload wins only if all conflicting workloads
        //   are listed; otherwise all candidates are dropped (see PriorityOrder)
        // - manual: Explicitly define overrides for all conflicts
        // +kubebuilder:validation:Enum=prefix;priority;manual
        // +kubebuilder:default=prefix
        // +optional
        ConflictResolution vmcp.ConflictResolutionStrategy `json:"conflictResolution" yaml:"conflictResolution"`

        // ConflictResolutionConfig provides configuration for the chosen strategy.
        // +optional
        ConflictResolutionConfig *ConflictResolutionConfig `json:"conflictResolutionConfig,omitempty" yaml:"conflictResolutionConfig,omitempty"` //nolint:lll

        // Tools defines per-workload tool filtering and overrides.
        // +optional
        Tools []*WorkloadToolConfig `json:"tools,omitempty" yaml:"tools,omitempty"`

        // ExcludeAllTools hides all backend tools from MCP clients when true.
        // Hidden tools are NOT advertised in tools/list responses, but they ARE
        // available in the routing table for composite tools to use.
        // This enables the use case where you want to hide raw backend tools from
        // direct client access while exposing curated composite tool workflows.
        // +optional
        ExcludeAllTools bool `json:"excludeAllTools,omitempty" yaml:"excludeAllTools,omitempty"`

        // NOTE (maintainers, not rendered into the CRD reference): there is
        // deliberately no kubebuilder default on DefaultToolVisibility. "" already
        // behaves as allow everywhere that reads this field, so defaulting would change
        // only the serialized bytes — and apiextensions applies structural defaults on
        // decode, so every existing VirtualMCPServer would come back with the field set,
        // changing config.yaml, its ConfigMap checksum, and the pod template that stamps
        // it. That restarts every vMCP deployment once for a no-op field.

        // DefaultToolVisibility controls whether a backend with NO entry in Tools has its
        // tools advertised to MCP clients.
        //   - allow (default): every tool from an unlisted backend is advertised, so
        //     adding a workload to the group exposes it without further configuration.
        //   - deny: an unlisted backend contributes no tools, so only backends named in
        //     Tools are advertised. Use this when the set of exposed tools must be
        //     enumerated deliberately rather than inherited from group membership.
        //
        // A backend that DOES have a Tools entry is unaffected by this setting: the
        // entry opts it in, and ExcludeAll/Filter on that entry decide the rest. Like
        // every other visibility setting here, this controls advertising only — hidden
        // tools remain in the routing table for composite tools (see the type doc).
        //
        // This gates TOOLS only. An unlisted backend's resources, resource templates,
        // and prompts are still advertised under deny; mergeResources/mergePrompts have
        // no equivalent check.
        // +kubebuilder:validation:Enum=allow;deny
        // +optional
        DefaultToolVisibility DefaultToolVisibility `json:"defaultToolVisibility,omitempty" yaml:"defaultToolVisibility,omitempty"`
}

// DefaultToolVisibility names the advertising default applied to backends with no
// per-workload Tools entry. The zero value is "" (unset), which behaves as
// DefaultToolVisibilityAllow so existing configs keep today's behavior.
// +gendoc
type DefaultToolVisibility string

const (
        // DefaultToolVisibilityAllow advertises every tool from a backend with no Tools
        // entry. This is the default and matches pre-DefaultToolVisibility behavior.
        DefaultToolVisibilityAllow DefaultToolVisibility = "allow"

        // DefaultToolVisibilityDeny advertises no tools from a backend with no Tools entry.
        DefaultToolVisibilityDeny DefaultToolVisibility = "deny"
)

// ConflictResolutionConfig provides configuration for conflict resolution strategies.
// +kubebuilder:object:generate=true
// +gendoc
type ConflictResolutionConfig struct {
        // PrefixFormat defines the prefix format for the "prefix" tool strategy
        // and for backend-prefixed prompt names (the default for every prompt;
        // under the "priority" strategy, backends listed in priorityOrder keep
        // their own prompt names).
        // Supports placeholders: {workload}, {workload}_, {workload}.
        // +kubebuilder:default="{workload}_"
        // +optional
        PrefixFormat string `json:"prefixFormat,omitempty" yaml:"prefixFormat,omitempty"`

        // PriorityOrder defines the workload priority order for the "priority" strategy
        // (first has highest priority). It is not an allowlist: unique tools remain
        // unchanged, including those from unlisted workloads. If all workloads in a
        // tool-name collision are listed, the highest-priority one wins and losers
        // are dropped. If any conflicting workload is unlisted, all candidates in that
        // collision are dropped from advertising and routing, with no prefix fallback,
        // and the collision is logged at ERROR.
        // Separately, listed workloads keep their own prompt names (unlisted workloads'
        // prompts stay backend-prefixed).
        // +optional
        PriorityOrder []string `json:"priorityOrder,omitempty" yaml:"priorityOrder,omitempty"`
}

// WorkloadToolConfig defines tool filtering and overrides for a specific workload.
// +kubebuilder:object:generate=true
// +gendoc
type WorkloadToolConfig struct {
        // Workload is the name of the backend MCPServer workload.
        // +kubebuilder:validation:Required
        Workload string `json:"workload" yaml:"workload"`

        // ToolConfigRef references an MCPToolConfig resource for tool filtering and renaming.
        // If specified, Filter and Overrides are ignored.
        // Only used when running in Kubernetes with the operator.
        // +optional
        ToolConfigRef *ToolConfigRef `json:"toolConfigRef,omitempty" yaml:"toolConfigRef,omitempty"`

        // Filter is an allow-list of tool names to advertise to MCP clients.
        // Tools NOT in this list are hidden from clients (not in tools/list response)
        // but remain available in the routing table for composite tools to use.
        // This enables selective exposure of backend tools while allowing composite
        // workflows to orchestrate all backend capabilities.
        // Only used if ToolConfigRef is not specified.
        // +optional
        Filter []string `json:"filter,omitempty" yaml:"filter,omitempty"`

        // Overrides is an inline map of tool overrides for renaming and description changes.
        // Overrides are applied to tools before conflict resolution and affect both
        // advertising and routing (the overridden name is used everywhere).
        // Only used if ToolConfigRef is not specified.
        // +optional
        Overrides map[string]*ToolOverride `json:"overrides,omitempty" yaml:"overrides,omitempty"`

        // ExcludeAll hides all tools from this workload from MCP clients when true.
        // Hidden tools are NOT advertised in tools/list responses, but they ARE
        // available in the routing table for composite tools to use.
        // This enables the use case where you want to hide raw backend tools from
        // direct client access while exposing curated composite tool workflows.
        // +optional
        ExcludeAll bool `json:"excludeAll,omitempty" yaml:"excludeAll,omitempty"`
}

// ToolConfigRef references an MCPToolConfig resource for tool filtering and renaming.
// Only used when running in Kubernetes with the operator.
// +kubebuilder:object:generate=true
// +gendoc
type ToolConfigRef struct {
        // Name is the name of the MCPToolConfig resource in the same namespace.
        // +kubebuilder:validation:Required
        Name string `json:"name" yaml:"name"`
}

// ToolAnnotationsOverride defines overrides for tool annotation fields.
// All fields use pointers so nil means "don't override" while zero values
// (empty string, false) mean "explicitly set to this value."
// +kubebuilder:object:generate=true
// +gendoc
type ToolAnnotationsOverride struct {
        // Title sets the human-readable title annotation.
        // For tool overrides this replaces the backend value; for composite tools
        // it is an explicit declaration merged over the derived safety floor.
        // +optional
        Title *string `json:"title,omitempty" yaml:"title,omitempty"`

        // ReadOnlyHint sets the read-only hint annotation.
        // For tool overrides this replaces the backend value; for composite tools
        // it is an explicit declaration merged over the derived safety floor.
        // +optional
        ReadOnlyHint *bool `json:"readOnlyHint,omitempty" yaml:"readOnlyHint,omitempty"`

        // DestructiveHint sets the destructive hint annotation.
        // For tool overrides this replaces the backend value; for composite tools
        // it is an explicit declaration merged over the derived safety floor.
        // +optional
        DestructiveHint *bool `json:"destructiveHint,omitempty" yaml:"destructiveHint,omitempty"`

        // IdempotentHint sets the idempotent hint annotation.
        // For tool overrides this replaces the backend value; for composite tools
        // it is an explicit declaration merged over the derived safety floor.
        // +optional
        IdempotentHint *bool `json:"idempotentHint,omitempty" yaml:"idempotentHint,omitempty"`

        // OpenWorldHint sets the open-world hint annotation.
        // For tool overrides this replaces the backend value; for composite tools
        // it is an explicit declaration merged over the derived safety floor.
        // +optional
        OpenWorldHint *bool `json:"openWorldHint,omitempty" yaml:"openWorldHint,omitempty"`
}

// ToolOverride defines tool name, description, and annotation overrides.
// +kubebuilder:object:generate=true
// +gendoc
type ToolOverride struct {
        // Name is the new tool name (for renaming).
        // +optional
        Name string `json:"name,omitempty" yaml:"name,omitempty"`

        // Description is the new tool description.
        // +optional
        Description string `json:"description,omitempty" yaml:"description,omitempty"`

        // Annotations overrides specific tool annotation fields.
        // Only specified fields are overridden; others pass through from the backend.
        // +optional
        Annotations *ToolAnnotationsOverride `json:"annotations,omitempty" yaml:"annotations,omitempty"`
}

// OperationalConfig contains operational settings.
// OperationalConfig defines operational settings like timeouts and health checks.
// +kubebuilder:object:generate=true
// +gendoc
type OperationalConfig struct {
        // LogLevel sets the logging level for the Virtual MCP server.
        // The only valid value is "debug" to enable debug logging.
        // When omitted or empty, the server uses info level logging.
        // +kubebuilder:validation:Enum=debug
        // +optional
        LogLevel string `json:"logLevel,omitempty" yaml:"logLevel,omitempty"`

        // Timeouts configures timeout settings.
        // +optional
        Timeouts *TimeoutConfig `json:"timeouts,omitempty" yaml:"timeouts,omitempty"`

        // FailureHandling configures failure handling behavior.
        // +optional
        FailureHandling *FailureHandlingConfig `json:"failureHandling,omitempty" yaml:"failureHandling,omitempty"`

        // ListChanged configures live list_changed propagation from backends.
        // +optional
        ListChanged *ListChangedConfig `json:"listChanged,omitempty" yaml:"listChanged,omitempty"`
}

// ListChangedConfig configures which backends vMCP subscribes to for
// list_changed notifications.
//
// Subscribing opens a standalone notification stream to the backend during
// session initialization. A backend that accepts that subscribe and then never
// services it stalls the handshake until the init deadline, and clients with
// their own connect timeout give up first. Excluding such a backend costs it
// live propagation only: its tools are still aggregated and callable, and they
// refresh on the next session.
// +kubebuilder:object:generate=true
// +gendoc
type ListChangedConfig struct {
        // Enabled turns propagation on or off for every backend. Defaults to true.
        // +optional
        Enabled *bool `json:"enabled,omitempty" yaml:"enabled,omitempty"`

        // DisabledWorkloads names backends to exclude while leaving the rest
        // subscribed. Prefer this over Enabled when a single backend misbehaves.
        // +optional
        // +listType=set
        DisabledWorkloads []string `json:"disabledWorkloads,omitempty" yaml:"disabledWorkloads,omitempty"`
}

// TimeoutConfig configures timeout settings.
// +kubebuilder:object:generate=true
// +gendoc
type TimeoutConfig struct {
        // Default is the default timeout for backend requests.
        // +kubebuilder:default="30s"
        // +optional
        Default Duration `json:"default,omitempty" yaml:"default,omitempty"`

        // PerWorkload defines per-workload timeout overrides.
        // +optional
        PerWorkload map[string]Duration `json:"perWorkload,omitempty" yaml:"perWorkload,omitempty"`

        // BackendInit caps how long session initialization waits for a single
        // backend to connect and hand back its capabilities. Unlike Default, which
        // only ever extends that deadline, an explicit BackendInit is authoritative:
        // it bounds session init even when a workload's request timeout is longer.
        //
        // Leave it unset unless a backend can stall the handshake. It exists for
        // backends that neither answer nor fail promptly: a stateless MCP server
        // that requires per-user auth cannot be classified by the unauthenticated
        // health probe, so vMCP opens the persistent connection its Modern skip
        // path would otherwise avoid, and that connection can hang until the
        // deadline. Clients with their own connect timeout give up first, and
        // because they never complete a call, the revision cache never warms and
        // the next session repeats it. A few seconds here lets session init fail
        // fast, which partialFailureMode: best_effort turns into a usable session.
        // +optional
        BackendInit Duration `json:"backendInit,omitempty" yaml:"backendInit,omitempty"`
}

// FailureHandlingConfig configures failure handling behavior.
// +kubebuilder:object:generate=true
// +gendoc
type FailureHandlingConfig struct {
        // HealthCheckInterval is the interval between health checks.
        // +kubebuilder:default="30s"
        // +optional
        HealthCheckInterval Duration `json:"healthCheckInterval,omitempty" yaml:"healthCheckInterval,omitempty"`

        // UnhealthyThreshold is the number of consecutive failures before marking unhealthy.
        // +kubebuilder:default=3
        // +optional
        UnhealthyThreshold int `json:"unhealthyThreshold,omitempty" yaml:"unhealthyThreshold,omitempty"`

        // HealthCheckTimeout is the maximum duration for a single health check operation.
        // Should be less than HealthCheckInterval to prevent checks from queuing up.
        // +kubebuilder:default="10s"
        // +optional
        HealthCheckTimeout Duration `json:"healthCheckTimeout,omitempty" yaml:"healthCheckTimeout,omitempty"`

        // StatusReportingInterval is the interval for reporting status updates to Kubernetes.
        // This controls how often the vMCP runtime reports backend health and phase changes.
        // Lower values provide faster status updates but increase API server load.
        // +kubebuilder:default="30s"
        // +optional
        StatusReportingInterval Duration `json:"statusReportingInterval,omitempty" yaml:"statusReportingInterval,omitempty"`

        // PartialFailureMode defines behavior when some backends are unavailable.
        // - fail: Fail entire request if any backend is unavailable
        // - best_effort: Continue with available backends
        // +kubebuilder:validation:Enum=fail;best_effort
        // +kubebuilder:default=fail
        // +optional
        PartialFailureMode string `json:"partialFailureMode,omitempty" yaml:"partialFailureMode,omitempty"`

        // CircuitBreaker configures circuit breaker behavior.
        // +optional
        CircuitBreaker *CircuitBreakerConfig `json:"circuitBreaker,omitempty" yaml:"circuitBreaker,omitempty"`
}

// CircuitBreakerConfig configures circuit breaker behavior.
// +kubebuilder:object:generate=true
// +gendoc
type CircuitBreakerConfig struct {
        // Enabled controls whether circuit breaker is enabled.
        // +kubebuilder:default=false
        // +optional
        Enabled bool `json:"enabled,omitempty" yaml:"enabled,omitempty"`

        // FailureThreshold is the number of failures before opening the circuit.
        // Must be >= 1.
        // +kubebuilder:default=5
        // +kubebuilder:validation:Minimum=1
        // +optional
        FailureThreshold int `json:"failureThreshold,omitempty" yaml:"failureThreshold,omitempty"`

        // Timeout is the duration to wait before attempting to close the circuit.
        // Must be >= 1s to prevent thrashing.
        // +kubebuilder:default="60s"
        // +kubebuilder:validation:XValidation:rule="self == '' || duration(self) >= duration('1s')",message="timeout must be >= 1s"
        // +optional
        Timeout Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"`
}

// CompositeToolConfig defines a composite tool workflow.
// This matches the YAML structure from the proposal (lines 173-255).
// +kubebuilder:object:generate=true
// +gendoc
type CompositeToolConfig struct {
        // Name is the workflow name (unique identifier).
        Name string `json:"name" yaml:"name"`

        // Description describes what the workflow does.
        Description string `json:"description,omitempty" yaml:"description,omitempty"`

        // Parameters defines input parameter schema in JSON Schema format.
        // Should be a JSON Schema object with "type": "object" and "properties".
        // Example:
        //   {
        //     "type": "object",
        //     "properties": {
        //       "param1": {"type": "string", "default": "value"},
        //       "param2": {"type": "integer"}
        //     },
        //     "required": ["param2"]
        //   }
        //
        // We use json.Map rather than a typed struct because JSON Schema is highly
        // flexible with many optional fields (default, enum, minimum, maximum, pattern,
        // items, additionalProperties, oneOf, anyOf, allOf, etc.). Using json.Map
        // allows full JSON Schema compatibility without needing to define every possible
        // field, and matches how the MCP SDK handles inputSchema.
        // +optional
        Parameters thvjson.Map `json:"parameters,omitempty" yaml:"parameters,omitempty"`

        // Timeout is the maximum workflow execution time.
        Timeout Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"`

        // Steps are the workflow steps to execute.
        Steps []WorkflowStepConfig `json:"steps" yaml:"steps"`

        // Output defines the structured output schema for this workflow.
        // If not specified, the workflow returns the last step's output (backward compatible).
        // +optional
        Output *OutputConfig `json:"output,omitempty" yaml:"output,omitempty"`

        // Annotations declares MCP tool annotations for the composite tool.
        //
        // Annotation derivation runs at ADVERTISE TIME (when tools/list is served
        // and the backend tools are aggregated), not at CRD admission — thv vmcp
        // validate does NOT check annotation contradictions. The derived floor is
        // fail-closed: when the workflow has one or more tool steps the floor is
        // always non-nil, and any step whose annotations are nil/unknown taints the
        // floor conservatively (readOnly=false, destructive=true, openWorld=true).
        // A workflow with no tool steps (e.g. only elicitation) has no floor.
        //
        // When nil, annotations are derived from the annotations of the backend tools
        // referenced by the workflow's steps (e.g. readOnlyHint is true only when every
        // step tool is read-only). When set, the values are an explicit author
        // declaration merged over the derived floor — subject to a safety-floor
        // guardrail that drops the composite tool (with a warning naming the offending
        // step tools) if an explicit hint would make the tool look safer than its
        // steps allow.
        // +optional
        Annotations *ToolAnnotationsOverride `json:"annotations,omitempty" yaml:"annotations,omitempty"`
}

// CompositeToolRef defines a reference to a VirtualMCPCompositeToolDefinition resource.
// The referenced resource must be in the same namespace as the VirtualMCPServer.
// +kubebuilder:object:generate=true
// +gendoc
type CompositeToolRef struct {
        // Name is the name of the VirtualMCPCompositeToolDefinition resource in the same namespace.
        // +kubebuilder:validation:Required
        Name string `json:"name" yaml:"name"`
}

// WorkflowStepConfig defines a single workflow step.
// This matches the proposal's step configuration (lines 180-255).
// +kubebuilder:object:generate=true
// +gendoc
type WorkflowStepConfig struct {
        // ID is the unique identifier for this step.
        // +kubebuilder:validation:Required
        ID string `json:"id" yaml:"id"`

        // Type is the step type (tool, elicitation, etc.)
        // +kubebuilder:validation:Enum=tool;elicitation;forEach
        // +kubebuilder:default=tool
        // +optional
        Type string `json:"type,omitempty" yaml:"type,omitempty"`

        // Tool is the tool to call (format: "workload.tool_name")
        // Only used when Type is "tool"
        // +optional
        Tool string `json:"tool,omitempty" yaml:"tool,omitempty"`

        // Arguments is a map of argument values with template expansion support.
        // Supports Go template syntax with .params and .steps for string values.
        // Non-string values (integers, booleans, arrays, objects) are passed as-is.
        // Note: the templating is only supported on the first level of the key-value pairs.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        Arguments thvjson.Map `json:"arguments,omitempty" yaml:"arguments,omitempty"`

        // Condition is a template expression that determines if the step should execute
        // +optional
        Condition string `json:"condition,omitempty" yaml:"condition,omitempty"`

        // DependsOn lists step IDs that must complete before this step
        // +optional
        DependsOn []string `json:"dependsOn,omitempty" yaml:"dependsOn,omitempty"`

        // OnError defines error handling behavior
        // +optional
        OnError *StepErrorHandling `json:"onError,omitempty" yaml:"onError,omitempty"`

        // Message is the elicitation message
        // Only used when Type is "elicitation"
        // +optional
        Message string `json:"message,omitempty" yaml:"message,omitempty"`

        // Schema defines the expected response schema for elicitation
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        Schema thvjson.Map `json:"schema,omitempty" yaml:"schema,omitempty"`

        // Timeout is the maximum execution time for this step
        // +optional
        Timeout Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"`

        // OnDecline defines the action to take when the user explicitly declines the elicitation
        // Only used when Type is "elicitation"
        // +optional
        OnDecline *ElicitationResponseConfig `json:"onDecline,omitempty" yaml:"onDecline,omitempty"`

        // OnCancel defines the action to take when the user cancels/dismisses the elicitation
        // Only used when Type is "elicitation"
        // +optional
        OnCancel *ElicitationResponseConfig `json:"onCancel,omitempty" yaml:"onCancel,omitempty"`

        // DefaultResults provides fallback output values when this step is skipped
        // (due to condition evaluating to false) or fails (when onError.action is "continue").
        // Each key corresponds to an output field name referenced by downstream steps.
        // Required if the step may be skipped AND downstream steps reference this step's output.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Schemaless
        DefaultResults thvjson.Map `json:"defaultResults,omitempty" yaml:"defaultResults,omitempty"`

        // Collection is a Go template expression that resolves to a JSON array or a slice.
        // Only used when Type is "forEach".
        // +optional
        Collection string `json:"collection,omitempty" yaml:"collection,omitempty"`

        // ItemVar is the variable name used to reference the current item in forEach templates.
        // Defaults to "item" if not specified.
        // Only used when Type is "forEach".
        // +optional
        ItemVar string `json:"itemVar,omitempty" yaml:"itemVar,omitempty"`

        // MaxParallel limits the number of concurrent iterations in a forEach step.
        // Defaults to the DAG executor's maxParallel (10).
        // Only used when Type is "forEach".
        // +optional
        MaxParallel int `json:"maxParallel,omitempty" yaml:"maxParallel,omitempty"`

        // MaxIterations limits the number of items that can be iterated over.
        // Defaults to 100, hard cap at 1000.
        // Only used when Type is "forEach".
        // +optional
        MaxIterations int `json:"maxIterations,omitempty" yaml:"maxIterations,omitempty"`

        // InnerStep defines the step to execute for each item in the collection.
        // Only used when Type is "forEach". Only tool-type inner steps are supported.
        // +optional
        // +kubebuilder:validation:Type=object
        // +kubebuilder:pruning:PreserveUnknownFields
        InnerStep *WorkflowStepConfig `json:"step,omitempty" yaml:"step,omitempty"`
}

// StepErrorHandling defines error handling behavior for workflow steps.
// +kubebuilder:object:generate=true
// +gendoc
type StepErrorHandling struct {
        // Action defines the action to take on error
        // +kubebuilder:validation:Enum=abort;continue;retry
        // +kubebuilder:default=abort
        // +optional
        Action string `json:"action,omitempty" yaml:"action,omitempty"`

        // RetryCount is the maximum number of retries
        // Only used when Action is "retry"
        // +optional
        RetryCount int `json:"retryCount,omitempty" yaml:"retryCount,omitempty"`

        // RetryDelay is the delay between retry attempts
        // Only used when Action is "retry"
        // +optional
        RetryDelay Duration `json:"retryDelay,omitempty" yaml:"retryDelay,omitempty"`
}

// ElicitationResponseConfig defines how to handle user responses to elicitation requests.
// +kubebuilder:object:generate=true
// +gendoc
type ElicitationResponseConfig struct {
        // Action defines the action to take when the user declines or cancels
        // - skip_remaining: Skip remaining steps in the workflow
        // - abort: Abort the entire workflow execution
        // - continue: Continue to the next step
        // +kubebuilder:validation:Enum=skip_remaining;abort;continue
        // +kubebuilder:default=abort
        // +optional
        Action string `json:"action,omitempty" yaml:"action,omitempty"`
}

// OutputConfig defines the structured output schema for a composite tool workflow.
// This follows the same pattern as the Parameters field, defining both the
// MCP output schema (type, description) and runtime value construction (value, default).
// +kubebuilder:object:generate=true
// +gendoc
type OutputConfig struct {
        // Properties defines the output properties.
        // Map key is the property name, value is the property definition.
        Properties map[string]OutputProperty `json:"properties" yaml:"properties"`

        // Required lists property names that must be present in the output.
        // +optional
        Required []string `json:"required,omitempty" yaml:"required,omitempty"`
}

// OutputProperty defines a single output property.
// For non-object types, Value is required.
// For object types, either Value or Properties must be specified (but not both).
// +kubebuilder:object:generate=true
// +gendoc
type OutputProperty struct {
        // Type is the JSON Schema type: "string", "integer", "number", "boolean", "object", "array"
        // +kubebuilder:validation:Required
        // +kubebuilder:validation:Enum=string;integer;number;boolean;object;array
        Type string `json:"type" yaml:"type"`

        // Description is a human-readable description exposed to clients and models
        // +optional
        Description string `json:"description" yaml:"description"`

        // Value is a template string for constructing the runtime value.
        // For object types, this can be a JSON string that will be deserialized.
        // Supports template syntax: {{.steps.step_id.output.field}}, {{.params.param_name}}
        // +optional
        Value string `json:"value,omitempty" yaml:"value,omitempty"`

        // Properties defines nested properties for object types.
        // Each nested property has full metadata (type, description, value/properties).
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Type=object
        // +kubebuilder:validation:Schemaless
        Properties map[string]OutputProperty `json:"properties,omitempty" yaml:"properties,omitempty"`

        // Default is the fallback value if template expansion fails.
        // Type coercion is applied to match the declared Type.
        // +optional
        // +kubebuilder:pruning:PreserveUnknownFields
        // +kubebuilder:validation:Schemaless
        Default thvjson.Any `json:"default,omitempty" yaml:"default,omitempty"`
}

// OptimizerConfig configures the MCP optimizer.
// When enabled, vMCP exposes only find_tool and call_tool operations to clients
// instead of all backend tools directly.
//
// +kubebuilder:object:generate=true
// +kubebuilder:validation:XValidation:rule="!has(self.embeddingHeaders) || (has(self.embeddingProvider) && self.embeddingProvider == 'openai')",message="embeddingHeaders is only supported when embeddingProvider is 'openai'"
// +kubebuilder:validation:XValidation:rule=`!has(self.embeddingHeaders) || self.embeddingHeaders.all(k, k.matches('^[!#$%&\\x27*+.^_\\x60|~0-9A-Za-z-]+$') && !(k.lowerAscii() in ['authorization', 'content-type']))`,message="embeddingHeaders names must be valid HTTP header names and must not include Authorization or Content-Type"
// +gendoc
//
//nolint:lll // CEL validation rules exceed line length limit
type OptimizerConfig struct {
        // EmbeddingService is the full base URL of the embedding service endpoint
        // (e.g., http://my-embedding.default.svc.cluster.local:8080) for semantic
        // tool discovery.
        //
        // In a Kubernetes environment, it is more convenient to use the
        // VirtualMCPServerSpec.EmbeddingServerRef field instead of setting this
        // directly. EmbeddingServerRef references an EmbeddingServer CRD by name,
        // and the operator automatically resolves the referenced resource's
        // Status.URL to populate this field. This provides managed lifecycle
        // (the operator watches the EmbeddingServer for readiness and URL changes)
        // and avoids hardcoding service URLs in the config. If both
        // EmbeddingServerRef and this field are set, EmbeddingServerRef takes
        // precedence and this value is overridden with a warning.
        // +optional
        EmbeddingService string `json:"embeddingService,omitempty" yaml:"embeddingService,omitempty"`

        // EmbeddingServiceTimeout is the HTTP request timeout for calls to the embedding service.
        // Defaults to 30s if not specified.
        // +kubebuilder:default="30s"
        // +optional
        EmbeddingServiceTimeout Duration `json:"embeddingServiceTimeout,omitempty" yaml:"embeddingServiceTimeout,omitempty"`

        // EmbeddingProvider selects the wire protocol used to talk to the embedding
        // service. "tei" speaks the HuggingFace Text Embeddings Inference API;
        // "openai" speaks the OpenAI-compatible /embeddings API, which lets the
        // optimizer use OpenAI, Azure OpenAI, or another OpenAI-compatible gateway.
        // Defaults to "tei" when empty.
        //
        // The "openai" provider reads EmbeddingService directly and cannot be combined
        // with EmbeddingServerRef, which provisions a managed TEI server; the operator
        // rejects that combination at admission.
        // +kubebuilder:validation:Enum=tei;openai
        // +kubebuilder:default="tei"
        // +optional
        EmbeddingProvider string `json:"embeddingProvider,omitempty" yaml:"embeddingProvider,omitempty"`

        // EmbeddingModel is the model name requested from the embedding service
        // (e.g. "text-embedding-3-small"). Required when EmbeddingProvider is
        // "openai". Ignored for the "tei" provider, where the model is fixed by the
        // running TEI container.
        //
        // The API key for an OpenAI-compatible service is not configured here: it is
        // read from the OPENAI_API_KEY environment variable so the secret never
        // lands in a CRD spec or ConfigMap. An empty key omits the Authorization
        // header, which supports keyless in-cluster gateways.
        // +optional
        EmbeddingModel string `json:"embeddingModel,omitempty" yaml:"embeddingModel,omitempty"`

        // EmbeddingHeaders holds additional HTTP headers sent with every embedding
        // request. Only supported when EmbeddingProvider is "openai". Values are
        // stored in plain text and must not contain secrets; Authorization
        // (derived from OPENAI_API_KEY) and Content-Type cannot be set.
        // +kubebuilder:validation:MaxProperties=32
        // +optional
        EmbeddingHeaders map[string]EmbeddingHeaderValue `json:"embeddingHeaders,omitempty" yaml:"embeddingHeaders,omitempty"`

        // MaxToolsToReturn is the maximum number of tool results returned by a search query.
        // Defaults to 8 if not specified or zero.
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:validation:Maximum=50
        // +optional
        MaxToolsToReturn int `json:"maxToolsToReturn,omitempty" yaml:"maxToolsToReturn,omitempty"`

        // HybridSearchSemanticRatio controls the balance between semantic (meaning-based)
        // and keyword search results. 0.0 = all keyword, 1.0 = all semantic.
        // Defaults to "0.5" if not specified or empty.
        // Serialized as a string because CRDs do not support float types portably.
        // +kubebuilder:validation:Pattern=`^([0-9]*[.])?[0-9]+$`
        // +optional
        HybridSearchSemanticRatio string `json:"hybridSearchSemanticRatio,omitempty" yaml:"hybridSearchSemanticRatio,omitempty"`

        // SemanticDistanceThreshold is the maximum distance for semantic search results.
        // Results exceeding this threshold are filtered out from semantic search.
        // This threshold does not apply to keyword search.
        // Range: 0 = identical, 2 = completely unrelated.
        // Defaults to "1.0" if not specified or empty.
        // Serialized as a string because CRDs do not support float types portably.
        // +kubebuilder:validation:Pattern=`^([0-9]*[.])?[0-9]+$`
        // +optional
        SemanticDistanceThreshold string `json:"semanticDistanceThreshold,omitempty" yaml:"semanticDistanceThreshold,omitempty"`
}

// EmbeddingHeaderValue is a custom embedding request header value: 1 to 8192
// characters with no control characters other than tab.
// +kubebuilder:validation:MinLength=1
// +kubebuilder:validation:MaxLength=8192
// +kubebuilder:validation:Pattern=`^[^\x00-\x08\x0A-\x1F\x7F]*$`
type EmbeddingHeaderValue string

// CodeModeConfig configures vMCP code mode (the execute_tool_script virtual tool).
// When enabled, agents can submit a Starlark script that calls multiple backend tools
// server-side — with loops, conditionals, and parallel() fan-out — and receive a single
// aggregated result, collapsing many tool-call round-trips into one.
// +kubebuilder:object:generate=true
// +gendoc
type CodeModeConfig struct {
        // Enabled turns code mode on. When false (the default), execute_tool_script is not
        // advertised in tools/list and scripts cannot be executed.
        // +optional
        Enabled bool `json:"enabled,omitempty" yaml:"enabled,omitempty"`

        // StepLimit is the maximum number of Starlark execution steps per script. It bounds
        // runaway loops and computation. Defaults to 100000 if unset or zero.
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:default=100000
        // +optional
        StepLimit int64 `json:"stepLimit,omitempty" yaml:"stepLimit,omitempty"`

        // ParallelMaxConcurrency caps the number of goroutines a script's parallel() builtin
        // may run concurrently. Defaults to 10 if unset or zero.
        // +kubebuilder:validation:Minimum=1
        // +kubebuilder:default=10
        // +optional
        ParallelMaxConcurrency int `json:"parallelMaxConcurrency,omitempty" yaml:"parallelMaxConcurrency,omitempty"`

        // ToolCallTimeout bounds each individual backend tool call made from within a script.
        // A call exceeding it is cancelled and surfaces a timeout error to the script.
        // Defaults to 30s if unset.
        // +kubebuilder:default="30s"
        // +optional
        ToolCallTimeout Duration `json:"toolCallTimeout,omitempty" yaml:"toolCallTimeout,omitempty"`
}

// SessionStorageConfig configures session storage for stateful horizontal scaling.
// The Redis password is not stored here; it is injected as the THV_SESSION_REDIS_PASSWORD
// environment variable by the operator when spec.sessionStorage.passwordRef is set.
// +kubebuilder:object:generate=true
// +gendoc
type SessionStorageConfig struct {
        // Provider is the session storage backend type.
        // +kubebuilder:validation:Enum=memory;redis
        // +kubebuilder:validation:Required
        Provider string `json:"provider" yaml:"provider"`

        // Address is the Redis server address (required when provider is redis).
        // +optional
        Address string `json:"address,omitempty" yaml:"address,omitempty"`

        // DB is the Redis database number.
        // +kubebuilder:validation:Minimum=0
        // +kubebuilder:default=0
        // +optional
        DB int32 `json:"db,omitempty" yaml:"db,omitempty"`

        // KeyPrefix is an optional prefix for all Redis keys used by ToolHive.
        // +optional
        KeyPrefix string `json:"keyPrefix,omitempty" yaml:"keyPrefix,omitempty"`

        // TLS enables TLS for session storage and rate limiting. An empty
        // configuration verifies the Redis server using system roots.
        // On Kubernetes, set spec.sessionStorage.tls instead: the operator derives this
        // field from it and rejects a value set under spec.config.sessionStorage.
        // +optional
        TLS *redisconfig.TLSConfig `json:"tls,omitempty" yaml:"tls,omitempty"`
}

// Validator validates configuration.
type Validator interface {
        // Validate checks if the configuration is valid.
        // Returns detailed validation errors.
        Validate(cfg *Config) error
}

// Loader loads configuration from a source.
type Loader interface {
        // Load loads configuration from the source.
        Load() (*Config, error)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package config provides the configuration model for Virtual MCP Server.
package config

import (
        "fmt"
        "time"

        "dario.cat/mergo"

        "github.com/stacklok/toolhive/pkg/authserver"
        authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
)

// Default constants for operational configuration.
// These values match the kubebuilder defaults in the VirtualMCPServer CRD.
const (
        // defaultHealthCheckInterval is the default interval between health checks.
        defaultHealthCheckInterval = 30 * time.Second

        // defaultUnhealthyThreshold is the default number of consecutive failures
        // before marking a backend as unhealthy.
        defaultUnhealthyThreshold = 3

        // defaultStatusReportingInterval is the default interval for reporting status updates.
        defaultStatusReportingInterval = 30 * time.Second

        // defaultPartialFailureMode defines the default behavior when some backends fail.
        // "fail" means the entire request fails if any backend is unavailable.
        defaultPartialFailureMode = "fail"

        // defaultTimeoutDefault is the default timeout for backend requests.
        defaultTimeoutDefault = 30 * time.Second

        // defaultCircuitBreakerFailureThreshold is the default number of failures
        // before opening the circuit breaker.
        defaultCircuitBreakerFailureThreshold = 5

        // defaultCircuitBreakerTimeout is the default duration to wait before
        // attempting to close the circuit.
        defaultCircuitBreakerTimeout = 60 * time.Second

        // defaultCircuitBreakerEnabled is the default state of the circuit breaker.
        defaultCircuitBreakerEnabled = false
)

// DefaultOperationalConfig returns a fully populated OperationalConfig with default values.
// This is the SINGLE SOURCE OF TRUTH for all operational defaults.
// This should be used when no operational config is provided.
func DefaultOperationalConfig() *OperationalConfig {
        return &OperationalConfig{
                Timeouts: &TimeoutConfig{
                        Default:     Duration(defaultTimeoutDefault),
                        PerWorkload: nil,
                },
                FailureHandling: &FailureHandlingConfig{
                        HealthCheckInterval:     Duration(defaultHealthCheckInterval),
                        UnhealthyThreshold:      defaultUnhealthyThreshold,
                        StatusReportingInterval: Duration(defaultStatusReportingInterval),
                        PartialFailureMode:      defaultPartialFailureMode,
                        CircuitBreaker: &CircuitBreakerConfig{
                                Enabled:          defaultCircuitBreakerEnabled,
                                FailureThreshold: defaultCircuitBreakerFailureThreshold,
                                Timeout:          Duration(defaultCircuitBreakerTimeout),
                        },
                },
        }
}

// EnsureOperationalDefaults ensures that the Config has a fully populated
// OperationalConfig with default values for any missing fields.
// If Operational is nil, it sets it to DefaultOperationalConfig().
// If Operational exists but has nil or zero-value nested fields, those fields
// are filled with defaults while preserving any user-provided values.
func (c *Config) EnsureOperationalDefaults() {
        if c == nil {
                return
        }

        if c.Operational == nil {
                c.Operational = DefaultOperationalConfig()
                return
        }

        // Merge defaults into target, only filling zero/nil values.
        // User-provided values are preserved.
        _ = mergo.Merge(c.Operational, DefaultOperationalConfig())
}

// InjectSubjectProviderNames auto-populates SubjectProviderName on every
// token_exchange, aws_sts, or xaa strategy in cfg.OutgoingAuth that has it
// unset, when an embedded auth server RunConfig is active.
//
// This is a defaulting operation: it ensures YAML-based vMCP deployments
// behave the same as the Kubernetes operator path. Without it a token_exchange
// strategy with no SubjectProviderName would silently fall back to
// identity.Token (the ToolHive-issued JWT), which the exchange endpoint rejects.
//
// When cfg or rc is nil, or rc has no upstreams, the call is a no-op. The
// provider name is derived via authserver.ResolveFirstUpstreamName over
// rc.Upstreams. For xaa, if more than
// one upstream is configured and SubjectProviderName is empty, defaulting is
// ambiguous and this returns an error wrapping
// authtypes.ErrAmbiguousSubjectProvider instead of silently defaulting;
// processing stops at the first error. Strategy structs are replaced via
// copy-then-assign rather than mutated in place, but cfg.OutgoingAuth itself
// is still updated by this call.
func InjectSubjectProviderNames(cfg *Config, rc *authserver.RunConfig) error {
        if cfg == nil || rc == nil || cfg.OutgoingAuth == nil {
                return nil
        }

        names := make([]string, len(rc.Upstreams))
        for i, u := range rc.Upstreams {
                names[i] = u.Name
        }
        if len(names) == 0 {
                return nil
        }
        providerName := authserver.ResolveFirstUpstreamName(names)
        hasMultipleUpstreams := len(names) > 1

        defaulted, err := authtypes.DefaultSubjectProviderName(cfg.OutgoingAuth.Default, providerName, hasMultipleUpstreams)
        if err != nil {
                return fmt.Errorf("default backend: %w", err)
        }
        cfg.OutgoingAuth.Default = defaulted

        for name, strategy := range cfg.OutgoingAuth.Backends {
                defaulted, err := authtypes.DefaultSubjectProviderName(strategy, providerName, hasMultipleUpstreams)
                if err != nil {
                        return fmt.Errorf("backend %q: %w", name, err)
                }
                cfg.OutgoingAuth.Backends[name] = defaulted
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "fmt"
        "net/http"
        "path/filepath"
        "slices"
        "strings"
        "time"

        httpval "github.com/stacklok/toolhive-core/validation/http"
        "github.com/stacklok/toolhive/pkg/authserver"
        "github.com/stacklok/toolhive/pkg/transport/middleware"
        "github.com/stacklok/toolhive/pkg/vmcp"
        authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
)

// Incoming auth type constants.
const (
        IncomingAuthTypeOIDC      = "oidc"
        IncomingAuthTypeLocal     = "local"
        IncomingAuthTypeAnonymous = "anonymous"
)

// defaultStrategyKey is the synthetic map key used for the default outgoing auth
// strategy in collectAllBackendStrategies. It is deliberately different from
// authserver.DefaultUpstreamName ("default") to avoid confusion with upstream
// provider names and to prevent key collisions with user-defined backend names.
const defaultStrategyKey = "<default-strategy>"

// DefaultValidator implements comprehensive configuration validation.
type DefaultValidator struct{}

// NewValidator creates a new configuration validator.
func NewValidator() *DefaultValidator {
        return &DefaultValidator{}
}

// Validate performs comprehensive validation of the configuration.
func (v *DefaultValidator) Validate(cfg *Config) error {
        if cfg == nil {
                return fmt.Errorf("%w: configuration is nil", vmcp.ErrInvalidConfig)
        }

        var errors []string

        // Validate basic fields
        if err := v.validateBasicFields(cfg); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate incoming authentication
        if err := v.validateIncomingAuth(cfg.IncomingAuth); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate outgoing authentication
        if err := v.validateOutgoingAuth(cfg.OutgoingAuth); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate aggregation configuration
        if err := v.validateAggregation(cfg.Aggregation); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate operational configuration
        if err := v.validateOperational(cfg.Operational); err != nil {
                errors = append(errors, err.Error())
        }

        if cfg.SessionStorage != nil && cfg.SessionStorage.TLS != nil && cfg.SessionStorage.Provider != "redis" {
                errors = append(errors, "session storage TLS requires provider redis")
        }

        // Validate static backends
        if err := v.validateStaticBackends(cfg.Backends); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate composite tools
        if err := v.validateCompositeTools(cfg.CompositeTools); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate composite tool references
        if err := v.validateCompositeToolRefs(cfg.CompositeToolRefs); err != nil {
                errors = append(errors, err.Error())
        }

        // Validate passthrough headers allowlist
        if err := v.validatePassthroughHeaders(cfg); err != nil {
                errors = append(errors, err.Error())
        }

        // Note: Optimizer validation is handled by optimizer.GetAndValidateConfig
        // in pkg/vmcp/optimizer/optimizer.go when the optimizer is constructed.

        if len(errors) > 0 {
                return fmt.Errorf("%w:\n  - %s", vmcp.ErrInvalidConfig, strings.Join(errors, "\n  - "))
        }

        return nil
}

func (*DefaultValidator) validateBasicFields(cfg *Config) error {
        if cfg.Name == "" {
                return fmt.Errorf("name is required")
        }

        if cfg.Group == "" {
                return fmt.Errorf("group reference is required")
        }

        return nil
}

func (*DefaultValidator) validateStaticBackends(backends []StaticBackendConfig) error {
        for i, b := range backends {
                // Validate type if specified
                if b.Type != "" && b.Type != string(vmcp.BackendTypeEntry) {
                        return fmt.Errorf("backends[%d].type must be empty or %q, got %q", i, vmcp.BackendTypeEntry, b.Type)
                }

                // CABundlePath is only valid for entry backends
                if b.CABundlePath != "" && b.Type != string(vmcp.BackendTypeEntry) {
                        return fmt.Errorf("backends[%d].caBundlePath is only valid when type is %q", i, vmcp.BackendTypeEntry)
                }

                // Validate CA bundle path: reject null bytes, path traversal, and relative paths
                if b.CABundlePath != "" {
                        if strings.ContainsRune(b.CABundlePath, 0) || strings.Contains(b.CABundlePath, "..") {
                                return fmt.Errorf("backends[%d].caBundlePath contains invalid path characters", i)
                        }
                        if !filepath.IsAbs(b.CABundlePath) {
                                return fmt.Errorf("backends[%d].caBundlePath must be an absolute path", i)
                        }
                }
        }
        return nil
}

func (v *DefaultValidator) validateIncomingAuth(auth *IncomingAuthConfig) error {
        if auth == nil {
                return fmt.Errorf("incomingAuth is required")
        }

        // Validate auth type
        validTypes := []string{IncomingAuthTypeOIDC, IncomingAuthTypeLocal, IncomingAuthTypeAnonymous}
        if !slices.Contains(validTypes, auth.Type) {
                return fmt.Errorf("incomingAuth.type must be one of: %s", strings.Join(validTypes, ", "))
        }

        // Validate OIDC configuration
        if auth.Type == IncomingAuthTypeOIDC {
                if auth.OIDC == nil {
                        return fmt.Errorf("incomingAuth.oidc is required when type is 'oidc'")
                }

                if auth.OIDC.Issuer == "" {
                        return fmt.Errorf("incomingAuth.oidc.issuer is required")
                }

                if auth.OIDC.Audience == "" {
                        return fmt.Errorf("incomingAuth.oidc.audience is required")
                }

                // ClientID is optional - only required for specific flows:
                // - Token introspection with client credentials
                // - Some OAuth flows requiring client identification
                // Not required for standard JWT validation using JWKS

                // ClientSecretEnv is optional - some OIDC flows don't require client secrets:
                // - PKCE flows (public clients)
                // - Token validation without introspection
                // - Kubernetes service account token validation

                // Validate CA bundle path: reject null bytes, path traversal, and relative paths.
                if auth.OIDC.CABundlePath != "" {
                        if strings.ContainsRune(auth.OIDC.CABundlePath, 0) || strings.Contains(auth.OIDC.CABundlePath, "..") {
                                return fmt.Errorf("incomingAuth.oidc.caBundlePath contains invalid path characters")
                        }
                        if !filepath.IsAbs(auth.OIDC.CABundlePath) {
                                return fmt.Errorf("incomingAuth.oidc.caBundlePath must be an absolute path")
                        }
                }
        }

        // Validate authorization configuration
        if auth.Authz != nil {
                if err := v.validateAuthz(auth.Authz); err != nil {
                        return fmt.Errorf("incomingAuth.authz: %w", err)
                }
        }

        return nil
}

func (*DefaultValidator) validateAuthz(authz *AuthzConfig) error {
        validTypes := []string{"cedar", "none"}
        if !slices.Contains(validTypes, authz.Type) {
                return fmt.Errorf("type must be one of: %s", strings.Join(validTypes, ", "))
        }

        if authz.Type == "cedar" && len(authz.Policies) == 0 {
                return fmt.Errorf("policies are required when type is 'cedar'")
        }

        return nil
}

func (v *DefaultValidator) validateOutgoingAuth(auth *OutgoingAuthConfig) error {
        if auth == nil {
                return fmt.Errorf("outgoingAuth is required")
        }

        // Validate source
        validSources := []string{"inline", "discovered"}
        if !slices.Contains(validSources, auth.Source) {
                return fmt.Errorf("outgoingAuth.source must be one of: %s", strings.Join(validSources, ", "))
        }

        // Validate default strategy
        if auth.Default != nil {
                if err := v.validateBackendAuthStrategy("default", auth.Default); err != nil {
                        return fmt.Errorf("outgoingAuth.default: %w", err)
                }
        }

        // Validate per-backend strategies
        for backendName, strategy := range auth.Backends {
                if err := v.validateBackendAuthStrategy(backendName, strategy); err != nil {
                        return fmt.Errorf("outgoingAuth.backends.%s: %w", backendName, err)
                }
        }

        return nil
}

//nolint:gocyclo // Strategy-specific validation requires checking multiple fields per type
func (*DefaultValidator) validateBackendAuthStrategy(_ string, strategy *authtypes.BackendAuthStrategy) error {
        if strategy == nil {
                return fmt.Errorf("strategy is nil")
        }

        validTypes := []string{
                authtypes.StrategyTypeUnauthenticated,
                authtypes.StrategyTypeHeaderInjection,
                authtypes.StrategyTypeTokenExchange,
                authtypes.StrategyTypeUpstreamInject,
                authtypes.StrategyTypeAwsSts,
                authtypes.StrategyTypeOBO,
                authtypes.StrategyTypeXAA,
        }
        if !slices.Contains(validTypes, strategy.Type) {
                return fmt.Errorf("type must be one of: %s", strings.Join(validTypes, ", "))
        }

        // Validate type-specific requirements
        switch strategy.Type {
        case authtypes.StrategyTypeTokenExchange:
                // Token exchange requires TokenExchange config with tokenUrl
                if strategy.TokenExchange == nil {
                        return fmt.Errorf("tokenExchange requires TokenExchange configuration")
                }
                if strategy.TokenExchange.TokenURL == "" {
                        return fmt.Errorf("tokenExchange requires tokenUrl field")
                }

        case authtypes.StrategyTypeHeaderInjection:
                // Header injection requires HeaderInjection config with header name and value
                if strategy.HeaderInjection == nil {
                        return fmt.Errorf("headerInjection requires HeaderInjection configuration")
                }
                if strategy.HeaderInjection.HeaderName == "" {
                        return fmt.Errorf("headerInjection requires headerName field")
                }
                if strategy.HeaderInjection.HeaderValue == "" {
                        return fmt.Errorf("headerInjection requires headerValue field")
                }

        case authtypes.StrategyTypeUpstreamInject:
                if strategy.UpstreamInject == nil {
                        return fmt.Errorf("upstream_inject requires UpstreamInject configuration")
                }
                // Note: empty ProviderName is allowed here; ValidateAuthServerIntegration
                // handles provider name resolution including the empty→"default" mapping.

        case authtypes.StrategyTypeAwsSts:
                if strategy.AwsSts == nil {
                        return fmt.Errorf("aws_sts requires AwsSts configuration")
                }
                if strategy.AwsSts.Region == "" {
                        return fmt.Errorf("aws_sts requires region field")
                }

        case authtypes.StrategyTypeOBO:
                if strategy.OBO == nil {
                        return fmt.Errorf("obo requires OBO configuration")
                }
                if strategy.OBO.TokenURL == "" {
                        return fmt.Errorf("obo requires tokenUrl field")
                }
                if strategy.OBO.ClientSecret != "" && strategy.OBO.ClientSecretEnv != "" {
                        return fmt.Errorf("obo: clientSecret and clientSecretEnv are mutually exclusive")
                }

        case authtypes.StrategyTypeXAA:
                if strategy.XAA == nil {
                        return fmt.Errorf("xaa requires XAA configuration")
                }
                if strategy.XAA.IDPTokenURL == "" {
                        return fmt.Errorf("xaa requires idpTokenUrl field")
                }
                if strategy.XAA.TargetTokenURL == "" {
                        return fmt.Errorf("xaa requires targetTokenUrl field")
                }
                if strategy.XAA.TargetAudience == "" {
                        return fmt.Errorf("xaa requires targetAudience field")
                }
                if strategy.XAA.SubjectTokenType != "" && strategy.XAA.SubjectTokenType != "urn:ietf:params:oauth:token-type:id_token" {
                        return fmt.Errorf("xaa: unsupported subjectTokenType %q; only %q is accepted",
                                strategy.XAA.SubjectTokenType, "urn:ietf:params:oauth:token-type:id_token")
                }
        }

        return nil
}

func (v *DefaultValidator) validateAggregation(agg *AggregationConfig) error {
        if agg == nil {
                return fmt.Errorf("aggregation is required")
        }

        // Validate conflict resolution strategy
        validStrategies := []vmcp.ConflictResolutionStrategy{
                vmcp.ConflictStrategyPrefix,
                vmcp.ConflictStrategyPriority,
                vmcp.ConflictStrategyManual,
        }
        if !slices.Contains(validStrategies, agg.ConflictResolution) {
                return fmt.Errorf("conflictResolution must be one of: prefix, priority, manual")
        }

        // Validate strategy-specific configuration
        if agg.ConflictResolutionConfig == nil {
                return fmt.Errorf("conflictResolutionConfig is required")
        }

        if err := v.validateConflictStrategy(agg); err != nil {
                return err
        }

        // Mirrors the CRD's enum for the CLI path, which has no API-server admission
        // to reject a bad value. "" is accepted and means allow, so configs written
        // before this setting existed stay valid.
        switch agg.DefaultToolVisibility {
        case "", DefaultToolVisibilityAllow, DefaultToolVisibilityDeny:
        default:
                return fmt.Errorf("defaultToolVisibility must be one of: allow, deny")
        }

        if err := validateDenyVisibilityPriorityOrder(agg); err != nil {
                return err
        }

        return v.validateToolConfigurations(agg.Tools)
}

// validateDenyVisibilityPriorityOrder rejects a priorityOrder entry that has no
// corresponding Tools entry when DefaultToolVisibility is deny.
//
// Conflict resolution runs BEFORE the advertising filter, and priorityOrder is
// independent of Tools. So an unlisted backend can win a name conflict — dropping
// the listed backend's version of that tool at resolution — and then be withheld
// by deny, leaving the tool advertised by nobody. That is strictly worse than
// allow, where the winner at least stays visible.
//
// Rejecting the combination up front is cheaper than reordering the pipeline,
// which would mean resolving conflicts twice: once over the visible candidates
// and once over the complete set that composite tools route against.
//
// validateAggregation guarantees ConflictResolutionConfig is non-nil before
// calling this, so it is safe to dereference below.
func validateDenyVisibilityPriorityOrder(agg *AggregationConfig) error {
        if agg.DefaultToolVisibility != DefaultToolVisibilityDeny ||
                agg.ConflictResolution != vmcp.ConflictStrategyPriority {
                return nil
        }

        listed := make(map[string]struct{}, len(agg.Tools))
        for _, wl := range agg.Tools {
                if wl != nil {
                        listed[wl.Workload] = struct{}{}
                }
        }

        for _, workload := range agg.ConflictResolutionConfig.PriorityOrder {
                if _, ok := listed[workload]; !ok {
                        return fmt.Errorf(
                                "defaultToolVisibility deny with priority strategy: priorityOrder entry %q has no tools entry; "+
                                        "it could win a name conflict and then be withheld, hiding the tool entirely. "+
                                        "Add a tools entry for %q or remove it from priorityOrder",
                                workload, workload)
                }
        }

        return nil
}

// validateConflictStrategy validates strategy-specific configuration
func (*DefaultValidator) validateConflictStrategy(agg *AggregationConfig) error {
        switch agg.ConflictResolution {
        case vmcp.ConflictStrategyPrefix:
                if agg.ConflictResolutionConfig.PrefixFormat == "" {
                        return fmt.Errorf("prefixFormat is required for prefix strategy")
                }

        case vmcp.ConflictStrategyPriority:
                if len(agg.ConflictResolutionConfig.PriorityOrder) == 0 {
                        return fmt.Errorf("priorityOrder is required for priority strategy")
                }

        case vmcp.ConflictStrategyManual:
                // Manual strategy requires explicit overrides
                if len(agg.Tools) == 0 {
                        return fmt.Errorf("tool overrides are required for manual strategy")
                }
        }

        return nil
}

// validateToolConfigurations validates tool override configurations
func (v *DefaultValidator) validateToolConfigurations(tools []*WorkloadToolConfig) error {
        workloadNames := make(map[string]bool)
        for i, tool := range tools {
                if tool.Workload == "" {
                        return fmt.Errorf("tools[%d].workload is required", i)
                }

                if workloadNames[tool.Workload] {
                        return fmt.Errorf("duplicate workload configuration: %s", tool.Workload)
                }
                workloadNames[tool.Workload] = true

                if err := v.validateToolOverrides(tool.Overrides, i); err != nil {
                        return err
                }
        }

        return nil
}

// validateToolOverrides validates individual tool overrides
func (*DefaultValidator) validateToolOverrides(overrides map[string]*ToolOverride, toolIndex int) error {
        for toolName, override := range overrides {
                if override.Name == "" && override.Description == "" {
                        return fmt.Errorf("tools[%d].overrides.%s: at least one of name or description must be specified", toolIndex, toolName)
                }
        }
        return nil
}

func (v *DefaultValidator) validateOperational(ops *OperationalConfig) error {
        if ops == nil {
                return nil // Operational config is optional (defaults apply)
        }

        // Validate timeouts
        if ops.Timeouts != nil {
                if ops.Timeouts.Default <= 0 {
                        return fmt.Errorf("operational.timeouts.default must be positive")
                }

                for workload, timeout := range ops.Timeouts.PerWorkload {
                        if timeout <= 0 {
                                return fmt.Errorf("operational.timeouts.perWorkload.%s must be positive", workload)
                        }
                }
        }

        // Validate failure handling
        if ops.FailureHandling != nil {
                if err := v.validateFailureHandling(ops.FailureHandling); err != nil {
                        return fmt.Errorf("operational.failureHandling: %w", err)
                }
        }

        return nil
}

func (*DefaultValidator) validateFailureHandling(fh *FailureHandlingConfig) error {
        if fh.HealthCheckInterval <= 0 {
                return fmt.Errorf("healthCheckInterval must be positive")
        }

        if fh.UnhealthyThreshold <= 0 {
                return fmt.Errorf("unhealthyThreshold must be positive")
        }

        // Validate health check timeout
        // Zero means no timeout (not recommended but valid)
        // Negative values are invalid
        healthCheckTimeout := time.Duration(fh.HealthCheckTimeout)
        if healthCheckTimeout < 0 {
                return fmt.Errorf("healthCheckTimeout must be >= 0 (zero means no timeout), got %v", healthCheckTimeout)
        }

        // If timeout is configured (non-zero), validate that it's less than interval
        if healthCheckTimeout > 0 {
                checkInterval := time.Duration(fh.HealthCheckInterval)

                // Validate that timeout is less than interval to prevent checks from queuing up
                if healthCheckTimeout >= checkInterval {
                        return fmt.Errorf("healthCheckTimeout (%v) must be less than healthCheckInterval (%v) to prevent checks from queuing up",
                                healthCheckTimeout, checkInterval)
                }
        }

        validModes := []string{"fail", "best_effort"}
        if !slices.Contains(validModes, fh.PartialFailureMode) {
                return fmt.Errorf("partialFailureMode must be one of: %s", strings.Join(validModes, ", "))
        }

        // Validate circuit breaker
        if fh.CircuitBreaker != nil && fh.CircuitBreaker.Enabled {
                if fh.CircuitBreaker.FailureThreshold < 1 {
                        return fmt.Errorf("circuitBreaker.failureThreshold must be >= 1, got %d",
                                fh.CircuitBreaker.FailureThreshold)
                }

                cbTimeout := time.Duration(fh.CircuitBreaker.Timeout)
                if cbTimeout <= 0 {
                        return fmt.Errorf("circuitBreaker.timeout must be > 0, got %v", cbTimeout)
                }

                if cbTimeout < time.Second {
                        return fmt.Errorf("circuitBreaker.timeout must be >= 1s to prevent thrashing, got %v",
                                cbTimeout)
                }
        }

        return nil
}

func (*DefaultValidator) validateCompositeTools(tools []CompositeToolConfig) error {
        if len(tools) == 0 {
                return nil // Composite tools are optional
        }

        toolNames := make(map[string]bool)

        for i := range tools {
                tool := &tools[i]

                // Check for duplicate tool names
                if toolNames[tool.Name] {
                        return fmt.Errorf("duplicate composite tool name: %s", tool.Name)
                }
                toolNames[tool.Name] = true

                // Use shared validation
                if err := ValidateCompositeToolConfig(fmt.Sprintf("compositeTools[%d]", i), tool); err != nil {
                        return err
                }
        }

        return nil
}

func (*DefaultValidator) validateCompositeToolRefs(refs []CompositeToolRef) error {
        if len(refs) == 0 {
                return nil // Composite tool references are optional
        }

        refNames := make(map[string]bool)

        for i := range refs {
                ref := &refs[i]

                if ref.Name == "" {
                        return fmt.Errorf("compositeToolRefs[%d].name is required", i)
                }

                if refNames[ref.Name] {
                        return fmt.Errorf("duplicate composite tool reference: %s", ref.Name)
                }
                refNames[ref.Name] = true
        }

        return nil
}

// The standalone header-forward middleware deliberately allows Authorization
// (an operator may legitimately forward it), but for vMCP passthrough the
// documented contract is stricter: Authorization and Cookie are rejected at
// startup because forwarding caller-supplied credentials verbatim to every
// backend is a credential-leak footgun, not a pass-through use case.
var vmcpRestrictedHeaders = map[string]bool{
        "Authorization": true,
        "Cookie":        true,
}

func (*DefaultValidator) validatePassthroughHeaders(cfg *Config) error {
        for i, name := range cfg.PassthroughHeaders {
                if name == "" {
                        return fmt.Errorf("passthroughHeaders[%d]: header name must not be empty", i)
                }

                canonical := http.CanonicalHeaderKey(name)

                if middleware.RestrictedHeaders[canonical] || vmcpRestrictedHeaders[canonical] {
                        return fmt.Errorf("passthroughHeaders[%d]: %q is a restricted header and cannot be forwarded", i, canonical)
                }

                if err := httpval.ValidateHeaderName(name); err != nil {
                        return fmt.Errorf("passthroughHeaders[%d]: invalid header name %q: %w", i, name, err)
                }
        }
        return nil
}

// Note: Workflow step validation is now handled by the shared ValidateWorkflowSteps function
// in composite_validation.go, which is called by ValidateCompositeToolConfig.

// ValidateAuthServerIntegration validates cross-cutting rules between the
// embedded auth server configuration and backend auth strategies.
// This is called separately from Validate() because it needs the runtime-only
// auth server RunConfig that is not part of the serializable Config.
func ValidateAuthServerIntegration(cfg *Config, rc *authserver.RunConfig) error {
        strategies := collectAllBackendStrategies(cfg)
        hasUpstreamInject := hasStrategyType(strategies, authtypes.StrategyTypeUpstreamInject)

        // Guard clause: nothing to validate if no auth server and no upstream_inject backends.
        if rc == nil && !hasUpstreamInject {
                return nil
        }

        // upstream_inject requires an auth server to obtain upstream tokens.
        if hasUpstreamInject && rc == nil {
                return fmt.Errorf("upstream_inject requires an embedded auth server (authServer must be configured)")
        }

        // Structural validation of the auth server RunConfig.
        if err := validateAuthServerRunConfig(rc); err != nil {
                return err
        }

        // Auth server requires OIDC incoming auth to validate issued tokens.
        if err := validateAuthServerRequiresOIDC(cfg); err != nil {
                return err
        }

        // Every upstream_inject providerName must reference an existing upstream.
        if err := validateUpstreamInjectProviders(rc, strategies); err != nil {
                return err
        }

        // Issuer and audience consistency between auth server and incoming auth.
        return validateAuthServerIncomingAuthConsistency(cfg, rc)
}

// validateAuthServerRunConfig performs lightweight structural validation of the
// auth server RunConfig (issuer, upstreams, allowed audiences).
func validateAuthServerRunConfig(rc *authserver.RunConfig) error {
        if rc == nil {
                return nil
        }
        if rc.Issuer == "" {
                return fmt.Errorf("auth server issuer is required")
        }
        if len(rc.Upstreams) == 0 && len(rc.DelegateClients) == 0 &&
                !authserver.JWTBearerGrantEnabled(rc.TrustedIssuers) {
                return fmt.Errorf("auth server requires at least one upstream unless delegate clients or a " +
                        "trusted issuer with JWT bearer grant is configured")
        }
        // AllowedAudiences is required for MCP compliance (RFC 8707).
        if len(rc.AllowedAudiences) == 0 {
                return fmt.Errorf("auth server requires at least one allowed audience (MCP clients must send RFC 8707 resource parameter)")
        }
        return nil
}

// validateUpstreamInjectProviders checks that every upstream_inject strategy
// references a provider that exists in the auth server upstreams.
func validateUpstreamInjectProviders(
        rc *authserver.RunConfig,
        strategies map[string]*authtypes.BackendAuthStrategy,
) error {
        if rc == nil {
                return nil
        }
        for name, strategy := range strategies {
                if strategy.Type != authtypes.StrategyTypeUpstreamInject || strategy.UpstreamInject == nil {
                        continue
                }
                if !upstreamExists(rc, strategy.UpstreamInject.ProviderName) {
                        return fmt.Errorf(
                                "backend %q: upstream_inject providerName %q not found in auth server upstreams",
                                name, strategy.UpstreamInject.ProviderName,
                        )
                }
        }
        return nil
}

// validateAuthServerIncomingAuthConsistency checks that the embedded auth server
// and the incoming OIDC middleware agree on issuer and audience.
//
// This is a general consistency check that applies whenever both the embedded AS
// and OIDC incoming auth are configured, regardless of which outgoing backend
// strategies (upstream_inject, token_exchange, etc.) are in use.
//
// The embedded AS issues tokens that the OIDC incoming auth middleware validates.
// If these two components disagree on issuer or audience, the middleware will
// reject every token the AS issues, and no authenticated request will succeed.
func validateAuthServerIncomingAuthConsistency(cfg *Config, rc *authserver.RunConfig) error {
        if !hasAuthServerWithOIDCIncoming(cfg, rc) {
                return nil
        }
        oidc := cfg.IncomingAuth.OIDC

        // The OIDC middleware validates the "iss" claim against incomingAuth.oidc.issuer.
        // If the AS uses a different issuer, every token it issues will fail validation.
        if rc.Issuer != oidc.Issuer {
                return fmt.Errorf(
                        "auth server issuer mismatch: auth server issuer %q != incomingAuth.oidc.issuer %q",
                        rc.Issuer, oidc.Issuer,
                )
        }

        // The embedded AS uses the RFC 8707 resource parameter value as the
        // token's aud claim (identity mapping). AllowedAudiences gates which resource
        // values the AS accepts. If incomingAuth expects an audience not in that list,
        // the AS will never issue a matching token.
        // Note: oidc.Audience is required when incomingAuth.type is "oidc" (enforced
        // by validateIncomingAuth), so the empty check is defensive for callers that
        // invoke ValidateAuthServerIntegration independently.
        if oidc.Audience != "" && !slices.Contains(rc.AllowedAudiences, oidc.Audience) {
                return fmt.Errorf(
                        "incomingAuth.oidc.audience %q not in auth server's allowed audiences %v",
                        oidc.Audience, rc.AllowedAudiences,
                )
        }

        return nil
}

// hasOIDCIncoming reports whether the config has OIDC incoming auth fully configured.
func hasOIDCIncoming(cfg *Config) bool {
        return cfg.IncomingAuth != nil &&
                cfg.IncomingAuth.Type == IncomingAuthTypeOIDC &&
                cfg.IncomingAuth.OIDC != nil
}

// validateAuthServerRequiresOIDC checks that when the auth server is configured,
// incomingAuth is OIDC. The AS issues tokens that the OIDC middleware
// validates; without OIDC incoming auth the entire OAuth flow is pointless.
func validateAuthServerRequiresOIDC(cfg *Config) error {
        if !hasOIDCIncoming(cfg) {
                return fmt.Errorf("embedded auth server requires OIDC incoming auth")
        }
        return nil
}

// hasAuthServerWithOIDCIncoming returns true when both the auth server and
// incoming OIDC auth are configured, enabling cross-cutting validation.
func hasAuthServerWithOIDCIncoming(cfg *Config, rc *authserver.RunConfig) bool {
        return rc != nil && hasOIDCIncoming(cfg)
}

// collectAllBackendStrategies returns all backend auth strategies from the config.
func collectAllBackendStrategies(cfg *Config) map[string]*authtypes.BackendAuthStrategy {
        result := make(map[string]*authtypes.BackendAuthStrategy)
        if cfg.OutgoingAuth == nil {
                return result
        }
        if cfg.OutgoingAuth.Default != nil {
                result[defaultStrategyKey] = cfg.OutgoingAuth.Default
        }
        for name, strategy := range cfg.OutgoingAuth.Backends {
                result[name] = strategy
        }
        return result
}

// hasStrategyType checks if any strategy in the map uses the given type.
func hasStrategyType(strategies map[string]*authtypes.BackendAuthStrategy, strategyType string) bool {
        for _, s := range strategies {
                if s.Type == strategyType {
                        return true
                }
        }
        return false
}

// upstreamExists checks if a provider name exists in the RunConfig's upstreams.
// Provider names and upstream names are resolved via authserver.ResolveUpstreamName
// before comparison to ensure consistent empty→"default" normalization.
func upstreamExists(rc *authserver.RunConfig, providerName string) bool {
        if rc == nil {
                return false
        }
        resolved := authserver.ResolveUpstreamName(providerName)
        for i := range rc.Upstreams {
                if authserver.ResolveUpstreamName(rc.Upstreams[i].Name) == resolved {
                        return true
                }
        }
        return false
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package config

import (
        "bytes"
        "fmt"
        "os"
        "time"

        "go.yaml.in/yaml/v3"

        "github.com/stacklok/toolhive-core/env"
        authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
)

// YAMLLoader loads configuration from a YAML file.
// This is the CLI-specific loader that parses the YAML format defined in the proposal.
type YAMLLoader struct {
        filePath  string
        envReader env.Reader
}

// NewYAMLLoader creates a new YAML configuration loader.
func NewYAMLLoader(filePath string, envReader env.Reader) *YAMLLoader {
        return &YAMLLoader{
                filePath:  filePath,
                envReader: envReader,
        }
}

// Load reads and parses the YAML configuration file.
// Uses strict unmarshalling to reject unknown fields.
func (l *YAMLLoader) Load() (*Config, error) {
        data, err := os.ReadFile(l.filePath)
        if err != nil {
                return nil, fmt.Errorf("failed to read config file: %w", err)
        }

        // Use yaml.Decoder with KnownFields for strict unmarshalling
        var cfg Config
        decoder := yaml.NewDecoder(bytes.NewReader(data))
        decoder.KnownFields(true) // Reject unknown fields

        if err := decoder.Decode(&cfg); err != nil {
                return nil, fmt.Errorf("failed to parse YAML: %w", err)
        }

        // Post-process the config
        if err := l.postProcess(&cfg); err != nil {
                return nil, fmt.Errorf("failed to process config: %w", err)
        }

        return &cfg, nil
}

// postProcess applies post-load processing to the config:
// - Resolves environment variables for secrets
// - Applies type inference for workflow steps
// - Sets default timeouts
// - Validates JSON schemas
func (l *YAMLLoader) postProcess(cfg *Config) error {
        // Process outgoing auth - resolve env vars
        if cfg.OutgoingAuth != nil {
                if err := l.processOutgoingAuth(cfg.OutgoingAuth); err != nil {
                        return fmt.Errorf("outgoingAuth: %w", err)
                }
        }

        // Process composite tools - type inference, defaults, validation
        for i := range cfg.CompositeTools {
                if err := l.processCompositeTool(&cfg.CompositeTools[i]); err != nil {
                        return fmt.Errorf("compositeTools[%d]: %w", i, err)
                }
        }

        // Apply operational defaults (fills missing values)
        cfg.EnsureOperationalDefaults()

        return nil
}

// processOutgoingAuth resolves environment variables for auth strategies.
func (l *YAMLLoader) processOutgoingAuth(auth *OutgoingAuthConfig) error {
        if auth.Default != nil {
                if err := l.processBackendAuthStrategy("default", auth.Default); err != nil {
                        return err
                }
        }

        for name, strategy := range auth.Backends {
                if err := l.processBackendAuthStrategy(name, strategy); err != nil {
                        return err
                }
        }

        return nil
}

// processBackendAuthStrategy resolves environment variables for a single auth strategy.
//
//nolint:gocyclo // Strategy-specific processing requires checking multiple fields
func (l *YAMLLoader) processBackendAuthStrategy(name string, strategy *authtypes.BackendAuthStrategy) error {
        switch strategy.Type {
        case authtypes.StrategyTypeHeaderInjection:
                if strategy.HeaderInjection == nil {
                        return fmt.Errorf("backend %s: headerInjection configuration is required", name)
                }

                hi := strategy.HeaderInjection
                hasValue := hi.HeaderValue != ""
                hasValueEnv := hi.HeaderValueEnv != ""

                if hasValue && hasValueEnv {
                        return fmt.Errorf("backend %s: only one of headerValue or headerValueEnv must be set", name)
                }
                if !hasValue && !hasValueEnv {
                        return fmt.Errorf("backend %s: either headerValue or headerValueEnv must be set", name)
                }

                // Resolve header value from environment if env var name is provided
                if hasValueEnv {
                        hi.HeaderValue = l.envReader.Getenv(hi.HeaderValueEnv)
                        if hi.HeaderValue == "" {
                                return fmt.Errorf("backend %s: environment variable %s not set or empty", name, hi.HeaderValueEnv)
                        }
                }

        case authtypes.StrategyTypeTokenExchange:
                if strategy.TokenExchange == nil {
                        return fmt.Errorf("backend %s: tokenExchange configuration is required", name)
                }

                te := strategy.TokenExchange
                if te.ClientSecretEnv != "" {
                        // Validate that the environment variable is set
                        resolvedSecret := l.envReader.Getenv(te.ClientSecretEnv)
                        if resolvedSecret == "" {
                                return fmt.Errorf("backend %s: environment variable %s not set", name, te.ClientSecretEnv)
                        }
                }

        case authtypes.StrategyTypeUnauthenticated:
                // No validation needed

        case authtypes.StrategyTypeAwsSts:
                if strategy.AwsSts == nil {
                        return fmt.Errorf("backend %s: aws_sts configuration is required", name)
                }
                if strategy.AwsSts.Region == "" {
                        return fmt.Errorf("backend %s: aws_sts requires region field", name)
                }

        case authtypes.StrategyTypeOBO:
                if strategy.OBO == nil {
                        return fmt.Errorf("backend %s: obo configuration is required", name)
                }
                obo := strategy.OBO
                if obo.ClientSecret != "" && obo.ClientSecretEnv != "" {
                        return fmt.Errorf("backend %s: only one of clientSecret or clientSecretEnv must be set", name)
                }
                if obo.ClientSecretEnv != "" {
                        resolvedSecret := l.envReader.Getenv(obo.ClientSecretEnv)
                        if resolvedSecret == "" {
                                return fmt.Errorf("backend %s: environment variable %s not set", name, obo.ClientSecretEnv)
                        }
                        obo.ClientSecret = resolvedSecret
                }

        default:
                // Unknown strategy type - let validation handle it
        }

        return nil
}

// processCompositeTool applies post-processing to a composite tool.
func (l *YAMLLoader) processCompositeTool(tool *CompositeToolConfig) error {
        // Validate parameters JSON Schema if present
        if !tool.Parameters.IsEmpty() {
                paramsMap, err := tool.Parameters.ToMap()
                if err != nil {
                        return fmt.Errorf("failed to unmarshal parameters for composite tool %s: %w", tool.Name, err)
                }
                if err := validateParametersJSONSchema(paramsMap, tool.Name); err != nil {
                        return err
                }
        }

        // Process each step
        for i := range tool.Steps {
                l.processWorkflowStep(&tool.Steps[i])
        }

        return nil
}

// processWorkflowStep applies post-processing to a workflow step.
func (*YAMLLoader) processWorkflowStep(step *WorkflowStepConfig) {
        // Apply type inference: if type is empty and tool field is present, infer as "tool"
        if step.Type == "" && step.Tool != "" {
                step.Type = "tool"
        }

        // Set default timeout for elicitation steps
        if step.Type == "elicitation" && step.Timeout == 0 {
                step.Timeout = Duration(5 * time.Minute)
        }
}

// validateParametersJSONSchema validates that parameters follows JSON Schema format.
// Per MCP specification, parameters should be a JSON Schema object with type "object".
//
// We enforce type="object" because MCP tools use named parameters (inputSchema.properties),
// and non-object types (e.g., type="string") would mean a tool takes a single unnamed value,
// which doesn't align with how MCP tool arguments work. The MCP SDK and specification
// expect tools to have named parameters accessible via inputSchema.properties.
func validateParametersJSONSchema(params map[string]any, toolName string) error {
        if len(params) == 0 {
                return nil
        }

        // Check if it has "type" field
        typeVal, hasType := params["type"]
        if !hasType {
                return fmt.Errorf("tool %s: parameters must have 'type' field (should be 'object' for JSON Schema)", toolName)
        }

        // Type must be a string
        typeStr, ok := typeVal.(string)
        if !ok {
                return fmt.Errorf("tool %s: parameters 'type' field must be a string", toolName)
        }

        // Type should be "object" for parameter schemas
        if typeStr != "object" {
                return fmt.Errorf("tool %s: parameters 'type' must be 'object' (got '%s')", toolName, typeStr)
        }

        // If properties exist, validate it's a map
        if properties, hasProps := params["properties"]; hasProps {
                if _, ok := properties.(map[string]any); !ok {
                        return fmt.Errorf("tool %s: parameters 'properties' must be an object", toolName)
                }
        }

        // If required exists, validate it's an array
        if required, hasRequired := params["required"]; hasRequired {
                if _, ok := required.([]any); !ok {
                        // Also accept []string which may come from YAML
                        if _, ok := required.([]string); !ok {
                                return fmt.Errorf("tool %s: parameters 'required' must be an array", toolName)
                        }
                }
        }

        return nil
}

//go:build !ignore_autogenerated

/*
Copyright 2025 Stacklok

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Code generated by controller-gen. DO NOT EDIT.

package config

import (
        "github.com/stacklok/toolhive/pkg/audit"
        "github.com/stacklok/toolhive/pkg/ratelimit/types"
        "github.com/stacklok/toolhive/pkg/redisconfig"
        "github.com/stacklok/toolhive/pkg/telemetry"
        authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
)

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AggregationConfig) DeepCopyInto(out *AggregationConfig) {
        *out = *in
        if in.ConflictResolutionConfig != nil {
                in, out := &in.ConflictResolutionConfig, &out.ConflictResolutionConfig
                *out = new(ConflictResolutionConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Tools != nil {
                in, out := &in.Tools, &out.Tools
                *out = make([]*WorkloadToolConfig, len(*in))
                for i := range *in {
                        if (*in)[i] != nil {
                                in, out := &(*in)[i], &(*out)[i]
                                *out = new(WorkloadToolConfig)
                                (*in).DeepCopyInto(*out)
                        }
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AggregationConfig.
func (in *AggregationConfig) DeepCopy() *AggregationConfig {
        if in == nil {
                return nil
        }
        out := new(AggregationConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *AuthzConfig) DeepCopyInto(out *AuthzConfig) {
        *out = *in
        if in.Policies != nil {
                in, out := &in.Policies, &out.Policies
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AuthzConfig.
func (in *AuthzConfig) DeepCopy() *AuthzConfig {
        if in == nil {
                return nil
        }
        out := new(AuthzConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *CircuitBreakerConfig) DeepCopyInto(out *CircuitBreakerConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new CircuitBreakerConfig.
func (in *CircuitBreakerConfig) DeepCopy() *CircuitBreakerConfig {
        if in == nil {
                return nil
        }
        out := new(CircuitBreakerConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *CodeModeConfig) DeepCopyInto(out *CodeModeConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new CodeModeConfig.
func (in *CodeModeConfig) DeepCopy() *CodeModeConfig {
        if in == nil {
                return nil
        }
        out := new(CodeModeConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *CompositeToolConfig) DeepCopyInto(out *CompositeToolConfig) {
        *out = *in
        in.Parameters.DeepCopyInto(&out.Parameters)
        if in.Steps != nil {
                in, out := &in.Steps, &out.Steps
                *out = make([]WorkflowStepConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.Output != nil {
                in, out := &in.Output, &out.Output
                *out = new(OutputConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Annotations != nil {
                in, out := &in.Annotations, &out.Annotations
                *out = new(ToolAnnotationsOverride)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new CompositeToolConfig.
func (in *CompositeToolConfig) DeepCopy() *CompositeToolConfig {
        if in == nil {
                return nil
        }
        out := new(CompositeToolConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *CompositeToolRef) DeepCopyInto(out *CompositeToolRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new CompositeToolRef.
func (in *CompositeToolRef) DeepCopy() *CompositeToolRef {
        if in == nil {
                return nil
        }
        out := new(CompositeToolRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *Config) DeepCopyInto(out *Config) {
        *out = *in
        if in.Backends != nil {
                in, out := &in.Backends, &out.Backends
                *out = make([]StaticBackendConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.IncomingAuth != nil {
                in, out := &in.IncomingAuth, &out.IncomingAuth
                *out = new(IncomingAuthConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.OutgoingAuth != nil {
                in, out := &in.OutgoingAuth, &out.OutgoingAuth
                *out = new(OutgoingAuthConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Aggregation != nil {
                in, out := &in.Aggregation, &out.Aggregation
                *out = new(AggregationConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.CompositeTools != nil {
                in, out := &in.CompositeTools, &out.CompositeTools
                *out = make([]CompositeToolConfig, len(*in))
                for i := range *in {
                        (*in)[i].DeepCopyInto(&(*out)[i])
                }
        }
        if in.CompositeToolRefs != nil {
                in, out := &in.CompositeToolRefs, &out.CompositeToolRefs
                *out = make([]CompositeToolRef, len(*in))
                copy(*out, *in)
        }
        if in.Operational != nil {
                in, out := &in.Operational, &out.Operational
                *out = new(OperationalConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Metadata != nil {
                in, out := &in.Metadata, &out.Metadata
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
        if in.Telemetry != nil {
                in, out := &in.Telemetry, &out.Telemetry
                *out = new(telemetry.Config)
                (*in).DeepCopyInto(*out)
        }
        if in.Audit != nil {
                in, out := &in.Audit, &out.Audit
                *out = new(audit.Config)
                (*in).DeepCopyInto(*out)
        }
        if in.Optimizer != nil {
                in, out := &in.Optimizer, &out.Optimizer
                *out = new(OptimizerConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.CodeMode != nil {
                in, out := &in.CodeMode, &out.CodeMode
                *out = new(CodeModeConfig)
                **out = **in
        }
        if in.SessionStorage != nil {
                in, out := &in.SessionStorage, &out.SessionStorage
                *out = new(SessionStorageConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.RateLimiting != nil {
                in, out := &in.RateLimiting, &out.RateLimiting
                *out = new(types.RateLimitConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.PassthroughHeaders != nil {
                in, out := &in.PassthroughHeaders, &out.PassthroughHeaders
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Config.
func (in *Config) DeepCopy() *Config {
        if in == nil {
                return nil
        }
        out := new(Config)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ConflictResolutionConfig) DeepCopyInto(out *ConflictResolutionConfig) {
        *out = *in
        if in.PriorityOrder != nil {
                in, out := &in.PriorityOrder, &out.PriorityOrder
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ConflictResolutionConfig.
func (in *ConflictResolutionConfig) DeepCopy() *ConflictResolutionConfig {
        if in == nil {
                return nil
        }
        out := new(ConflictResolutionConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ElicitationResponseConfig) DeepCopyInto(out *ElicitationResponseConfig) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ElicitationResponseConfig.
func (in *ElicitationResponseConfig) DeepCopy() *ElicitationResponseConfig {
        if in == nil {
                return nil
        }
        out := new(ElicitationResponseConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *FailureHandlingConfig) DeepCopyInto(out *FailureHandlingConfig) {
        *out = *in
        if in.CircuitBreaker != nil {
                in, out := &in.CircuitBreaker, &out.CircuitBreaker
                *out = new(CircuitBreakerConfig)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new FailureHandlingConfig.
func (in *FailureHandlingConfig) DeepCopy() *FailureHandlingConfig {
        if in == nil {
                return nil
        }
        out := new(FailureHandlingConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *IncomingAuthConfig) DeepCopyInto(out *IncomingAuthConfig) {
        *out = *in
        if in.OIDC != nil {
                in, out := &in.OIDC, &out.OIDC
                *out = new(OIDCConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.Authz != nil {
                in, out := &in.Authz, &out.Authz
                *out = new(AuthzConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new IncomingAuthConfig.
func (in *IncomingAuthConfig) DeepCopy() *IncomingAuthConfig {
        if in == nil {
                return nil
        }
        out := new(IncomingAuthConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ListChangedConfig) DeepCopyInto(out *ListChangedConfig) {
        *out = *in
        if in.Enabled != nil {
                in, out := &in.Enabled, &out.Enabled
                *out = new(bool)
                **out = **in
        }
        if in.DisabledWorkloads != nil {
                in, out := &in.DisabledWorkloads, &out.DisabledWorkloads
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ListChangedConfig.
func (in *ListChangedConfig) DeepCopy() *ListChangedConfig {
        if in == nil {
                return nil
        }
        out := new(ListChangedConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OIDCConfig) DeepCopyInto(out *OIDCConfig) {
        *out = *in
        if in.Scopes != nil {
                in, out := &in.Scopes, &out.Scopes
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OIDCConfig.
func (in *OIDCConfig) DeepCopy() *OIDCConfig {
        if in == nil {
                return nil
        }
        out := new(OIDCConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OperationalConfig) DeepCopyInto(out *OperationalConfig) {
        *out = *in
        if in.Timeouts != nil {
                in, out := &in.Timeouts, &out.Timeouts
                *out = new(TimeoutConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.FailureHandling != nil {
                in, out := &in.FailureHandling, &out.FailureHandling
                *out = new(FailureHandlingConfig)
                (*in).DeepCopyInto(*out)
        }
        if in.ListChanged != nil {
                in, out := &in.ListChanged, &out.ListChanged
                *out = new(ListChangedConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OperationalConfig.
func (in *OperationalConfig) DeepCopy() *OperationalConfig {
        if in == nil {
                return nil
        }
        out := new(OperationalConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OptimizerConfig) DeepCopyInto(out *OptimizerConfig) {
        *out = *in
        if in.EmbeddingHeaders != nil {
                in, out := &in.EmbeddingHeaders, &out.EmbeddingHeaders
                *out = make(map[string]EmbeddingHeaderValue, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OptimizerConfig.
func (in *OptimizerConfig) DeepCopy() *OptimizerConfig {
        if in == nil {
                return nil
        }
        out := new(OptimizerConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OutgoingAuthConfig) DeepCopyInto(out *OutgoingAuthConfig) {
        *out = *in
        if in.Default != nil {
                in, out := &in.Default, &out.Default
                *out = new(authtypes.BackendAuthStrategy)
                (*in).DeepCopyInto(*out)
        }
        if in.Backends != nil {
                in, out := &in.Backends, &out.Backends
                *out = make(map[string]*authtypes.BackendAuthStrategy, len(*in))
                for key, val := range *in {
                        var outVal *authtypes.BackendAuthStrategy
                        if val == nil {
                                (*out)[key] = nil
                        } else {
                                inVal := (*in)[key]
                                in, out := &inVal, &outVal
                                *out = new(authtypes.BackendAuthStrategy)
                                (*in).DeepCopyInto(*out)
                        }
                        (*out)[key] = outVal
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OutgoingAuthConfig.
func (in *OutgoingAuthConfig) DeepCopy() *OutgoingAuthConfig {
        if in == nil {
                return nil
        }
        out := new(OutgoingAuthConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OutputConfig) DeepCopyInto(out *OutputConfig) {
        *out = *in
        if in.Properties != nil {
                in, out := &in.Properties, &out.Properties
                *out = make(map[string]OutputProperty, len(*in))
                for key, val := range *in {
                        (*out)[key] = *val.DeepCopy()
                }
        }
        if in.Required != nil {
                in, out := &in.Required, &out.Required
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OutputConfig.
func (in *OutputConfig) DeepCopy() *OutputConfig {
        if in == nil {
                return nil
        }
        out := new(OutputConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *OutputProperty) DeepCopyInto(out *OutputProperty) {
        *out = *in
        if in.Properties != nil {
                in, out := &in.Properties, &out.Properties
                *out = make(map[string]OutputProperty, len(*in))
                for key, val := range *in {
                        (*out)[key] = *val.DeepCopy()
                }
        }
        in.Default.DeepCopyInto(&out.Default)
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OutputProperty.
func (in *OutputProperty) DeepCopy() *OutputProperty {
        if in == nil {
                return nil
        }
        out := new(OutputProperty)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *SessionStorageConfig) DeepCopyInto(out *SessionStorageConfig) {
        *out = *in
        if in.TLS != nil {
                in, out := &in.TLS, &out.TLS
                *out = new(redisconfig.TLSConfig)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new SessionStorageConfig.
func (in *SessionStorageConfig) DeepCopy() *SessionStorageConfig {
        if in == nil {
                return nil
        }
        out := new(SessionStorageConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *StaticBackendConfig) DeepCopyInto(out *StaticBackendConfig) {
        *out = *in
        if in.Metadata != nil {
                in, out := &in.Metadata, &out.Metadata
                *out = make(map[string]string, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new StaticBackendConfig.
func (in *StaticBackendConfig) DeepCopy() *StaticBackendConfig {
        if in == nil {
                return nil
        }
        out := new(StaticBackendConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *StepErrorHandling) DeepCopyInto(out *StepErrorHandling) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new StepErrorHandling.
func (in *StepErrorHandling) DeepCopy() *StepErrorHandling {
        if in == nil {
                return nil
        }
        out := new(StepErrorHandling)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *TimeoutConfig) DeepCopyInto(out *TimeoutConfig) {
        *out = *in
        if in.PerWorkload != nil {
                in, out := &in.PerWorkload, &out.PerWorkload
                *out = make(map[string]Duration, len(*in))
                for key, val := range *in {
                        (*out)[key] = val
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TimeoutConfig.
func (in *TimeoutConfig) DeepCopy() *TimeoutConfig {
        if in == nil {
                return nil
        }
        out := new(TimeoutConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ToolAnnotationsOverride) DeepCopyInto(out *ToolAnnotationsOverride) {
        *out = *in
        if in.Title != nil {
                in, out := &in.Title, &out.Title
                *out = new(string)
                **out = **in
        }
        if in.ReadOnlyHint != nil {
                in, out := &in.ReadOnlyHint, &out.ReadOnlyHint
                *out = new(bool)
                **out = **in
        }
        if in.DestructiveHint != nil {
                in, out := &in.DestructiveHint, &out.DestructiveHint
                *out = new(bool)
                **out = **in
        }
        if in.IdempotentHint != nil {
                in, out := &in.IdempotentHint, &out.IdempotentHint
                *out = new(bool)
                **out = **in
        }
        if in.OpenWorldHint != nil {
                in, out := &in.OpenWorldHint, &out.OpenWorldHint
                *out = new(bool)
                **out = **in
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ToolAnnotationsOverride.
func (in *ToolAnnotationsOverride) DeepCopy() *ToolAnnotationsOverride {
        if in == nil {
                return nil
        }
        out := new(ToolAnnotationsOverride)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ToolConfigRef) DeepCopyInto(out *ToolConfigRef) {
        *out = *in
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ToolConfigRef.
func (in *ToolConfigRef) DeepCopy() *ToolConfigRef {
        if in == nil {
                return nil
        }
        out := new(ToolConfigRef)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *ToolOverride) DeepCopyInto(out *ToolOverride) {
        *out = *in
        if in.Annotations != nil {
                in, out := &in.Annotations, &out.Annotations
                *out = new(ToolAnnotationsOverride)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ToolOverride.
func (in *ToolOverride) DeepCopy() *ToolOverride {
        if in == nil {
                return nil
        }
        out := new(ToolOverride)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *WorkflowStepConfig) DeepCopyInto(out *WorkflowStepConfig) {
        *out = *in
        in.Arguments.DeepCopyInto(&out.Arguments)
        if in.DependsOn != nil {
                in, out := &in.DependsOn, &out.DependsOn
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.OnError != nil {
                in, out := &in.OnError, &out.OnError
                *out = new(StepErrorHandling)
                **out = **in
        }
        in.Schema.DeepCopyInto(&out.Schema)
        if in.OnDecline != nil {
                in, out := &in.OnDecline, &out.OnDecline
                *out = new(ElicitationResponseConfig)
                **out = **in
        }
        if in.OnCancel != nil {
                in, out := &in.OnCancel, &out.OnCancel
                *out = new(ElicitationResponseConfig)
                **out = **in
        }
        in.DefaultResults.DeepCopyInto(&out.DefaultResults)
        if in.InnerStep != nil {
                in, out := &in.InnerStep, &out.InnerStep
                *out = new(WorkflowStepConfig)
                (*in).DeepCopyInto(*out)
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new WorkflowStepConfig.
func (in *WorkflowStepConfig) DeepCopy() *WorkflowStepConfig {
        if in == nil {
                return nil
        }
        out := new(WorkflowStepConfig)
        in.DeepCopyInto(out)
        return out
}

// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
func (in *WorkloadToolConfig) DeepCopyInto(out *WorkloadToolConfig) {
        *out = *in
        if in.ToolConfigRef != nil {
                in, out := &in.ToolConfigRef, &out.ToolConfigRef
                *out = new(ToolConfigRef)
                **out = **in
        }
        if in.Filter != nil {
                in, out := &in.Filter, &out.Filter
                *out = make([]string, len(*in))
                copy(*out, *in)
        }
        if in.Overrides != nil {
                in, out := &in.Overrides, &out.Overrides
                *out = make(map[string]*ToolOverride, len(*in))
                for key, val := range *in {
                        var outVal *ToolOverride
                        if val == nil {
                                (*out)[key] = nil
                        } else {
                                inVal := (*in)[key]
                                in, out := &inVal, &outVal
                                *out = new(ToolOverride)
                                (*in).DeepCopyInto(*out)
                        }
                        (*out)[key] = outVal
                }
        }
}

// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new WorkloadToolConfig.
func (in *WorkloadToolConfig) DeepCopy() *WorkloadToolConfig {
        if in == nil {
                return nil
        }
        out := new(WorkloadToolConfig)
        in.DeepCopyInto(out)
        return out
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package conversion provides utilities for converting between MCP SDK types and vmcp wrapper types.
// This package centralizes conversion logic to ensure consistency and eliminate duplication.
package conversion

import (
        "encoding/json"
        "fmt"
        "log/slog"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive/pkg/vmcp"
)

// ConvertMCPAnnotations converts mcp.Annotations to vmcp.ContentAnnotations.
// Returns nil if the input is nil or all fields are zero-valued.
func ConvertMCPAnnotations(ann *mcp.Annotations) *vmcp.ContentAnnotations {
        if ann == nil {
                return nil
        }
        // Convert []mcp.Role to []string for the ACL boundary.
        var audience []string
        if len(ann.Audience) > 0 {
                audience = make([]string, len(ann.Audience))
                for i, r := range ann.Audience {
                        audience[i] = string(r)
                }
        }
        if len(audience) == 0 && ann.Priority == nil && ann.LastModified == "" {
                return nil
        }
        var priority *float64
        if ann.Priority != nil {
                p := *ann.Priority
                priority = &p
        }
        return &vmcp.ContentAnnotations{
                Audience:     audience,
                Priority:     priority,
                LastModified: ann.LastModified,
        }
}

// ToMCPAnnotations converts vmcp.ContentAnnotations to mcp.Annotations.
// Returns nil if the input is nil or all fields are zero-valued.
func ToMCPAnnotations(ann *vmcp.ContentAnnotations) *mcp.Annotations {
        if ann == nil {
                return nil
        }
        var audience []mcp.Role
        if len(ann.Audience) > 0 {
                audience = make([]mcp.Role, len(ann.Audience))
                for i, a := range ann.Audience {
                        audience[i] = mcp.Role(a)
                }
        }
        if len(audience) == 0 && ann.Priority == nil && ann.LastModified == "" {
                return nil
        }
        var priority *float64
        if ann.Priority != nil {
                p := *ann.Priority
                priority = &p
        }
        return &mcp.Annotations{
                Audience:     audience,
                Priority:     priority,
                LastModified: ann.LastModified,
        }
}

// ConvertMCPContent converts a single mcp.Content item to vmcp.Content.
// Unknown content types are returned as vmcp.Content{Type: "unknown"}.
func ConvertMCPContent(content mcp.Content) vmcp.Content {
        if text, ok := mcp.AsTextContent(content); ok {
                return vmcp.Content{Type: vmcp.ContentTypeText, Text: text.Text, Annotations: ConvertMCPAnnotations(text.Annotations)}
        }
        if img, ok := mcp.AsImageContent(content); ok {
                return vmcp.Content{
                        Type: vmcp.ContentTypeImage, Data: img.Data,
                        MimeType: img.MIMEType, Annotations: ConvertMCPAnnotations(img.Annotations),
                }
        }
        if audio, ok := mcp.AsAudioContent(content); ok {
                return vmcp.Content{
                        Type: vmcp.ContentTypeAudio, Data: audio.Data,
                        MimeType: audio.MIMEType, Annotations: ConvertMCPAnnotations(audio.Annotations),
                }
        }
        if res, ok := mcp.AsEmbeddedResource(content); ok {
                ann := ConvertMCPAnnotations(res.Annotations)
                if textRes, ok := mcp.AsTextResourceContents(res.Resource); ok {
                        return vmcp.Content{
                                Type: vmcp.ContentTypeResource, Text: textRes.Text,
                                URI: textRes.URI, MimeType: textRes.MIMEType, Annotations: ann,
                        }
                }
                if blobRes, ok := mcp.AsBlobResourceContents(res.Resource); ok {
                        return vmcp.Content{
                                Type: vmcp.ContentTypeResource, Data: blobRes.Blob,
                                URI: blobRes.URI, MimeType: blobRes.MIMEType, Annotations: ann,
                        }
                }
                slog.Debug("Embedded resource has unknown resource contents type", "type", fmt.Sprintf("%T", res.Resource))
                return vmcp.Content{Type: vmcp.ContentTypeResource, Annotations: ann}
        }
        if link, ok := content.(mcp.ResourceLink); ok {
                return vmcp.Content{
                        Type:        vmcp.ContentTypeLink,
                        URI:         link.URI,
                        Name:        link.Name,
                        Description: link.Description,
                        MimeType:    link.MIMEType,
                        Annotations: ConvertMCPAnnotations(link.Annotations),
                }
        }
        slog.Debug("Encountered unknown MCP content type", "type", fmt.Sprintf("%T", content))
        return vmcp.Content{Type: "unknown"}
}

// ConvertMCPContents converts a slice of mcp.Content to []vmcp.Content.
// Returns an empty (non-nil) slice for a nil or empty input.
func ConvertMCPContents(contents []mcp.Content) []vmcp.Content {
        result := make([]vmcp.Content, 0, len(contents))
        for _, c := range contents {
                result = append(result, ConvertMCPContent(c))
        }
        return result
}

// ToMCPContent converts a single vmcp.Content item to mcp.Content.
// Unknown content types are converted to empty text with a warning.
func ToMCPContent(content vmcp.Content) mcp.Content {
        ann := toAnnotated(content.Annotations)

        switch content.Type {
        case vmcp.ContentTypeText:
                tc := mcp.NewTextContent(content.Text)
                tc.Annotated = ann
                return tc
        case vmcp.ContentTypeImage:
                ic := mcp.NewImageContent(content.Data, content.MimeType)
                ic.Annotated = ann
                return ic
        case vmcp.ContentTypeAudio:
                ac := mcp.NewAudioContent(content.Data, content.MimeType)
                ac.Annotated = ann
                return ac
        case vmcp.ContentTypeResource:
                // Reconstruct embedded resource from vmcp.Content fields.
                // Text content takes precedence over blob content when both are present.
                if content.Text != "" {
                        er := mcp.NewEmbeddedResource(mcp.TextResourceContents{
                                URI:      content.URI,
                                MIMEType: content.MimeType,
                                Text:     content.Text,
                        })
                        er.Annotated = ann
                        return er
                }
                if content.Data != "" {
                        er := mcp.NewEmbeddedResource(mcp.BlobResourceContents{
                                URI:      content.URI,
                                MIMEType: content.MimeType,
                                Blob:     content.Data,
                        })
                        er.Annotated = ann
                        return er
                }
                // Empty resource content — preserve resource wrapper and metadata with empty contents.
                slog.Warn("converting empty resource content to empty embedded resource - no Text or Data field present")
                er := mcp.NewEmbeddedResource(mcp.TextResourceContents{
                        URI:      content.URI,
                        MIMEType: content.MimeType,
                        Text:     "",
                })
                er.Annotated = ann
                return er
        case vmcp.ContentTypeLink:
                // Reconstruct a ResourceLink from vmcp.Content fields.
                rl := mcp.NewResourceLink(content.URI, content.Name, content.Description, content.MimeType)
                rl.Annotated = ann
                return rl
        default:
                slog.Warn("converting unknown content type to empty text - this may cause data loss", "type", content.Type)
                tc := mcp.NewTextContent("")
                tc.Annotated = ann
                return tc
        }
}

// toAnnotated converts vmcp.ContentAnnotations to mcp.Annotated.
// Returns a zero-valued mcp.Annotated if annotations is nil.
func toAnnotated(ann *vmcp.ContentAnnotations) mcp.Annotated {
        mcpAnn := ToMCPAnnotations(ann)
        if mcpAnn == nil {
                return mcp.Annotated{}
        }
        return mcp.Annotated{Annotations: mcpAnn}
}

// ToMCPContents converts a slice of vmcp.Content to []mcp.Content.
func ToMCPContents(contents []vmcp.Content) []mcp.Content {
        result := make([]mcp.Content, 0, len(contents))
        for _, c := range contents {
                result = append(result, ToMCPContent(c))
        }
        return result
}

// ConvertMCPResourceContents converts []mcp.ResourceContents to []vmcp.ResourceContent,
// preserving the text vs blob distinction and per-item metadata.
func ConvertMCPResourceContents(contents []mcp.ResourceContents) []vmcp.ResourceContent {
        result := make([]vmcp.ResourceContent, 0, len(contents))
        for _, c := range contents {
                if textRes, ok := mcp.AsTextResourceContents(c); ok {
                        result = append(result, vmcp.ResourceContent{
                                URI:      textRes.URI,
                                MimeType: textRes.MIMEType,
                                Text:     textRes.Text,
                        })
                } else if blobRes, ok := mcp.AsBlobResourceContents(c); ok {
                        result = append(result, vmcp.ResourceContent{
                                URI:      blobRes.URI,
                                MimeType: blobRes.MIMEType,
                                Blob:     blobRes.Blob,
                        })
                } else {
                        // Warn rather than debug: an unrecognized resource type likely
                        // indicates a protocol change or bug, and silently dropping data
                        // should be visible to operators.
                        slog.Warn("Skipping unknown resource contents type", "type", fmt.Sprintf("%T", c))
                }
        }
        return result
}

// ToMCPResourceContents converts []vmcp.ResourceContent to []mcp.ResourceContents,
// reconstructing the text vs blob distinction.
func ToMCPResourceContents(contents []vmcp.ResourceContent) []mcp.ResourceContents {
        result := make([]mcp.ResourceContents, 0, len(contents))
        for _, c := range contents {
                // Blob takes precedence: a non-empty Blob field means this is a blob resource.
                // If both Text and Blob are set the Text field is ignored.
                if c.Blob != "" {
                        result = append(result, mcp.BlobResourceContents{
                                URI:      c.URI,
                                MIMEType: c.MimeType,
                                Blob:     c.Blob,
                        })
                } else {
                        result = append(result, mcp.TextResourceContents{
                                URI:      c.URI,
                                MIMEType: c.MimeType,
                                Text:     c.Text,
                        })
                }
        }
        return result
}

// ConvertToolInputSchema converts a mcp.ToolInputSchema to map[string]any via a
// JSON round-trip, capturing all fields (type, properties, required, $defs,
// additionalProperties, etc.) without enumerating them manually. Falls back to
// {type: schema.Type} if marshalling fails.
func ConvertToolInputSchema(schema mcp.ToolInputSchema) map[string]any {
        result := make(map[string]any)
        b, err := json.Marshal(schema)
        if err != nil {
                return map[string]any{"type": schema.Type}
        }
        if err := json.Unmarshal(b, &result); err != nil {
                return map[string]any{"type": schema.Type}
        }
        return result
}

// ConvertMCPPromptMessages converts []mcp.PromptMessage to []vmcp.PromptMessage,
// preserving individual message roles and content types.
func ConvertMCPPromptMessages(messages []mcp.PromptMessage) []vmcp.PromptMessage {
        result := make([]vmcp.PromptMessage, 0, len(messages))
        for _, msg := range messages {
                result = append(result, vmcp.PromptMessage{
                        Role:    string(msg.Role),
                        Content: ConvertMCPContent(msg.Content),
                })
        }
        return result
}

// ToMCPPromptMessages converts []vmcp.PromptMessage to []mcp.PromptMessage.
func ToMCPPromptMessages(messages []vmcp.PromptMessage) []mcp.PromptMessage {
        result := make([]mcp.PromptMessage, 0, len(messages))
        for _, msg := range messages {
                result = append(result, mcp.PromptMessage{
                        Role:    mcp.Role(msg.Role),
                        Content: ToMCPContent(msg.Content),
                })
        }
        return result
}

// ConvertPromptArguments converts map[string]any to map[string]string by
// formatting each value with fmt.Sprintf("%v", v). Required by the MCP
// GetPrompt API which accepts only string-typed arguments.
func ConvertPromptArguments(arguments map[string]any) map[string]string {
        result := make(map[string]string, len(arguments))
        for k, v := range arguments {
                result[k] = fmt.Sprintf("%v", v)
        }
        return result
}

// ConvertToolAnnotations converts mcp.ToolAnnotation to *vmcp.ToolAnnotations.
// Returns nil if all fields are zero-valued (empty Title, all hint pointers nil).
func ConvertToolAnnotations(ann mcp.ToolAnnotation) *vmcp.ToolAnnotations {
        if ann.Title == "" && ann.ReadOnlyHint == nil && ann.DestructiveHint == nil &&
                ann.IdempotentHint == nil && ann.OpenWorldHint == nil {
                return nil
        }
        return &vmcp.ToolAnnotations{
                Title:           ann.Title,
                ReadOnlyHint:    ann.ReadOnlyHint,
                DestructiveHint: ann.DestructiveHint,
                IdempotentHint:  ann.IdempotentHint,
                OpenWorldHint:   ann.OpenWorldHint,
        }
}

// ConvertToolOutputSchema converts a mcp.ToolOutputSchema to map[string]any via a
// JSON round-trip, same pattern as ConvertToolInputSchema.
// Returns nil if the schema has no meaningful type (empty Type field).
func ConvertToolOutputSchema(schema mcp.ToolOutputSchema) map[string]any {
        // A zero-valued ToolOutputSchema has Type="" — this means the backend
        // did not provide an output schema. Return nil to distinguish from a
        // schema that was explicitly set.
        if schema.Type == "" {
                return nil
        }
        b, err := json.Marshal(schema)
        if err != nil {
                return nil
        }
        result := make(map[string]any)
        if err := json.Unmarshal(b, &result); err != nil {
                return nil
        }
        if len(result) == 0 {
                return nil
        }
        return result
}

// ToMCPToolAnnotations converts *vmcp.ToolAnnotations back to mcp.ToolAnnotation.
// Returns a zero-valued mcp.ToolAnnotation if annotations is nil.
func ToMCPToolAnnotations(annotations *vmcp.ToolAnnotations) mcp.ToolAnnotation {
        if annotations == nil {
                return mcp.ToolAnnotation{}
        }
        return mcp.ToolAnnotation{
                Title:           annotations.Title,
                ReadOnlyHint:    annotations.ReadOnlyHint,
                DestructiveHint: annotations.DestructiveHint,
                IdempotentHint:  annotations.IdempotentHint,
                OpenWorldHint:   annotations.OpenWorldHint,
        }
}

// ContentArrayToMap converts a vmcp.Content array to a map for template variable substitution.
// This is used by composite tool workflows and backend result handling.
//
// Conversion rules:
//   - First text content: key="text"
//   - Subsequent text content: key="text_1", "text_2", etc.
//   - Image content: key="image_0", "image_1", etc.
//   - First resource content: key="resource" (text resources use .Text, blob resources use .Data)
//   - Subsequent resource content: key="resource_1", "resource_2", etc.
//   - Audio content: ignored (not supported for template substitution)
//   - Resource links: ignored (not supported for template substitution)
//   - Unknown content types: ignored (warnings logged at conversion boundaries)
//
// This ensures consistent behavior between client response handling and workflow step output processing.
func ContentArrayToMap(content []vmcp.Content) map[string]any {
        result := make(map[string]any)
        if len(content) == 0 {
                return result
        }

        textIndex := 0
        imageIndex := 0
        resourceIndex := 0

        for _, item := range content {
                switch item.Type {
                case vmcp.ContentTypeText:
                        key := string(vmcp.ContentTypeText)
                        if textIndex > 0 {
                                key = fmt.Sprintf("text_%d", textIndex)
                        }
                        result[key] = item.Text
                        textIndex++

                case vmcp.ContentTypeImage:
                        key := fmt.Sprintf("image_%d", imageIndex)
                        result[key] = item.Data
                        imageIndex++

                case vmcp.ContentTypeResource:
                        key := "resource"
                        if resourceIndex > 0 {
                                key = fmt.Sprintf("resource_%d", resourceIndex)
                        }
                        // Text resources use .Text, blob resources use .Data
                        value := item.Text
                        if value == "" {
                                value = item.Data
                        }
                        result[key] = value
                        resourceIndex++

                case vmcp.ContentTypeAudio, vmcp.ContentTypeLink:
                        // Purposely ignored for template substitution:
                        // - Audio content is ignored (not supported for template substitution)
                        // - Resource links are ignored (not supported for template substitution)
                default:
                        // Unknown content types are ignored (warnings logged at conversion boundaries)
                }
        }

        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package conversion

import (
        "errors"

        sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        thvmcp "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/vmcp"
)

// ErrorToToolResult converts domain/tool execution errors into MCP tool results.
//
// mcp-go tool handlers cannot reliably carry custom JSON-RPC codes by returning
// Go errors: the SDK maps them to generic internal errors. For domain errors
// that opt into CodedError, preserve the code/data in StructuredContent instead.
func ErrorToToolResult(err error) *sdkmcp.CallToolResult {
        // Denial first (matches writeModernDispatchError): an error that is both a
        // CodedError and wraps ErrAuthorizationFailed must render as the denial,
        // never as retry-shaped coded data (the sets are disjoint today; this
        // ordering is the invariant).
        if errors.Is(err, vmcp.ErrAuthorizationFailed) {
                return sdkmcp.NewToolResultError(vmcp.DenyMessageToolCall)
        }
        var coded thvmcp.CodedError
        if errors.As(err, &coded) {
                return CodedErrorResult(err, coded)
        }
        return sdkmcp.NewToolResultError(err.Error())
}

// CodedErrorResult builds an IsError tool result carrying stable machine-readable
// error details in StructuredContent. message uses err.Error() so callers can wrap
// a coded error with context using %w without losing the code or data.
func CodedErrorResult(err error, coded thvmcp.CodedError) *sdkmcp.CallToolResult {
        result := sdkmcp.NewToolResultError(err.Error())
        structured := map[string]any{
                "code":    coded.Code(),
                "message": err.Error(),
        }
        if data := coded.Data(); len(data) > 0 {
                structured["data"] = data
        }
        result.StructuredContent = structured
        return result
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package conversion

import (
        "maps"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        mcpparser "github.com/stacklok/toolhive/pkg/mcp"
)

// FromMCPMeta converts MCP SDK meta to map[string]any for vmcp wrapper types.
// This preserves the _meta field from backend MCP server responses.
//
// Returns nil if meta is nil or empty, following the MCP specification that
// _meta is optional and should be omitted when empty.
func FromMCPMeta(meta *mcp.Meta) map[string]any {
        if meta == nil {
                return nil
        }

        result := make(map[string]any)

        // Merge additional fields first (custom metadata like trace context)
        maps.Copy(result, meta.AdditionalFields)

        // Set progressToken last to ensure it takes precedence over any
        // progressToken key in AdditionalFields (prevents malicious/incorrect overrides)
        if meta.ProgressToken != nil {
                result["progressToken"] = meta.ProgressToken
        }

        // Return nil if the map is empty (no metadata to preserve)
        if len(result) == 0 {
                return nil
        }

        return result
}

// ToMCPMeta converts vmcp meta map to MCP SDK meta for forwarding to clients.
// This reconstructs the _meta field when sending responses back through the MCP protocol.
//
// Reserved io.modelcontextprotocol/* keys are stripped first: vMCP, not the
// backend, is the client's MCP peer, so a backend must not speak for it. This is
// the single chokepoint for every Legacy egress that carries backend _meta
// (serve_handlers, sessionmanager, the elicitation adapter); the Modern path's
// mirror is newModernResultMeta, and both call the same mcpparser helper so the
// two revisions cannot drift.
//
// Note the Legacy/Modern asymmetry: Modern re-stamps its own
// io.modelcontextprotocol/serverInfo after stripping, and Legacy deliberately
// does not — the 2025-11-25 revision has no serverInfo _meta key at all. The
// missing Legacy stamp is correct, not an oversight.
//
// Returns nil if meta is nil or empty, following the MCP specification that
// _meta is optional and should be omitted when empty. A map consisting only of
// reserved keys therefore collapses to nil rather than an empty _meta object.
func ToMCPMeta(meta map[string]any) *mcp.Meta {
        meta = mcpparser.StripReservedMeta(meta)
        if len(meta) == 0 {
                return nil
        }

        result := &mcp.Meta{
                AdditionalFields: make(map[string]any),
        }

        for k, v := range meta {
                if k == "progressToken" {
                        result.ProgressToken = v
                } else {
                        result.AdditionalFields[k] = v
                }
        }

        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package vmcp

import (
        "errors"
        "strings"
)

// Common domain errors used across vmcp subpackages.
// Following DDD principles, domain errors are defined at the package root.
// These errors should be checked using errors.Is().

var (
        // ErrNotFound indicates a requested resource (tool, resource, prompt, workflow) was not found.
        // Wrapping errors should provide specific details about what was not found.
        ErrNotFound = errors.New("not found")

        // ErrInvalidConfig indicates invalid configuration was provided.
        // Wrapping errors should provide specific details about what is invalid.
        ErrInvalidConfig = errors.New("invalid configuration")

        // ErrAuthenticationFailed indicates authentication failure.
        // Wrapping errors should include the underlying authentication error.
        ErrAuthenticationFailed = errors.New("authentication failed")

        // ErrAuthorizationFailed indicates authorization failure.
        // Wrapping errors should include the policy or permission that was denied.
        ErrAuthorizationFailed = errors.New("authorization failed")

        // ErrWorkflowFailed indicates workflow execution failed.
        // Wrapping errors should include the step ID and failure reason.
        ErrWorkflowFailed = errors.New("workflow execution failed")

        // ErrTimeout indicates an operation timed out.
        // Wrapping errors should include the operation type and timeout duration.
        ErrTimeout = errors.New("operation timed out")

        // ErrCancelled indicates an operation was cancelled.
        // Context cancellation should wrap this error with context.Cause().
        ErrCancelled = errors.New("operation cancelled")

        // ErrInvalidInput indicates invalid input parameters.
        // Wrapping errors should specify which parameter is invalid and why.
        ErrInvalidInput = errors.New("invalid input")

        // ErrUnsupportedTransport indicates an unsupported MCP transport type.
        // Wrapping errors should specify which transport type is not supported.
        ErrUnsupportedTransport = errors.New("unsupported transport type")

        // ErrToolExecutionFailed indicates an MCP tool execution failed (domain error).
        // This represents the tool running but returning an error result (IsError=true in MCP).
        // These errors should be forwarded to the client transparently as the LLM needs to see them.
        // Wrapping errors should include the tool name and error message from MCP.
        ErrToolExecutionFailed = errors.New("tool execution failed")

        // ErrBackendUnavailable indicates a backend MCP server is unreachable (operational error).
        // This represents infrastructure issues (network down, server not responding, etc.).
        // These errors may be retried, circuit-broken, or handled differently from domain errors.
        // Wrapping errors should include the backend ID and underlying cause.
        ErrBackendUnavailable = errors.New("backend unavailable")

        // ErrToolNameConflict indicates a composite tool name conflicts with a backend tool name.
        // This prevents ambiguity in routing/execution where the same name could refer to
        // either a backend tool or a composite workflow tool.
        // Wrapping errors should list the conflicting tool names.
        ErrToolNameConflict = errors.New("tool name conflict")
)

// Authorization-denial messages, one per admission kind.
//
// These are kind-only: each names the capability class, never the specific
// tool/resource/prompt or whether it is advertised vs. nonexistent, so a denial
// can never be used to enumerate what exists (the "no oracle" property). They
// live here at the domain root and are referenced by every site that emits a
// denial — conversion.ErrorToToolResult, the Serve resources/read handler, the
// codemode script closure, and the pre-dispatch call gate — so a denial has one
// wording per kind by reference, not by copied literals that can silently drift
// apart. Test assertions deliberately keep the literal to pin the exact wording.
const (
        // DenyMessageToolCall is emitted when a tools/call is denied by policy.
        DenyMessageToolCall = "call denied by authorization policy"
        // DenyMessageResourceRead is emitted when a resources/read is denied by policy.
        DenyMessageResourceRead = "read denied by authorization policy"
        // DenyMessagePromptGet is emitted when a prompts/get is denied by policy.
        DenyMessagePromptGet = "prompt denied by authorization policy"
)

// Error Categorization Helpers
//
// These functions categorize errors by examining error message strings.
// They serve as a fallback mechanism for error detection when:
//
// 1. Errors come from external libraries that use their own error types and formats
// 2. Legacy code paths don't wrap errors with sentinel errors
// 3. Backwards compatibility is needed for error detection
//
// Note: BackendClient now wraps all errors with appropriate sentinel errors
// (ErrAuthenticationFailed, ErrTimeout, ErrBackendUnavailable). Health monitoring
// code should prefer errors.Is() checks over these string-based functions.
// These functions remain for backwards compatibility and as a fallback mechanism.

// IsAuthenticationError checks if an error message indicates an authentication failure.
// Uses case-insensitive pattern matching to detect various auth error formats from
// HTTP libraries, MCP protocol errors, and authentication middleware.
func IsAuthenticationError(err error) bool {
        if err == nil {
                return false
        }

        errLower := strings.ToLower(err.Error())

        // Check for explicit authentication failure messages
        if strings.Contains(errLower, "authentication failed") ||
                strings.Contains(errLower, "authentication error") {
                return true
        }

        // Check for HTTP 401/403 status codes with context.
        // Match patterns like "401 Unauthorized", "HTTP 401", "status code 401".
        // Also match mcp-go's ErrUnauthorized = "unauthorized (401)" which uses
        // reversed order compared to the "401 unauthorized" pattern above.
        if strings.Contains(errLower, "401 unauthorized") ||
                strings.Contains(errLower, "unauthorized (401)") ||
                strings.Contains(errLower, "403 forbidden") ||
                strings.Contains(errLower, "http 401") ||
                strings.Contains(errLower, "http 403") ||
                strings.Contains(errLower, "status code 401") ||
                strings.Contains(errLower, "status code 403") {
                return true
        }

        // Check for explicit unauthenticated/unauthorized errors
        if strings.Contains(errLower, "request unauthenticated") ||
                strings.Contains(errLower, "request unauthorized") ||
                strings.Contains(errLower, "access denied") {
                return true
        }

        // mcp-go ErrAuthorizationRequired ("authorization required") and
        // ErrOAuthAuthorizationRequired ("no valid token available, authorization
        // required") string forms. Used as a fallback when the typed sentinel is
        // not preserved through the error chain (issue #5223).
        if strings.Contains(errLower, "authorization required") {
                return true
        }

        return false
}

// IsTimeoutError checks if an error message indicates a timeout.
// Detects various timeout formats from context deadlines, HTTP timeouts,
// and network timeout errors.
func IsTimeoutError(err error) bool {
        if err == nil {
                return false
        }

        errLower := strings.ToLower(err.Error())
        return strings.Contains(errLower, "timeout") ||
                strings.Contains(errLower, "deadline exceeded") ||
                strings.Contains(errLower, "context deadline exceeded")
}

// IsConnectionError checks if an error message indicates a connection failure.
// Detects network-level errors like connection refused, reset, unreachable, etc.
// Also detects broken pipes, EOF errors, and HTTP 5xx server errors that indicate
// backend unavailability.
func IsConnectionError(err error) bool {
        if err == nil {
                return false
        }

        errStr := err.Error()
        errLower := strings.ToLower(errStr)

        // Check against list of known connection error patterns
        networkPatterns := []string{
                "connection refused", "connection reset", "no route to host",
                "network is unreachable", "broken pipe", "connection closed",
        }
        for _, pattern := range networkPatterns {
                if strings.Contains(errLower, pattern) {
                        return true
                }
        }

        // EOF errors (be specific - check exact case to avoid false positives)
        if strings.Contains(errStr, "EOF") {
                return true
        }

        // HTTP 5xx server errors
        httpErrorPatterns := []string{
                "500 internal server error", "502 bad gateway",
                "503 service unavailable", "504 gateway timeout",
                "status code 500", "status code 502",
                "status code 503", "status code 504",
        }
        for _, pattern := range httpErrorPatterns {
                if strings.Contains(errLower, pattern) {
                        return true
                }
        }

        return false
}

// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package vmcp

import (
        "encoding/json"
        "errors"
        "fmt"
)

// InputRequiredResult carries a Modern (2026-07-28) backend's
// resultType:"input_required" envelope — a Multi Round-Trip Request round
// (SEP-2322): the backend cannot complete the call until the caller fulfills
// InputRequests and retries the original request echoing RequestState.
//
// This is the domain-typed egress surface of MRTR (the seam the client-edge
// limitation in docs/arch/10-virtual-mcp-architecture.md names): the backend
// client decodes the envelope into this type and surfaces it through
// InputRequiredError, so upper layers can relay a round to a Modern
// downstream client or fulfill it in-request for a Legacy one. Full flow
// design in docs/arch/16-vmcp-mrtr.md.
type InputRequiredResult struct {
        // InputRequests maps the backend's server-assigned keys to raw request
        // objects (ElicitRequest, CreateMessageRequest, or ListRootsRequest —
        // schema/draft's InputRequest union). Values are deliberately opaque
        // json.RawMessage: on the pass-through path vMCP relays them verbatim and
        // must not reinterpret them. Use Methods to inspect only the wire method
        // names (e.g. for capability gating). Nil on a load-shedding round
        // (RequestState only).
        InputRequests map[string]json.RawMessage

        // RequestState is the backend's opaque round-trip state. Per SEP-2322 the
        // client "MUST echo back the exact value" on retry — vMCP relays it
        // verbatim and never interprets it. Nil when the backend's envelope had no
        // requestState field, which the retry must preserve: client requirement 2
        // says that if the result does not contain the field, the client "MUST NOT
        // include one in the retry". The schema puts no minimum length on the
        // field, so present-and-empty ("") is legal, distinct from absent, and
        // must be echoed as the empty string.
        RequestState *string
}

// InputRequiredError is the typed error that carries an input_required round
// through the backend client's error return, so every existing
// complete-result path stays untouched. Result is non-nil only for a valid,
// drivable round (see the extraction rules on InputRequiredFromError); for
// any other non-"complete" envelope only the sentinel classification remains.
//
// Sentinel is the classification error this unwraps to — the backend client
// sets its errModernInputRequired so every existing errors.Is check keeps
// working unchanged. The fields are exported so tests (e.g. against
// mocks.MockBackendClient) can construct the error without going through the
// real client's envelope decode.
type InputRequiredError struct {
        // ResultType is the envelope's raw resultType. Backend-controlled text:
        // Error() truncates it, and it must never be interpolated elsewhere
        // unbounded.
        ResultType string

        // Result is the decoded round, nil unless the envelope was a valid
        // input_required round on a method that may carry one.
        Result *InputRequiredResult

        // Sentinel is the error Unwrap returns; required. It leads the Error()
        // message, so the rendered text stays byte-identical to a plain
        // fmt.Errorf("%w: ...", Sentinel, ...) wrapping. Because the field is
        // exported for direct construction, nothing enforces "required" at compile
        // time: a literal that omits it unwraps to nil and renders
        // missingSentinelText in place of the classification.
        Sentinel error
}

// maxResultTypeInError bounds how much of the backend-controlled resultType
// Error() interpolates. Real resultTypes are short tokens ("complete",
// "input_required", "task"); anything longer is hostile or broken, and this
// error's text reaches the downstream client via writeModernDispatchError —
// report a prefix and the length instead of splicing the payload (the #6079 /
// #6066 rule).
const maxResultTypeInError = 64

// missingSentinelText leads Error() when Sentinel is nil, which a direct
// construction that omits the required field can produce. Naming the mistake
// beats letting fmt splice "%!s(<nil>)" into a message that reaches the
// downstream client.
const missingSentinelText = "input_required error with no classification sentinel"

func (e *InputRequiredError) Error() string {
        sentinel := func() string {
                if e.Sentinel == nil {
                        return missingSentinelText
                }
                return e.Sentinel.Error()
        }()
        if len(e.ResultType) > maxResultTypeInError {
                return fmt.Sprintf("%s: resultType=%q... (%d bytes)",
                        sentinel, e.ResultType[:maxResultTypeInError], len(e.ResultType))
        }
        return fmt.Sprintf("%s: resultType=%q", sentinel, e.ResultType)
}

// Unwrap returns the classification sentinel, keeping every existing
// errors.Is check on the untyped sentinel working unchanged.
func (e *InputRequiredError) Unwrap() error { return e.Sentinel }

// InputRequiredFromError extracts the SEP-2322 input_required payload from an
// error chain, however deeply the backend client wrapped it. It reports false
// for errors that are not a drivable input_required round — including the
// unrecognized-resultType, disallowed-method, and invalid-payload variants of
// the same sentinel — so callers can use it as the single MRTR branch point:
//
//        if round, ok := vmcp.InputRequiredFromError(err); ok { ... relay round ... }
//
// This is the seam the ingress half (dispatchModern's input_required envelope
// and the Legacy-client bridge; docs/arch/16-vmcp-mrtr.md slices 2-4) consumes.
func InputRequiredFromError(err error) (*InputRequiredResult, bool) {
        var ire *InputRequiredError
        if errors.As(err, &ire) && ire.Result != nil {
                return ire.Result, true
        }
        return nil, false
}

// Methods returns the JSON-RPC method name of each input request, keyed by
// the backend's request key. An entry whose value does not decode as an
// object with a "method" field maps to the empty string, so callers gating on
// method (capability checks) fail closed on malformed entries instead of
// skipping them.
func (r *InputRequiredResult) Methods() map[string]string {
        if len(r.InputRequests) == 0 {
                return nil
        }
        methods := make(map[string]string, len(r.InputRequests))
        for key, raw := range r.InputRequests {
                var probe struct {
                        Method string `json:"method"`
                }
                // A failed decode leaves probe.Method empty — the fail-closed value.
                _ = json.Unmarshal(raw, &probe)
                methods[key] = probe.Method
        }
        return methods
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package similarity provides vector distance functions for semantic search.
package similarity

import (
        "fmt"
        "math"
)

// CosineSimilarity computes the cosine similarity between two vectors.
// Returns a value in [-1, 1] where 1 means identical direction,
// 0 means orthogonal, and -1 means opposite direction.
//
// Vectors of different lengths return an error. The guard lives here rather
// than at the call sites because the loop below indexes both slices
// positionally: a shorter b would panic and a longer b would silently ignore
// its tail, and a caller that forgets its own check gets one or the other.
func CosineSimilarity(a, b []float32) (float64, error) {
        if len(a) != len(b) {
                return 0, fmt.Errorf("vectors have different dimensions: %d and %d", len(a), len(b))
        }

        var dot, normA, normB float64
        for i := range a {
                ai := float64(a[i])
                bi := float64(b[i])
                dot += ai * bi
                normA += ai * ai
                normB += bi * bi
        }

        denom := math.Sqrt(normA) * math.Sqrt(normB)
        if denom == 0 {
                return 0, nil
        }
        return dot / denom, nil
}

// CosineDistance computes the cosine distance between two vectors.
// Returns a value in [0, 2] where 0 means identical direction and 2 means
// opposite direction. Lower values indicate more similar vectors.
// Vectors of different lengths return an error (see CosineSimilarity).
func CosineDistance(a, b []float32) (float64, error) {
        sim, err := CosineSimilarity(a, b)
        if err != nil {
                return 0, err
        }
        return 1 - sim, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package similarity

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types"
)

// NewEmbeddingClient creates an EmbeddingClient from the given optimizer
// configuration, selecting the backend implementation from EmbeddingProvider.
// It returns (nil, nil) if cfg is nil or no embedding service URL is configured,
// meaning semantic search will be disabled.
func NewEmbeddingClient(cfg *types.OptimizerConfig) (types.EmbeddingClient, error) {
        if cfg == nil || cfg.EmbeddingService == "" {
                return nil, nil
        }

        switch cfg.EmbeddingProvider {
        case "", types.EmbeddingProviderTEI:
                return newTEIClient(cfg.EmbeddingService, cfg.EmbeddingServiceTimeout)
        case types.EmbeddingProviderOpenAI:
                return newOpenAIClient(cfg.EmbeddingService, cfg.EmbeddingModel, cfg.EmbeddingAPIKey,
                        cfg.EmbeddingHeaders, cfg.EmbeddingServiceTimeout)
        default:
                return nil, fmt.Errorf("unsupported embedding provider %q (supported: %q, %q)",
                        cfg.EmbeddingProvider, types.EmbeddingProviderTEI, types.EmbeddingProviderOpenAI)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package similarity

import (
        "bytes"
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "maps"
        "net/http"
        "strings"
        "time"
)

const (
        embeddingsPath = "/embeddings"

        // openAIMaxBatchSize is the OpenAI cap on inputs per /embeddings request;
        // compatible gateways generally honor the same limit.
        openAIMaxBatchSize = 2048
)

// openAIClient implements types.EmbeddingClient against an OpenAI-compatible
// /embeddings API (OpenAI, Azure OpenAI, or another OpenAI-compatible gateway).
type openAIClient struct {
        baseURL      string
        apiKey       string
        model        string
        headers      map[string]string
        httpClient   *http.Client
        maxBatchSize int
}

// newOpenAIClient creates a client that POSTs to baseURL+"/embeddings" using the
// given model. A non-empty apiKey is sent as a Bearer token; an empty apiKey
// omits the Authorization header so keyless endpoints work. headers are set on
// every request but cannot override Content-Type or Authorization. Zero timeout
// uses defaultTimeout.
func newOpenAIClient(baseURL, model, apiKey string, headers map[string]string, timeout time.Duration) (*openAIClient, error) {
        if baseURL == "" {
                return nil, fmt.Errorf("OpenAI embedding base URL is required")
        }
        if model == "" {
                return nil, fmt.Errorf("OpenAI embedding model is required")
        }
        baseURL = strings.TrimSuffix(baseURL, "/")

        if timeout == 0 {
                timeout = defaultTimeout
        }

        slog.Debug("OpenAI embedding client created",
                "base_url", baseURL, "model", model, "timeout", timeout, "custom_headers", len(headers))

        return &openAIClient{
                baseURL:      baseURL,
                apiKey:       apiKey,
                model:        model,
                headers:      maps.Clone(headers),
                httpClient:   &http.Client{Timeout: timeout},
                maxBatchSize: openAIMaxBatchSize,
        }, nil
}

type openAIEmbedRequest struct {
        Model string   `json:"model"`
        Input []string `json:"input"`
        // EncodingFormat pins the response to float arrays, since we decode into
        // []float32; without it a compatible server may return base64.
        EncodingFormat string `json:"encoding_format"`
}

type openAIEmbedResponse struct {
        Data []openAIEmbedding `json:"data"`
}

type openAIEmbedding struct {
        Index     int       `json:"index"`
        Embedding []float32 `json:"embedding"`
}

// Embed returns a vector embedding for the given text.
func (c *openAIClient) Embed(ctx context.Context, text string) ([]float32, error) {
        results, err := c.EmbedBatch(ctx, []string{text})
        if err != nil {
                return nil, err
        }
        if len(results) == 0 {
                return nil, fmt.Errorf("OpenAI returned empty response for single input")
        }
        return results[0], nil
}

// EmbedBatch returns embeddings for multiple texts, chunking to respect the
// OpenAI /embeddings input batch size.
func (c *openAIClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) {
        if len(texts) == 0 {
                return nil, nil
        }

        allEmbeddings := make([][]float32, 0, len(texts))

        for start := 0; start < len(texts); start += c.maxBatchSize {
                end := min(start+c.maxBatchSize, len(texts))
                embeddings, err := c.embedChunk(ctx, texts[start:end])
                if err != nil {
                        return nil, err
                }
                allEmbeddings = append(allEmbeddings, embeddings...)
        }

        slog.Debug("OpenAI embedding batch completed",
                "inputs", len(texts), "chunks", (len(texts)+c.maxBatchSize-1)/c.maxBatchSize,
                "dimensions", len(allEmbeddings[0]))

        return allEmbeddings, nil
}

// embedChunk sends one batch to the /embeddings endpoint and returns the
// embeddings ordered to match texts.
func (c *openAIClient) embedChunk(ctx context.Context, texts []string) ([][]float32, error) {
        bodyBytes, err := json.Marshal(openAIEmbedRequest{Model: c.model, Input: texts, EncodingFormat: "float"})
        if err != nil {
                return nil, fmt.Errorf("failed to marshal OpenAI request: %w", err)
        }

        url := c.baseURL + embeddingsPath
        req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyBytes))
        if err != nil {
                return nil, fmt.Errorf("failed to create OpenAI request: %w", err)
        }
        for name, value := range c.headers {
                req.Header.Set(name, value)
        }
        // Set after the custom headers so they can never be overridden.
        req.Header.Set("Content-Type", "application/json")
        if c.apiKey != "" {
                req.Header.Set("Authorization", "Bearer "+c.apiKey)
        }

        resp, err := c.httpClient.Do(req) // #nosec G704 -- URL is built from the configured embedding base URL
        if err != nil {
                return nil, fmt.Errorf("OpenAI request failed: %w", err)
        }
        defer func() {
                _, _ = io.Copy(io.Discard, resp.Body)
                _ = resp.Body.Close()
        }()

        if resp.StatusCode != http.StatusOK {
                body, _ := io.ReadAll(resp.Body)
                return nil, fmt.Errorf("OpenAI returned status %d: %s", resp.StatusCode, string(body))
        }

        var embedResp openAIEmbedResponse
        if err := json.NewDecoder(resp.Body).Decode(&embedResp); err != nil {
                return nil, fmt.Errorf("failed to decode OpenAI response: %w", err)
        }

        if len(embedResp.Data) != len(texts) {
                return nil, fmt.Errorf("OpenAI returned %d embeddings for %d inputs", len(embedResp.Data), len(texts))
        }

        // Place each embedding at its reported index; the API is free to return
        // entries out of order.
        embeddings := make([][]float32, len(texts))
        for _, d := range embedResp.Data {
                if d.Index < 0 || d.Index >= len(texts) {
                        return nil, fmt.Errorf("OpenAI returned out-of-range embedding index %d for %d inputs", d.Index, len(texts))
                }
                embeddings[d.Index] = d.Embedding
        }

        return embeddings, nil
}

// ModelID returns the configured model name. The OpenAI /embeddings API
// selects the model per request from this same field, so unlike TEI there is
// no server-side state that could drift from it.
func (c *openAIClient) ModelID(context.Context) (string, error) {
        return c.model, nil
}

// Close is a no-op for the OpenAI client.
func (*openAIClient) Close() error {
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package similarity

import (
        "bytes"
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "net/http"
        "time"
)

const (
        // defaultTimeout is the default HTTP timeout for TEI requests.
        defaultTimeout = 30 * time.Second

        // embedPath is the TEI endpoint path for generating embeddings.
        embedPath = "/embed"

        // infoPath is the TEI endpoint that returns server metadata including max batch size.
        infoPath = "/info"

        // defaultMaxBatchSize is used when the TEI /info endpoint does not report a max batch size.
        defaultMaxBatchSize = 32
)

// teiClient implements types.EmbeddingClient by calling the HuggingFace
// Text Embeddings Inference (TEI) HTTP API.
type teiClient struct {
        baseURL      string
        httpClient   *http.Client
        maxBatchSize int
}

// newTEIClient creates a new TEI embedding client that calls the specified endpoint.
// It queries the TEI /info endpoint to discover the server's maximum batch size.
func newTEIClient(baseURL string, timeout time.Duration) (*teiClient, error) {
        if baseURL == "" {
                return nil, fmt.Errorf("TEI BaseURL is required")
        }

        if timeout == 0 {
                timeout = defaultTimeout
        }

        httpClient := &http.Client{Timeout: timeout}

        maxBatch, err := fetchMaxBatchSize(baseURL, httpClient)
        if err != nil {
                slog.Warn("failed to query TEI /info, using default max batch size",
                        "error", err, "default", defaultMaxBatchSize)
                maxBatch = defaultMaxBatchSize
        }

        slog.Debug("TEI embedding client created",
                "base_url", baseURL, "timeout", timeout, "max_batch_size", maxBatch)

        return &teiClient{
                baseURL:      baseURL,
                httpClient:   httpClient,
                maxBatchSize: maxBatch,
        }, nil
}

// teiInfoResponse is a subset of the TEI /info endpoint response.
type teiInfoResponse struct {
        ModelID            string `json:"model_id"`
        MaxClientBatchSize int    `json:"max_client_batch_size"`
}

// fetchInfo queries the TEI /info endpoint.
func fetchInfo(ctx context.Context, baseURL string, httpClient *http.Client) (teiInfoResponse, error) {
        var info teiInfoResponse

        req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+infoPath, nil)
        if err != nil {
                return info, fmt.Errorf("failed to create TEI /info request: %w", err)
        }

        resp, err := httpClient.Do(req) // #nosec G704 -- URL is built from the configured TEI base URL
        if err != nil {
                return info, fmt.Errorf("TEI /info request failed: %w", err)
        }
        defer func() { _ = resp.Body.Close() }()

        if resp.StatusCode != http.StatusOK {
                return info, fmt.Errorf("TEI /info returned status %d", resp.StatusCode)
        }

        if err := json.NewDecoder(resp.Body).Decode(&info); err != nil {
                return info, fmt.Errorf("failed to decode TEI /info response: %w", err)
        }

        return info, nil
}

// fetchMaxBatchSize queries the TEI /info endpoint and returns the max client batch size.
func fetchMaxBatchSize(baseURL string, httpClient *http.Client) (int, error) {
        info, err := fetchInfo(context.Background(), baseURL, httpClient)
        if err != nil {
                return 0, err
        }

        if info.MaxClientBatchSize <= 0 {
                return defaultMaxBatchSize, nil
        }

        return info.MaxClientBatchSize, nil
}

// ModelID returns the id of the model the TEI server is currently running,
// read from /info on every call. The model is a property of the running
// container, not of this client's configuration, so it is deliberately not
// cached: the point is letting a caller observe a redeploy that swapped the
// model behind an unchanged URL.
func (c *teiClient) ModelID(ctx context.Context) (string, error) {
        info, err := fetchInfo(ctx, c.baseURL, c.httpClient)
        if err != nil {
                return "", err
        }
        if info.ModelID == "" {
                return "", fmt.Errorf("TEI /info reported no model_id")
        }
        return info.ModelID, nil
}

// embedRequest is the JSON body sent to the TEI /embed endpoint.
type embedRequest struct {
        Inputs []string `json:"inputs"`
        // Truncate tells the TEI server to silently truncate input texts that
        // exceed the model's maximum token length instead of returning an error.
        // We always set this to true because tool descriptions may exceed model
        // limits and we prefer embedding a truncated description over a request failure.
        Truncate bool `json:"truncate"`
}

// Embed returns a vector embedding for the given text.
func (c *teiClient) Embed(ctx context.Context, text string) ([]float32, error) {
        results, err := c.EmbedBatch(ctx, []string{text})
        if err != nil {
                return nil, err
        }
        if len(results) == 0 {
                return nil, fmt.Errorf("TEI returned empty response for single input")
        }
        return results[0], nil
}

// EmbedBatch returns vector embeddings for multiple texts, automatically
// chunking requests to respect the TEI server's maximum batch size.
func (c *teiClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) {
        if len(texts) == 0 {
                return nil, nil
        }

        allEmbeddings := make([][]float32, 0, len(texts))

        for start := 0; start < len(texts); start += c.maxBatchSize {
                end := min(start+c.maxBatchSize, len(texts))
                chunk := texts[start:end]

                embeddings, err := c.embedChunk(ctx, chunk)
                if err != nil {
                        return nil, err
                }
                allEmbeddings = append(allEmbeddings, embeddings...)
        }

        slog.Debug("TEI embedding batch completed",
                "inputs", len(texts), "chunks", (len(texts)+c.maxBatchSize-1)/c.maxBatchSize,
                "dimensions", len(allEmbeddings[0]))

        return allEmbeddings, nil
}

// embedChunk sends a single batch of texts to the TEI /embed endpoint.
func (c *teiClient) embedChunk(ctx context.Context, texts []string) ([][]float32, error) {
        reqBody := embedRequest{
                Inputs:   texts,
                Truncate: true,
        }
        bodyBytes, err := json.Marshal(reqBody)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal TEI request: %w", err)
        }

        url := c.baseURL + embedPath
        req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyBytes))
        if err != nil {
                return nil, fmt.Errorf("failed to create TEI request: %w", err)
        }
        req.Header.Set("Content-Type", "application/json")

        resp, err := c.httpClient.Do(req) // #nosec G704 -- URL is built from the configured TEI base URL
        if err != nil {
                return nil, fmt.Errorf("TEI request failed: %w", err)
        }
        defer func() { _ = resp.Body.Close() }()

        if resp.StatusCode != http.StatusOK {
                body, _ := io.ReadAll(resp.Body)
                return nil, fmt.Errorf("TEI returned status %d: %s", resp.StatusCode, string(body))
        }

        var embeddings [][]float32
        if err := json.NewDecoder(resp.Body).Decode(&embeddings); err != nil {
                return nil, fmt.Errorf("failed to decode TEI response: %w", err)
        }

        if len(embeddings) != len(texts) {
                return nil, fmt.Errorf("TEI returned %d embeddings for %d inputs", len(embeddings), len(texts))
        }

        return embeddings, nil
}

// Close is a no-op for the TEI client.
func (*teiClient) Close() error {
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package tokencounter provides token estimation for MCP tool definitions.
package tokencounter

import (
        "encoding/json"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
)

// Counter estimates the number of tokens a tool definition would consume
// when sent to an LLM. Implementations may use character-based heuristics or
// real tokenizers.
type Counter interface {
        CountTokens(tool mcp.Tool) int
}

// JSONByteDivisionCounter estimates token count by serialising the full mcp.Tool
// to JSON and dividing the byte length by a configurable divisor.
type JSONByteDivisionCounter struct {
        Divisor int
}

// CountTokens returns len(json(tool)) / divisor.
// Returns 0 if the divisor is zero or serialisation fails.
func (c JSONByteDivisionCounter) CountTokens(tool mcp.Tool) int {
        if c.Divisor <= 0 {
                return 0
        }
        data, err := json.Marshal(tool)
        if err != nil {
                return 0
        }
        return len(data) / c.Divisor
}

// NewJSONByteCounter returns a JSONByteDivisionCounter with a divisor of 4,
// which is a reasonable approximation for most LLM tokenizers.
func NewJSONByteCounter() Counter {
        return JSONByteDivisionCounter{Divisor: 4}
}

// TokenMetrics provides information about token usage optimization.
type TokenMetrics struct {
        // BaselineTokens is the estimated tokens if all tools were sent.
        BaselineTokens int `json:"baseline_tokens"`

        // ReturnedTokens is the actual tokens for the returned tools.
        ReturnedTokens int `json:"returned_tokens"`

        // SavingsPercent is the percentage of tokens saved.
        SavingsPercent float64 `json:"savings_percent"`
}

// ComputeTokenMetrics calculates token savings by comparing the precomputed
// baseline (all tools) against only the matched tool names.
func ComputeTokenMetrics(baselineTokens int, tokenCounts map[string]int, matchedToolNames []string) TokenMetrics {
        if baselineTokens == 0 {
                return TokenMetrics{}
        }

        var returnedTokens int
        for _, name := range matchedToolNames {
                returnedTokens += tokenCounts[name]
        }

        savingsPercent := float64(baselineTokens-returnedTokens) / float64(baselineTokens) * 100

        return TokenMetrics{
                BaselineTokens: baselineTokens,
                ReturnedTokens: returnedTokens,
                SavingsPercent: savingsPercent,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package toolstore implements a SQLite-based ToolStore for search over
// MCP tool metadata. It uses FTS5 for full-text search and optional
// embedding-based semantic search for hybrid retrieval.
package toolstore

import (
        "context"
        "crypto/sha256"
        "database/sql"
        _ "embed"
        "encoding/binary"
        "encoding/hex"
        "encoding/json"
        "errors"
        "fmt"
        "log/slog"
        "math"
        "sort"
        "strings"
        "sync/atomic"

        "golang.org/x/sync/errgroup"
        _ "modernc.org/sqlite" // registers the "sqlite" database/sql driver

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive-core/mcpcompat/server"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types"
)

// Default values for configurable search parameters.
const (
        // DefaultMaxToolsToReturn is the maximum number of results returned to the caller.
        DefaultMaxToolsToReturn = 8

        // DefaultHybridSemanticToolsRatio controls the proportion of semantic vs FTS5
        // results in hybrid mode: 0 = all FTS5, 1 = all semantic.
        DefaultHybridSemanticToolsRatio = 0.5

        // DefaultSemanticDistanceThreshold is the maximum cosine distance for semantic search results.
        // Results with distance > threshold are filtered out in searchSemantic only.
        // Cosine distance: 0 = identical, 2 = opposite.
        DefaultSemanticDistanceThreshold = 1.0
)

//go:embed schema.sql
var schemaSQL string

// sqliteToolStore implements a tool store using SQLite with FTS5 for full-text search
// and optional vector embedding-based semantic search.
// It satisfies the types.ToolStore interface.
type sqliteToolStore struct {
        db *sql.DB

        // keepAlive is a dedicated connection held open for the store's whole lifetime.
        //
        // The database is in-memory with a shared cache, so it exists only while at
        // least one connection to it is open. If the database/sql pool ever drops to
        // zero connections, SQLite discards the entire database — table, FTS5 index
        // and triggers — and nothing recreates it, because the schema is executed
        // only in the constructor. Every later operation then fails with
        // "no such table: llm_capabilities" until the process restarts.
        //
        // The pool does reach zero in practice. modernc.org/sqlite implements context
        // cancellation with sqlite3_interrupt and reports an interrupted connection as
        // unusable from conn.IsValid, so database/sql discards a connection whose
        // statement was canceled instead of returning it to the pool. A single
        // canceled request while the pool holds one connection is enough to destroy
        // the database for the life of the process (#5889).
        //
        // This connection is acquired with a background context and never used to run
        // statements, so it can never be interrupted and never discarded. It exists
        // solely to keep the database alive.
        keepAlive *sql.Conn

        embeddingClient           types.EmbeddingClient // nil = FTS5-only
        maxToolsToReturn          int
        hybridSemanticRatio       float64
        semanticDistanceThreshold float64

        // configIdentity digests the configured embedding provider, endpoint and
        // model. The model id read live from the backend on every build is folded
        // in next to it (see embeddingIdentity), so vectors are not reused across
        // any provider, endpoint, or model change the identity can see — including
        // a TEI container redeployed with a different model behind an unchanged
        // URL, which the config alone cannot. When the id cannot be read the
        // identity falls back and reuse proceeds unverified, within the bounds
        // documented on embeddingIdentity. Immutable after construction.
        configIdentity string

        // lastModelID is the most recent model id successfully read from the
        // embedding backend, or nil before the first successful read. It is the
        // fallback for a failed read: keeping the last-seen id keeps cache keys
        // stable across a transient failure, where switching to a different
        // identity would force a full re-embed and a second one on recovery.
        // Held by pointer because the store is used by value.
        lastModelID *atomic.Pointer[string]

        // embeddingDim is the vector width most recently observed from the
        // embedding backend, or 0 before any embedding call has succeeded. It
        // bounds which stored vectors may be reused, as defense in depth behind
        // the model id in the cache key. Held by pointer because the store is
        // used by value.
        embeddingDim *atomic.Int64
}

// NewSQLiteToolStore creates a new ToolStore backed by a shared in-memory
// SQLite database. All callers of this constructor share the same database,
// which is the intended production behavior (one shared store per server).
// If embeddingClient is non-nil, semantic search is enabled alongside FTS5.
// If cfg is non-nil, its search parameters override the defaults; nil values use defaults.
func NewSQLiteToolStore(embeddingClient types.EmbeddingClient, cfg *types.OptimizerConfig) (types.ToolStore, error) {
        return newSQLiteToolStore("file:memdb?mode=memory&cache=shared", embeddingClient, cfg)
}

// newSQLiteToolStore creates a tool store backed by a database described
// in the connectionString. It is useful for tests, where we want multiple
// isolated (non-shared) databases.
func newSQLiteToolStore(
        connectionString string, embeddingClient types.EmbeddingClient, cfg *types.OptimizerConfig,
) (sqliteToolStore, error) {
        db, err := sql.Open("sqlite", connectionString)
        if err != nil {
                return sqliteToolStore{}, fmt.Errorf("failed to open sqlite database: %w", err)
        }

        // Pin the keep-alive connection before anything else, so the database cannot
        // be discarded from this point on. See the keepAlive field for why.
        keepAlive, err := db.Conn(context.Background())
        if err != nil {
                _ = db.Close()
                return sqliteToolStore{}, fmt.Errorf("failed to acquire keep-alive connection: %w", err)
        }

        // Execute schema
        if _, err := db.Exec(schemaSQL); err != nil {
                _ = keepAlive.Close()
                _ = db.Close()
                return sqliteToolStore{}, fmt.Errorf("failed to initialize schema: %w", err)
        }

        maxTools := DefaultMaxToolsToReturn
        hybridRatio := DefaultHybridSemanticToolsRatio
        semanticThreshold := DefaultSemanticDistanceThreshold
        if cfg != nil {
                if cfg.MaxToolsToReturn != nil {
                        maxTools = *cfg.MaxToolsToReturn
                }
                if cfg.HybridSemanticRatio != nil {
                        hybridRatio = *cfg.HybridSemanticRatio
                }
                if cfg.SemanticDistanceThreshold != nil {
                        semanticThreshold = *cfg.SemanticDistanceThreshold
                }
        }

        store := sqliteToolStore{
                db:                        db,
                keepAlive:                 keepAlive,
                embeddingClient:           embeddingClient,
                maxToolsToReturn:          maxTools,
                hybridSemanticRatio:       hybridRatio,
                semanticDistanceThreshold: semanticThreshold,
                configIdentity:            configIdentity(cfg),
                lastModelID:               &atomic.Pointer[string]{},
                embeddingDim:              &atomic.Int64{},
        }

        slog.Debug("optimizer tool store created",
                "max_tools_to_return", maxTools,
                "hybrid_semantic_ratio", hybridRatio,
                "semantic_distance_threshold", semanticThreshold,
                "semantic_search_enabled", embeddingClient != nil,
        )

        return store, nil
}

// UpsertTools adds or updates tools in the store.
func (s sqliteToolStore) UpsertTools(ctx context.Context, tools []server.ServerTool) (retErr error) {
        // Resolve embeddings before opening the write transaction, so no lock is
        // held across the multi-second embedding call. SQLite in shared-cache mode
        // takes table-level locks: a read inside the transaction would pin a read
        // lock on llm_capabilities, and concurrent builds would then fail to take
        // the write lock with SQLITE_LOCKED, which the busy handler does not retry.
        // See TestSQLiteToolStore_UpsertTools_ConcurrentBuilds.
        embBlobs, hashes, err := s.resolveEmbeddings(ctx, tools)
        if err != nil {
                return err
        }

        tx, err := s.db.BeginTx(ctx, nil)
        if err != nil {
                return fmt.Errorf("failed to begin transaction: %w", err)
        }
        defer func() {
                if retErr != nil {
                        _ = tx.Rollback()
                }
        }()

        stmt, err := tx.PrepareContext(ctx,
                "INSERT OR REPLACE INTO llm_capabilities (name, description, embedding, content_hash) VALUES (?, ?, ?, ?)")
        if err != nil {
                return fmt.Errorf("failed to prepare statement: %w", err)
        }
        defer func() { _ = stmt.Close() }()

        for i, tool := range tools {
                if _, err := stmt.ExecContext(ctx, tool.Tool.Name, tool.Tool.Description, embBlobs[i], hashes[i]); err != nil {
                        return fmt.Errorf("failed to upsert tool %s: %w", tool.Tool.Name, err)
                }
        }

        slog.Debug("upserted tools into store", "count", len(tools))

        return tx.Commit()
}

// maxResolveAttempts bounds how many times one build may restart because the
// embedding model changed under it. Two is enough for the legitimate case of
// a single redeploy; a backend that swaps models on consecutive builds is an
// operational problem no retry count fixes.
const maxResolveAttempts = 2

// resolveEmbeddings returns an encoded embedding blob and a content hash for
// each tool, embedding only the tools whose hash is not already stored.
//
// The Serve path rebuilds a per-session optimizer on every session registration
// and every cross-pod rehydration, each upserting the session's whole tool set.
// Embedding all of it every time costs O(tools x sessions) on the client's
// initialize round-trip; reuse makes it O(tools whose text changed). See
// stacklok/toolhive#5847.
//
// The backend identity is read at the start of an attempt and re-read after
// the embedding batch. If the two differ, the batch spanned a model swap and
// none of its vectors is attributable to either model — each chunk may have
// hit either side of the swap — so the attempt is discarded and re-run under
// the new identity, where the stale rows simply miss by key. Without the
// re-read, vectors produced by the new model would be committed under keys
// naming the old one: wrong vectors under valid-looking keys, which no later
// build would ever re-check. Reused blobs need no such guard: a blob is only
// found under the identity it was committed with, and only identities that
// were verified when the blob was committed ever carry a hash — vectors
// embedded while the id could not be read are committed hashless, searchable
// but never reusable (see the note at the fill loop below).
//
// With no embedding client it returns nil blobs and hashes (FTS5-only mode).
func (s sqliteToolStore) resolveEmbeddings(
        ctx context.Context, tools []server.ServerTool,
) ([][]byte, []sql.NullString, error) {
        blobs := make([][]byte, len(tools))
        hashes := make([]sql.NullString, len(tools))

        if s.embeddingClient == nil {
                return blobs, hashes, nil
        }

        texts := make([]string, len(tools))
        for i, tool := range tools {
                texts[i] = embeddedText(tool.Tool.Name, tool.Tool.Description)
        }

        keys := make([]string, len(tools))
        for attempt := 1; ; attempt++ {
                identity, verified := s.embeddingIdentity(ctx)

                for i := range tools {
                        keys[i] = embeddingCacheKey(identity, texts[i])
                        hashes[i] = sql.NullString{String: keys[i], Valid: true}
                }

                cached, err := s.cachedEmbeddings(ctx, keys)
                if err != nil {
                        return nil, nil, err
                }

                missIndexByKey, missTexts, reused := partitionMisses(blobs, keys, texts, cached)

                slog.Debug("resolved tool embeddings",
                        "tools", len(tools), "reused", reused, "embedded", len(missTexts))

                // All cache hits: nothing was embedded, so there is no batch to have
                // spanned a swap. The reused vectors are correct for the keys they
                // carry by construction.
                if len(missTexts) == 0 {
                        return blobs, hashes, nil
                }

                embeddings, err := s.embedTexts(ctx, missTexts)
                if err != nil {
                        return nil, nil, err
                }

                after, afterVerified := s.embeddingIdentity(ctx)
                if after != identity {
                        if attempt >= maxResolveAttempts {
                                return nil, nil, fmt.Errorf(
                                        "embedding model changed during %d consecutive build attempts; giving up", attempt)
                        }
                        slog.Warn("embedding model changed during the batch; discarding it and re-embedding under the new identity")
                        continue
                }
                verified = verified && afterVerified

                // Published only for a batch that will commit: a discarded attempt's
                // width must not make concurrent builds reject their own valid rows.
                if n := len(embeddings[0]); n > 0 {
                        s.embeddingDim.Store(int64(n))
                }

                if err := fillMisses(blobs, hashes, keys, tools, missIndexByKey, embeddings, verified); err != nil {
                        return nil, nil, err
                }

                return blobs, hashes, nil
        }
}

// fillMisses encodes the freshly embedded vectors into their tools' slots.
// When the attempt's identity was not verified on both sides of the batch,
// the fresh rows are committed hashless — searchable but never reusable. A
// hashed row here would be permanent poison if the backend is later rolled
// back to the fallback model: the next verified build would derive the very
// identity these keys name, cache-hit the mislabelled vectors, and re-key
// nothing.
func fillMisses(
        blobs [][]byte, hashes []sql.NullString, keys []string,
        tools []server.ServerTool, missIndexByKey map[string]int, embeddings [][]float32, verified bool,
) error {
        for i, key := range keys {
                if blobs[i] != nil {
                        continue
                }
                idx, ok := missIndexByKey[key]
                if !ok {
                        // Unreachable, but a missing key would otherwise index 0 and store
                        // another tool's vector under this name.
                        return fmt.Errorf("no embedding resolved for tool %s", tools[i].Tool.Name)
                }
                blobs[i] = encodeEmbedding(embeddings[idx])
                if !verified {
                        hashes[i] = sql.NullString{}
                }
        }
        return nil
}

// partitionMisses fills blobs with the cached vectors and returns the texts
// still to embed, deduplicated by key — the same tool may appear twice in one
// batch — along with the number of reused entries. Entries with no cached
// vector are reset to nil: a retry after a model swap must not carry blobs
// reused under the previous attempt's identity.
func partitionMisses(
        blobs [][]byte, keys, texts []string, cached map[string][]byte,
) (map[string]int, []string, int) {
        missIndexByKey := make(map[string]int, len(keys))
        var missTexts []string
        reused := 0
        for i, key := range keys {
                if blob, ok := cached[key]; ok {
                        blobs[i] = blob
                        reused++
                        continue
                }
                blobs[i] = nil
                if _, seen := missIndexByKey[key]; !seen {
                        missIndexByKey[key] = len(missTexts)
                        missTexts = append(missTexts, texts[i])
                }
        }
        return missIndexByKey, missTexts, reused
}

// embedTexts runs one embedding batch, returning exactly one vector per text.
func (s sqliteToolStore) embedTexts(ctx context.Context, texts []string) ([][]float32, error) {
        embeddings, err := s.embeddingClient.EmbedBatch(ctx, texts)
        if err != nil {
                return nil, fmt.Errorf("failed to generate embeddings: %w", err)
        }
        if len(embeddings) != len(texts) {
                return nil, fmt.Errorf("embedding client returned %d embeddings for %d inputs",
                        len(embeddings), len(texts))
        }
        return embeddings, nil
}

// embeddingIdentity derives the backend identity mixed into every cache key
// for one build, folding the live model id into the configured identity.
//
// The id is read from the backend on every build rather than once at
// construction: the store lives for the whole process and the embedding
// service is addressed by a stable URL, so a TEI container can be redeployed
// with a different model behind an unchanged URL — invisible to the config,
// and to the dimension check when the widths match. A live id turns that
// swap into different cache keys, so stale rows simply stop being found;
// there is nothing to detect and nothing to discard.
//
// Every failure path keeps the previous identity and reports it unverified.
// A read that fails says the backend is unreachable, not that it changed —
// and an unreachable backend cannot re-embed the catalogue either, so
// flipping the identity would turn a working build into a full re-embed now
// and a second one on recovery. Before any successful read it degrades to
// the configured identity alone. Previously verified rows stay reusable
// under an unverified identity; what an unverified identity must never do is
// attribute NEW vectors (see resolveEmbeddings).
func (s sqliteToolStore) embeddingIdentity(ctx context.Context) (identity string, verified bool) {
        modelID, err := s.embeddingClient.ModelID(ctx)
        verified = err == nil && modelID != ""
        if !verified {
                modelID = ""
                if last := s.lastModelID.Load(); last != nil {
                        modelID = *last
                }
                slog.Warn("could not read the embedding model id; proceeding with the last known identity unverified",
                        "error", err, "assumed_model_id", modelID)
        } else {
                s.lastModelID.Store(&modelID)
        }
        // Digested rather than joined so the identity is fixed-length and cannot
        // shift the field boundaries of the cache key it is folded into.
        return hashParts(s.configIdentity, modelID), verified
}

// cachedEmbeddings returns the reusable stored embeddings among the given
// content hashes, keyed by hash. Hashes with no usable vector are absent.
//
// Matching on content_hash rather than tool name confirms in one lookup that
// the embedded text and the producing backend are both unchanged; a changed
// description or a repointed backend cannot quietly reuse a vector. A rename
// changes the text, so it re-embeds. Runs outside any transaction — see the
// lock note in UpsertTools.
//
// A stale-width vector is treated as a miss rather than reused: reuse would be
// permanent, since it is handed back and re-stored on every rebuild while
// search discards it, silently dropping the tool from semantic results.
func (s sqliteToolStore) cachedEmbeddings(ctx context.Context, keys []string) (map[string][]byte, error) {
        keysJSON, err := json.Marshal(keys)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal content hashes: %w", err)
        }

        queryStr := `SELECT content_hash, embedding
                FROM llm_capabilities
                WHERE embedding IS NOT NULL
                  AND content_hash IN (SELECT value FROM json_each(?))`

        rows, err := s.db.QueryContext(ctx, queryStr, string(keysJSON))
        if err != nil {
                return nil, fmt.Errorf("embedding cache lookup failed: %w", err)
        }
        defer func() { _ = rows.Close() }()

        wantBytes := int(s.embeddingDim.Load()) * 4
        cached := make(map[string][]byte, len(keys))
        var unusable int
        for rows.Next() {
                var hash string
                var blob []byte
                if err := rows.Scan(&hash, &blob); err != nil {
                        return nil, fmt.Errorf("failed to scan cached embedding: %w", err)
                }
                if len(blob) == 0 || (wantBytes > 0 && len(blob) != wantBytes) {
                        unusable++
                        continue
                }
                cached[hash] = blob
        }
        if err := rows.Err(); err != nil {
                return nil, err
        }

        if unusable > 0 {
                slog.Warn("ignoring unusable stored embeddings, they will be recomputed",
                        "count", unusable, "expected_bytes", wantBytes)
        }

        return cached, nil
}

// embeddedText builds the exact string sent to the embedding backend. The
// content hash is only meaningful if it covers precisely what was embedded, so
// callers must not construct this string themselves.
func embeddedText(name, description string) string {
        return fmt.Sprintf("name: %s description: %s", name, description)
}

// cacheKeyVersion invalidates every stored key if the embedded-text format changes.
const cacheKeyVersion = "v1"

// embeddingCacheKey hashes the embedded text together with the identity of the
// backend that produced it.
//
// The backend identity is what makes reuse safe: an embedding is interchangeable
// only with one from the same provider, endpoint, and model. Without it,
// repointing the service would silently serve vectors from a different semantic
// space.
func embeddingCacheKey(identity, text string) string {
        return hashParts(cacheKeyVersion, identity, text)
}

// hashParts hashes an ordered list of strings so that distinct lists can never
// produce the same digest.
//
// Each part is length-prefixed rather than delimiter-joined. A delimiter alone
// is ambiguous — moving the separator between two adjacent parts yields an
// identical byte stream, so ("a\x00b", "c") and ("a", "b\x00c") would collide.
// That matters because tool names and descriptions reach this function from
// aggregated backend servers, so a backend could otherwise craft a description
// that collides with another tool's key and take over its stored embedding.
func hashParts(parts ...string) string {
        h := sha256.New()
        var length [8]byte
        for _, part := range parts {
                binary.LittleEndian.PutUint64(length[:], uint64(len(part)))
                h.Write(length[:])
                h.Write([]byte(part))
        }
        return hex.EncodeToString(h.Sum(nil))
}

// configIdentity digests the configured half of the backend identity; the
// model id read live per build completes it (see embeddingIdentity).
func configIdentity(cfg *types.OptimizerConfig) string {
        if cfg == nil {
                return ""
        }
        // Digested rather than joined so the identity is fixed-length and cannot
        // shift the field boundaries of the cache key it is folded into.
        return hashParts(cfg.EmbeddingProvider, cfg.EmbeddingService, cfg.EmbeddingModel)
}

// Search finds tools matching q using FTS5 full-text search and optional
// semantic search when an embedding client is configured. The lexical arm
// prefers q.Keywords, falling back to q.Description when Keywords is empty;
// the semantic arm always embeds q.Description.
// The allowedTools parameter limits results to only tools with names in the given set.
// If allowedTools is empty, no results are returned (empty = no access).
// Returns matches ranked by relevance.
func (s sqliteToolStore) Search(ctx context.Context, q types.SearchQuery, allowedTools []string) ([]mcp.Tool, error) {
        if len(allowedTools) == 0 {
                slog.Debug("search skipped, no allowed tools")
                return nil, nil
        }

        // The lexical arm prefers explicit keywords and falls back to the words of
        // the description, either when no keywords were supplied or when every
        // keyword was dropped as too common to discriminate. The semantic arm
        // always embeds the description.
        ftsExpr := sanitizeFTS5Terms(q.Keywords)
        if ftsExpr == "" {
                ftsExpr = sanitizeFTS5Terms(strings.Fields(q.Description))
        }

        // FTS5-only path (no embedding client)
        if s.embeddingClient == nil {
                if ftsExpr == "" {
                        slog.Debug("search skipped, empty FTS5 expression", "query", q.Description, "keywords", q.Keywords)
                        return nil, nil
                }
                results, err := s.searchFTS5(ctx, ftsExpr, allowedTools, s.maxToolsToReturn)
                if err != nil {
                        return nil, err
                }
                slog.Debug("search completed (FTS5-only)",
                        "query", q.Description, "keywords", q.Keywords, "results", len(results), "matched_tools", matchNames(results))
                return results, nil
        }

        // Hybrid search: derive per-method limits from the ratio.
        ftsLimit, semanticLimit := hybridSearchLimits(s.maxToolsToReturn, s.hybridSemanticRatio)

        g, gCtx := errgroup.WithContext(ctx)

        var ftsResults []mcp.Tool
        if ftsExpr != "" && ftsLimit > 0 {
                g.Go(func() error {
                        var err error
                        ftsResults, err = s.searchFTS5(gCtx, ftsExpr, allowedTools, ftsLimit)
                        return err
                })
        }

        var semanticResults []mcp.Tool
        if semanticLimit > 0 {
                g.Go(func() error {
                        var err error
                        semanticResults, err = s.searchSemantic(gCtx, q.Description, allowedTools, semanticLimit)
                        return err
                })
        }

        if err := g.Wait(); err != nil {
                return nil, err
        }

        merged := mergeResults(ftsResults, semanticResults, s.maxToolsToReturn)

        slog.Debug("search completed (hybrid)",
                "query", q.Description,
                "keywords", q.Keywords,
                "fts5_results", len(ftsResults),
                "semantic_results", len(semanticResults),
                "merged_results", len(merged),
                "matched_tools", matchNames(merged),
        )

        return merged, nil
}

// Close releases the underlying database connections. Releasing the keep-alive
// connection drops the in-memory database, so it happens only here, on the way
// to closing the pool itself.
//
// Close may be called more than once: a keep-alive connection that was already
// released reports sql.ErrConnDone, which is not a failure.
func (s sqliteToolStore) Close() error {
        var embErr error
        if s.embeddingClient != nil {
                embErr = s.embeddingClient.Close()
        }

        var keepAliveErr error
        if s.keepAlive != nil {
                if err := s.keepAlive.Close(); err != nil && !errors.Is(err, sql.ErrConnDone) {
                        keepAliveErr = err
                }
        }

        dbErr := s.db.Close()
        return errors.Join(embErr, keepAliveErr, dbErr)
}

// searchFTS5 performs a full-text search using FTS5 MATCH with BM25 ranking.
// It uses json_each() to pass the allowed tool names as a single JSON array
// parameter, avoiding manual placeholder construction.
//
// The limit parameter caps results per this method. In hybrid mode, FTS5 and
// semantic search each independently return their top-k results (split by
// hybridSemanticToolsRatio). A tool with a low BM25 rank won't be missed if
// it has high cosine similarity, because the semantic query runs separately
// and will surface it.
//
// The ftsExpr is produced by sanitizeFTS5Terms and is always passed as a
// parameterized ? value, never interpolated into SQL.
func (s sqliteToolStore) searchFTS5(
        ctx context.Context, ftsExpr string, allowedTools []string, limit int,
) ([]mcp.Tool, error) {
        allowedJSON, err := json.Marshal(allowedTools)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal allowed tools: %w", err)
        }

        queryStr := `SELECT t.name, t.description, rank
                FROM llm_capabilities_fts fts
                JOIN llm_capabilities t ON t.rowid = fts.rowid
                WHERE llm_capabilities_fts MATCH ?
                  AND t.name IN (SELECT value FROM json_each(?))
                ORDER BY rank
                LIMIT ?`

        rows, err := s.db.QueryContext(ctx, queryStr, ftsExpr, string(allowedJSON), limit)
        if err != nil {
                return nil, fmt.Errorf("FTS5 query failed: %w", err)
        }
        defer func() { _ = rows.Close() }()

        var matches []mcp.Tool
        for rows.Next() {
                var name, description string
                var rank float64
                if err := rows.Scan(&name, &description, &rank); err != nil {
                        return nil, fmt.Errorf("failed to scan row: %w", err)
                }
                matches = append(matches, mcp.Tool{
                        Name:        name,
                        Description: description,
                })
        }

        if err := rows.Err(); err != nil {
                return nil, err
        }

        slog.Debug("FTS5 search completed",
                "fts_expression", ftsExpr,
                "allowed_tools", len(allowedTools),
                "limit", limit,
                "results", len(matches),
                "matched_tools", matchNames(matches),
        )

        return matches, nil
}

// searchSemantic performs embedding-based semantic search.
// It embeds the query, loads all candidate embeddings from the database,
// computes cosine distance, and returns the closest matches.
//
// This runs as a separate query from searchFTS5 because BM25 rank and cosine
// similarity are fundamentally different metrics that cannot be meaningfully
// combined in a single SQL query. BM25 rank is a hidden FTS5 column computed
// on-the-fly from term frequency, while cosine similarity requires loading
// embedding blobs and computing distances in Go. Merging happens afterward
// in mergeResults, which deduplicates and keeps the best score per tool.
//
//nolint:unparam // limit kept for API consistency with searchFTS5
func (s sqliteToolStore) searchSemantic(
        ctx context.Context, query string, allowedTools []string, limit int,
) ([]mcp.Tool, error) {
        queryVec, err := s.embeddingClient.Embed(ctx, query)
        if err != nil {
                return nil, fmt.Errorf("failed to embed query: %w", err)
        }
        // A build whose tools are all cache hits never embeds anything, so the query
        // round-trip is the only place a model swap is still observable.
        if len(queryVec) > 0 {
                s.embeddingDim.Store(int64(len(queryVec)))
        }

        allowedJSON, err := json.Marshal(allowedTools)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal allowed tools: %w", err)
        }

        queryStr := `SELECT name, description, embedding
                FROM llm_capabilities
                WHERE embedding IS NOT NULL
                  AND name IN (SELECT value FROM json_each(?))`

        rows, err := s.db.QueryContext(ctx, queryStr, string(allowedJSON))
        if err != nil {
                return nil, fmt.Errorf("semantic query failed: %w", err)
        }
        defer func() { _ = rows.Close() }()

        type rankedMatch struct {
                name        string
                description string
                dist        float64
        }

        var ranked []rankedMatch
        var candidatesEvaluated, dimensionMismatches int
        for rows.Next() {
                var name, description string
                var embBlob []byte
                if err := rows.Scan(&name, &description, &embBlob); err != nil {
                        return nil, fmt.Errorf("failed to scan row: %w", err)
                }

                candidatesEvaluated++
                emb := decodeEmbedding(embBlob)

                // A width mismatch means the vector survived a model change (see
                // embeddingIdentity); skip it rather than fail the whole search.
                dist, err := similarity.CosineDistance(queryVec, emb)
                if err != nil {
                        dimensionMismatches++
                        continue
                }

                // Filter by semantic distance threshold.
                // This is meaningful only for cosine distance (semantic search).
                // FTS5 ranks are normalized BM25 scores, not true distance measures.
                if dist > s.semanticDistanceThreshold {
                        continue
                }

                ranked = append(ranked, rankedMatch{
                        name:        name,
                        description: description,
                        dist:        dist,
                })
        }

        if err := rows.Err(); err != nil {
                return nil, err
        }

        // Sort by distance ascending (lower = better match)
        sort.Slice(ranked, func(i, j int) bool {
                return ranked[i].dist < ranked[j].dist
        })

        if len(ranked) > limit {
                ranked = ranked[:limit]
        }

        matches := make([]mcp.Tool, len(ranked))
        for i, r := range ranked {
                matches[i] = mcp.Tool{
                        Name:        r.name,
                        Description: r.description,
                }
        }

        if dimensionMismatches > 0 {
                slog.Warn("skipped stored embeddings with a mismatched dimension",
                        "count", dimensionMismatches, "query_dimensions", len(queryVec))
        }

        slog.Debug("semantic search completed",
                "allowed_tools", len(allowedTools),
                "limit", limit,
                "candidates_evaluated", candidatesEvaluated,
                "dimension_mismatches", dimensionMismatches,
                "results", len(matches),
                "matched_tools", matchNames(matches),
        )

        return matches, nil
}

// mergeResults combines semantic and FTS5 results, deduplicating by name.
// Semantic results are listed first (preserving their distance-based order),
// followed by FTS5 results not already present, and truncated to maxResults.
func mergeResults(fts, semantic []mcp.Tool, maxResults int) []mcp.Tool {
        seen := make(map[string]struct{}, len(fts)+len(semantic))
        merged := make([]mcp.Tool, 0, len(fts)+len(semantic))

        // Semantic results first.
        for _, m := range semantic {
                if _, ok := seen[m.Name]; ok {
                        continue
                }
                seen[m.Name] = struct{}{}
                merged = append(merged, m)
        }

        // Then FTS5 results not already seen.
        for _, m := range fts {
                if _, ok := seen[m.Name]; ok {
                        continue
                }
                seen[m.Name] = struct{}{}
                merged = append(merged, m)
        }

        if len(merged) > maxResults {
                merged = merged[:maxResults]
        }

        return merged
}

// matchNames extracts tool names from a slice of ToolMatch results for logging.
func matchNames(matches []mcp.Tool) []string {
        names := make([]string, len(matches))
        for i, m := range matches {
                names[i] = m.Name
        }
        return names
}

// problematicWords contains words that FTS5 interprets as operators or that
// are too common in tool metadata to be useful search terms.
//
// sanitizeFTS5Terms drops them. The alternative, falling back to a phrase
// search over the whole query when one of these appears, demands exact token
// adjacency and so reliably matches nothing; dropping the noise word and
// OR-joining the rest avoids the result flood without going empty-handed.
var problematicWords = map[string]struct{}{
        "name": {}, "description": {}, "schema": {}, "input": {},
        "output": {}, "type": {}, "properties": {}, "required": {},
        "title": {}, "id": {}, "tool": {}, "server": {},
        "meta": {}, "data": {}, "content": {}, "text": {},
        "value": {}, "field": {}, "column": {}, "table": {},
        "index": {}, "key": {}, "primary": {},
}

// sanitizeFTS5Terms builds an FTS5 MATCH expression from a list of search terms.
//
// Each entry is split on whitespace (a single term may arrive multi-word), words
// in problematicWords are dropped, the remainder is deduplicated case-insensitively,
// and what is left is OR-joined. Returns "" when no usable term remains, letting
// the caller fall back to another source of terms.
//
// The returned string is designed to be passed as a single ? parameter to
// QueryContext. FTS5 MATCH requires a single string operand containing the full
// query expression (e.g., "read" OR "write"). Individual terms cannot be separate
// ? SQL parameters because the OR/AND operators are part of the FTS5 query
// language, not SQL.
// See: https://sqlite.org/fts5.html#full_text_query_syntax
//
// Safety:
//   - SQL injection is prevented because the expression is always bound via ?.
//   - FTS5 operator injection is prevented by double-quoting each term and
//     escaping embedded double-quotes (standard FTS5 escaping).
func sanitizeFTS5Terms(terms []string) string {
        seen := make(map[string]struct{}, len(terms))
        var quoted []string
        for _, term := range terms {
                for _, word := range strings.Fields(term) {
                        lower := strings.ToLower(word)
                        if _, ok := problematicWords[lower]; ok {
                                continue
                        }
                        if _, ok := seen[lower]; ok {
                                continue
                        }
                        seen[lower] = struct{}{}
                        escaped := strings.ReplaceAll(word, `"`, `""`)
                        quoted = append(quoted, `"`+escaped+`"`)
                }
        }

        return strings.Join(quoted, " OR ")
}

// hybridSearchLimits computes the per-method result limits for hybrid search
// from the total limit and the semantic ratio (0 = all FTS5, 1 = all semantic).
func hybridSearchLimits(total int, semanticRatio float64) (ftsLimit, semanticLimit int) {
        semanticLimit = int(math.Round(float64(total) * semanticRatio))
        ftsLimit = total - semanticLimit
        return ftsLimit, semanticLimit
}

// encodeEmbedding serializes a float32 slice to a little-endian byte slice.
func encodeEmbedding(vec []float32) []byte {
        buf := make([]byte, len(vec)*4)
        for i, v := range vec {
                binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(v))
        }
        return buf
}

// decodeEmbedding deserializes a little-endian byte slice to a float32 slice.
func decodeEmbedding(buf []byte) []float32 {
        vec := make([]float32, len(buf)/4)
        for i := range vec {
                vec[i] = math.Float32frombits(binary.LittleEndian.Uint32(buf[i*4:]))
        }
        return vec
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package optimizer provides the Optimizer interface for intelligent tool discovery
// and invocation in the Virtual MCP Server.
//
// When the optimizer is enabled, vMCP exposes only two tools to clients:
//   - find_tool: Semantic search over available tools
//   - call_tool: Dynamic invocation of any backend tool
//
// This reduces token usage by avoiding the need to send all tool definitions
// to the LLM, instead allowing it to discover relevant tools on demand.
package optimizer

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "maps"
        "os"
        "slices"
        "strconv"
        "strings"
        "time"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive-core/mcpcompat/server"
        httpval "github.com/stacklok/toolhive-core/validation/http"
        vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/tokencounter"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/toolstore"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types"
)

// embeddingAPIKeyEnvVar holds the bearer token for an OpenAI-compatible
// embedding service. It is an env var, not a config field, so the secret never
// lands in a CRD spec or ConfigMap.
// #nosec G101 -- This is an environment variable name, not a hardcoded credential
const embeddingAPIKeyEnvVar = "OPENAI_API_KEY"

// Config defines configuration options for the Optimizer.
// It is defined in the internal/types package and aliased here so that
// external consumers continue to use optimizer.Config.
type Config = types.OptimizerConfig

// GetAndValidateConfig validates the CRD-compatible OptimizerConfig and converts it
// to the internal optimizer.Config with parsed, typed values.
// Returns (nil, nil) if cfg is nil.
func GetAndValidateConfig(cfg *vmcpconfig.OptimizerConfig) (*Config, error) {
        if cfg == nil {
                return nil, nil
        }

        optCfg := &Config{
                EmbeddingService:        cfg.EmbeddingService,
                EmbeddingServiceTimeout: time.Duration(cfg.EmbeddingServiceTimeout),
                EmbeddingProvider:       cfg.EmbeddingProvider,
                EmbeddingModel:          cfg.EmbeddingModel,
                EmbeddingHeaders:        convertEmbeddingHeaders(cfg.EmbeddingHeaders),
        }

        if err := resolveEmbeddingProvider(optCfg); err != nil {
                return nil, err
        }

        if cfg.MaxToolsToReturn != 0 {
                if cfg.MaxToolsToReturn < 1 || cfg.MaxToolsToReturn > 50 {
                        return nil, fmt.Errorf("optimizer.maxToolsToReturn must be between 1 and 50, got %d", cfg.MaxToolsToReturn)
                }
                optCfg.MaxToolsToReturn = &cfg.MaxToolsToReturn
        }

        if cfg.HybridSearchSemanticRatio != "" {
                ratio, err := strconv.ParseFloat(cfg.HybridSearchSemanticRatio, 64)
                if err != nil {
                        return nil, fmt.Errorf("optimizer.hybridSearchSemanticRatio must be a valid number: %w", err)
                }
                if ratio < 0 || ratio > 1 {
                        return nil, fmt.Errorf(
                                "optimizer.hybridSearchSemanticRatio must be between 0.0 and 1.0, got %s",
                                cfg.HybridSearchSemanticRatio,
                        )
                }
                optCfg.HybridSemanticRatio = &ratio
        }

        if cfg.SemanticDistanceThreshold != "" {
                threshold, err := strconv.ParseFloat(cfg.SemanticDistanceThreshold, 64)
                if err != nil {
                        return nil, fmt.Errorf("optimizer.semanticDistanceThreshold must be a valid number: %w", err)
                }
                if threshold < 0 || threshold > 2 {
                        return nil, fmt.Errorf(
                                "optimizer.semanticDistanceThreshold must be between 0.0 and 2.0, got %s",
                                cfg.SemanticDistanceThreshold,
                        )
                }
                optCfg.SemanticDistanceThreshold = &threshold
        }

        return optCfg, nil
}

// resolveEmbeddingProvider normalizes and validates the embedding provider on
// optCfg in place. An empty provider defaults to TEI so existing configs keep
// working; the OpenAI provider requires a service and model, reads its API
// key from the environment, and is the only provider that accepts custom
// embedding headers.
func resolveEmbeddingProvider(optCfg *Config) error {
        switch optCfg.EmbeddingProvider {
        case "":
                optCfg.EmbeddingProvider = types.EmbeddingProviderTEI
        case types.EmbeddingProviderTEI:
        case types.EmbeddingProviderOpenAI:
                if optCfg.EmbeddingService == "" {
                        return fmt.Errorf("optimizer.embeddingService is required when optimizer.embeddingProvider is %q",
                                types.EmbeddingProviderOpenAI)
                }
                if optCfg.EmbeddingModel == "" {
                        return fmt.Errorf("optimizer.embeddingModel is required when optimizer.embeddingProvider is %q",
                                types.EmbeddingProviderOpenAI)
                }
                if err := validateEmbeddingHeaders(optCfg.EmbeddingHeaders); err != nil {
                        return err
                }
                optCfg.EmbeddingAPIKey = os.Getenv(embeddingAPIKeyEnvVar)
        default:
                return fmt.Errorf("optimizer.embeddingProvider must be %q or %q, got %q",
                        types.EmbeddingProviderTEI, types.EmbeddingProviderOpenAI, optCfg.EmbeddingProvider)
        }

        // Defense in depth: mirrors the CEL rule on config.OptimizerConfig,
        // covering config sources with no admission validation.
        if optCfg.EmbeddingProvider != types.EmbeddingProviderOpenAI && len(optCfg.EmbeddingHeaders) > 0 {
                return fmt.Errorf("optimizer.embeddingHeaders is only supported when optimizer.embeddingProvider is %q",
                        types.EmbeddingProviderOpenAI)
        }

        return nil
}

// convertEmbeddingHeaders converts the config header map to the internal
// plain-string representation, returning nil for an empty map.
func convertEmbeddingHeaders(headers map[string]vmcpconfig.EmbeddingHeaderValue) map[string]string {
        if len(headers) == 0 {
                return nil
        }
        out := make(map[string]string, len(headers))
        for name, value := range headers {
                out[name] = string(value)
        }
        return out
}

// validateEmbeddingHeaders rejects custom embedding headers with invalid
// RFC 7230 names or values, and headers the OpenAI client sets itself:
// Content-Type is always application/json, and Authorization is derived from
// the OPENAI_API_KEY environment variable so the token never lands in config.
// Reserved names are compared case-insensitively, matching HTTP semantics.
// Mirrors the CEL rule on config.OptimizerConfig as defense in depth.
func validateEmbeddingHeaders(headers map[string]string) error {
        for name, value := range headers {
                if err := httpval.ValidateHeaderName(name); err != nil {
                        return fmt.Errorf("optimizer.embeddingHeaders: invalid header name %q: %w", name, err)
                }
                if err := httpval.ValidateHeaderValue(value); err != nil {
                        return fmt.Errorf("optimizer.embeddingHeaders[%q]: %w", name, err)
                }
                switch strings.ToLower(name) {
                case "authorization":
                        return fmt.Errorf("optimizer.embeddingHeaders must not set %q: the Authorization header is derived "+
                                "from the %s environment variable", name, embeddingAPIKeyEnvVar)
                case "content-type":
                        return fmt.Errorf("optimizer.embeddingHeaders must not set %q: the Content-Type header is always "+
                                "application/json", name)
                }
        }
        return nil
}

// Optimizer defines the interface for intelligent tool discovery and invocation.
//
// The default implementation delegates search to a ToolStore (SQLite FTS5 with
// optional embedding-based semantic search) and scopes results to the tools
// registered for each session.
type Optimizer interface {
        // FindTool searches for tools matching the given description and keywords.
        // Returns matching tools ranked by relevance.
        FindTool(ctx context.Context, input FindToolInput) (*FindToolOutput, error)

        // CallTool invokes a tool by name with the given parameters.
        // Returns the tool's result or an error if the tool is not found or execution fails.
        // Returns the MCP CallToolResult directly from the underlying tool handler.
        CallTool(ctx context.Context, input CallToolInput) (*mcp.CallToolResult, error)
}

// FindToolInput contains the parameters for finding tools.
type FindToolInput struct {
        // ToolDescription is a natural language description of the tool to find.
        //nolint:lll // Long description tag provides essential context for LLM tool usage.
        ToolDescription string `json:"tool_description" description:"Description of the task or capability needed (e.g. 'web search', 'analyze CSV file', 'send an email'). This is used for semantic similarity matching against available tools."`

        // ToolKeywords is an optional list of keywords to narrow the search.
        //nolint:lll // Long description tag provides essential context for LLM tool usage.
        ToolKeywords []string `json:"tool_keywords,omitempty" description:"Optional keywords driving the BM25 keyword-search arm (e.g. ['list', 'issues', 'github'] or ['SQL', 'query', 'postgres']). Semantic matching always uses tool_description, so provide a complete description even when supplying keywords."`
}

// FindToolOutput contains the results of a tool search.
type FindToolOutput struct {
        // Tools contains the matching tools, ranked by relevance.
        Tools []mcp.Tool `json:"tools"`

        // TokenMetrics provides information about token savings from using the optimizer.
        TokenMetrics TokenMetrics `json:"token_metrics"`
}

// TokenMetrics provides information about token usage optimization.
// It is defined in the internal/tokencounter package and aliased here so that
// external consumers continue to use optimizer.TokenMetrics.
type TokenMetrics = tokencounter.TokenMetrics

// CallToolInput contains the parameters for calling a tool.
type CallToolInput struct {
        // ToolName is the name of the tool to invoke.
        //nolint:lll // Long description tag provides essential context for LLM tool usage.
        ToolName string `json:"tool_name" description:"The name of the tool to execute (obtain this from find_tool results - it is the tool's name field)"`

        // Parameters are the arguments to pass to the tool.
        //nolint:lll // Long description tag provides essential context for LLM tool usage.
        Parameters map[string]any `json:"parameters" description:"Dictionary of arguments required by the tool. The structure must match the tool's input schema as returned by find_tool."`
}

// callToolArgToolName is the parameters key resolveCallToolTarget hoists a nested
// name out of. It must match the json tag on CallToolInput.ToolName, since the
// nested lookup is a map index while the top-level one goes through encoding/json;
// TestCallToolArgToolNameMatchesStructTag guards the pair against drift.
const callToolArgToolName = "tool_name"

// resolveCallToolTarget resolves the tool a call_tool invocation targets,
// accepting the common LLM malformation where tool_name is nested inside
// parameters instead of sitting alongside it. A top-level name always wins, so a
// backend tool with its own tool_name argument still works.
//
// It is deliberately unexported: decoding a payload into a CallToolInput is the
// only supported way to learn which tool a call_tool request names. Reading the
// name out of a raw arguments map instead misses the case-variant keys
// encoding/json accepts, and a target that authorization and dispatch disagree on
// is a tool executing under a policy decision made for a different name.
//
// params is never modified; a copy is returned when a nested name is hoisted.
func resolveCallToolTarget(name string, params map[string]any) (string, map[string]any) {
        if name != "" {
                return name, params
        }
        nested, ok := params[callToolArgToolName].(string)
        if !ok || nested == "" {
                return name, params
        }
        hoisted := maps.Clone(params)
        delete(hoisted, callToolArgToolName)
        return nested, hoisted
}

// UnmarshalJSON hoists a nested tool_name so dispatch targets the same tool
// authorization approved. See resolveCallToolTarget.
func (in *CallToolInput) UnmarshalJSON(data []byte) error {
        type rawCallToolInput CallToolInput // drops the method set to avoid recursion
        var raw rawCallToolInput
        if err := json.Unmarshal(data, &raw); err != nil {
                return err
        }
        raw.ToolName, raw.Parameters = resolveCallToolTarget(raw.ToolName, raw.Parameters)
        *in = CallToolInput(raw)
        return nil
}

// NewOptimizerFactory creates the embedding client and SQLite tool store from
// the given OptimizerConfig, then returns an OptimizerFactory and a cleanup
// function that closes the store. The caller must invoke the cleanup function
// during shutdown to release resources.
func NewOptimizerFactory(cfg *Config) (
        func(context.Context, []server.ServerTool) (Optimizer, error),
        func(context.Context) error,
        error,
) {
        embClient, err := similarity.NewEmbeddingClient(cfg)
        if err != nil {
                return nil, nil, fmt.Errorf("failed to create embedding client: %w", err)
        }

        store, err := toolstore.NewSQLiteToolStore(embClient, cfg)
        if err != nil {
                return nil, nil, fmt.Errorf("failed to create optimizer store: %w", err)
        }

        factory := newOptimizerFactoryWithStore(store, tokencounter.NewJSONByteCounter())
        cleanup := func(_ context.Context) error {
                return store.Close()
        }

        slog.Debug("optimizer factory created",
                "embedding_service", cfg.EmbeddingService,
                "semantic_search_enabled", embClient != nil,
        )

        return factory, cleanup, nil
}

// toolOptimizer implements the Optimizer interface using a shared ToolStore
// for search and a local handler map for tool invocation.
//
// It delegates search to the ToolStore (which uses SQLite FTS5 with optional
// embedding-based semantic search) and scopes results to only the tools this
// instance was created with.
type toolOptimizer struct {
        // store is the shared tool store used for search.
        store types.ToolStore

        // tools contains all available tools indexed by name.
        tools map[string]server.ServerTool

        // toolNames is the precomputed list of tool names from the tools map.
        // Immutable after construction; avoids re-allocation on every FindTool call.
        toolNames []string

        // tokenCounts holds precomputed per-tool token estimates, indexed by tool name.
        // Immutable after construction: token counts are computed once in newToolOptimizer
        // and never modified. The tools are fixed per session (one optimizer per session),
        // and the tokencounter.Counter is set at configuration time, so counts cannot change at runtime.
        tokenCounts map[string]int

        // baselineTokens is the precomputed sum of all per-tool token counts.
        // Immutable after construction; used as the denominator for savings metrics.
        baselineTokens int
}

// newToolOptimizer creates a new toolOptimizer backed by the given ToolStore.
//
// The tools slice should contain all backend tools (as ServerTool with handlers).
// Tools are upserted into the shared store and scoped for this optimizer instance.
// Token counts are precomputed using the provided counter for metrics calculation.
func newToolOptimizer(
        ctx context.Context, store types.ToolStore, counter tokencounter.Counter, tools []server.ServerTool,
) (Optimizer, error) {
        toolMap := make(map[string]server.ServerTool, len(tools))
        names := make([]string, 0, len(tools))
        tokenCounts := make(map[string]int, len(tools))
        var baselineTokens int
        for _, tool := range tools {
                toolMap[tool.Tool.Name] = tool
                names = append(names, tool.Tool.Name)
                tc := counter.CountTokens(tool.Tool)
                tokenCounts[tool.Tool.Name] = tc
                baselineTokens += tc
        }

        if err := store.UpsertTools(ctx, tools); err != nil {
                return nil, fmt.Errorf("failed to upsert tools into store: %w", err)
        }

        slog.Debug("optimizer session created",
                "tools", len(tools),
                "baseline_tokens", baselineTokens,
        )

        return &toolOptimizer{
                store:          store,
                tools:          toolMap,
                toolNames:      names,
                tokenCounts:    tokenCounts,
                baselineTokens: baselineTokens,
        }, nil
}

// FindTool searches for tools using the shared ToolStore, scoped to this instance's tools.
//
// TokenMetrics quantify the token savings from returning only matching tools
// instead of the full set of available tools.
func (d *toolOptimizer) FindTool(ctx context.Context, input FindToolInput) (*FindToolOutput, error) {
        if input.ToolDescription == "" {
                return nil, fmt.Errorf("tool_description is required")
        }

        matches, err := d.store.Search(ctx, types.SearchQuery{
                Description: input.ToolDescription,
                Keywords:    input.ToolKeywords,
        }, d.toolNames)
        if err != nil {
                return nil, fmt.Errorf("tool search failed: %w", err)
        }

        // Enrich each match with the full tool from the in-memory map.
        // The store only returns Name and Description; replacing with the full
        // mcp.Tool gives us InputSchema, OutputSchema, Annotations, etc.
        for i, m := range matches {
                if tool, ok := d.tools[m.Name]; ok {
                        matches[i] = tool.Tool
                }
        }

        matchedNames := make([]string, len(matches))
        for i, m := range matches {
                matchedNames[i] = m.Name
        }
        metrics := tokencounter.ComputeTokenMetrics(d.baselineTokens, d.tokenCounts, matchedNames)

        slog.Debug("find_tool completed",
                "query", input.ToolDescription,
                "keywords", input.ToolKeywords,
                "results", len(matches),
                "baseline_tokens", metrics.BaselineTokens,
                "returned_tokens", metrics.ReturnedTokens,
                "savings_percent", metrics.SavingsPercent,
        )

        return &FindToolOutput{
                Tools:        matches,
                TokenMetrics: metrics,
        }, nil
}

// CallTool invokes a tool by name using its registered handler.
//
// The tool is looked up by exact name match. If found, the handler
// is invoked directly with the given parameters.
func (d *toolOptimizer) CallTool(ctx context.Context, input CallToolInput) (*mcp.CallToolResult, error) {
        if input.ToolName == "" {
                return nil, fmt.Errorf(
                        `tool_name is required: call_tool expects {"tool_name": "<name from find_tool>", `+
                                `"parameters": {<tool arguments>}}, got parameters keys %v`,
                        slices.Sorted(maps.Keys(input.Parameters)))
        }

        // Verify the tool exists
        tool, exists := d.tools[input.ToolName]
        if !exists {
                slog.Debug("call_tool failed, tool not found", "tool", input.ToolName)
                return mcp.NewToolResultError(fmt.Sprintf("tool not found: %s", input.ToolName)), nil
        }

        slog.Debug("call_tool invoking backend tool", "tool", input.ToolName)

        // Build the MCP request
        request := mcp.CallToolRequest{}
        request.Params.Name = input.ToolName
        request.Params.Arguments = input.Parameters

        // Call the tool handler directly
        return tool.Handler(ctx, request)
}

// newOptimizerFactoryWithStore returns an OptimizerFactory that creates
// toolOptimizer instances backed by the given ToolStore. All optimizers created
// by the returned factory share the same store, enabling cross-session search.
func newOptimizerFactoryWithStore(
        store types.ToolStore, counter tokencounter.Counter,
) func(context.Context, []server.ServerTool) (Optimizer, error) {
        return func(ctx context.Context, tools []server.ServerTool) (Optimizer, error) {
                return newToolOptimizer(ctx, store, counter, tools)
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package vmcp

import (
        "context"
        "fmt"
        "sync"
)

// BackendRegistry provides thread-safe access to discovered backends.
// This is a shared kernel interface used across vmcp bounded contexts
// (aggregator, router, health monitoring).
//
// The registry serves as the single source of truth for backend information
// during the lifecycle of a virtual MCP server instance. It supports both
// immutable (Phase 1) and mutable (future phases) implementations.
//
// Design Philosophy:
//   - Phase 1: Immutable registry (backends discovered once, never change)
//   - Future: Mutable registry with health monitoring and dynamic updates
//   - Thread-safe for concurrent reads across all implementations
//   - Implementations may support concurrent writes with appropriate locking
//
//go:generate mockgen -destination=mocks/mock_registry.go -package=mocks -source=registry.go BackendRegistry DynamicRegistry
type BackendRegistry interface {
        // Get retrieves a backend by ID.
        // Returns nil if the backend is not found.
        // This method is safe for concurrent reads.
        //
        // Example:
        //   backend := registry.Get(ctx, "github-mcp")
        //   if backend == nil {
        //       return fmt.Errorf("backend not found")
        //   }
        Get(ctx context.Context, backendID string) *Backend

        // List returns all registered backends.
        // The returned slice is a fresh, non-nil snapshot the caller owns: safe to
        // iterate without additional locking and safe to mutate (implementations must
        // not return an internal or shared backing array). Order is not guaranteed
        // unless specified by the implementation.
        //
        // Example:
        //   backends := registry.List(ctx)
        //   for _, backend := range backends {
        //       fmt.Printf("Backend: %s\n", backend.Name)
        //   }
        List(ctx context.Context) []Backend

        // Count returns the number of registered backends.
        // This is more efficient than len(List()) for large registries.
        Count() int
}

// immutableRegistry is a Phase 1 implementation that stores a static list
// of backends discovered at startup. It's thread-safe for concurrent reads
// and never changes after construction.
//
// Use NewImmutableRegistry() to create instances.
type immutableRegistry struct {
        // backends maps backend ID to backend information.
        // This map is built once at construction and never modified.
        backends map[string]Backend
}

// NewImmutableRegistry creates a registry from a static list of backends.
//
// This implementation is used in Phase 1 where backends are discovered once
// at startup and don't change during the virtual MCP server's lifetime.
// The registry is thread-safe for concurrent reads.
//
// Parameters:
//   - backends: List of discovered backends to register
//
// Returns:
//   - BackendRegistry: An immutable registry instance
//
// Example:
//
//        backends := discoverer.Discover(ctx, "engineering-team")
//        registry := vmcp.NewImmutableRegistry(backends)
//        backend := registry.Get(ctx, "github-mcp")
func NewImmutableRegistry(backends []Backend) BackendRegistry {
        reg := &immutableRegistry{
                backends: make(map[string]Backend, len(backends)),
        }
        for _, b := range backends {
                reg.backends[b.ID] = b
        }
        return reg
}

// Get retrieves a backend by ID from the immutable registry.
// Returns nil if the backend is not found.
func (r *immutableRegistry) Get(_ context.Context, backendID string) *Backend {
        if b, exists := r.backends[backendID]; exists {
                // Return a copy to prevent external modifications
                return &b
        }
        return nil
}

// List returns all registered backends as a slice.
// The order is not guaranteed. The returned slice is a copy and safe to modify.
func (r *immutableRegistry) List(_ context.Context) []Backend {
        backends := make([]Backend, 0, len(r.backends))
        for _, b := range r.backends {
                backends = append(backends, b)
        }
        return backends
}

// Count returns the number of registered backends.
func (r *immutableRegistry) Count() int {
        return len(r.backends)
}

// DynamicRegistry extends BackendRegistry with mutable operations.
// This implementation supports thread-safe backend updates with version-based
// cache invalidation for dynamic discovery mode.
//
// The registry maintains a monotonic version counter that increments on every
// mutation (Upsert/Remove). This enables lazy cache invalidation in the
// discovery manager without thundering herd problems.
//
// Design Philosophy:
//   - Thread-safe for concurrent reads and writes using RWMutex
//   - Idempotent operations (Upsert/Remove safe to call multiple times)
//   - Version-based cache invalidation (no event callbacks needed)
//   - Used in dynamic discovery mode for K8s-aware vMCP servers
type DynamicRegistry interface {
        BackendRegistry

        // Upsert adds or updates a backend atomically.
        // Idempotent - calling with the same backend multiple times is safe.
        // Increments Version() on every call.
        //
        // Parameters:
        //   - backend: The backend to add or update
        //
        // Returns:
        //   - error: Returns error if backend has empty ID
        //
        // Example:
        //   err := registry.Upsert(vmcp.Backend{
        //       ID: "github-mcp",
        //       Name: "GitHub MCP",
        //       BaseURL: "http://github-mcp.default.svc.cluster.local:8080",
        //   })
        Upsert(backend Backend) error

        // Remove deletes a backend by ID.
        // Idempotent - removing non-existent backends returns nil.
        // Increments Version() on every call.
        //
        // Parameters:
        //   - backendID: The ID of the backend to remove
        //
        // Returns:
        //   - error: Always returns nil (operation is always successful)
        //
        // Example:
        //   err := registry.Remove("github-mcp")
        Remove(backendID string) error

        // Version returns the current registry version.
        // Increments on every mutation (Upsert/Remove).
        // Used for cache invalidation in discovery manager.
        //
        // Returns:
        //   - uint64: Monotonic version counter
        //
        // Example:
        //   version := registry.Version()
        //   // Cache entries tagged with this version
        Version() uint64
}

// dynamicRegistry is a mutable implementation that supports thread-safe
// backend updates with version tracking for cache invalidation.
//
// Use NewDynamicRegistry() to create instances.
type dynamicRegistry struct {
        mu       sync.RWMutex
        backends map[string]Backend
        version  uint64
}

// NewDynamicRegistry creates a new mutable registry with optional initial backends.
//
// This implementation is used in dynamic discovery mode where backends can change
// during the virtual MCP server's lifetime (e.g., K8s-aware vMCP servers).
// The registry is thread-safe for concurrent reads and writes.
//
// Version Tracking:
// The registry starts at version 0 regardless of the number of initial backends.
// Initial backends are considered the baseline state, not mutations. This ensures:
//   - Cache coherence: Discovery manager caches capabilities with version 0 at startup
//   - Consistency: All servers starting with same backends have same initial version
//   - Predictability: Version only increments for runtime changes (Upsert/Remove)
//
// Version increments only occur when backends are added/removed AFTER initialization,
// which triggers cache invalidation in the discovery manager.
//
// Parameters:
//   - backends: Optional list of initial backends to register
//
// Returns:
//   - DynamicRegistry: A mutable registry instance starting at version 0
//
// Example:
//
//        // Start with 2 backends, version = 0
//        registry := vmcp.NewDynamicRegistry([]Backend{backend1, backend2})
//        // Add a third backend, version = 1 (triggers cache invalidation)
//        err := registry.Upsert(backend3)
func NewDynamicRegistry(backends []Backend) DynamicRegistry {
        reg := &dynamicRegistry{
                backends: make(map[string]Backend),
                version:  0, // Initial state is version 0, regardless of backend count
        }
        for _, b := range backends {
                // Store by value - Go will automatically make a copy when storing in the map
                reg.backends[b.ID] = b
        }
        return reg
}

// Get retrieves a backend by ID from the dynamic registry.
// Returns nil if the backend is not found.
// Thread-safe for concurrent access.
func (r *dynamicRegistry) Get(_ context.Context, backendID string) *Backend {
        r.mu.RLock()
        defer r.mu.RUnlock()

        if b, exists := r.backends[backendID]; exists {
                // Go automatically makes a copy when retrieving from map of values
                return &b
        }
        return nil
}

// List returns all registered backends as a slice.
// The order is not guaranteed. The returned slice is a copy and safe to modify.
// Thread-safe for concurrent access.
func (r *dynamicRegistry) List(_ context.Context) []Backend {
        r.mu.RLock()
        defer r.mu.RUnlock()

        backends := make([]Backend, 0, len(r.backends))
        for _, b := range r.backends {
                // Go automatically makes a copy when iterating over map of values
                backends = append(backends, b)
        }
        return backends
}

// Count returns the number of registered backends.
// Thread-safe for concurrent access.
func (r *dynamicRegistry) Count() int {
        r.mu.RLock()
        defer r.mu.RUnlock()

        return len(r.backends)
}

// Upsert adds or updates a backend atomically.
//
// Idempotent: Calling with the same backend multiple times is safe (no error).
// Version Behavior: Increments version on EVERY call, even if backend data is identical.
//
// Design Trade-off:
// This implementation prioritizes simplicity over cache efficiency by always incrementing
// the version. This means identical updates will trigger cache invalidation.
//
// Rationale:
//   - Simpler implementation (no deep equality checks)
//   - Conservative correctness (won't miss actual changes)
//   - Version semantics: tracks mutations (calls), not changes (data diffs)
//
// Impact on Cache Efficiency:
// High-frequency updates with identical data (e.g., health status polling) will cause
// cache invalidation even when nothing changed. For such scenarios, consider:
//  1. Track dynamic state (like health) separately from backend membership
//  2. Implement equality check before incrementing version (if cache efficiency is critical)
//  3. Rate-limit updates at the source to reduce mutation frequency
//
// Example:
//
//        registry.Upsert(backend1)  // version = 1
//        registry.Upsert(backend1)  // identical data, but version = 2 (cache invalidated)
func (r *dynamicRegistry) Upsert(backend Backend) error {
        if backend.ID == "" {
                return fmt.Errorf("backend ID cannot be empty")
        }

        r.mu.Lock()
        defer r.mu.Unlock()

        // Go automatically makes a copy when passing by value and storing in the map
        r.backends[backend.ID] = backend
        r.version++ // Always increment - see design trade-off in godoc

        return nil
}

// Remove deletes a backend by ID atomically.
//
// Idempotent: Removing non-existent backends returns nil (no error).
// Version Behavior: Increments version on EVERY call, even if the backend doesn't exist.
//
// Design Trade-off:
// This implementation prioritizes simplicity and consistency over cache efficiency by always
// incrementing the version. This means duplicate removals will trigger cache invalidation.
//
// Rationale:
//   - Consistent semantics with Upsert (both always increment)
//   - Simpler implementation (no existence checks needed)
//   - Conservative correctness (won't miss actual changes)
//   - Predictable behavior (version always tracks operation calls, not state changes)
//
// Impact on Cache Efficiency:
// If the K8s watcher or reconciliation loop calls Remove multiple times for the same backend
// (e.g., due to event replays or duplicate notifications), each call will increment the version
// and invalidate cached capabilities even though the registry state hasn't changed.
//
// For scenarios where duplicate removals are common, consider:
//  1. Deduplicate remove operations at the watcher level
//  2. Check existence before calling Remove (if cache efficiency is critical)
//  3. Implement equality check before incrementing version (adds complexity)
//
// Example:
//
//        registry.Remove("backend1")  // backend exists, version = 1
//        registry.Remove("backend1")  // backend already gone, but version = 2 (cache invalidated)
func (r *dynamicRegistry) Remove(backendID string) error {
        r.mu.Lock()
        defer r.mu.Unlock()

        delete(r.backends, backendID)
        r.version++

        return nil
}

// Version returns the current registry version.
// Increments on every mutation (Upsert/Remove).
// Thread-safe for concurrent access.
func (r *dynamicRegistry) Version() uint64 {
        r.mu.RLock()
        defer r.mu.RUnlock()

        return r.version
}

// BackendToTarget converts a Backend to a BackendTarget for routing.
// This helper is used when populating routing tables during capability aggregation.
//
// The BackendTarget contains all information needed to forward requests to
// a specific backend workload, including authentication strategy and metadata.
func BackendToTarget(backend *Backend) *BackendTarget {
        if backend == nil {
                return nil
        }

        return &BackendTarget{
                WorkloadID:      backend.ID,
                WorkloadName:    backend.Name,
                BaseURL:         backend.BaseURL,
                TransportType:   backend.TransportType,
                CABundlePath:    backend.CABundlePath,
                CABundleData:    backend.CABundleData,
                AuthConfig:      backend.AuthConfig,
                SessionAffinity: false, // TODO: Add session affinity support in future phases
                HealthStatus:    backend.HealthStatus,
                HeaderForward:   backend.HeaderForward,
                Metadata:        backend.Metadata,
        }
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package schema

// arraySchema represents an array type with item schema.
type arraySchema struct {
        items TypeCoercer
}

// makeArraySchema creates an arraySchema from a raw JSON Schema map.
func makeArraySchema(raw map[string]any) arraySchema {
        schema := arraySchema{}

        items, ok := raw["items"].(map[string]any)
        if ok {
                schema.items = MakeSchema(items)
        } else {
                schema.items = passthroughSchema{}
        }

        return schema
}

// TryCoerce coerces array elements to their expected types.
// Each element is coerced independently according to the items schema.
//
// Partial coercion behavior: If coercion fails for one element, other
// elements are still coerced. Failed elements retain their original value.
//
// Returns a new slice; the original is not modified.
func (s arraySchema) TryCoerce(value any) any {
        arr, ok := value.([]any)
        if !ok {
                return value
        }

        result := make([]any, len(arr))
        for i, elem := range arr {
                result[i] = s.items.TryCoerce(elem)
        }

        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package schema

// objectSchema represents a JSON Schema object type.
type objectSchema struct {
        properties map[string]TypeCoercer
}

// makeObjectSchema creates an objectSchema from a raw JSON Schema map.
func makeObjectSchema(raw map[string]any) objectSchema {
        schema := objectSchema{
                properties: make(map[string]TypeCoercer),
        }

        properties, ok := raw["properties"].(map[string]any)
        if !ok {
                return schema
        }

        for propName, propSchema := range properties {
                propMap, ok := propSchema.(map[string]any)
                if !ok {
                        continue
                }

                // Recursively create schema for this property
                schema.properties[propName] = MakeSchema(propMap)

        }

        return schema
}

// TryCoerce coerces object properties to their expected types.
// Each property is coerced independently according to its schema.
// Properties without a schema pass through unchanged.
//
// Partial coercion behavior: If coercion fails for one property, other
// properties are still coerced. Failed properties retain their original value.
// For example, given {"count": "abc", "enabled": "true"} with integer/boolean
// schema, the result is {"count": "abc", "enabled": true} - count stays as
// string (invalid integer) while enabled is coerced to boolean.
//
// Returns a new map; the original is not modified.
func (s objectSchema) TryCoerce(value any) any {
        obj, ok := value.(map[string]any)
        if !ok {
                return value
        }

        result := make(map[string]any, len(obj))
        for key, val := range obj {
                if propSchema, exists := s.properties[key]; exists {
                        result[key] = propSchema.TryCoerce(val)
                } else {
                        result[key] = val
                }
        }

        return result
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package schema

import (
        "log/slog"
        "strconv"
)

// primitiveSchema represents string/integer/number/boolean types.
type primitiveSchema struct {
        schemaType string
}

// makePrimitiveSchema creates a primitiveSchema from a raw JSON Schema map.
func makePrimitiveSchema(schemaType string) primitiveSchema {
        schema := primitiveSchema{
                schemaType: schemaType,
        }

        return schema
}

// TryCoerce coerces a value to the expected primitive type.
// String values are converted to integer/number/boolean as specified.
// Non-string values pass through unchanged.
// Returns the original value if coercion fails.
func (s primitiveSchema) TryCoerce(value any) any {
        str, ok := value.(string)
        if !ok {
                // Non-string values pass through unchanged
                return value
        }

        switch s.schemaType {
        case "string":
                return value

        case "integer":
                v, err := strconv.ParseInt(str, 10, 64)
                if err != nil {
                        slog.Debug("failed to coerce to integer", "value", str, "error", err)
                        return value
                }
                return v

        case "number":
                v, err := strconv.ParseFloat(str, 64)
                if err != nil {
                        slog.Debug("failed to coerce to number", "value", str, "error", err)
                        return value
                }
                return v

        case "boolean":
                b, err := strconv.ParseBool(str)
                if err != nil {
                        slog.Debug("failed to coerce to boolean", "value", str, "error", err)
                        return value
                }
                return b
        }

        return value
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package schema

import (
        "encoding/json"
        "fmt"
        "reflect"
        "strings"
)

// GenerateSchema generates a JSON Schema from a Go struct type using reflection.
//
// The function inspects struct tags to determine:
//   - json: Field name in the schema (uses json tag name)
//   - description: Field description (from `description` tag)
//   - omitempty: Whether the field is optional (not in required array)
//
// Supported types:
//   - string -> {"type": "string"}
//   - int, int64, etc. -> {"type": "integer"}
//   - float64, float32 -> {"type": "number"}
//   - bool -> {"type": "boolean"}
//   - []T -> {"type": "array", "items": {...}}
//   - map[string]any -> {"type": "object"}
//   - struct -> {"type": "object", "properties": {...}}
//
// Example:
//
//        type FindToolInput struct {
//            ToolDescription string   `json:"tool_description" description:"Natural language description"`
//            ToolKeywords    []string `json:"tool_keywords,omitempty" description:"Optional keywords"`
//        }
//        schema := GenerateSchema[FindToolInput]()
func GenerateSchema[T any]() (map[string]any, error) {
        var zero T
        t := reflect.TypeOf(zero)
        return generateSchemaForType(t)
}

// generateSchemaForType generates schema for a reflect.Type.
func generateSchemaForType(t reflect.Type) (map[string]any, error) {
        // Handle pointer types
        if t.Kind() == reflect.Pointer {
                t = t.Elem()
        }

        switch t.Kind() {
        case reflect.Struct:
                return generateObjectSchema(t)
        case reflect.String:
                return map[string]any{"type": "string"}, nil
        case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
                reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
                reflect.Uintptr:
                return map[string]any{"type": "integer"}, nil
        case reflect.Float32, reflect.Float64:
                return map[string]any{"type": "number"}, nil
        case reflect.Bool:
                return map[string]any{"type": "boolean"}, nil
        case reflect.Slice, reflect.Array:
                items, err := generateSchemaForType(t.Elem())
                if err != nil {
                        return nil, err
                }
                return map[string]any{
                        "type":  "array",
                        "items": items,
                }, nil
        case reflect.Map:
                // For map[string]any, just return object type
                return map[string]any{"type": "object"}, nil
        case reflect.Pointer, reflect.Interface,
                reflect.UnsafePointer, reflect.Chan, reflect.Func,
                reflect.Complex64, reflect.Complex128,
                reflect.Invalid:
                return nil, fmt.Errorf("unsupported type: %s", t.Kind())
        default:
                // Should never happen, but appease the linter.
                return nil, fmt.Errorf("unsupported type: %s", t.Kind())
        }
}

// generateObjectSchema generates schema for a struct type.
func generateObjectSchema(t reflect.Type) (map[string]any, error) {
        properties := make(map[string]any)
        var required []string

        for i := range t.NumField() {
                field := t.Field(i)

                // Skip unexported fields
                if !field.IsExported() {
                        continue
                }

                // Get JSON tag for field name
                jsonTag := field.Tag.Get("json")
                if jsonTag == "-" {
                        continue
                }

                // Parse json tag (name,omitempty)
                jsonName, isOptional := parseJSONTag(jsonTag)
                if jsonName == "" {
                        jsonName = field.Name
                }

                // Generate schema for field type
                fieldSchema, err := generateSchemaForType(field.Type)
                if err != nil {
                        return nil, err
                }

                // Add description if present
                if desc := field.Tag.Get("description"); desc != "" {
                        fieldSchema["description"] = desc
                }

                properties[jsonName] = fieldSchema

                // Add to required if not optional
                if !isOptional {
                        required = append(required, jsonName)
                }
        }

        schema := map[string]any{
                "type":       "object",
                "properties": properties,
        }

        if len(required) > 0 {
                schema["required"] = required
        }

        return schema, nil
}

// parseJSONTag parses a json struct tag and returns the field name and whether it's optional.
func parseJSONTag(tag string) (name string, optional bool) {
        if tag == "" {
                return "", false
        }

        parts := strings.Split(tag, ",")
        name = parts[0]

        for _, part := range parts[1:] {
                if part == "omitempty" {
                        optional = true
                }
        }

        return name, optional
}

// Translate converts an untyped input (typically map[string]any from MCP request arguments)
// to a typed struct using JSON marshalling/unmarshalling.
//
// This provides a simple, reliable way to convert MCP tool arguments to typed Go structs
// without manual field-by-field extraction.
//
// Example:
//
//        args := request.Params.Arguments // map[string]any
//        input, err := Translate[FindToolInput](args)
//        if err != nil {
//            return nil, fmt.Errorf("invalid arguments: %w", err)
//        }
func Translate[T any](input any) (T, error) {
        var result T

        // Marshal to JSON
        data, err := json.Marshal(input)
        if err != nil {
                return result, fmt.Errorf("failed to marshal input: %w", err)
        }

        // Unmarshal to typed struct
        if err := json.Unmarshal(data, &result); err != nil {
                return result, fmt.Errorf("failed to unmarshal to %T: %w", result, err)
        }

        return result, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package schema provides typed JSON Schema structures for type coercion and defaults.
//
// This package wraps raw JSON Schema maps (map[string]any) into typed structures
// that provide type-safe operations like coercion and default value application.
package schema

// TypeCoercer coerces values to their expected types.
// Implementations return the coerced value, or the original if coercion fails.
type TypeCoercer interface {
        TryCoerce(value any) any
}

// MakeSchema parses a raw JSON Schema map into typed Schema structures.
// Always returns a valid TypeCoercer; callers do not need to nil-check.
// For nil, empty, or unknown schema types, returns a passthrough coercer
// that returns values unchanged.
func MakeSchema(raw map[string]any) TypeCoercer {
        if len(raw) == 0 {
                return passthroughSchema{}
        }

        schemaType, _ := raw["type"].(string)

        switch schemaType {
        case "object":
                return makeObjectSchema(raw)
        case "array":
                return makeArraySchema(raw)
        case "string", "integer", "number", "boolean":
                return makePrimitiveSchema(schemaType)
        default:
                return passthroughSchema{}
        }
}

// passthroughSchema is a no-op schema that returns values unchanged.
type passthroughSchema struct{}

// TryCoerce returns the value unchanged.
func (passthroughSchema) TryCoerce(value any) any {
        return value
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package binding preserves the vMCP binding API, delegating the canonical
// format to pkg/auth/sessionbinding. An identity binding encodes the principal that
// created a session as a single opaque string suitable for storage in a
// key/value store such as Redis or Valkey.
//
// # Format
//
// A bound identity is encoded as iss + "\x00" + sub. NUL is rejected from
// either half by Format and Parse: OIDC Core does not formally forbid NUL in
// sub, but no real-world issuer emits one, and accepting it would let a
// corrupted or adversarial value re-split during Parse.
//
// Sessions created before any auth middleware ran use the literal
// "unauthenticated" sentinel. The sentinel must not contain '\x00' so it
// cannot collide with any bound form; any future format change must keep
// the sentinel and the bound-form value sets non-overlapping.
//
// # Trust boundary
//
// Bindings are stored plaintext at rest. They are PII but not credentials —
// they identify but do not authenticate. They carry no freshness signal (no
// exp, no nonce) and are NOT a substitute for token validation: callers must
// only compare a stored binding against a freshly-validated (iss, sub) pair
// from the current request's token.
package binding

import "github.com/stacklok/toolhive/pkg/auth/sessionbinding"

// UnauthenticatedSentinel is the binding value stored for sessions that were
// created without an authenticated identity (auth middleware not present or
// identity nil).
const UnauthenticatedSentinel = sessionbinding.UnauthenticatedSentinel

// ErrInvalidBinding is returned by Format when either input is empty or
// contains a NUL byte.
var ErrInvalidBinding = sessionbinding.ErrInvalidBinding

// Format returns the canonical on-the-wire form of an identity binding:
// iss + "\x00" + sub. Returns ErrInvalidBinding when either input is empty
// or contains a NUL byte.
func Format(iss, sub string) (string, error) {
        return sessionbinding.Format(iss, sub)
}

// Parse splits an on-the-wire binding into its (iss, sub) components.
// Returns ok=true only when s contains exactly one NUL and both halves are
// non-empty. Returns ok=false for the unauthenticated sentinel, for malformed
// input, and for empty strings. Callers must check ok; the empty-string
// return values are not meaningful when ok=false.
func Parse(s string) (iss, sub string, ok bool) {
        return sessionbinding.Parse(s)
}

// IsUnauthenticated reports whether s is the literal unauthenticated sentinel.
func IsUnauthenticated(s string) bool {
        return s == UnauthenticatedSentinel
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package optimizerdec provides a MultiSession decorator that replaces the
// full tool list with two optimizer tools: find_tool and call_tool.
package optimizerdec

import (
        "context"
        "encoding/json"
        "fmt"

        "github.com/stacklok/toolhive-core/mcpcompat/mcp"
        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/vmcp"
        "github.com/stacklok/toolhive/pkg/vmcp/conversion"
        "github.com/stacklok/toolhive/pkg/vmcp/optimizer"
        "github.com/stacklok/toolhive/pkg/vmcp/schema"
        sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types"
)

const (
        // FindToolName is the tool name for semantic tool discovery.
        FindToolName = "find_tool"
        // CallToolName is the tool name for routing a call to any backend tool.
        CallToolName = "call_tool"
)

// Pre-generated schemas for find_tool and call_tool, computed at init time.
var (
        findToolInputSchema = mustGenerateSchema[optimizer.FindToolInput]()
        callToolInputSchema = mustGenerateSchema[optimizer.CallToolInput]()
)

// optimizerDecorator wraps a MultiSession to expose only find_tool and call_tool.
// Tools() returns only those two tools. CallTool("find_tool") routes through the
// optimizer's FindTool; CallTool("call_tool") routes through the optimizer's
// CallTool so that all optimizer telemetry (traces, metrics) is recorded.
type optimizerDecorator struct {
        sessiontypes.MultiSession
        opt            optimizer.Optimizer
        optimizerTools []vmcp.Tool
}

// NewDecorator wraps sess with optimizer mode. Only find_tool and call_tool are
// exposed via Tools(). find_tool calls opt.FindTool. call_tool calls opt.CallTool,
// which routes through the instrumented optimizer (telemetry, traces, metrics).
func NewDecorator(sess sessiontypes.MultiSession, opt optimizer.Optimizer) sessiontypes.MultiSession {
        return &optimizerDecorator{
                MultiSession:   sess,
                opt:            opt,
                optimizerTools: OptimizerTools(),
        }
}

// OptimizerTools returns the find_tool and call_tool meta-tool definitions (name,
// description, input schema) that replace the full backend tool list in optimizer
// mode. The definitions are shared so that the legacy MultiSession decorator (this
// package) and the Serve-path optimizer wiring (pkg/vmcp/server) advertise an
// identical pair; each consumer wires its own handlers around these definitions.
// A fresh slice is returned on every call so callers cannot mutate shared state.
func OptimizerTools() []vmcp.Tool {
        return []vmcp.Tool{
                {
                        Name: FindToolName,
                        Description: "Find and return tools that can help accomplish the user's request. " +
                                "This searches available MCP server tools using semantic and keyword-based matching. " +
                                "Use this function when you need to: " +
                                "(1) discover what tools are available for a specific task, " +
                                "(2) find the right tool(s) before attempting to solve a problem, " +
                                "(3) check if required functionality exists in the current environment. " +
                                "Returns matching tools ranked by relevance including their names, descriptions, " +
                                "required parameters and schemas, plus token efficiency metrics showing " +
                                "baseline_tokens, returned_tokens, and savings_percent. " +
                                "Always call this before call_tool to discover the correct tool name and parameter schema.",
                        InputSchema: findToolInputSchema,
                },
                {
                        Name: CallToolName,
                        Description: "Execute a specific tool with the provided parameters. " +
                                "Use this function to run a tool after identifying it with find_tool. " +
                                "Important: always use find_tool first to get the correct tool_name " +
                                "and parameter schema before calling this function.",
                        InputSchema: callToolInputSchema,
                },
        }
}

// Tools returns only find_tool and call_tool, replacing the full backend tool list.
// A defensive copy is returned so callers cannot mutate the decorator's internal slice.
func (d *optimizerDecorator) Tools() []vmcp.Tool {
        result := make([]vmcp.Tool, len(d.optimizerTools))
        copy(result, d.optimizerTools)
        return result
}

// CallTool handles find_tool and call_tool. Both route through the optimizer so
// that all optimizer telemetry is recorded. Any other tool name returns an error.
func (d *optimizerDecorator) CallTool(
        ctx context.Context,
        _ *auth.Identity,
        toolName string,
        arguments map[string]any,
        _ map[string]any,
) (*vmcp.ToolCallResult, error) {
        switch toolName {
        case FindToolName:
                return d.handleFindTool(ctx, arguments)
        case CallToolName:
                return d.handleCallTool(ctx, arguments)
        default:
                return nil, fmt.Errorf("tool not found: %s", toolName)
        }
}

func (d *optimizerDecorator) handleFindTool(ctx context.Context, arguments map[string]any) (*vmcp.ToolCallResult, error) {
        input, err := schema.Translate[optimizer.FindToolInput](arguments)
        if err != nil {
                return errorResult(fmt.Sprintf("invalid arguments: %v", err)), nil
        }

        output, err := d.opt.FindTool(ctx, input)
        if err != nil {
                return errorResult(fmt.Sprintf("find_tool failed: %v", err)), nil
        }
        if output == nil {
                return errorResult("find_tool: optimizer returned nil result"), nil
        }

        jsonBytes, err := json.Marshal(output)
        if err != nil {
                return errorResult(fmt.Sprintf("failed to marshal find_tool output: %v", err)), nil
        }

        var structured map[string]any
        // Unmarshal cannot fail: jsonBytes was just produced by json.Marshal above.
        _ = json.Unmarshal(jsonBytes, &structured)

        return &vmcp.ToolCallResult{
                Content:           []vmcp.Content{{Type: "text", Text: string(jsonBytes)}},
                StructuredContent: structured,
        }, nil
}

func (d *optimizerDecorator) handleCallTool(
        ctx context.Context,
        arguments map[string]any,
) (*vmcp.ToolCallResult, error) {
        input, err := schema.Translate[optimizer.CallToolInput](arguments)
        if err != nil {
                return errorResult(fmt.Sprintf("invalid arguments: %v", err)), nil
        }

        mcpResult, err := d.opt.CallTool(ctx, input)
        if err != nil {
                return errorResult(fmt.Sprintf("call_tool failed: %v", err)), nil
        }
        if mcpResult == nil {
                return errorResult("call_tool: optimizer returned nil result"), nil
        }

        return mcpResultToVMCPResult(mcpResult), nil
}

// mcpResultToVMCPResult converts an MCP SDK CallToolResult to the vmcp domain type.
func mcpResultToVMCPResult(r *mcp.CallToolResult) *vmcp.ToolCallResult {
        structured, _ := r.StructuredContent.(map[string]any)
        return &vmcp.ToolCallResult{
                Content:           conversion.ConvertMCPContents(r.Content),
                StructuredContent: structured,
                IsError:           r.IsError,
        }
}

func errorResult(msg string) *vmcp.ToolCallResult {
        return &vmcp.ToolCallResult{
                Content: []vmcp.Content{{Type: "text", Text: msg}},
                IsError: true,
        }
}

func mustGenerateSchema[T any]() map[string]any {
        s, err := schema.GenerateSchema[T]()
        if err != nil {
                panic(fmt.Sprintf("optimizerdec: failed to generate schema: %v", err))
        }
        return s
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package types defines shared session interfaces for the vmcp/session package
// hierarchy. Placing the common types here allows both the internal backend
// package and the top-level session package to share a definition without
// introducing an import cycle.
package types

//go:generate mockgen -destination=mocks/mock_session.go -package=mocks github.com/stacklok/toolhive/pkg/vmcp/session/types MultiSession

import (
        "context"
        "errors"
        "log/slog"

        "github.com/stacklok/toolhive/pkg/auth"
        transportsession "github.com/stacklok/toolhive/pkg/transport/session"
        "github.com/stacklok/toolhive/pkg/vmcp"
        "github.com/stacklok/toolhive/pkg/vmcp/session/binding"
)

// Caller represents the ability to invoke MCP protocol operations against a
// backend. It is the common subset shared by both a single-backend
// [backend.Session] and the multi-backend [session.MultiSession].
//
// Implementations must be safe for concurrent use.
type Caller interface {
        // CallTool invokes toolName on the backend.
        //
        // caller identifies the requesting user/service. For bound sessions, caller
        // must be non-nil and its identity must match the session creator. For
        // anonymous sessions, caller may be nil.
        //
        // Returns:
        //   - ErrNilCaller if caller is nil for a bound session
        //   - ErrUnauthorizedCaller if the caller identity does not match the session owner
        //
        // arguments contains the tool input parameters.
        // meta contains protocol-level metadata (_meta) forwarded from the client.
        CallTool(
                ctx context.Context,
                caller *auth.Identity,
                toolName string,
                arguments map[string]any,
                meta map[string]any,
        ) (*vmcp.ToolCallResult, error)

        // ReadResource retrieves the resource identified by uri from the backend.
        //
        // caller identifies the requesting user/service. For bound sessions, caller
        // must be non-nil and its identity must match the session creator. For
        // anonymous sessions, caller may be nil.
        //
        // Returns:
        //   - ErrNilCaller if caller is nil for a bound session
        //   - ErrUnauthorizedCaller if the caller identity does not match the session owner
        ReadResource(ctx context.Context, caller *auth.Identity, uri string) (*vmcp.ResourceReadResult, error)

        // GetPrompt retrieves the named prompt from the backend.
        //
        // caller identifies the requesting user/service. For bound sessions, caller
        // must be non-nil and its identity must match the session creator. For
        // anonymous sessions, caller may be nil.
        //
        // Returns:
        //   - ErrNilCaller if caller is nil for a bound session
        //   - ErrUnauthorizedCaller if the caller identity does not match the session owner
        //
        // arguments contains the prompt input parameters.
        GetPrompt(
                ctx context.Context,
                caller *auth.Identity,
                name string,
                arguments map[string]any,
        ) (*vmcp.PromptGetResult, error)

        // Close releases all resources held by this caller. Implementations must
        // be idempotent: calling Close multiple times returns nil.
        Close() error
}

// MultiSession is the vMCP domain session interface. It extends the
// transport-layer Session with behaviour: capability access and session-scoped
// backend routing across multiple backend connections.
//
// A MultiSession is a "session of sessions": each backend contributes its own
// persistent connection (see [backend.Session] in pkg/vmcp/session/internal/backend),
// and the MultiSession aggregates them behind a single routing table.
//
// # Distributed deployment note
//
// Because MCP clients cannot be serialised, horizontal scaling requires sticky
// sessions (session affinity at the load balancer). Without sticky sessions, a
// request routed to a different vMCP instance must recreate backend clients
// (one-time cost per re-route). This is an accepted trade-off.
//
// # Storage
//
// A MultiSession uses a two-layer storage model:
//
//   - Runtime layer (in-process only): backend HTTP connections, routing
//     table, and capability lists. These cannot be serialized and are lost
//     when the process exits. Sessions are therefore node-local.
//
//   - Metadata layer (serializable): identity subject and connected backend
//     IDs are written to the embedded transportsession.Session so that
//     pluggable transportsession.Storage backends (e.g. Redis) can persist
//     them. This enables auditing and future session reconstruction, but
//     does not make the session itself portable — the runtime layer must
//     be rebuilt from scratch on a different node.
type MultiSession interface {
        transportsession.Session
        Caller

        // Tools returns the advertised tools available in this session (shown to MCP clients).
        // The list is built once at session creation and is read-only thereafter.
        Tools() []vmcp.Tool

        // AllTools returns all resolved tools in this session, including tools that are
        // excluded from advertising to MCP clients via excludeAll or filter configuration.
        // Used by the workflow engine for argument type coercion via InputSchema lookup.
        AllTools() []vmcp.Tool

        // Resources returns the resolved resources available in this session.
        Resources() []vmcp.Resource

        // Prompts returns the resolved prompts available in this session.
        Prompts() []vmcp.Prompt

        // BackendSessions returns a snapshot of the backend-assigned session IDs,
        // keyed by backend workload ID. The backend session ID is assigned by the
        // backend MCP server and is used to correlate vMCP sessions with backend
        // sessions for debugging and auditing.
        BackendSessions() map[string]string

        // GetRoutingTable returns the session's routing table.
        // Used by the discovery middleware to inject DiscoveredCapabilities into the
        // request context so composite tool workflow steps can route backend tool calls.
        GetRoutingTable() *vmcp.RoutingTable
}

const (
        // MetadataKeyTokenHash held hex(HMAC-SHA256(bearerToken)) for authenticated
        // sessions and "" for anonymous sessions.
        //
        // Legacy: superseded by MetadataKeyIdentityBinding (#5306); sessions written
        // under this key are invalidated on read. Constant removed in the follow-up PR.
        MetadataKeyTokenHash = "vmcp.token.hash" //nolint:gosec // This is a metadata key name, not a credential.

        // MetadataKeyTokenSalt held the hex-encoded per-session salt used for
        // HMAC-SHA256 token hashing.
        //
        // Legacy: superseded by MetadataKeyIdentityBinding (#5306); removed in the follow-up PR.
        MetadataKeyTokenSalt = "vmcp.token.salt" //nolint:gosec // This is a metadata key name, not a credential.

        // MetadataKeyIdentityBinding is the session metadata key that holds the
        // identity-binding string in the format defined by
        // [pkg/vmcp/session/binding]. Bound sessions store binding.Format(iss, sub);
        // unauthenticated sessions store binding.UnauthenticatedSentinel.
        //
        // Storage is plaintext PII — the Redis/Valkey instance must be access-controlled.
        MetadataKeyIdentityBinding = "vmcp.identity.binding"
)

// ShouldAllowAnonymous reports whether a session has no per-user binding and
// should be treated as anonymous. The identity is bound when Token is set, or
// when Claims yields a (iss, sub) pair accepted by binding.Format.
//
// Fail-closed: a Claims["iss"] or Claims["sub"] that is present but not a
// string (a misbehaving validator) is treated as bound, with a WARN logged.
//
// Contract note: this function answers "should this session be CREATED as
// anonymous?" using Token presence as a fast path. It does NOT guarantee that
// a non-anonymous identity will produce a successful binding: an identity with
// Token != "" but missing iss/sub claims passes this check as bound, then
// fails in BindSession with an extraction error. In practice this is
// pathological — all shipping middlewares (JWT validator, LocalUserMiddleware,
// AnonymousMiddleware) populate Claims. Callers who need "will this identity
// actually produce a binding?" must use extractBindingID directly, which is
// deliberately kept internal to the security decorator.
func ShouldAllowAnonymous(identity *auth.Identity) bool {
        if identity == nil {
                return true
        }
        if identity.Token != "" {
                return false
        }
        iss, issOK := claimString(identity.Claims, "iss")
        sub, subOK := claimString(identity.Claims, "sub")
        if !issOK || !subOK {
                slog.Warn("auth identity has present-but-non-string iss/sub claim; treating as bound")
                return false
        }
        if _, err := binding.Format(iss, sub); err == nil {
                return false
        }
        return true
}

// claimString returns (value, ok) for a string claim. ok is true when the
// claim is missing entirely (absence is benign — the caller proceeds to
// binding.Format which rejects empty halves) or when the claim is present
// and a string (including the empty string). ok is false only when the
// claim is present but not a string — the caller must treat that as a
// misconfigured validator and fail closed.
func claimString(claims map[string]any, key string) (string, bool) {
        v, present := claims[key]
        if !present {
                return "", true
        }
        s, isString := v.(string)
        if !isString {
                return "", false
        }
        return s, true
}

// Identity binding errors returned by Caller methods when caller identity
// validation fails.
var (
        // ErrUnauthorizedCaller is returned when the caller identity does not
        // match the session owner's identity (identity binding mismatch).
        ErrUnauthorizedCaller = errors.New("caller identity does not match session owner")

        // ErrNilCaller is returned when a bound session receives a nil caller.
        // Bound sessions require explicit caller identity on every method call.
        ErrNilCaller = errors.New("caller identity is required for bound sessions")

        // ErrSessionOwnerUnknown is returned when the session has no bound identity
        // but is configured to require one. This indicates a configuration error.
        ErrSessionOwnerUnknown = errors.New("session has no bound identity")
)

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package vmcp

import (
        "context"
        "time"

        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

        "github.com/stacklok/toolhive/pkg/mcp"
        authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
)

// This file contains shared domain types used across multiple vmcp subpackages.
// Following DDD principles, these are core domain concepts that cross bounded contexts.

// HeaderForwardConfig configures HTTP headers injected into outbound requests
// from the vMCP runtime to a static backend. AddPlaintextHeaders carries literal
// values; AddHeadersFromSecret maps header names to secret identifiers resolved
// via secrets.EnvironmentProvider at request time (env var TOOLHIVE_SECRET_<identifier>).
//
// Secret values MUST NOT appear in this struct — only identifiers. The operator
// injects actual Secret values into the vMCP pod as env vars via
// valueFrom.secretKeyRef at Deployment rendering time. Plaintext values are
// also injected via env vars by the operator (literal values, not secret refs)
// and ingested at vMCP startup; the type is therefore purely a runtime
// representation and is not part of any CRD schema.
type HeaderForwardConfig struct {
        // AddPlaintextHeaders is a map of canonical HTTP header name to literal value.
        AddPlaintextHeaders map[string]string `json:"addPlaintextHeaders,omitempty" yaml:"addPlaintextHeaders,omitempty"`

        // AddHeadersFromSecret maps canonical HTTP header name to secret identifier.
        // The vMCP runtime resolves each identifier via secrets.EnvironmentProvider,
        // which reads TOOLHIVE_SECRET_<identifier> from the pod environment.
        AddHeadersFromSecret map[string]string `json:"addHeadersFromSecret,omitempty" yaml:"addHeadersFromSecret,omitempty"`
}

// BackendTarget identifies a specific backend workload and provides
// the information needed to forward requests to it.
type BackendTarget struct {
        // WorkloadID is the unique identifier for the backend workload.
        WorkloadID string

        // WorkloadName is the human-readable name of the workload.
        WorkloadName string

        // BaseURL is the base URL for the backend's MCP server.
        // For local deployments: http://localhost:PORT
        // For Kubernetes: http://service-name.namespace.svc.cluster.local:PORT
        BaseURL string

        // TransportType specifies the MCP transport protocol.
        // Supported: "stdio", "http", "sse", "streamable-http"
        TransportType string

        // CABundlePath is the file path to a custom CA certificate bundle for TLS verification.
        // When set, the HTTP client uses this CA (appended to system roots) for backend connections.
        // Only applicable to entry-type backends with self-signed or internal certificates.
        // Used in static mode where CA bundles are volume-mounted by the operator.
        CABundlePath string

        // CABundleData contains raw CA certificate PEM bytes for TLS verification.
        // Used in dynamic mode where CA bundles are fetched from K8s ConfigMaps at
        // discovery time (no volume mount available). Takes precedence over CABundlePath.
        CABundleData []byte

        // OriginalCapabilityName is the original name of the capability (tool/resource/prompt)
        // as known by the backend. This is used when forwarding requests to the backend.
        //
        // When conflict resolution renames capabilities, this field preserves the original name:
        // - Prefix strategy: "fetch" → "fetch_fetch" (OriginalCapabilityName="fetch")
        // - Priority strategy: usually unchanged (OriginalCapabilityName="tool_name")
        // - Manual strategy: "fetch" → "custom_name" (OriginalCapabilityName="fetch")
        //
        // If empty, the resolved name is used when forwarding to the backend.
        //
        // IMPORTANT: Do NOT access this field directly when forwarding requests to backends.
        // Use GetBackendCapabilityName() method instead, which handles both renamed and
        // non-renamed capabilities correctly. Direct access can lead to incorrect behavior
        // when capabilities are not renamed (OriginalCapabilityName will be empty).
        //
        // Example (WRONG):
        //   client.CallTool(ctx, target, target.OriginalCapabilityName, args) // BUG: fails when empty
        //
        // Example (CORRECT):
        //   client.CallTool(ctx, target, target.GetBackendCapabilityName(toolName), args)
        OriginalCapabilityName string

        // AuthConfig contains the typed authentication configuration for this backend.
        // The actual authentication is handled by OutgoingAuthRegistry interface.
        // If nil, the backend requires no authentication.
        AuthConfig *authtypes.BackendAuthStrategy

        // SessionAffinity indicates if requests from the same session
        // must be routed to this specific backend instance.
        SessionAffinity bool

        // HealthStatus indicates the current health of the backend.
        HealthStatus BackendHealthStatus

        // HeaderForward carries per-backend HTTP header injection configuration
        // applied by the vMCP client to every outbound request targeting this backend
        // (list, call, health-check). Nil when no headers are configured.
        HeaderForward *HeaderForwardConfig

        // Metadata stores additional backend-specific information.
        Metadata map[string]string
}

// GetBackendCapabilityName returns the name to use when forwarding a request to the backend.
// If conflict resolution renamed the capability, this returns the original name that the backend expects.
// Otherwise, it returns the resolved name as-is.
//
// This method encapsulates the name translation logic for all capability types (tools, resources, prompts).
//
// ALWAYS use this method when forwarding capability calls to backends. Do NOT access
// OriginalCapabilityName directly, as it may be empty when no renaming occurred.
//
// Usage example:
//
//        target, _ := router.RouteTool(ctx, "fetch_fetch")  // Prefixed name from client
//        backendName := target.GetBackendCapabilityName("fetch_fetch")  // Returns "fetch"
//        client.CallTool(ctx, target, backendName, args)  // Backend receives original name
//
// This ensures correct behavior regardless of conflict resolution strategy:
//   - Prefix strategy: "fetch_fetch" → "fetch" (renamed, uses OriginalCapabilityName)
//   - Priority strategy: "list_issues" → "list_issues" (not renamed, returns resolvedName)
//   - Manual strategy: "custom_fetch" → "fetch" (renamed, uses OriginalCapabilityName)
func (t *BackendTarget) GetBackendCapabilityName(resolvedName string) string {
        if t.OriginalCapabilityName != "" {
                return t.OriginalCapabilityName
        }
        return resolvedName
}

// BackendType represents the type of backend workload.
type BackendType string

const (
        // BackendTypeContainer indicates a container-based backend managed by ToolHive.
        // Currently represented by an empty Type field for backward compatibility.
        BackendTypeContainer BackendType = "container"

        // BackendTypeProxy indicates a proxy-based backend (MCPRemoteProxy).
        // Currently represented by an empty Type field for backward compatibility.
        BackendTypeProxy BackendType = "proxy"

        // BackendTypeEntry indicates an external MCP server declared via MCPServerEntry.
        BackendTypeEntry BackendType = "entry"
)

// BackendHealthStatus represents the health state of a backend.
type BackendHealthStatus string

const (
        // BackendHealthy indicates the backend is healthy and accepting requests.
        BackendHealthy BackendHealthStatus = "healthy"

        // BackendDegraded indicates the backend is operational but experiencing issues.
        // This occurs when:
        // - Health checks succeed but response times exceed the degraded threshold (slow but working)
        // - Backend just recovered from failures and is in a stabilizing state
        // - Background OAuth token refresh is failing transiently while the
        //   workload monitor retries (auth_retrying workload status)
        BackendDegraded BackendHealthStatus = "degraded"

        // BackendUnhealthy indicates the backend is not responding to health checks.
        BackendUnhealthy BackendHealthStatus = "unhealthy"

        // BackendUnknown indicates the backend health status is unknown.
        BackendUnknown BackendHealthStatus = "unknown"

        // BackendUnauthenticated indicates the backend returned an authentication error
        // (HTTP 401/403) while the backend target had no outgoing auth strategy
        // configured (AuthConfig nil or StrategyTypeUnauthenticated). This signals
        // operator misconfiguration: the backend requires authentication but no
        // auth strategy was configured on the backend target.
        //
        // Note: a 401/403 from a backend with a configured outgoing auth strategy is
        // treated as BackendHealthy, because health probes deliberately do not carry
        // user credentials and the backend's challenge proves reachability and a
        // working auth layer.
        BackendUnauthenticated BackendHealthStatus = "unauthenticated"
)

// ToCRDStatus converts BackendHealthStatus to CRD-friendly status string.
// This maps internal health states to user-facing status values:
//   - healthy → ready
//   - degraded → degraded
//   - unhealthy → unavailable
//   - unauthenticated → unauthenticated (misconfig: backend requires auth but none configured)
//   - unknown → unknown
func (s BackendHealthStatus) ToCRDStatus() string {
        switch s {
        case BackendHealthy:
                return "ready"
        case BackendDegraded:
                return "degraded"
        case BackendUnhealthy:
                return "unavailable"
        case BackendUnauthenticated:
                return "unauthenticated"
        case BackendUnknown:
                return "unknown"
        default:
                return "unknown"
        }
}

// Condition represents a specific aspect of vMCP server status.
type Condition = metav1.Condition

// Phase represents the operational lifecycle phase of a vMCP server.
type Phase string

// Phase constants for vMCP server lifecycle.
const (
        PhasePending  Phase = "Pending"
        PhaseReady    Phase = "Ready"
        PhaseDegraded Phase = "Degraded"
        PhaseFailed   Phase = "Failed"
)

// Condition type constants for common vMCP conditions.
const (
        ConditionTypeBackendsDiscovered = "BackendsDiscovered"
        ConditionTypeReady              = "Ready"
        ConditionTypeAuthConfigured     = "AuthConfigured"
)

// Reason constants for condition reasons.
const (
        ReasonBackendDiscoverySucceeded = "BackendDiscoverySucceeded"
        ReasonBackendDiscoveryFailed    = "BackendDiscoveryFailed"
        ReasonServerReady               = "ServerReady"
        ReasonServerStarting            = "ServerStarting"
        ReasonServerDegraded            = "ServerDegraded"
        ReasonServerFailed              = "ServerFailed"
)

// DiscoveredBackend represents a backend server discovered by vMCP runtime.
// This type is shared with the Kubernetes operator CRD (VirtualMCPServer.Status.DiscoveredBackends).
// +gendoc
type DiscoveredBackend struct {
        // Name is the name of the backend MCPServer
        Name string `json:"name"`

        // URL is the URL of the backend MCPServer
        // +optional
        URL string `json:"url,omitempty"`

        // Status is the current status of the backend (ready, degraded, unavailable, unauthenticated, unknown).
        // Use BackendHealthStatus.ToCRDStatus() to populate this field.
        // +optional
        Status string `json:"status,omitempty"`

        // AuthConfigRef is the name of the discovered MCPExternalAuthConfig (if any)
        // +optional
        AuthConfigRef string `json:"authConfigRef,omitempty"`

        // AuthType is the type of authentication configured
        // +optional
        AuthType string `json:"authType,omitempty"`

        // LastHealthCheck is the timestamp of the last health check
        // +optional
        LastHealthCheck metav1.Time `json:"lastHealthCheck,omitempty"`

        // Message provides additional information about the backend status
        // +optional
        Message string `json:"message,omitempty"`

        // CircuitBreakerState is the current circuit breaker state (closed, open, half-open).
        // Empty when circuit breaker is disabled or not configured.
        // +optional
        // +kubebuilder:validation:Enum=closed;open;half-open
        CircuitBreakerState string `json:"circuitBreakerState,omitempty"`

        // CircuitLastChanged is the timestamp when the circuit breaker state last changed.
        // Empty when circuit breaker is disabled or has never changed state.
        // +optional
        CircuitLastChanged metav1.Time `json:"circuitLastChanged,omitempty"`

        // ConsecutiveFailures is the current count of consecutive health check failures.
        // Resets to 0 when the backend becomes healthy again.
        // +optional
        ConsecutiveFailures int `json:"consecutiveFailures,omitempty"`

        // MCPRevision is the backend's negotiated MCP protocol revision
        // ("2026-07-28" or "2025-11-25"). Empty when the backend has not been probed.
        // +optional
        MCPRevision string `json:"mcpRevision,omitempty"`
}

// DeepCopyInto copies the receiver into out. Required for Kubernetes CRD types.
func (in *DiscoveredBackend) DeepCopyInto(out *DiscoveredBackend) {
        *out = *in
        in.LastHealthCheck.DeepCopyInto(&out.LastHealthCheck)
}

// DeepCopy creates a deep copy of DiscoveredBackend. Required for Kubernetes CRD types.
func (in *DiscoveredBackend) DeepCopy() *DiscoveredBackend {
        if in == nil {
                return nil
        }
        out := new(DiscoveredBackend)
        in.DeepCopyInto(out)
        return out
}

// Status represents the runtime status of a vMCP server.
type Status struct {
        Phase              Phase               `json:"phase"`
        Message            string              `json:"message,omitempty"`
        Conditions         []Condition         `json:"conditions,omitempty"`
        DiscoveredBackends []DiscoveredBackend `json:"discoveredBackends,omitempty"`
        BackendCount       int32               `json:"backendCount,omitempty"`
        ObservedGeneration int64               `json:"observedGeneration,omitempty"`
        Timestamp          time.Time           `json:"timestamp"`
}

// Backend represents a discovered backend MCP server workload.
type Backend struct {
        // ID is the unique identifier for this backend.
        ID string

        // Name is the human-readable name.
        Name string

        // BaseURL is the backend's MCP server URL.
        BaseURL string

        // TransportType is the MCP transport protocol.
        TransportType string

        // Type is the backend workload type (container, proxy, entry).
        // Empty string is treated as container/proxy for backward compatibility.
        Type BackendType

        // CABundlePath is the file path to a custom CA certificate bundle for TLS verification.
        // Only applicable to entry-type backends with self-signed or internal certificates.
        // Used in static mode where CA bundles are volume-mounted by the operator.
        CABundlePath string

        // CABundleData contains raw CA certificate PEM bytes for TLS verification.
        // Used in dynamic mode where CA bundles are fetched from K8s ConfigMaps at
        // discovery time (no volume mount available). Takes precedence over CABundlePath.
        CABundleData []byte

        // HealthStatus is the current health state.
        HealthStatus BackendHealthStatus

        // AuthConfig contains the typed authentication configuration for this backend.
        // The actual authentication is handled by OutgoingAuthRegistry interface.
        // If nil, the backend requires no authentication.
        AuthConfig *authtypes.BackendAuthStrategy

        // AuthConfigRef is the name of the MCPExternalAuthConfig resource (if any).
        // This field is populated during backend discovery and is useful for
        // debugging and status reporting.
        // +optional
        AuthConfigRef string

        // HeaderForward carries per-backend HTTP header injection configuration.
        // Only populated for entry-type backends whose MCPServerEntry declares
        // spec.headerForward. Nil when the entry has no header forwarding configured.
        HeaderForward *HeaderForwardConfig

        // Metadata stores additional backend information.
        Metadata map[string]string
}

// Tool represents an MCP tool capability.
type Tool struct {
        // Name is the tool name (may conflict with other backends).
        Name string

        // Description describes what the tool does.
        Description string

        // InputSchema is the JSON Schema for tool parameters.
        InputSchema map[string]any

        // OutputSchema is the JSON Schema for tool output (optional).
        // Per MCP specification, this describes the structure of the tool's response.
        OutputSchema map[string]any

        // Annotations describes behavioral hints for the tool (optional).
        // Per MCP specification, these include readOnlyHint, destructiveHint, etc.
        Annotations *ToolAnnotations

        // BackendID identifies the backend that provides this tool.
        BackendID string
}

// ToolAnnotations describes behavioral hints for a tool.
// These are the vmcp-domain equivalents of mcp.ToolAnnotation, following the
// Anti-Corruption Layer pattern (vmcp types are decoupled from mcp-go).
type ToolAnnotations struct {
        // Title is a human-readable title for the tool.
        Title string `json:"title,omitempty"`
        // ReadOnlyHint indicates whether the tool does not modify its environment.
        ReadOnlyHint *bool `json:"readOnlyHint,omitempty"`
        // DestructiveHint indicates whether the tool may perform destructive updates.
        DestructiveHint *bool `json:"destructiveHint,omitempty"`
        // IdempotentHint indicates whether repeated calls with same args have no additional effect.
        IdempotentHint *bool `json:"idempotentHint,omitempty"`
        // OpenWorldHint indicates whether the tool interacts with external entities.
        OpenWorldHint *bool `json:"openWorldHint,omitempty"`
}

// Resource represents an MCP resource capability.
type Resource struct {
        // URI is the resource URI (should be globally unique).
        URI string

        // Name is a human-readable name.
        Name string

        // Description describes the resource.
        Description string

        // MimeType is the resource's MIME type (optional).
        MimeType string

        // BackendID identifies the backend that provides this resource.
        BackendID string
}

// ResourceTemplate represents an MCP resource template capability.
// Unlike a Resource, whose URI identifies a single concrete resource, a
// ResourceTemplate declares an RFC 6570 URI template that expands to a family
// of resource URIs (e.g. "file:///logs/{date}.txt"). Reading an expanded URI is
// served through the existing resource-read path: the router matches the URI
// against these templates when the exact-URI lookup misses.
type ResourceTemplate struct {
        // URITemplate is the RFC 6570 URI template (should be globally unique).
        URITemplate string

        // Name is a human-readable name.
        Name string

        // Description describes the resource template.
        Description string

        // MimeType is the MIME type shared by resources matching this template (optional).
        MimeType string

        // BackendID identifies the backend that provides this resource template.
        BackendID string
}

// Prompt represents an MCP prompt capability.
type Prompt struct {
        // Name is the prompt name (may conflict with other backends).
        Name string

        // Description describes the prompt.
        Description string

        // Arguments are the prompt parameters.
        Arguments []PromptArgument

        // BackendID identifies the backend that provides this prompt.
        BackendID string
}

// PromptArgument represents a prompt parameter.
type PromptArgument struct {
        // Name is the argument name.
        Name string

        // Description describes the argument.
        Description string

        // Required indicates if the argument is mandatory.
        Required bool
}

// ContentType represents the type of content in an MCP message.
type ContentType string

const (
        // ContentTypeText represents text content.
        ContentTypeText ContentType = "text"
        // ContentTypeImage represents image content.
        ContentTypeImage ContentType = "image"
        // ContentTypeAudio represents audio content.
        ContentTypeAudio ContentType = "audio"
        // ContentTypeResource represents embedded resource content.
        ContentTypeResource ContentType = "resource"
        // ContentTypeLink represents a resource link.
        ContentTypeLink ContentType = "resource_link"
)

// ContentAnnotations describes per-content metadata annotations.
// These are the vmcp-domain equivalents of mcp.Annotations, following the
// Anti-Corruption Layer pattern (vmcp types are decoupled from mcp-go).
type ContentAnnotations struct {
        // Audience describes who the content is intended for.
        // Valid values are the mcp.Role constants: "user" and "assistant".
        Audience []string `json:"audience,omitempty"`
        // Priority is a hint for display ordering in the closed interval [0.0, 1.0]
        // per the MCP spec, where 0.0 is least important and 1.0 is most important.
        // Nil means no priority hint was provided.
        Priority *float64 `json:"priority,omitempty"`
        // LastModified is an ISO 8601 timestamp (e.g., "2025-01-12T15:00:58Z").
        LastModified string `json:"lastModified,omitempty"`
}

// Content represents MCP content (text, image, audio, embedded resource, resource link).
// This is used by ToolCallResult to preserve the full content structure from backends.
type Content struct {
        // Type indicates the content type.
        Type ContentType

        // Text is the content text (for TextContent)
        Text string

        // Data is the base64-encoded data (for ImageContent/AudioContent)
        Data string

        // MimeType is the MIME type (for ImageContent/AudioContent/ResourceLink)
        MimeType string

        // URI is the resource URI (for EmbeddedResource/ResourceLink)
        URI string

        // Name is the resource name (for ResourceLink)
        Name string

        // Description is the resource description (for ResourceLink)
        Description string

        // Annotations contains per-content metadata (audience, priority, lastModified).
        // Nil means no annotations were provided by the backend.
        Annotations *ContentAnnotations
}

// ToolCallResult wraps a tool call response with metadata.
// This preserves both the tool output AND the _meta field from the backend MCP server.
type ToolCallResult struct {
        // Content is the tool output (text, image, etc.)
        // This is the array of content items returned by the backend.
        Content []Content

        // StructuredContent is structured output (preferred for composite tools and workflows).
        // If the backend MCP server provides StructuredContent, it is used directly.
        // Otherwise, this is populated by converting the Content array to a map:
        //   - First text item: key="text"
        //   - Additional text items: key="text_1", "text_2", etc.
        //   - Image items: key="image_0", "image_1", etc.
        // This allows templates to access fields via {{.steps.stepID.output.text}}.
        // Note: No JSON parsing is performed - backends must provide structured data explicitly.
        StructuredContent map[string]any

        // IsError indicates if the tool call failed.
        IsError bool

        // Meta contains protocol-level metadata from the backend (_meta field).
        // This includes progressToken, trace context, and custom backend metadata.
        // Per MCP specification, this field is optional and may be nil.
        Meta map[string]any

        // BackendID is the logical backend that served the call, set by the core
        // after routing. Empty for a composite tool (no single serving backend).
        // Audit-only: never serialized to the client.
        BackendID string `json:"-"`
}

// ResourceContent represents a single resource content item,
// preserving the text vs blob distinction from the MCP protocol.
type ResourceContent struct {
        // URI is the resource URI.
        URI string
        // MimeType is the content type of this resource item.
        MimeType string
        // Text is the text content (non-empty for text resources).
        Text string
        // Blob is the base64-encoded binary content (non-empty for blob resources).
        // Exactly one of Text or Blob should be set; Blob takes precedence in ToMCPResourceContents.
        Blob string
}

// ResourceReadResult wraps a resource read response with metadata.
// This preserves both the resource data AND the _meta field from the backend MCP server.
type ResourceReadResult struct {
        // Contents preserves individual resource content items with their
        // per-item URIs, MIME types, and text/blob distinction.
        Contents []ResourceContent

        // Meta contains protocol-level metadata from the backend (_meta field).
        // Populated on both the Legacy and Modern client paths; forwarded to the
        // client by the vMCP Serve handlers (non-reserved keys only).
        Meta map[string]any

        // BackendID is the logical backend that served the read, set by the core
        // after routing. Audit-only: never serialized to the client.
        BackendID string `json:"-"`
}

// PromptMessage represents a single message in a prompt response,
// preserving the role and content structure from the backend.
type PromptMessage struct {
        // Role is the message role. The MCP spec defines "user" and "assistant";
        // backends may also send other values which are relayed as-is.
        Role string
        // Content is the message content, supporting all MCP content types.
        Content Content
}

// PromptGetResult wraps a prompt response with metadata.
// This preserves both the prompt messages AND the _meta field from the backend MCP server.
type PromptGetResult struct {
        // Messages preserves individual prompt messages with their roles
        // and full content structure (text, images, audio, resources).
        Messages []PromptMessage

        // Description is an optional description of the prompt.
        Description string

        // Meta contains protocol-level metadata from the backend (_meta field).
        // This includes progressToken, trace context, and custom backend metadata.
        // Per MCP specification, this field is optional and may be nil.
        Meta map[string]any

        // BackendID is the logical backend that served the get, set by the core
        // after routing. Audit-only: never serialized to the client.
        BackendID string `json:"-"`
}

// Completion reference type constants, matching the MCP completion/complete
// "ref" discriminator values.
const (
        // CompletionRefTypePrompt identifies a prompt-argument completion reference.
        CompletionRefTypePrompt = "ref/prompt"
        // CompletionRefTypeResource identifies a resource(-template) completion reference.
        CompletionRefTypeResource = "ref/resource"
)

// CompletionRef identifies what a completion/complete request is completing:
// either a prompt argument (Type == CompletionRefTypePrompt, Name set) or a
// resource-template variable (Type == CompletionRefTypeResource, URI set). It is
// the domain equivalent of the MCP PromptReference/ResourceReference shapes,
// following the Anti-Corruption Layer pattern (no mcp-go types cross the core
// boundary, vmcp anti-pattern #5).
type CompletionRef struct {
        // Type is the reference discriminator: CompletionRefTypePrompt or
        // CompletionRefTypeResource.
        Type string
        // Name is the prompt name (set only for CompletionRefTypePrompt).
        Name string
        // URI is the resource URI or URI template (set only for
        // CompletionRefTypeResource).
        URI string
}

// CompletionResult holds argument-completion candidates returned by a backend.
// It is the domain equivalent of the MCP CompletionResultDetails.
type CompletionResult struct {
        // Values are the completion candidate strings (MCP caps this at 100 items).
        Values []string
        // Total is the total number of options available, which may exceed len(Values).
        Total int
        // HasMore indicates additional options exist beyond those in Values.
        HasMore bool
}

// MaxResourceURILength caps the length, in bytes, of a client-supplied resource
// URI that is matched against RFC 6570 resource templates.
//
// Template matching is not linear: its cost is quadratic in the number of
// separators a URI places inside a template variable, and it takes no context,
// so a closed connection or request deadline cannot stop it. The request body
// cap is 8MB and only CallTool is rate-limited, so without this cap a handful
// of resources/read, resources/subscribe or completion/complete requests could
// pin a CPU for minutes. Exact concrete-resource and template-string lookups are
// map lookups and are not subject to the cap. 2048 bytes is the conventional
// practical URL limit and far beyond any resource URI a backend advertises.
const MaxResourceURILength = 2048

// RoutingTable contains the mappings from capability names to backend targets.
// This is the output of the aggregation phase and input to the router.
// Placed in vmcp root package to avoid circular dependencies between
// aggregator and router packages.
//
// Note: Composite tools are NOT included here. They are executed by the composer
// package and do not route to a single backend.
type RoutingTable struct {
        // Tools maps tool names to their backend targets.
        // After conflict resolution, tool names are unique.
        Tools map[string]*BackendTarget

        // Resources maps resource URIs to their backend targets.
        Resources map[string]*BackendTarget

        // ResourceTemplates maps resource URI-template strings to their backend targets.
        // The router consults this when an exact Resources lookup misses: it matches
        // the requested URI against each template and routes to the first match.
        ResourceTemplates map[string]*BackendTarget

        // Prompts maps prompt names to their backend targets.
        Prompts map[string]*BackendTarget
}

// ConflictResolutionStrategy defines how to handle capability name conflicts.
// Placed in vmcp root package to be shared by config and aggregator packages.
// +gendoc
type ConflictResolutionStrategy string

const (
        // ConflictStrategyPrefix prefixes all tools with workload identifier.
        ConflictStrategyPrefix ConflictResolutionStrategy = "prefix"

        // ConflictStrategyPriority uses explicit priority ordering (first wins).
        ConflictStrategyPriority ConflictResolutionStrategy = "priority"

        // ConflictStrategyManual requires explicit overrides for all conflicts.
        ConflictStrategyManual ConflictResolutionStrategy = "manual"
)

// HealthChecker performs health checks on backend MCP servers.
type HealthChecker interface {
        // CheckHealth checks if a backend is healthy and responding.
        // Returns the current health status and any error encountered.
        CheckHealth(ctx context.Context, target *BackendTarget) (BackendHealthStatus, error)
}

// BackendClient abstracts MCP protocol communication with backend servers.
// This interface handles the protocol-level details of calling backend MCP servers,
// supporting multiple transport types (HTTP, SSE, stdio, streamable-http).
//
// All methods return wrapper types that preserve the _meta field from backend
// MCP server responses. Protocol-level metadata (progress tokens, trace context,
// custom metadata) is forwarded to clients where supported (tools and prompts).
// Note: Resource _meta forwarding is not currently supported due to MCP SDK handler
// signature limitations; the Meta field is preserved for future SDK improvements.
//
//go:generate mockgen -destination=mocks/mock_backend_client.go -package=mocks -source=types.go BackendClient HealthChecker
type BackendClient interface {
        // CallTool invokes a tool on the backend MCP server.
        // The meta parameter contains _meta fields from the client request that should be forwarded to the backend.
        // Returns the complete tool result including _meta field from the backend response.
        //
        // paramHeaders carries the SEP-2243 Mcp-Param-* headers mirrored from the tool's
        // x-mcp-header-designated arguments, keyed by full header name. It is supplied by
        // the caller because deriving it needs the tool's inputSchema, which lives in the
        // aggregated capability view rather than on BackendTarget; passing it explicitly
        // keeps this client free of a schema cache and its staleness window. nil or empty
        // means nothing to mirror, which is the common case. It applies only to a Modern
        // (2026-07-28) backend hop -- SEP-2243 is part of that revision -- and is ignored
        // for a Legacy backend.
        CallTool(
                ctx context.Context, target *BackendTarget, toolName string, arguments map[string]any,
                meta map[string]any, paramHeaders map[string]string,
        ) (*ToolCallResult, error)

        // ReadResource retrieves a resource from the backend MCP server.
        // Returns the complete resource result including _meta field.
        ReadResource(ctx context.Context, target *BackendTarget, uri string) (*ResourceReadResult, error)

        // GetPrompt retrieves a prompt from the backend MCP server.
        // Returns the complete prompt result including _meta field.
        GetPrompt(ctx context.Context, target *BackendTarget, name string, arguments map[string]any) (*PromptGetResult, error)

        // Complete requests argument-completion candidates from the backend MCP server
        // (the completion/complete method). ref identifies the prompt or resource being
        // completed; argName/argValue are the argument name and its partial value;
        // contextArgs carries previously-resolved argument values for multi-argument
        // completion (may be nil). For a CompletionRefTypePrompt ref, the prompt name is
        // translated to the backend's own name via target.GetBackendCapabilityName.
        //
        // If the backend does not advertise the completions capability, an empty (non-nil)
        // CompletionResult is returned without error, matching the MCP spec's lenient
        // completion semantics.
        Complete(
                ctx context.Context, target *BackendTarget, ref CompletionRef,
                argName, argValue string, contextArgs map[string]string,
        ) (*CompletionResult, error)

        // ListCapabilities queries a backend for its capabilities.
        // Returns tools, resources, and prompts exposed by the backend.
        ListCapabilities(ctx context.Context, target *BackendTarget) (*CapabilityList, error)
}

// RevisionReporter reports a backend's resolved MCP protocol revision, keyed by
// workload ID. The second return distinguishes "known" from "not probed yet" —
// callers must not read the Revision when it is false, since RevisionLegacy is
// the zero value.
//
// Deliberately NOT part of [BackendClient]: it is a read-model for
// status/telemetry and for deciding whether a Legacy session connect is worth
// attempting, not a protocol operation, and every BackendClient implementation
// (including the generated mocks) would otherwise have to satisfy it. Consumers
// type-assert to it and degrade gracefully when the assertion fails.
//
// Because that assertion is the failure mode — a wrapper that forgets to forward
// CachedRevision silently disables the behaviour with no compile error —
// implementations should pin themselves with a compile-time assertion, e.g.
// `var _ vmcp.RevisionReporter = (*httpBackendClient)(nil)`.
type RevisionReporter interface {
        CachedRevision(workloadID string) (mcp.Revision, bool)
}

// CapabilityList contains the capabilities from a backend's MCP server.
// This is returned by BackendClient.ListCapabilities().
type CapabilityList struct {
        // Tools available on this backend.
        Tools []Tool

        // Resources available on this backend.
        Resources []Resource

        // ResourceTemplates available on this backend.
        ResourceTemplates []ResourceTemplate

        // Prompts available on this backend.
        Prompts []Prompt

        // SupportsLogging indicates if the backend supports MCP logging.
        SupportsLogging bool

        // SupportsSampling indicates if the backend supports MCP sampling.
        SupportsSampling bool
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package webhook

import (
        "bytes"
        "context"
        "crypto/tls"
        "crypto/x509"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "net"
        "net/http"
        "os"
        "path/filepath"
        "strconv"
        "time"

        "github.com/stacklok/toolhive/pkg/networking"
)

// Client is an HTTP client for calling webhook endpoints.
type Client struct {
        httpClient *http.Client
        config     Config
        hmacSecret []byte
        // TODO: webhookType will be used by a future Send() method that dispatches
        // to Call or CallMutating based on type. For now callers pick the method directly.
        webhookType Type
}

// NewClient creates a new webhook Client from the given configuration.
// The hmacSecret parameter is the resolved secret bytes for HMAC signing;
// pass nil if signing is not configured.
func NewClient(cfg Config, webhookType Type, hmacSecret []byte) (*Client, error) {
        if err := cfg.Validate(); err != nil {
                return nil, fmt.Errorf("invalid webhook config: %w", err)
        }
        if cfg.HMACSecretRef != "" && len(hmacSecret) == 0 {
                return nil, fmt.Errorf("webhook %q has HMAC configured but resolved secret is empty", cfg.Name)
        }

        timeout := cfg.Timeout
        if timeout == 0 {
                timeout = DefaultTimeout
        }

        transport, err := buildTransport(cfg.TLSConfig)
        if err != nil {
                return nil, fmt.Errorf("failed to build HTTP transport: %w", err)
        }

        return &Client{
                httpClient: &http.Client{
                        Transport: transport,
                        Timeout:   timeout,
                },
                config:      cfg,
                hmacSecret:  hmacSecret,
                webhookType: webhookType,
        }, nil
}

// Call sends a request to a validating webhook and returns its response.
func (c *Client) Call(ctx context.Context, req *Request) (*Response, error) {
        body, err := json.Marshal(req)
        if err != nil {
                return nil, NewInvalidResponseError(c.config.Name, fmt.Errorf("failed to marshal request: %w", err), 0)
        }

        respBody, err := c.doHTTPCall(ctx, body)
        if err != nil {
                return nil, err
        }

        var resp Response
        if err := json.Unmarshal(respBody, &resp); err != nil {
                return nil, NewInvalidResponseError(c.config.Name, fmt.Errorf("failed to unmarshal response: %w", err), 0)
        }

        return &resp, nil
}

// CallMutating sends a request to a mutating webhook and returns its response.
func (c *Client) CallMutating(ctx context.Context, req *Request) (*MutatingResponse, error) {
        body, err := json.Marshal(req)
        if err != nil {
                return nil, NewInvalidResponseError(c.config.Name, fmt.Errorf("failed to marshal request: %w", err), 0)
        }

        respBody, err := c.doHTTPCall(ctx, body)
        if err != nil {
                return nil, err
        }

        var resp MutatingResponse
        if err := json.Unmarshal(respBody, &resp); err != nil {
                return nil, NewInvalidResponseError(c.config.Name, fmt.Errorf("failed to unmarshal mutating response: %w", err), 0)
        }

        return &resp, nil
}

// doHTTPCall performs the HTTP POST to the webhook endpoint, handling signing,
// error classification, and response size limiting.
func (c *Client) doHTTPCall(ctx context.Context, body []byte) ([]byte, error) {
        httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.config.URL, bytes.NewReader(body))
        if err != nil {
                return nil, NewNetworkError(c.config.Name, fmt.Errorf("failed to create HTTP request: %w", err))
        }
        httpReq.Header.Set("Content-Type", "application/json")

        hmacSecret, err := c.hmacSecretForRequest(ctx)
        if err != nil {
                return nil, NewNetworkError(c.config.Name, fmt.Errorf("failed to resolve HMAC secret: %w", err))
        }

        // Apply HMAC signing if configured.
        if len(hmacSecret) > 0 {
                timestamp := time.Now().Unix()
                signature := SignPayload(hmacSecret, timestamp, body)
                httpReq.Header.Set(SignatureHeader, signature)
                httpReq.Header.Set(TimestampHeader, strconv.FormatInt(timestamp, 10))
        }

        // #nosec G704 -- URL is validated in Config.Validate and we use ValidatingTransport for SSRF protection.
        resp, err := c.httpClient.Do(httpReq)
        if err != nil {
                return nil, classifyError(c.config.Name, err)
        }
        defer func() {
                _ = resp.Body.Close()
        }()

        // Enforce response size limit.
        limitedReader := io.LimitReader(resp.Body, MaxResponseSize+1)
        respBody, err := io.ReadAll(limitedReader)
        if err != nil {
                return nil, NewNetworkError(c.config.Name, fmt.Errorf("failed to read response body: %w", err))
        }
        if int64(len(respBody)) > MaxResponseSize {
                return nil, NewInvalidResponseError(c.config.Name,
                        fmt.Errorf("response body exceeds maximum size of %d bytes", MaxResponseSize), 0)
        }

        // 5xx errors indicate webhook operational failures.
        if resp.StatusCode >= http.StatusInternalServerError {
                // Body preview is logged at debug level so operators can troubleshoot,
                // but is kept out of the returned error chain to avoid surfacing
                // potentially sensitive bytes (e.g. from an internal service reached
                // via a misconfigured URL) into higher-level error logs.
                slog.Debug("webhook returned server error",
                        "webhook", c.config.Name,
                        "url", c.config.URL,
                        "status_code", resp.StatusCode,
                        "body_preview", truncateBody(respBody),
                )
                return nil, NewNetworkError(c.config.Name,
                        fmt.Errorf("webhook returned HTTP %d", resp.StatusCode))
        }

        // Non-200 responses (excluding 5xx handled above) are treated as invalid.
        // The StatusCode is surfaced so callers can distinguish HTTP 422 (RFC always-deny)
        // from other non-2xx codes that may follow the failure policy.
        if resp.StatusCode != http.StatusOK {
                slog.Debug("webhook returned non-2xx response",
                        "webhook", c.config.Name,
                        "url", c.config.URL,
                        "status_code", resp.StatusCode,
                        "body_preview", truncateBody(respBody),
                )
                return nil, NewInvalidResponseError(c.config.Name,
                        fmt.Errorf("webhook returned HTTP %d", resp.StatusCode),
                        resp.StatusCode)
        }

        return respBody, nil
}

func (c *Client) hmacSecretForRequest(ctx context.Context) ([]byte, error) {
        if c.config.HMACSecretRef == "" {
                return c.hmacSecret, nil
        }
        if !filepath.IsAbs(c.config.HMACSecretRef) {
                return c.hmacSecret, nil
        }

        secret, err := ResolveSecret(ctx, c.config.HMACSecretRef)
        if err != nil {
                return nil, err
        }
        if len(secret) == 0 {
                return nil, fmt.Errorf("resolved HMAC secret is empty")
        }
        return secret, nil
}

// buildTransport creates an http.RoundTripper with the specified TLS configuration,
// always wrapped in ValidatingTransport for SSRF protection.
func buildTransport(tlsCfg *TLSConfig) (http.RoundTripper, error) {
        transport := &http.Transport{
                TLSHandshakeTimeout:   10 * time.Second,
                ResponseHeaderTimeout: 10 * time.Second,
                MaxIdleConns:          100,
                MaxIdleConnsPerHost:   10,
                IdleConnTimeout:       90 * time.Second,
        }

        // allowHTTP is true when InsecureSkipVerify is set, which also covers in-cluster
        allowHTTP := tlsCfg != nil && tlsCfg.InsecureSkipVerify

        if tlsCfg != nil {
                tlsConfig := &tls.Config{
                        MinVersion: tls.VersionTLS12,
                }

                // Load CA bundle if provided.
                if tlsCfg.CABundlePath != "" {
                        caCert, err := os.ReadFile(tlsCfg.CABundlePath)
                        if err != nil {
                                return nil, fmt.Errorf("failed to read CA bundle: %w", err)
                        }
                        caCertPool := x509.NewCertPool()
                        if !caCertPool.AppendCertsFromPEM(caCert) {
                                return nil, fmt.Errorf("failed to parse CA certificate bundle")
                        }
                        tlsConfig.RootCAs = caCertPool
                }

                // Load client certificate for mTLS if provided.
                if tlsCfg.ClientCertPath != "" && tlsCfg.ClientKeyPath != "" {
                        cert, err := tls.LoadX509KeyPair(tlsCfg.ClientCertPath, tlsCfg.ClientKeyPath)
                        if err != nil {
                                return nil, fmt.Errorf("failed to load client certificate: %w", err)
                        }
                        tlsConfig.Certificates = []tls.Certificate{cert}
                }

                if tlsCfg.InsecureSkipVerify {
                        //#nosec G402 -- InsecureSkipVerify is intentionally user-configurable for development/testing only.
                        tlsConfig.InsecureSkipVerify = true
                }

                transport.TLSClientConfig = tlsConfig
        }

        // Always wrap in ValidatingTransport for SSRF protection, even without TLS config.
        return &networking.ValidatingTransport{
                Transport:         transport,
                InsecureAllowHTTP: allowHTTP,
        }, nil
}

// classifyError examines an HTTP client error and returns an appropriately
// typed webhook error (TimeoutError or NetworkError).
func classifyError(webhookName string, err error) error {
        // Check for context cancellation/deadline first, as these may not wrap net.Error.
        if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
                return NewTimeoutError(webhookName, err)
        }
        // Check for timeout errors via net.Error interface (e.g., dial timeout).
        var netErr net.Error
        if errors.As(err, &netErr) && netErr.Timeout() {
                return NewTimeoutError(webhookName, err)
        }
        return NewNetworkError(webhookName, err)
}

// truncateBody returns a preview of the response body for error messages.
func truncateBody(body []byte) string {
        const maxPreview = 256
        if len(body) <= maxPreview {
                return string(body)
        }
        return string(body[:maxPreview]) + "..."
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package webhook

import (
        "encoding/json"
        "errors"
        "fmt"
        "os"
        "path/filepath"
        "strings"

        "go.yaml.in/yaml/v3"
)

// FileConfig is the top-level structure for a webhook configuration file.
// It supports both YAML and JSON formats.
//
// Example YAML:
//
//        validating:
//          - name: policy-check
//            url: https://policy.example.com/validate
//            timeout: 5s
//            failure_policy: fail
//
//        mutating:
//          - name: hr-enrichment
//            url: https://hr-api.example.com/enrich
//            timeout: 3s
//            failure_policy: ignore
type FileConfig struct {
        // Validating is the list of validating webhook configurations.
        Validating []Config `yaml:"validating" json:"validating"`
        // Mutating is the list of mutating webhook configurations.
        Mutating []Config `yaml:"mutating" json:"mutating"`
}

// LoadConfig reads and parses a webhook configuration file.
// The format is auto-detected by file extension: ".json" uses JSON decoding;
// all other extensions (including ".yaml" and ".yml") use YAML decoding.
func LoadConfig(path string) (*FileConfig, error) {
        data, err := os.ReadFile(path) // #nosec G304 -- path is caller-supplied
        if err != nil {
                return nil, fmt.Errorf("webhook config file not found: %s", path)
        }

        var cfg FileConfig
        ext := strings.ToLower(filepath.Ext(path))
        if ext == ".json" {
                if err := json.Unmarshal(data, &cfg); err != nil {
                        return nil, fmt.Errorf("failed to parse webhook config %s as JSON: %w", path, err)
                }
        } else {
                if err := yaml.Unmarshal(data, &cfg); err != nil {
                        return nil, fmt.Errorf("failed to parse webhook config %s as YAML: %w", path, err)
                }
        }

        normalizeConfig(cfg)

        return &cfg, nil
}

// normalizeConfig applies effective defaults after parsing so validation sees
// the same values the runtime will use.
func normalizeConfig(cfg FileConfig) {
        for i := range cfg.Validating {
                if cfg.Validating[i].Timeout == 0 {
                        cfg.Validating[i].Timeout = DefaultTimeout
                }
        }
        for i := range cfg.Mutating {
                if cfg.Mutating[i].Timeout == 0 {
                        cfg.Mutating[i].Timeout = DefaultTimeout
                }
        }
}

// MergeConfigs merges multiple FileConfigs into one.
// Webhooks with the same name are de-duplicated: entries from later configs
// override entries from earlier ones (last-writer-wins per webhook name).
// The resulting Validating and Mutating slices preserve the order in which
// unique names were first seen and apply overrides in place.
func MergeConfigs(configs ...*FileConfig) *FileConfig {
        merged := &FileConfig{}

        validatingIndex := make(map[string]int) // name -> index in merged.Validating
        mutatingIndex := make(map[string]int)   // name -> index in merged.Mutating

        for _, cfg := range configs {
                if cfg == nil {
                        continue
                }
                for _, wh := range cfg.Validating {
                        if idx, exists := validatingIndex[wh.Name]; exists {
                                merged.Validating[idx] = wh
                        } else {
                                validatingIndex[wh.Name] = len(merged.Validating)
                                merged.Validating = append(merged.Validating, wh)
                        }
                }
                for _, wh := range cfg.Mutating {
                        if idx, exists := mutatingIndex[wh.Name]; exists {
                                merged.Mutating[idx] = wh
                        } else {
                                mutatingIndex[wh.Name] = len(merged.Mutating)
                                merged.Mutating = append(merged.Mutating, wh)
                        }
                }
        }

        return merged
}

// ValidateConfig validates all webhook configurations in a FileConfig,
// collecting all validation errors before returning.
func ValidateConfig(cfg *FileConfig) error {
        if cfg == nil {
                return nil
        }

        var errs []error
        for i, wh := range cfg.Validating {
                if err := wh.Validate(); err != nil {
                        errs = append(errs, fmt.Errorf("validating webhook[%d] %q: %w", i, wh.Name, err))
                }
        }
        for i, wh := range cfg.Mutating {
                if err := wh.Validate(); err != nil {
                        errs = append(errs, fmt.Errorf("mutating webhook[%d] %q: %w", i, wh.Name, err))
                }
        }

        return errors.Join(errs...)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package webhook

import (
        "errors"
        "fmt"
        "net/http"
)

// WebhookError is the base error type for all webhook-related errors.
//
//nolint:revive // WebhookError is the canonical name; renaming to Error conflicts with Error() method.
type WebhookError struct {
        // WebhookName is the name of the webhook that caused the error.
        WebhookName string
        // Err is the underlying error.
        Err error
}

// Error implements the error interface.
func (e *WebhookError) Error() string {
        return fmt.Sprintf("webhook %q: %v", e.WebhookName, e.Err)
}

// Unwrap returns the underlying error for errors.Is/errors.As support.
func (e *WebhookError) Unwrap() error {
        return e.Err
}

// TimeoutError indicates that a webhook call timed out.
type TimeoutError struct {
        WebhookError
}

// Error implements the error interface.
func (e *TimeoutError) Error() string {
        return fmt.Sprintf("webhook %q: timeout: %v", e.WebhookName, e.Err)
}

// NetworkError indicates a network-level failure when calling a webhook.
type NetworkError struct {
        WebhookError
}

// Error implements the error interface.
func (e *NetworkError) Error() string {
        return fmt.Sprintf("webhook %q: network error: %v", e.WebhookName, e.Err)
}

// InvalidResponseError indicates that a webhook returned an unparsable or invalid response.
type InvalidResponseError struct {
        WebhookError
        // StatusCode is the HTTP status code returned by the webhook, if applicable.
        // A value of 0 means no HTTP response was received (e.g., JSON decode error).
        StatusCode int
}

// Error implements the error interface.
func (e *InvalidResponseError) Error() string {
        if e.StatusCode != 0 {
                return fmt.Sprintf("webhook %q: invalid response (HTTP %d): %v", e.WebhookName, e.StatusCode, e.Err)
        }
        return fmt.Sprintf("webhook %q: invalid response: %v", e.WebhookName, e.Err)
}

// NewTimeoutError creates a new TimeoutError.
func NewTimeoutError(webhookName string, err error) *TimeoutError {
        return &TimeoutError{
                WebhookError: WebhookError{
                        WebhookName: webhookName,
                        Err:         err,
                },
        }
}

// NewNetworkError creates a new NetworkError.
func NewNetworkError(webhookName string, err error) *NetworkError {
        return &NetworkError{
                WebhookError: WebhookError{
                        WebhookName: webhookName,
                        Err:         err,
                },
        }
}

// NewInvalidResponseError creates a new InvalidResponseError.
// statusCode is the HTTP status code from the webhook response (0 if not applicable).
func NewInvalidResponseError(webhookName string, err error, statusCode int) *InvalidResponseError {
        return &InvalidResponseError{
                WebhookError: WebhookError{
                        WebhookName: webhookName,
                        Err:         err,
                },
                StatusCode: statusCode,
        }
}

// IsAlwaysDenyError reports whether the webhook error should deny the request
// regardless of the configured failure policy.
func IsAlwaysDenyError(err error) bool {
        var invalidRespErr *InvalidResponseError
        return errors.As(err, &invalidRespErr) && invalidRespErr.StatusCode == http.StatusUnprocessableEntity
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package mutating implements a mutating webhook middleware for ToolHive.
// It calls external HTTP services to transform MCP requests using JSONPatch (RFC 6902).
package mutating

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/webhook"
)

// MiddlewareParams holds the configuration parameters for the mutating webhook middleware.
type MiddlewareParams struct {
        // Webhooks is the list of mutating webhook configurations to call.
        // Webhooks are called in configuration order; each webhook receives the output
        // of the previous mutation. All patches are applied sequentially.
        Webhooks []webhook.Config `json:"webhooks"`
}

// Validate checks that the MiddlewareParams are valid.
func (p *MiddlewareParams) Validate() error {
        if len(p.Webhooks) == 0 {
                return fmt.Errorf("mutating webhook middleware requires at least one webhook")
        }
        for i, wh := range p.Webhooks {
                if err := wh.Validate(); err != nil {
                        return fmt.Errorf("webhook[%d] (%q): %w", i, wh.Name, err)
                }
        }
        return nil
}

// FactoryMiddlewareParams extends MiddlewareParams with context for the factory.
type FactoryMiddlewareParams struct {
        MiddlewareParams
        // ServerName is the name of the ToolHive instance.
        ServerName string `json:"server_name"`
        // Transport is the transport type (e.g., sse, stdio).
        Transport string `json:"transport"`
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mutating

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

        "github.com/google/uuid"
        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/webhook"
)

// MiddlewareType is the type constant for the mutating webhook middleware.
const MiddlewareType = "mutating-webhook"

// Middleware wraps mutating webhook functionality for the factory pattern.
type Middleware struct {
        handler types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *Middleware) Handler() types.MiddlewareFunction {
        return m.handler
}

// Close cleans up any resources used by the middleware.
func (*Middleware) Close() error {
        return nil
}

type clientExecutor struct {
        client *webhook.Client
        config webhook.Config
}

// CreateMiddleware is the factory function for mutating webhook middleware.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params FactoryMiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal mutating webhook middleware parameters: %w", err)
        }

        if err := params.Validate(); err != nil {
                return fmt.Errorf("invalid mutating webhook configuration: %w", err)
        }

        // Create clients for each webhook.
        var executors []clientExecutor
        for i, whCfg := range params.Webhooks {
                hmacSecret, err := webhook.ResolveSecret(context.Background(), whCfg.HMACSecretRef)
                if err != nil {
                        return fmt.Errorf("failed to resolve HMAC secret for webhook[%d] (%q): %w", i, whCfg.Name, err)
                }

                client, err := webhook.NewClient(whCfg, webhook.TypeMutating, hmacSecret)
                if err != nil {
                        return fmt.Errorf("failed to create client for webhook[%d] (%q): %w", i, whCfg.Name, err)
                }
                executors = append(executors, clientExecutor{client: client, config: whCfg})
        }

        mw := &Middleware{
                handler: createMutatingHandler(executors, params.ServerName, params.Transport),
        }
        runner.AddMiddleware(MiddlewareType, mw)
        return nil
}

func createMutatingHandler(executors []clientExecutor, serverName, transport string) types.MiddlewareFunction {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // The configured chain MUST place MCP parsing before this middleware.
                        // A nil POST context may be an admitted response or a missing-parser
                        // misconfiguration. These controls do not apply to non-POST traffic.
                        parsedMCP := mcp.GetParsedMCPRequest(r.Context())
                        if parsedMCP == nil {
                                next.ServeHTTP(w, r)
                                return
                        }

                        // Read the request body to get the raw MCP request, capped at MaxRequestSize
                        // to prevent unbounded memory consumption from oversized inbound requests.
                        r.Body = http.MaxBytesReader(w, r.Body, webhook.MaxRequestSize)
                        bodyBytes, err := io.ReadAll(r.Body)
                        if err != nil {
                                var maxErr *http.MaxBytesError
                                if errors.As(err, &maxErr) {
                                        sendErrorResponse(w, http.StatusRequestEntityTooLarge, "Request body exceeds maximum size", parsedMCP.ID)
                                        return
                                }
                                sendErrorResponse(w, http.StatusInternalServerError, "Failed to read request body", parsedMCP.ID)
                                return
                        }
                        // Restore the request body immediately; we will replace it after mutations.
                        r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))

                        // currentMCPBody is the MCP JSON-RPC body we thread through the webhook chain.
                        // Each successful mutation replaces this with the patched version.
                        currentMCPBody := bodyBytes

                        // Build the base webhook request context (reused across all webhooks).
                        reqContext := &webhook.RequestContext{
                                ServerName: serverName,
                                SourceIP:   readSourceIP(r),
                                Transport:  transport,
                        }

                        // Resolve principal once (same for all webhooks in this chain).
                        var principal *auth.PrincipalInfo
                        if identity, ok := auth.IdentityFromContext(r.Context()); ok {
                                principal = identity.GetPrincipalInfo()
                        }

                        // Execute the webhook chain to apply mutations.
                        mutatedBody, err := executeMutations(r.Context(), executors, currentMCPBody, reqContext, principal, parsedMCP.ID, w)
                        if err != nil {
                                // executeMutations handles writing the error to the response implicitly when it returns an error.
                                return
                        }

                        // A mutating webhook rewrote the body, so the parse cached by ParsingMiddleware
                        // now describes a request the backend will not execute. Refresh it for downstream
                        // consumers and publish it through the holder so outer audit and telemetry
                        // wrappers observe the final request after the inner chain returns.
                        if !bytes.Equal(bodyBytes, mutatedBody) {
                                republished, err := mcp.RepublishParsedMCPRequest(r, mutatedBody)
                                if err != nil {
                                        var batchErr *mcp.BatchUnsupportedError
                                        if errors.As(err, &batchErr) {
                                                mcp.WriteBatchUnsupportedError(w)
                                                return
                                        }
                                        slog.Error("Mutating webhook produced an unparsable MCP request", "error", err)
                                        sendErrorResponse(w, http.StatusInternalServerError, "Webhook produced an invalid MCP request", parsedMCP.ID)
                                        return
                                }
                                r = republished
                        }

                        // Replace the request body with the (potentially mutated) MCP body for downstream handlers.
                        // ContentLength must be updated alongside it: a patch almost always changes the body
                        // length, and the reverse proxy rejects a forwarded request whose declared length
                        // disagrees with the body it can actually read.
                        r.Body = io.NopCloser(bytes.NewBuffer(mutatedBody))
                        r.ContentLength = int64(len(mutatedBody))
                        next.ServeHTTP(w, r)
                })
        }
}

// executeMutations runs the chain of mutating webhooks sequentially.
// It returns the final mutated body, or an error if the chain was aborted.
// If an error occurs that should abort the request, this function writes the error response.
func executeMutations(
        ctx context.Context,
        executors []clientExecutor,
        initialBody []byte,
        reqContext *webhook.RequestContext,
        principal *auth.PrincipalInfo,
        msgID interface{},
        w http.ResponseWriter,
) ([]byte, error) {
        currentBody := initialBody

        for _, exec := range executors {
                mutatedBody, err := executeSingleMutation(ctx, exec, currentBody, reqContext, principal, msgID, w)
                if err != nil {
                        return nil, err
                }
                currentBody = mutatedBody
        }

        return currentBody, nil
}

// executeSingleMutation applies a single mutating webhook.
func executeSingleMutation(
        ctx context.Context,
        exec clientExecutor,
        currentBody []byte,
        reqContext *webhook.RequestContext,
        principal *auth.PrincipalInfo,
        msgID interface{},
        w http.ResponseWriter,
) ([]byte, error) {
        whName := exec.config.Name

        whReq := &webhook.Request{
                Version:    webhook.APIVersion,
                UID:        uuid.New().String(),
                Timestamp:  time.Now().UTC(),
                MCPRequest: json.RawMessage(currentBody),
                Context:    reqContext,
                Principal:  principal,
        }

        resp, err := exec.client.CallMutating(ctx, whReq)
        if err != nil {
                if webhook.IsAlwaysDenyError(err) {
                        slog.Info("Mutating webhook denied request due to HTTP 422 response", "webhook", whName, "error", err)
                        sendErrorResponse(w, http.StatusUnprocessableEntity, "Request denied by webhook policy", msgID)
                        return nil, err
                }

                if exec.config.FailurePolicy == webhook.FailurePolicyIgnore {
                        slog.Warn("Mutating webhook error ignored due to fail-open policy", "webhook", whName, "error", err)
                        return currentBody, nil
                }
                slog.Error("Mutating webhook error caused request denial", "webhook", whName, "error", err)
                sendErrorResponse(w, http.StatusInternalServerError, "Webhook error", msgID)
                return nil, err
        }

        // Explicit denial from a mutating webhook is always honored, regardless of failure policy.
        // This differs from operational errors (network, timeout) where the failure policy applies.
        if !resp.Allowed {
                slog.Info("Mutating webhook denied request", "webhook", whName, "reason", resp.Reason)
                sendErrorResponse(w, http.StatusForbidden, "Request denied by webhook policy", msgID)
                return nil, fmt.Errorf("webhook denied request")
        }

        if resp.PatchType == "" || len(resp.Patch) == 0 {
                return currentBody, nil
        }

        if resp.PatchType != patchTypeJSONPatch {
                slog.Error("Mutating webhook returned unsupported patch type", "webhook", whName, "patch_type", resp.PatchType)
                if exec.config.FailurePolicy == webhook.FailurePolicyIgnore {
                        return currentBody, nil
                }
                sendErrorResponse(w, http.StatusInternalServerError, "Unsupported patch type from webhook", msgID)
                return nil, fmt.Errorf("unsupported patch type")
        }

        return applyMutationPatch(resp, whName, exec.config.FailurePolicy, currentBody, msgID, w)
}

func applyMutationPatch(
        resp *webhook.MutatingResponse,
        whName string,
        failurePolicy webhook.FailurePolicy,
        currentBody []byte,
        msgID interface{},
        w http.ResponseWriter,
) ([]byte, error) {
        var patchOps []JSONPatchOp
        if err := json.Unmarshal(resp.Patch, &patchOps); err != nil {
                slog.Error("Mutating webhook returned malformed patch", "webhook", whName, "error", err)
                if failurePolicy == webhook.FailurePolicyIgnore {
                        return currentBody, nil
                }
                sendErrorResponse(w, http.StatusInternalServerError, "Malformed patch from webhook", msgID)
                return nil, err
        }

        if err := ValidatePatch(patchOps); err != nil {
                slog.Error("Mutating webhook patch failed validation", "webhook", whName, "error", err)
                if failurePolicy == webhook.FailurePolicyIgnore {
                        return currentBody, nil
                }
                sendErrorResponse(w, http.StatusInternalServerError, "Invalid patch from webhook", msgID)
                return nil, err
        }

        if !IsPatchScopedToMCPRequest(patchOps) {
                slog.Error("Mutating webhook patch targets fields outside mcp_request — rejected", "webhook", whName)
                if failurePolicy == webhook.FailurePolicyIgnore {
                        return currentBody, nil
                }
                sendErrorResponse(w, http.StatusInternalServerError, "Patch must be scoped to mcp_request", msgID)
                return nil, fmt.Errorf("patch scope violation")
        }

        envelopeJSON, err := json.Marshal(struct {
                MCPRequest json.RawMessage `json:"mcp_request"`
        }{
                MCPRequest: json.RawMessage(currentBody),
        })
        if err != nil {
                slog.Error("Failed to marshal webhook request envelope", "webhook", whName, "error", err)
                if failurePolicy == webhook.FailurePolicyIgnore {
                        return currentBody, nil
                }
                sendErrorResponse(w, http.StatusInternalServerError, "Internal error applying patch", msgID)
                return nil, err
        }

        patchedEnvelope, err := ApplyPatch(envelopeJSON, patchOps)
        if err != nil {
                slog.Error("Mutating webhook patch application failed", "webhook", whName, "error", err)
                if failurePolicy == webhook.FailurePolicyIgnore {
                        return currentBody, nil
                }
                sendErrorResponse(w, http.StatusInternalServerError, "Failed to apply patch from webhook", msgID)
                return nil, err
        }

        mutatedMCPBody, err := extractMCPRequest(patchedEnvelope)
        if err != nil {
                slog.Error("Failed to extract mcp_request", "webhook", whName, "error", err)
                if failurePolicy == webhook.FailurePolicyIgnore {
                        return currentBody, nil
                }
                sendErrorResponse(w, http.StatusInternalServerError, "Internal error extracting patched request", msgID)
                return nil, err
        }

        slog.Debug("Mutating webhook applied patch successfully", "webhook", whName)
        return mutatedMCPBody, nil
}

// extractMCPRequest extracts the raw mcp_request bytes from a patched webhook envelope.
func extractMCPRequest(envelope []byte) ([]byte, error) {
        var env struct {
                MCPRequest json.RawMessage `json:"mcp_request"`
        }
        if err := json.Unmarshal(envelope, &env); err != nil {
                return nil, fmt.Errorf("failed to unmarshal patched envelope: %w", err)
        }
        if len(env.MCPRequest) == 0 {
                return nil, fmt.Errorf("mcp_request field missing or empty in patched envelope")
        }
        return env.MCPRequest, nil
}

func readSourceIP(r *http.Request) string {
        return r.RemoteAddr
}

func sendErrorResponse(w http.ResponseWriter, statusCode int, message string, msgID interface{}) {
        id, err := mcp.ConvertToJSONRPC2ID(msgID)
        if err != nil {
                id = jsonrpc2.ID{} // Use empty ID if conversion fails.
        }
        code := mcp.JSONRPCCodeForStatus(statusCode)

        // Return a JSON-RPC 2.0 error so MCP clients can parse the denial.
        errResp := &jsonrpc2.Response{
                ID:    id,
                Error: jsonrpc2.NewError(code, message),
        }

        // The helper encodes before writing any header, so an encode failure never
        // leaves a half-written response (e.g. a status header followed by a
        // second write). Nothing to do on a write error for a denial body.
        _ = mcp.WriteJSONRPCError(w, statusCode, errResp)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package mutating

import (
        "encoding/json"
        "fmt"
        "strings"

        jsonpatch "github.com/evanphx/json-patch/v5"
)

// patchTypeJSONPatch is the patch_type value for RFC 6902 JSON Patch.
const patchTypeJSONPatch = "json_patch"

// mcpRequestPathPrefix is the required prefix for all patch paths.
// Patches are scoped to the mcp_request container only.
const mcpRequestPathPrefix = "/mcp_request/"

// validOps is the set of valid RFC 6902 operations.
var validOps = map[string]bool{
        "add":     true,
        "remove":  true,
        "replace": true,
        "copy":    true,
        "move":    true,
        "test":    true,
}

// JSONPatchOp represents a single RFC 6902 JSON Patch operation.
type JSONPatchOp struct {
        // Op is the patch operation type (add, remove, replace, copy, move, test).
        Op string `json:"op"`
        // Path is the JSON Pointer (RFC 6901) path to apply the operation to.
        Path string `json:"path"`
        // Value is the value to use for add, replace, and test operations.
        Value json.RawMessage `json:"value,omitempty"`
        // From is the source path for copy and move operations.
        From string `json:"from,omitempty"`
}

// ValidatePatch checks that all operations in the patch are well-formed.
// It validates that all operations are supported RFC 6902 types and paths are non-empty.
func ValidatePatch(patch []JSONPatchOp) error {
        for i, op := range patch {
                if !validOps[op.Op] {
                        return fmt.Errorf("patch[%d]: unsupported operation %q (valid ops: add, remove, replace, copy, move, test)", i, op.Op)
                }
                if op.Path == "" {
                        return fmt.Errorf("patch[%d]: path is required", i)
                }
                // copy and move also require a From field.
                if (op.Op == "copy" || op.Op == "move") && op.From == "" {
                        return fmt.Errorf("patch[%d]: %q operation requires a 'from' field", i, op.Op)
                }
        }
        return nil
}

// IsPatchScopedToMCPRequest returns true if all patch operations target paths
// within the mcp_request container. This prevents webhooks from accidentally
// or maliciously modifying principal, context, or other immutable envelope fields.
// The root "/mcp_request" path is intentionally rejected so webhooks must make
// granular changes beneath the MCP request instead of replacing it wholesale.
func IsPatchScopedToMCPRequest(patch []JSONPatchOp) bool {
        for _, op := range patch {
                if !strings.HasPrefix(op.Path, mcpRequestPathPrefix) {
                        return false
                }
                // For copy/move, also check the From path.
                if (op.Op == "copy" || op.Op == "move") && op.From != "" {
                        if !strings.HasPrefix(op.From, mcpRequestPathPrefix) {
                                return false
                        }
                }
        }
        return true
}

// ApplyPatch applies a set of RFC 6902 JSON Patch operations to the original JSON document.
// Returns the patched JSON document. The patch operations are applied in order.
func ApplyPatch(original []byte, patch []JSONPatchOp) ([]byte, error) {
        // Marshal the patch ops to JSON so the library can parse them.
        patchJSON, err := json.Marshal(patch)
        if err != nil {
                return nil, fmt.Errorf("failed to marshal patch operations: %w", err)
        }

        jp, err := jsonpatch.DecodePatch(patchJSON)
        if err != nil {
                return nil, fmt.Errorf("failed to decode JSON patch: %w", err)
        }

        patched, err := jp.Apply(original)
        if err != nil {
                return nil, fmt.Errorf("failed to apply JSON patch: %w", err)
        }

        return patched, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package webhook

import (
        "context"
        "fmt"
        "os"
        "path/filepath"

        "github.com/stacklok/toolhive/pkg/config"
        "github.com/stacklok/toolhive/pkg/secrets"
)

// ResolveSecret resolves a webhook secret reference using the active secrets provider.
// It accepts either a raw secret name or CLI-style "<name>,target=<target>" format.
func ResolveSecret(ctx context.Context, ref string) ([]byte, error) {
        if ref == "" {
                return nil, nil
        }

        if filepath.IsAbs(ref) {
                // #nosec G304 -- absolute refs are generated by the operator for mounted webhook Secret files.
                value, err := os.ReadFile(ref)
                if err != nil {
                        return nil, fmt.Errorf("failed to read webhook secret file %q: %w", ref, err)
                }
                return value, nil
        }

        secretName := ref
        // Raw secret names are also valid references, so parse errors intentionally fall back to ref as-is.
        if secretParam, err := secrets.ParseSecretParameter(ref); err == nil {
                secretName = secretParam.Name
        }

        provider, err := resolveSecretProvider()
        if err != nil {
                return nil, fmt.Errorf("failed to initialize webhook secrets provider: %w", err)
        }

        value, err := provider.GetSecret(ctx, secretName)
        if err != nil {
                return nil, fmt.Errorf("failed to resolve webhook secret %q: %w", secretName, err)
        }

        return []byte(value), nil
}

func resolveSecretProvider() (secrets.Provider, error) {
        if providerType := os.Getenv(secrets.ProviderEnvVar); providerType != "" {
                return secrets.CreateSecretProvider(secrets.ProviderType(providerType))
        }

        cfg := config.NewDefaultProvider().GetConfig()
        if !cfg.Secrets.SetupCompleted {
                return nil, secrets.ErrSecretsNotSetup
        }

        providerType, err := cfg.Secrets.GetProviderType()
        if err != nil {
                return nil, fmt.Errorf("failed to get configured secrets provider type: %w", err)
        }

        return secrets.CreateProvider(providerType, secrets.WithUserFacing())
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package webhook

import (
        "crypto/hmac"
        "crypto/sha256"
        "encoding/hex"
        "fmt"
        "strings"
)

// Header names for webhook HMAC signing.
const (
        // SignatureHeader is the HTTP header containing the HMAC signature.
        SignatureHeader = "X-ToolHive-Signature"
        // TimestampHeader is the HTTP header containing the Unix timestamp.
        TimestampHeader = "X-ToolHive-Timestamp"
)

// signaturePrefix is the prefix for the HMAC-SHA256 signature value.
const signaturePrefix = "sha256="

// SignPayload computes an HMAC-SHA256 signature over the given timestamp and
// payload. The signature is computed over the string "timestamp.payload" and
// returned in the format "sha256=<hex-encoded-signature>".
func SignPayload(secret []byte, timestamp int64, payload []byte) string {
        mac := hmac.New(sha256.New, secret)
        // Write the message: "timestamp.payload"
        msg := fmt.Sprintf("%d.", timestamp)
        mac.Write([]byte(msg))
        mac.Write(payload)
        return signaturePrefix + hex.EncodeToString(mac.Sum(nil))
}

// VerifySignature verifies an HMAC-SHA256 signature against the given timestamp
// and payload. The signature should be in the format "sha256=<hex-encoded-signature>".
// Comparison is done in constant time to prevent timing attacks.
//
// Note: This function only verifies cryptographic correctness. Callers should
// independently verify that the timestamp is recent (e.g., within 5 minutes)
// to prevent replay attacks.
func VerifySignature(secret []byte, timestamp int64, payload []byte, signature string) bool {
        if !strings.HasPrefix(signature, signaturePrefix) {
                return false
        }

        sigBytes, err := hex.DecodeString(strings.TrimPrefix(signature, signaturePrefix))
        if err != nil {
                return false
        }

        mac := hmac.New(sha256.New, secret)
        msg := fmt.Sprintf("%d.", timestamp)
        mac.Write([]byte(msg))
        mac.Write(payload)

        return hmac.Equal(mac.Sum(nil), sigBytes)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package webhook implements the core types, HTTP client, HMAC signing,
// and error handling for ToolHive's dynamic webhook middleware system.
package webhook

import (
        "encoding/json"
        "fmt"
        "net/url"
        "os"
        "time"

        "github.com/stacklok/toolhive/pkg/auth"
)

// APIVersion is the version of the webhook API protocol.
const APIVersion = "v0.1.0"

// DefaultTimeout is the default timeout for webhook HTTP calls.
const DefaultTimeout = 10 * time.Second

// MaxTimeout is the maximum allowed timeout for webhook HTTP calls.
const MaxTimeout = 30 * time.Second

// MinTimeout is the minimum allowed timeout for webhook HTTP calls.
const MinTimeout = 1 * time.Second

// MaxResponseSize is the maximum allowed size in bytes for webhook responses (1 MB).
const MaxResponseSize = 1 << 20

// MaxRequestSize is the maximum allowed size in bytes for inbound webhook
// middleware request bodies (1 MB). Requests exceeding this size are
// rejected with HTTP 413 before the body is buffered or forwarded.
const MaxRequestSize = 1 << 20

// Type indicates whether a webhook is validating or mutating.
type Type string

const (
        // TypeValidating indicates a validating webhook that accepts or denies requests.
        TypeValidating Type = "validating"
        // TypeMutating indicates a mutating webhook that transforms requests.
        TypeMutating Type = "mutating"
)

// FailurePolicy defines how webhook errors are handled.
type FailurePolicy string

const (
        // FailurePolicyFail denies the request on webhook error (fail-closed).
        FailurePolicyFail FailurePolicy = "fail"
        // FailurePolicyIgnore allows the request on webhook error (fail-open).
        FailurePolicyIgnore FailurePolicy = "ignore"
)

// TLSConfig holds TLS-related configuration for webhook HTTP communication.
type TLSConfig struct {
        // CABundlePath is the path to a CA certificate bundle for server verification.
        CABundlePath string `json:"ca_bundle_path,omitempty" yaml:"ca_bundle_path,omitempty"`
        // ClientCertPath is the path to a client certificate for mTLS.
        ClientCertPath string `json:"client_cert_path,omitempty" yaml:"client_cert_path,omitempty"`
        // ClientKeyPath is the path to a client key for mTLS.
        ClientKeyPath string `json:"client_key_path,omitempty" yaml:"client_key_path,omitempty"`
        // InsecureSkipVerify disables server certificate verification.
        // WARNING: This should only be used for development/testing.
        InsecureSkipVerify bool `json:"insecure_skip_verify,omitempty" yaml:"insecure_skip_verify,omitempty"`
}

// Config holds the configuration for a single webhook.
type Config struct {
        // Name is a unique identifier for this webhook.
        Name string `json:"name" yaml:"name"`
        // URL is the HTTPS endpoint to call.
        URL string `json:"url" yaml:"url"`
        // Timeout is the maximum time to wait for a webhook response.
        Timeout time.Duration `json:"timeout" yaml:"timeout" swaggertype:"primitive,integer"`
        // FailurePolicy determines behavior when the webhook call fails.
        FailurePolicy FailurePolicy `json:"failure_policy" yaml:"failure_policy"`
        // TLSConfig holds optional TLS configuration (CA bundles, client certs).
        TLSConfig *TLSConfig `json:"tls_config,omitempty" yaml:"tls_config,omitempty"`
        // HMACSecretRef is an optional reference to an HMAC secret for payload signing.
        HMACSecretRef string `json:"hmac_secret_ref,omitempty" yaml:"hmac_secret_ref,omitempty"`
}

// ValidateDefinition checks semantic validity of a webhook config without touching the filesystem.
// This is useful when a higher layer is responsible for mounting referenced files later.
func (c *Config) ValidateDefinition() error {
        if c.Name == "" {
                return fmt.Errorf("webhook name is required")
        }
        if c.URL == "" {
                return fmt.Errorf("webhook URL is required")
        }
        parsed, err := url.ParseRequestURI(c.URL)
        if err != nil {
                return fmt.Errorf("webhook URL is invalid: %w", err)
        }
        // Enforce HTTPS unless InsecureSkipVerify is explicitly set (for in-cluster HTTP endpoints).
        insecureHTTPAllowed := c.TLSConfig != nil && c.TLSConfig.InsecureSkipVerify
        if parsed.Scheme != "https" && !insecureHTTPAllowed {
                return fmt.Errorf("webhook URL must use HTTPS (set insecure_skip_verify to allow HTTP for development/in-cluster use)")
        }
        if c.FailurePolicy != FailurePolicyFail && c.FailurePolicy != FailurePolicyIgnore {
                return fmt.Errorf("webhook failure_policy must be %q or %q, got %q",
                        FailurePolicyFail, FailurePolicyIgnore, c.FailurePolicy)
        }
        if c.Timeout != 0 && c.Timeout < MinTimeout {
                return fmt.Errorf("webhook timeout must be between %v and %v", MinTimeout, MaxTimeout)
        }
        if c.Timeout > MaxTimeout {
                return fmt.Errorf("webhook timeout %v exceeds maximum %v", c.Timeout, MaxTimeout)
        }
        if c.TLSConfig != nil {
                if err := validateTLSConfigShape(c.TLSConfig); err != nil {
                        return fmt.Errorf("webhook TLS config: %w", err)
                }
        }
        return nil
}

// Validate checks that the WebhookConfig has valid required fields and referenced files exist.
func (c *Config) Validate() error {
        if err := c.ValidateDefinition(); err != nil {
                return err
        }

        if c.TLSConfig != nil {
                if err := validateTLSConfigFiles(c.TLSConfig); err != nil {
                        return fmt.Errorf("webhook TLS config: %w", err)
                }
        }
        return nil
}

// UnmarshalJSON accepts webhook timeout values as either strings (for example "5s")
// or numeric nanoseconds, while keeping the rest of the struct on the standard JSON path.
func (c *Config) UnmarshalJSON(data []byte) error {
        var raw struct {
                Name          string          `json:"name"`
                URL           string          `json:"url"`
                Timeout       json.RawMessage `json:"timeout"`
                FailurePolicy FailurePolicy   `json:"failure_policy"`
                TLSConfig     *TLSConfig      `json:"tls_config,omitempty"`
                HMACSecretRef string          `json:"hmac_secret_ref,omitempty"`
        }
        if err := json.Unmarshal(data, &raw); err != nil {
                return err
        }

        *c = Config{
                Name:          raw.Name,
                URL:           raw.URL,
                FailurePolicy: raw.FailurePolicy,
                TLSConfig:     raw.TLSConfig,
                HMACSecretRef: raw.HMACSecretRef,
        }

        if len(raw.Timeout) == 0 || string(raw.Timeout) == "null" {
                c.Timeout = DefaultTimeout
                return nil
        }

        if raw.Timeout[0] == '"' {
                var timeoutStr string
                if err := json.Unmarshal(raw.Timeout, &timeoutStr); err != nil {
                        return fmt.Errorf("invalid timeout value: %w", err)
                }
                d, err := time.ParseDuration(timeoutStr)
                if err != nil {
                        return fmt.Errorf("invalid timeout value %q: %w", timeoutStr, err)
                }
                c.Timeout = d
                return nil
        }

        var timeoutNanos int64
        if err := json.Unmarshal(raw.Timeout, &timeoutNanos); err != nil {
                return fmt.Errorf("invalid timeout value: %w", err)
        }
        c.Timeout = time.Duration(timeoutNanos)
        return nil
}

// Request is the payload sent to webhook endpoints.
type Request struct {
        // Version is the webhook API protocol version.
        Version string `json:"version"`
        // UID is a unique identifier for this request, used for idempotency.
        UID string `json:"uid"`
        // Timestamp is when the request was created.
        Timestamp time.Time `json:"timestamp"`
        // Principal contains the authenticated user's identity information.
        // Uses PrincipalInfo (not Identity) so credentials never enter the webhook payload.
        Principal *auth.PrincipalInfo `json:"principal"`
        // MCPRequest is the raw MCP JSON-RPC request.
        MCPRequest json.RawMessage `json:"mcp_request"`
        // Context provides additional metadata about the request origin.
        Context *RequestContext `json:"context"`
}

// RequestContext provides metadata about the request origin and environment.
type RequestContext struct {
        // ServerName is the ToolHive/vMCP instance name handling the request.
        ServerName string `json:"server_name"`
        // BackendServer is the actual MCP server being proxied (when using vMCP).
        BackendServer string `json:"backend_server,omitempty"`
        // Namespace is the Kubernetes namespace, if applicable.
        Namespace string `json:"namespace,omitempty"`
        // SourceIP is the client's IP address.
        SourceIP string `json:"source_ip"`
        // Transport is the connection transport type (e.g., "sse", "stdio").
        Transport string `json:"transport"`
}

// Response is the response from a validating webhook.
type Response struct {
        // Version is the webhook API protocol version.
        Version string `json:"version"`
        // UID is the unique request identifier, echoed back for correlation.
        UID string `json:"uid"`
        // Allowed indicates whether the request is permitted.
        Allowed bool `json:"allowed"`
        // Code is an optional HTTP status code for denied requests.
        Code int `json:"code,omitempty"`
        // Message is an optional human-readable explanation.
        Message string `json:"message,omitempty"`
        // Reason is an optional machine-readable denial reason.
        Reason string `json:"reason,omitempty"`
        // Details contains optional structured information about the denial.
        Details map[string]string `json:"details,omitempty"`
}

// MutatingResponse is the response from a mutating webhook.
type MutatingResponse struct {
        Response
        // PatchType indicates the type of patch (e.g., "json_patch").
        PatchType string `json:"patch_type,omitempty"`
        // Patch contains the JSON Patch operations to apply.
        Patch json.RawMessage `json:"patch,omitempty"`
}

// validateTLSConfigShape validates TLS configuration for consistency.
func validateTLSConfigShape(cfg *TLSConfig) error {
        // If one of client cert/key is provided, both must be present.
        if (cfg.ClientCertPath == "") != (cfg.ClientKeyPath == "") {
                return fmt.Errorf("both client_cert_path and client_key_path must be provided for mTLS")
        }
        return nil
}

// validateTLSConfigFiles validates referenced TLS files exist.
func validateTLSConfigFiles(cfg *TLSConfig) error {
        if cfg.CABundlePath != "" {
                if _, err := os.Stat(cfg.CABundlePath); err != nil {
                        return fmt.Errorf("ca_bundle_path %q not found: %w", cfg.CABundlePath, err)
                }
        }
        if cfg.ClientCertPath != "" {
                if _, err := os.Stat(cfg.ClientCertPath); err != nil {
                        return fmt.Errorf("client_cert_path %q not found: %w", cfg.ClientCertPath, err)
                }
        }
        if cfg.ClientKeyPath != "" {
                if _, err := os.Stat(cfg.ClientKeyPath); err != nil {
                        return fmt.Errorf("client_key_path %q not found: %w", cfg.ClientKeyPath, err)
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package validating implements a validating webhook middleware for ToolHive.
// It calls external HTTP services to approve or deny MCP requests.
package validating

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/webhook"
)

// MiddlewareParams holds the configuration parameters for the validating webhook middleware.
type MiddlewareParams struct {
        // Webhooks is the list of validating webhook configurations to call.
        // Webhooks are called in configuration order; if any webhook denies the request,
        // the request is rejected. All webhooks must allow the request for it to proceed.
        Webhooks []webhook.Config `json:"webhooks"`
}

// Validate checks that the MiddlewareParams are valid.
func (p *MiddlewareParams) Validate() error {
        if len(p.Webhooks) == 0 {
                return fmt.Errorf("validating webhook middleware requires at least one webhook")
        }
        for i, wh := range p.Webhooks {
                if err := wh.Validate(); err != nil {
                        return fmt.Errorf("webhook[%d] (%q): %w", i, wh.Name, err)
                }
        }
        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package validating

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

        "github.com/google/uuid"
        "golang.org/x/exp/jsonrpc2"

        "github.com/stacklok/toolhive/pkg/auth"
        "github.com/stacklok/toolhive/pkg/mcp"
        "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/webhook"
)

// MiddlewareType is the type constant for the validating webhook middleware.
const MiddlewareType = "validating-webhook"

// FactoryMiddlewareParams extends MiddlewareParams with context for the factory.
type FactoryMiddlewareParams struct {
        MiddlewareParams
        // ServerName is the name of the ToolHive instance.
        ServerName string `json:"server_name"`
        // Transport is the transport type (e.g., sse, stdio).
        Transport string `json:"transport"`
}

// Middleware wraps validating webhook functionality for the factory pattern.
type Middleware struct {
        handler types.MiddlewareFunction
}

// Handler returns the middleware function used by the proxy.
func (m *Middleware) Handler() types.MiddlewareFunction {
        return m.handler
}

// Close cleans up any resources used by the middleware.
func (*Middleware) Close() error {
        return nil
}

type clientExecutor struct {
        client *webhook.Client
        config webhook.Config
}

// CreateMiddleware is the factory function for validating webhook middleware.
func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRunner) error {
        var params FactoryMiddlewareParams
        if err := json.Unmarshal(config.Parameters, &params); err != nil {
                return fmt.Errorf("failed to unmarshal validating webhook middleware parameters: %w", err)
        }

        if err := params.Validate(); err != nil {
                return fmt.Errorf("invalid validating webhook configuration: %w", err)
        }

        // Create clients for each webhook
        var executors []clientExecutor
        for i, whCfg := range params.Webhooks {
                hmacSecret, err := webhook.ResolveSecret(context.Background(), whCfg.HMACSecretRef)
                if err != nil {
                        return fmt.Errorf("failed to resolve HMAC secret for webhook[%d] (%q): %w", i, whCfg.Name, err)
                }

                client, err := webhook.NewClient(whCfg, webhook.TypeValidating, hmacSecret)
                if err != nil {
                        return fmt.Errorf("failed to create client for webhook[%d] (%q): %w", i, whCfg.Name, err)
                }
                executors = append(executors, clientExecutor{client: client, config: whCfg})
        }

        mw := &Middleware{
                handler: createValidatingHandler(executors, params.ServerName, params.Transport),
        }
        runner.AddMiddleware(MiddlewareType, mw)
        return nil
}

func createValidatingHandler(executors []clientExecutor, serverName, transport string) types.MiddlewareFunction {
        return func(next http.Handler) http.Handler {
                return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                        // The configured chain MUST place MCP parsing before this middleware.
                        // A nil POST context may be an admitted response or a missing-parser
                        // misconfiguration. These controls do not apply to non-POST traffic.
                        parsedMCP := mcp.GetParsedMCPRequest(r.Context())
                        if parsedMCP == nil {
                                next.ServeHTTP(w, r)
                                return
                        }

                        // Read the request body to get the raw MCP request, capped at MaxRequestSize
                        // to prevent unbounded memory consumption from oversized inbound requests.
                        r.Body = http.MaxBytesReader(w, r.Body, webhook.MaxRequestSize)
                        bodyBytes, err := io.ReadAll(r.Body)
                        if err != nil {
                                var maxErr *http.MaxBytesError
                                if errors.As(err, &maxErr) {
                                        sendErrorResponse(w, http.StatusRequestEntityTooLarge, "Request body exceeds maximum size", parsedMCP.ID)
                                        return
                                }
                                sendErrorResponse(w, http.StatusInternalServerError, "Failed to read request body", parsedMCP.ID)
                                return
                        }
                        // Restore the request body for downstream handlers
                        r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))

                        // Build the webhook request payload
                        reqUID := uuid.New().String()
                        whReq := &webhook.Request{
                                Version:    webhook.APIVersion,
                                UID:        reqUID,
                                Timestamp:  time.Now().UTC(),
                                MCPRequest: json.RawMessage(bodyBytes),
                                Context: &webhook.RequestContext{
                                        ServerName: serverName,
                                        SourceIP:   readSourceIP(r),
                                        Transport:  transport,
                                },
                        }

                        // Add Principal if authenticated
                        if identity, ok := auth.IdentityFromContext(r.Context()); ok {
                                whReq.Principal = identity.GetPrincipalInfo()
                        }

                        // Call each webhook in order
                        for _, exec := range executors {
                                whName := exec.config.Name

                                resp, err := exec.client.Call(r.Context(), whReq)
                                if err != nil {
                                        if webhook.IsAlwaysDenyError(err) {
                                                slog.Info("Validating webhook denied request due to HTTP 422 response",
                                                        "webhook", whName, "error", err)
                                                sendErrorResponse(w, http.StatusForbidden, "Request denied by policy", parsedMCP.ID)
                                                return
                                        }

                                        // Handle error based on failure policy
                                        if exec.config.FailurePolicy == webhook.FailurePolicyIgnore {
                                                slog.Warn("Validating webhook error ignored due to fail-open policy",
                                                        "webhook", whName, "error", err)
                                                continue
                                        }

                                        slog.Error("Validating webhook error caused request denial",
                                                "webhook", whName, "error", err)
                                        sendErrorResponse(w, http.StatusForbidden, "Request denied by policy", parsedMCP.ID)
                                        return
                                }

                                if !resp.Allowed {
                                        slog.Info("Validating webhook denied request", "webhook", whName, "reason", resp.Reason, "message", resp.Message)

                                        // Prevent information leaks by ignoring the webhook's message
                                        sendErrorResponse(w, http.StatusForbidden, "Request denied by policy", parsedMCP.ID)
                                        return
                                }
                        }

                        // All webhooks allowed or ignored errors
                        next.ServeHTTP(w, r)
                })
        }
}

func readSourceIP(r *http.Request) string {
        // Let runner handle X-Forwarded-For if TrustProxyHeaders is set.
        // For now, simple RemoteAddr.
        return r.RemoteAddr
}

func sendErrorResponse(w http.ResponseWriter, statusCode int, message string, msgID interface{}) {
        id, err := mcp.ConvertToJSONRPC2ID(msgID)
        if err != nil {
                id = jsonrpc2.ID{} // Use empty ID if conversion fails
        }
        code := mcp.JSONRPCCodeForStatus(statusCode)

        // Return a JSON-RPC 2.0 error so MCP clients can parse the denial.
        // The HTTP status code signals the error at the transport level; the JSON-RPC body carries the detail.
        errResp := &jsonrpc2.Response{
                ID:    id,
                Error: jsonrpc2.NewError(code, message),
        }

        // The helper encodes before writing any header, so an encode failure never
        // leaves a half-written response (e.g. a status header followed by a
        // second write). Nothing to do on a write error for a denial body.
        _ = mcp.WriteJSONRPCError(w, statusCode, errResp)
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package statuses

import (
        "context"
        "encoding/json"
        "fmt"
        "log/slog"
        "os"
        "path/filepath"
        "strings"
        "time"

        "github.com/adrg/xdg"

        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/core"
        "github.com/stacklok/toolhive/pkg/fileutils"
        "github.com/stacklok/toolhive/pkg/labels"
        "github.com/stacklok/toolhive/pkg/lockfile"
        "github.com/stacklok/toolhive/pkg/process"
        "github.com/stacklok/toolhive/pkg/state"
        "github.com/stacklok/toolhive/pkg/transport"
        transporttypes "github.com/stacklok/toolhive/pkg/transport/types"
        "github.com/stacklok/toolhive/pkg/workloads/types"
)

const (
        // statusesPrefix is the prefix used for status files in the XDG data directory
        statusesPrefix = "toolhive/statuses"
        // lockTimeout is the maximum time to wait for a file lock
        lockTimeout = 1 * time.Second
        // lockRetryInterval is the interval between lock attempts
        lockRetryInterval = 100 * time.Millisecond
)

// NewFileStatusManager creates a new file-based StatusManager.
// Status files will be stored in the XDG data directory under "statuses/".
func NewFileStatusManager(runtime rt.Runtime) (StatusManager, error) {
        // Get the base directory using XDG data directory
        baseDir, err := xdg.DataFile(statusesPrefix)
        if err != nil {
                return nil, fmt.Errorf("failed to get directory for status files: %w", err)
        }

        // Ensure the base directory exists (equivalent to mkdir -p)
        if err := os.MkdirAll(baseDir, 0o750); err != nil {
                return nil, fmt.Errorf("failed to create status directory %s: %w", baseDir, err)
        }

        // Create run config store for accessing run configurations
        runConfigStore, err := state.NewRunConfigStore(state.DefaultAppName)
        if err != nil {
                return nil, fmt.Errorf("failed to create run config store: %w", err)
        }

        return &fileStatusManager{
                baseDir:        baseDir,
                runtime:        runtime,
                runConfigStore: runConfigStore,
        }, nil
}

// fileStatusManager is an implementation of StatusManager that persists
// workload status to files on disk with JSON serialization and file locking
// to prevent concurrent access issues.
type fileStatusManager struct {
        baseDir string
        runtime rt.Runtime
        // runConfigStore is used to access run configurations without import cycles
        // TODO: This is a temporary solution to check if a workload is remote
        runConfigStore state.Store
}

// isRemoteWorkload checks if a workload is remote by attempting to load its run configuration
// and checking if it has a RemoteURL field set.
// TODO: This is a temporary solution to check if a workload is remote
// because of the import cycle between this package and the runconfig package.
// We can easily load run config and check if it has a RemoteURL field set when we resolve the import cycle.
func (f *fileStatusManager) isRemoteWorkload(ctx context.Context, workloadName string) (bool, error) {
        // Check if the run configuration exists
        exists, err := f.runConfigStore.Exists(ctx, workloadName)
        if err != nil {
                return false, err
        }

        if !exists {
                return false, rt.ErrWorkloadNotFound
        }

        // Get a reader for the run configuration
        reader, err := f.runConfigStore.GetReader(ctx, workloadName)
        if err != nil {
                return false, err
        }
        defer func() {
                if err := reader.Close(); err != nil {
                        slog.Warn("failed to close reader", "error", err)
                }
        }()

        // Parse the JSON to check for remote_url field
        var config struct {
                RemoteURL string `json:"remote_url"`
        }
        decoder := json.NewDecoder(reader)
        if err := decoder.Decode(&config); err != nil {
                return false, err
        }

        // Check if the remote_url field is set
        return strings.TrimSpace(config.RemoteURL) != "", nil
}

// remoteWorkloadConfig is a minimal struct to parse only the fields we need from RunConfig
// without importing the runner package (which would create a circular dependency).
type remoteWorkloadConfig struct {
        RemoteURL       string            `json:"remote_url"`
        Port            int               `json:"port"`
        Transport       string            `json:"transport"` // Parse as string, then convert to transport type
        ProxyMode       string            `json:"proxy_mode"`
        Group           string            `json:"group"`
        ContainerLabels map[string]string `json:"container_labels"`
}

// populateRemoteWorkloadData populates a workload with data from the run config for remote workloads.
// This includes URL, port, transport type, and other fields that are stored in the run config.
func (f *fileStatusManager) populateRemoteWorkloadData(ctx context.Context, workload *core.Workload) error {
        // Load the run configuration JSON
        reader, err := f.runConfigStore.GetReader(ctx, workload.Name)
        if err != nil {
                return fmt.Errorf("failed to load run config for remote workload %s: %w", workload.Name, err)
        }
        defer func() {
                if err := reader.Close(); err != nil {
                        slog.Warn("failed to close reader", "error", err)
                }
        }()

        // Parse only the fields we need
        var config remoteWorkloadConfig
        decoder := json.NewDecoder(reader)
        if err := decoder.Decode(&config); err != nil {
                return fmt.Errorf("failed to decode run config for remote workload %s: %w", workload.Name, err)
        }

        if config.RemoteURL == "" {
                return fmt.Errorf("workload %s does not have a remote URL", workload.Name)
        }

        // Parse the transport type from string
        transportType, err := transporttypes.ParseTransportType(config.Transport)
        if err != nil {
                return fmt.Errorf("failed to parse transport type for remote workload %s: %w", workload.Name, err)
        }

        proxyURL := ""
        if config.Port > 0 {
                proxyURL = transport.GenerateMCPServerURL(
                        transportType.String(),
                        config.ProxyMode,
                        transport.LocalhostIPv4,
                        config.Port,
                        workload.Name,
                        config.RemoteURL,
                )
        }

        effectiveProxyMode := types.GetEffectiveProxyMode(transportType, config.ProxyMode)

        workload.Package = "remote"
        workload.URL = proxyURL
        workload.Port = config.Port
        workload.TransportType = transportType
        workload.ProxyMode = effectiveProxyMode
        workload.Group = config.Group
        workload.Labels = config.ContainerLabels
        workload.Remote = true

        return nil
}

// workloadStatusFile represents the JSON structure stored on disk
type workloadStatusFile struct {
        Status        rt.WorkloadStatus `json:"status"`
        StatusContext string            `json:"status_context,omitempty"`
        CreatedAt     time.Time         `json:"created_at"`
        UpdatedAt     time.Time         `json:"updated_at"`
        ProcessID     int               `json:"process_id"`
}

// GetWorkload retrieves the status of a workload by its name.
func (f *fileStatusManager) GetWorkload(ctx context.Context, workloadName string) (core.Workload, error) {
        var pid int
        result := core.Workload{Name: workloadName}
        fileFound := false

        err := f.withFileReadLock(ctx, workloadName, func(statusFilePath string) error {
                // Check if file exists
                if _, err := os.Stat(statusFilePath); os.IsNotExist(err) {
                        // File doesn't exist, we'll fall back to runtime check
                        return nil
                } else if err != nil {
                        return fmt.Errorf("failed to check status file for workload %s: %w", workloadName, err)
                }

                statusFile, err := f.readStatusFile(statusFilePath)
                if err != nil {
                        return fmt.Errorf("failed to read status for workload %s: %w", workloadName, err)
                }

                result.Status = statusFile.Status
                result.StatusContext = statusFile.StatusContext
                result.CreatedAt = statusFile.CreatedAt

                fileFound = true

                // Check if PID migration is needed
                if statusFile.Status == rt.WorkloadStatusRunning && statusFile.ProcessID == 0 {
                        // Try PID migration - the migration function will handle cases
                        // where container info is not available gracefully
                        if migratedPID, wasMigrated := f.migratePIDFromFile(workloadName, nil); wasMigrated {
                                // Update the status file with the migrated PID
                                statusFile.ProcessID = migratedPID
                                statusFile.UpdatedAt = time.Now()
                                if err := f.writeStatusFile(statusFilePath, *statusFile); err != nil {
                                        slog.Warn("failed to write migrated PID", "workload", workloadName, "error", err)
                                } else {
                                        slog.Debug("successfully migrated PID to status file", "pid", migratedPID, "workload", workloadName)
                                }
                        }
                }

                pid = statusFile.ProcessID

                return nil
        })
        if err != nil {
                return core.Workload{}, err
        }

        // If file was found, check if this is a remote workload
        if fileFound {
                // Check if this is a remote workload using the state package
                remote, err := f.isRemoteWorkload(ctx, workloadName)
                if err != nil {
                        // error is expected
                        slog.Debug("failed to check if workload is remote", "workload", workloadName, "error", err)
                }
                if remote {
                        // Populate remote workload data from run config
                        if err := f.populateRemoteWorkloadData(ctx, &result); err != nil {
                                slog.Warn("failed to populate remote workload data", "workload", workloadName, "error", err)
                                // Mark as remote even if we couldn't load full data
                                result.Remote = true
                                result.Package = "remote (data failed to load)"
                        }
                }

                // If workload is running, validate against runtime
                if result.Status == rt.WorkloadStatusRunning {
                        return f.validateRunningWorkload(ctx, workloadName, result, pid)
                }

                // Return file data
                return result, nil
        }

        // File not found, fall back to runtime check
        return f.getWorkloadFromRuntime(ctx, workloadName)
}

func (f *fileStatusManager) ListWorkloads(ctx context.Context, listAll bool, labelFilters []string) ([]core.Workload, error) {
        // Parse the filters into a format we can use for matching.
        parsedFilters, err := types.ParseLabelFilters(labelFilters)
        if err != nil {
                return nil, fmt.Errorf("failed to parse label filters: %w", err)
        }

        // Get workloads from files BEFORE snapshotting the runtime. The order is
        // load-bearing: a workload's status file only flips to `running` after its
        // container has been created and started, so a file that says `running` in
        // a snapshot taken at time T implies the container existed before T. If a
        // later runtime snapshot lacks it, the container is genuinely gone and
        // handleRuntimeMissing's unhealthy verdict is correct. With the snapshots
        // the other way round, a workload that starts between the runtime snapshot
        // and the file read looks like "file running, runtime missing" and gets its
        // status file permanently overwritten as unhealthy — a healthy, freshly
        // started workload poisoned by a concurrent `thv list` (see #4432 for the
        // broader write-on-read hazard in this path).
        fileWorkloadsWithPID, err := f.getWorkloadsFromFiles()
        if err != nil {
                return nil, fmt.Errorf("failed to get workloads from files: %w", err)
        }

        // Get workloads from runtime
        runtimeContainers, err := f.runtime.ListWorkloads(ctx)
        if err != nil {
                return nil, fmt.Errorf("failed to list workloads from runtime: %w", err)
        }

        // TODO: Fetch the runconfig if present to populate additional fields like package, tool type, group etc.
        // There's currently an import cycle between this package and the runconfig package

        for _, fileWorkload := range fileWorkloadsWithPID {
                if fileWorkload.workload.Remote { // Remote workloads are not managed by the container runtime
                        delete(fileWorkloadsWithPID, fileWorkload.workload.Name) // Skip remote workloads here, we add them in workload manager
                }
        }

        // Create a map of runtime workloads by name for easy lookup
        workloadMap := f.mergeRuntimeAndFileWorkloads(ctx, runtimeContainers, fileWorkloadsWithPID)

        // Convert map to slice and apply filters
        var workloads []core.Workload
        for _, workload := range workloadMap {
                // Apply listAll filter. AuthRetrying is intentionally surfaced
                // alongside Running because the workload is still self-recovering
                // without user intervention on a longer cadence — every other
                // non-Running status is terminal until external intervention,
                // and is hidden.
                if !listAll &&
                        workload.Status != rt.WorkloadStatusRunning &&
                        workload.Status != rt.WorkloadStatusAuthRetrying {
                        continue
                }

                // Apply label filters
                if len(parsedFilters) > 0 {
                        if !types.MatchesLabelFilters(workload.Labels, parsedFilters) {
                                continue
                        }
                }

                workloads = append(workloads, workload)
        }

        return workloads, nil
}

// setWorkloadStatusInternal handles the core logic for updating workload status files.
// pidPtr controls PID behavior: nil means preserve existing PID, non-nil means set to provided value.
func (f *fileStatusManager) setWorkloadStatusInternal(
        ctx context.Context,
        workloadName string,
        status rt.WorkloadStatus,
        contextMsg string,
        pidPtr *int,
) error {
        err := f.withFileLock(ctx, workloadName, func(statusFilePath string) error {
                // Check if file exists
                fileExists := true
                if _, err := os.Stat(statusFilePath); os.IsNotExist(err) {
                        fileExists = false
                } else if err != nil {
                        return fmt.Errorf("failed to check status file for workload %s: %w", workloadName, err)
                }

                var statusFile *workloadStatusFile
                var err error
                now := time.Now()

                if fileExists {
                        // Read existing file to preserve created_at timestamp and other fields
                        statusFile, err = f.readStatusFile(statusFilePath)
                        if err != nil {
                                return fmt.Errorf("failed to read existing status for workload %s: %w", workloadName, err)
                        }
                } else {
                        // Create new status file with CreatedAt set
                        statusFile = &workloadStatusFile{
                                CreatedAt: now,
                        }
                }

                // Update status, context, and optionally PID
                statusFile.Status = status
                statusFile.StatusContext = contextMsg
                statusFile.UpdatedAt = now

                // Only update PID if pidPtr is provided
                if pidPtr != nil {
                        statusFile.ProcessID = *pidPtr
                }

                if err = f.writeStatusFile(statusFilePath, *statusFile); err != nil {
                        return fmt.Errorf("failed to write updated status for workload %s: %w", workloadName, err)
                }

                // Log with appropriate message based on whether PID was set
                if pidPtr != nil {
                        slog.Debug("workload status set with PID", "workload", workloadName, "status", status, "pid", *pidPtr, "context", contextMsg)
                } else {
                        slog.Debug("workload status set", "workload", workloadName, "status", status, "context", contextMsg)
                }
                return nil
        })
        if err != nil {
                if pidPtr != nil {
                        slog.Error("error updating workload status and PID", "workload", workloadName, "error", err)
                } else {
                        slog.Error("error updating workload status", "workload", workloadName, "error", err)
                }
        }
        return err
}

// SetWorkloadStatus sets the status of a workload by its name.
func (f *fileStatusManager) SetWorkloadStatus(
        ctx context.Context,
        workloadName string,
        status rt.WorkloadStatus,
        contextMsg string,
) error {
        return f.setWorkloadStatusInternal(ctx, workloadName, status, contextMsg, nil)
}

// DeleteWorkloadStatus removes the status of a workload by its name.
func (f *fileStatusManager) DeleteWorkloadStatus(ctx context.Context, workloadName string) error {
        return f.withFileLock(ctx, workloadName, func(statusFilePath string) error {
                // Remove status file
                if err := os.Remove(statusFilePath); err != nil && !os.IsNotExist(err) {
                        return fmt.Errorf("failed to delete status file for workload %s: %w", workloadName, err)
                }

                // Remove lock file (best effort) - done by withFileLock after this function returns
                slog.Debug("workload status deleted", "workload", workloadName)
                return nil
        })
}

// SetWorkloadPID sets the PID of a workload by its name.
// This method will do nothing if the workload does not exist.
func (f *fileStatusManager) SetWorkloadPID(ctx context.Context, workloadName string, pid int) error {
        err := f.withFileLock(ctx, workloadName, func(statusFilePath string) error {
                // Check if file exists
                if _, err := os.Stat(statusFilePath); os.IsNotExist(err) {
                        // File doesn't exist, nothing to do
                        slog.Debug("workload does not exist, skipping PID update", "workload", workloadName)
                        return nil
                } else if err != nil {
                        return fmt.Errorf("failed to check status file for workload %s: %w", workloadName, err)
                }

                // Read existing file
                statusFile, err := f.readStatusFile(statusFilePath)
                if err != nil {
                        return fmt.Errorf("failed to read existing status for workload %s: %w", workloadName, err)
                }

                // Update only the PID and UpdatedAt timestamp
                statusFile.ProcessID = pid
                statusFile.UpdatedAt = time.Now()

                if err = f.writeStatusFile(statusFilePath, *statusFile); err != nil {
                        return fmt.Errorf("failed to write updated PID for workload %s: %w", workloadName, err)
                }

                slog.Debug("workload PID set", "workload", workloadName, "pid", pid)
                return nil
        })
        if err != nil {
                slog.Error("error updating workload PID", "workload", workloadName, "error", err)
        }
        return err
}

// ResetWorkloadPID resets the PID of a workload to 0.
// This method will do nothing if the workload does not exist.
func (f *fileStatusManager) ResetWorkloadPID(ctx context.Context, workloadName string) error {
        // As a side effect, get rid of the PID file if any exists
        err := removePIDFile(workloadName)
        if err != nil {
                // This is an expected error in most cases.
                slog.Debug("no PID for workload was removed", "workload", workloadName)
        }

        return f.SetWorkloadPID(ctx, workloadName, 0)
}

// ResetWorkloadPIDIfMatch resets the PID of a workload to 0 only if the
// current PID in the status file matches expectedPID. This prevents a dying
// process from clobbering a PID written by a replacement process.
func (f *fileStatusManager) ResetWorkloadPIDIfMatch(ctx context.Context, workloadName string, expectedPID int) error {
        // As a side effect, get rid of the PID file if any exists
        if err := removePIDFile(workloadName); err != nil {
                slog.Debug("no PID for workload was removed", "workload", workloadName)
        }

        err := f.withFileLock(ctx, workloadName, func(statusFilePath string) error {
                if _, err := os.Stat(statusFilePath); os.IsNotExist(err) {
                        return nil
                } else if err != nil {
                        return fmt.Errorf("failed to check status file for workload %s: %w", workloadName, err)
                }

                statusFile, err := f.readStatusFile(statusFilePath)
                if err != nil {
                        return fmt.Errorf("failed to read status for workload %s: %w", workloadName, err)
                }

                if statusFile.ProcessID != expectedPID {
                        slog.Debug("skipping PID reset: current PID does not match",
                                "workload", workloadName,
                                "current_pid", statusFile.ProcessID,
                                "expected_pid", expectedPID)
                        return nil
                }

                statusFile.ProcessID = 0
                statusFile.UpdatedAt = time.Now()
                return f.writeStatusFile(statusFilePath, *statusFile)
        })
        if err != nil {
                slog.Error("error resetting workload PID", "workload", workloadName, "error", err)
        }
        return err
}

// GetWorkloadPID retrieves the PID of a workload from its status file.
func (f *fileStatusManager) GetWorkloadPID(ctx context.Context, workloadName string) (int, error) {
        var pid int

        err := f.withFileReadLock(ctx, workloadName, func(statusFilePath string) error {
                // Check if file exists
                if _, err := os.Stat(statusFilePath); os.IsNotExist(err) {
                        // File doesn't exist, return 0
                        pid = 0
                        return nil
                } else if err != nil {
                        return fmt.Errorf("failed to check status file for workload %s: %w", workloadName, err)
                }

                statusFile, err := f.readStatusFile(statusFilePath)
                if err != nil {
                        return fmt.Errorf("failed to read status file for workload %s: %w", workloadName, err)
                }

                pid = statusFile.ProcessID
                return nil
        })
        if err != nil {
                return 0, err
        }

        slog.Debug("workload PID retrieved", "workload", workloadName, "pid", pid)
        return pid, nil
}

// migratePIDFromFile migrates PID from legacy PID file to status file if needed.
// This is called when the status is running and ProcessID is 0.
// Returns (migratedPID, wasUpdated) where wasUpdated indicates if the PID was successfully migrated
func (*fileStatusManager) migratePIDFromFile(workloadName string, containerInfo *rt.ContainerInfo) (int, bool) {
        // Get the base name from container labels
        var baseName string
        if containerInfo != nil {
                baseName = labels.GetContainerBaseName(containerInfo.Labels)
        } else {
                // If we don't have container info, try using workload name as base name
                baseName = workloadName
        }

        if baseName == "" {
                slog.Debug("no base name available for workload, skipping PID migration", "workload", workloadName)
                return 0, false
        }

        // Try to read PID from PID file
        // The readPIDFile function handles checking both old and new locations
        pid, err := readPIDFile(baseName)
        if err != nil {
                slog.Debug("failed to read PID file for workload", "workload", workloadName, "baseName", baseName, "error", err)
                return 0, false
        }
        slog.Debug("found PID in PID file for workload, will update status file", "pid", pid, "workload", workloadName)

        // Delete the PID file after successful migration
        if err := removePIDFile(baseName); err != nil {
                slog.Warn("failed to remove PID file for workload", "workload", workloadName, "baseName", baseName, "error", err)
                // Don't return false here - the migration succeeded, cleanup just failed
        }

        return pid, true
}

// getStatusFilePath returns the file path for a given workload's status file.
func (f *fileStatusManager) getStatusFilePath(workloadName string) string {
        return filepath.Join(f.baseDir, fmt.Sprintf("%s.json", workloadName))
}

// getLockFilePath returns the lock file path for a given workload.
func (f *fileStatusManager) getLockFilePath(workloadName string) string {
        return filepath.Join(f.baseDir, fmt.Sprintf("%s.lock", workloadName))
}

// ensureBaseDir creates the base directory if it doesn't exist.
func (f *fileStatusManager) ensureBaseDir() error {
        return os.MkdirAll(f.baseDir, 0o750)
}

// TODO: This can probably be de-duped with withFileReadLock
// withFileLock executes the provided function while holding a write lock on the workload's lock file.
func (f *fileStatusManager) withFileLock(ctx context.Context, workloadName string, fn func(string) error) error {
        // Remove any slashes from the workload name to avoid problems.
        workloadName = strings.ReplaceAll(workloadName, "/", "-")

        // Validate workload name for safe path construction
        if err := fileutils.ValidateWorkloadNameForPath(workloadName); err != nil {
                return fmt.Errorf("invalid workload name '%s': %w", workloadName, err)
        }
        if err := f.ensureBaseDir(); err != nil {
                return fmt.Errorf("failed to create base directory: %w", err)
        }

        statusFilePath := f.getStatusFilePath(workloadName)
        lockFilePath := f.getLockFilePath(workloadName)

        // Create file lock
        fileLock := lockfile.NewTrackedLock(lockFilePath)
        defer lockfile.ReleaseTrackedLock(lockFilePath, fileLock)

        // Create context with timeout
        lockCtx, cancel := context.WithTimeout(ctx, lockTimeout)
        defer cancel()

        // Acquire lock with context
        locked, err := fileLock.TryLockContext(lockCtx, lockRetryInterval)
        if err != nil {
                return fmt.Errorf("failed to acquire lock for workload %s: %w", workloadName, err)
        }
        if !locked {
                return fmt.Errorf("could not acquire lock for workload %s: timeout after %v", workloadName, lockTimeout)
        }

        return fn(statusFilePath)
}

// withFileReadLock executes the provided function while holding a read lock on the workload's lock file.
func (f *fileStatusManager) withFileReadLock(ctx context.Context, workloadName string, fn func(string) error) error {
        // Remove any slashes from the workload name to avoid problems.
        workloadName = strings.ReplaceAll(workloadName, "/", "-")

        // Validate workload name for safe path construction
        if err := fileutils.ValidateWorkloadNameForPath(workloadName); err != nil {
                return fmt.Errorf("invalid workload name '%s': %w", workloadName, err)
        }
        if err := f.ensureBaseDir(); err != nil {
                return fmt.Errorf("failed to create base directory: %w", err)
        }
        statusFilePath := f.getStatusFilePath(workloadName)
        lockFilePath := f.getLockFilePath(workloadName)

        // Create file lock
        fileLock := lockfile.NewTrackedLock(lockFilePath)
        defer lockfile.ReleaseTrackedLock(lockFilePath, fileLock)

        // Create context with timeout
        lockCtx, cancel := context.WithTimeout(ctx, lockTimeout)
        defer cancel()

        // Acquire read lock with context
        locked, err := fileLock.TryRLockContext(lockCtx, lockRetryInterval)
        if err != nil {
                return fmt.Errorf("failed to acquire read lock for workload %s: %w", workloadName, err)
        }
        if !locked {
                return fmt.Errorf("could not acquire read lock for workload %s: timeout after %v", workloadName, lockTimeout)
        }

        return fn(statusFilePath)
}

// readStatusFile reads and parses a workload status file from disk.
// If the file is corrupted, it attempts recovery using various strategies.
func (f *fileStatusManager) readStatusFile(statusFilePath string) (*workloadStatusFile, error) {
        data, err := os.ReadFile(statusFilePath) //nolint:gosec // file path is constructed by our own function
        if err != nil {
                return nil, fmt.Errorf("failed to read status file: %w", err)
        }

        // Validate file content before parsing
        if len(data) == 0 {
                return nil, fmt.Errorf("status file is empty")
        }

        // Attempt to parse the JSON
        var statusFile workloadStatusFile
        parseErr := json.Unmarshal(data, &statusFile)

        // If parsing succeeded, validate and return
        if parseErr == nil {
                if statusFile.Status == "" {
                        return nil, fmt.Errorf("status file missing required 'status' field")
                }
                if statusFile.CreatedAt.IsZero() {
                        return nil, fmt.Errorf("status file missing or invalid 'created_at' field")
                }
                return &statusFile, nil
        }

        // Parsing failed - check if JSON is valid
        if json.Valid(data) {
                // JSON is structurally valid but unmarshal failed - this is unexpected
                return nil, fmt.Errorf("failed to unmarshal valid JSON: %w", parseErr)
        }

        // JSON is invalid - attempt recovery
        slog.Warn("status file contains invalid JSON, attempting recovery", "path", statusFilePath)

        recoveredFile, recoveryErr := f.attemptJSONRecovery(statusFilePath, data)
        if recoveryErr != nil {
                // Recovery failed - back up the corrupted file
                backupPath := statusFilePath + ".corrupted"
                //nolint:gosec // G703 - path derived from trusted status file with fixed suffix
                if backupErr := os.WriteFile(backupPath, data, 0o600); backupErr == nil {
                        slog.Warn("backed up corrupted status file", "path", backupPath)
                }
                return nil, fmt.Errorf("failed to parse status file (original error: %v, recovery failed: %v)", parseErr, recoveryErr)
        }

        slog.Info("successfully recovered corrupted status file", "path", statusFilePath)

        // Auto-repair: write the recovered file back atomically
        if repairErr := f.writeStatusFile(statusFilePath, *recoveredFile); repairErr != nil {
                slog.Warn("recovered status file but failed to auto-repair", "error", repairErr)
                // Don't fail - we successfully recovered the data
        } else {
                slog.Debug("auto-repaired status file", "path", statusFilePath)
        }

        return recoveredFile, nil
}

// attemptJSONRecovery tries multiple strategies to recover corrupted JSON data.
//
//nolint:gocyclo // Multiple recovery strategies require conditional logic
func (*fileStatusManager) attemptJSONRecovery(statusFilePath string, data []byte) (*workloadStatusFile, error) {
        str := string(data)
        var statusFile workloadStatusFile

        // Strategy 1: Remove extra closing braces
        openBraces := strings.Count(str, "{")
        closeBraces := strings.Count(str, "}")

        if closeBraces > openBraces {
                // Remove extra closing braces from the end
                trimmed := strings.TrimRight(str, "}")
                reconstructed := trimmed + strings.Repeat("}", openBraces)

                if err := json.Unmarshal([]byte(reconstructed), &statusFile); err == nil {
                        if statusFile.Status != "" && !statusFile.CreatedAt.IsZero() {
                                slog.Debug("recovered by removing extra closing braces", "path", statusFilePath, "count", closeBraces-openBraces)
                                return &statusFile, nil
                        }
                }
        }

        // Strategy 2: Add missing closing braces (truncated file)
        if closeBraces < openBraces {
                augmented := str + strings.Repeat("}", openBraces-closeBraces)

                if err := json.Unmarshal([]byte(augmented), &statusFile); err == nil {
                        if statusFile.Status != "" && !statusFile.CreatedAt.IsZero() {
                                slog.Debug("recovered by adding missing closing braces", "path", statusFilePath, "count", openBraces-closeBraces)
                                return &statusFile, nil
                        }
                }
        }

        // Strategy 3: Try trimming whitespace and control characters
        cleaned := strings.TrimSpace(str)
        // Remove any null bytes or other control characters that might have been introduced
        cleaned = strings.Map(func(r rune) rune {
                if r == 0 || (r < 32 && r != '\n' && r != '\r' && r != '\t') {
                        return -1 // Remove character
                }
                return r
        }, cleaned)

        if err := json.Unmarshal([]byte(cleaned), &statusFile); err == nil {
                if statusFile.Status != "" && !statusFile.CreatedAt.IsZero() {
                        slog.Debug("recovered by cleaning whitespace/control characters", "path", statusFilePath)
                        return &statusFile, nil
                }
        }

        return nil, fmt.Errorf("all recovery strategies failed")
}

// writeStatusFile writes a workload status file to disk with proper formatting.
// Uses atomic file writes to prevent corruption from interrupted writes.
func (*fileStatusManager) writeStatusFile(statusFilePath string, statusFile workloadStatusFile) error {
        data, err := json.MarshalIndent(statusFile, "", "  ")
        if err != nil {
                return fmt.Errorf("failed to marshal status file: %w", err)
        }

        if err := fileutils.AtomicWriteFile(statusFilePath, data, 0o600); err != nil {
                return fmt.Errorf("failed to write status file: %w", err)
        }

        return nil
}

// getWorkloadFromRuntime retrieves workload information from the runtime.
func (f *fileStatusManager) getWorkloadFromRuntime(ctx context.Context, workloadName string) (core.Workload, error) {
        info, err := f.runtime.GetWorkloadInfo(ctx, workloadName)
        if err != nil {
                return core.Workload{}, fmt.Errorf("failed to get workload info from runtime: %w", err)
        }

        return types.WorkloadFromContainerInfo(&info, f.runConfigStore)
}

// workloadWithPID holds a workload and its associated PID for internal processing
type workloadWithPID struct {
        workload core.Workload
        pid      int
}

// getWorkloadsFromFiles retrieves all workloads from status files.
func (f *fileStatusManager) getWorkloadsFromFiles() (map[string]workloadWithPID, error) {
        // Ensure base directory exists
        if err := f.ensureBaseDir(); err != nil {
                return nil, fmt.Errorf("failed to ensure base directory: %w", err)
        }

        // List all .json files in the base directory
        files, err := filepath.Glob(filepath.Join(f.baseDir, "*.json"))
        if err != nil {
                return nil, fmt.Errorf("failed to list status files: %w", err)
        }

        workloads := make(map[string]workloadWithPID)
        ctx := context.Background() // Create context for file locking

        for _, file := range files {
                // Extract workload name from filename (remove .json extension)
                workloadName := strings.TrimSuffix(filepath.Base(file), ".json")

                // Use write lock since we may need to update the file for PID migration
                err := f.withFileLock(ctx, workloadName, func(statusFilePath string) error {
                        // Check if file exists first
                        if _, err := os.Stat(statusFilePath); os.IsNotExist(err) {
                                slog.Debug("status file for workload no longer exists, skipping", "workload", workloadName)
                                return nil // Not an error, file was removed
                        } else if err != nil {
                                return fmt.Errorf("failed to check status file: %w", err)
                        }

                        // Read the status file with proper error handling
                        statusFile, err := f.readStatusFile(statusFilePath)
                        if err != nil {
                                // Distinguish between different types of errors
                                if os.IsPermission(err) {
                                        return fmt.Errorf("permission denied reading status file: %w", err)
                                }
                                // For JSON parsing errors or corrupted files, log details
                                slog.Error("failed to read or parse status file", "path", statusFilePath, "workload", workloadName, "error", err)
                                return fmt.Errorf("corrupted or invalid status file: %w", err)
                        }

                        // Create workload from file data
                        workload := core.Workload{
                                Name:          workloadName,
                                Status:        statusFile.Status,
                                StatusContext: statusFile.StatusContext,
                                CreatedAt:     statusFile.CreatedAt,
                        }

                        // Check if this is a remote workload using the state package
                        remote, err := f.isRemoteWorkload(ctx, workloadName)
                        if err != nil {
                                // This error is expected
                                slog.Debug("failed to check if workload is remote", "workload", workloadName, "error", err)
                        }
                        if remote {
                                workload.Remote = true
                        }

                        // Check if PID migration is needed
                        pid := statusFile.ProcessID
                        if statusFile.Status == rt.WorkloadStatusRunning && statusFile.ProcessID == 0 {
                                // Try PID migration - the migration function will handle cases
                                // where container info is not available gracefully
                                if migratedPID, wasMigrated := f.migratePIDFromFile(workloadName, nil); wasMigrated {
                                        // Update the status file with the migrated PID
                                        statusFile.ProcessID = migratedPID
                                        statusFile.UpdatedAt = time.Now()
                                        pid = migratedPID
                                        if err := f.writeStatusFile(statusFilePath, *statusFile); err != nil {
                                                slog.Warn("failed to write migrated PID", "workload", workloadName, "error", err)
                                        } else {
                                                slog.Debug("successfully migrated PID to status file", "pid", migratedPID, "workload", workloadName)
                                        }
                                }
                        }

                        workloads[workloadName] = workloadWithPID{
                                workload: workload,
                                pid:      pid,
                        }
                        return nil
                })
                if err != nil {
                        // Log the specific error but continue processing other workloads
                        // This maintains the existing behavior but with better diagnostics
                        slog.Warn("failed to process status file for workload", "workload", workloadName, "error", err)
                        continue
                }
        }

        return workloads, nil
}

// validateRunningWorkload validates that a workload marked as running in the file
// is actually running in the runtime and has a healthy proxy process if applicable.
func (f *fileStatusManager) validateRunningWorkload(
        ctx context.Context, workloadName string, result core.Workload, pid int,
) (core.Workload, error) {
        // For remote workloads, we don't need to validate against the container runtime
        // since they don't have containers
        if result.Remote {
                return result, nil
        }

        // Get raw container info from runtime (before label filtering)
        containerInfo, err := f.runtime.GetWorkloadInfo(ctx, workloadName)
        if err != nil {
                return core.Workload{}, err
        }

        // Check if runtime status matches file status
        if containerInfo.State != rt.WorkloadStatusRunning {
                return f.handleRuntimeMismatch(ctx, workloadName, result, containerInfo)
        }

        // Check if proxy process is running when workload is running
        if unhealthyWorkload, isUnhealthy := f.isProxyUnhealthy(ctx, workloadName, result, containerInfo, pid); isUnhealthy {
                return unhealthyWorkload, nil
        }

        // Runtime and proxy confirm workload is healthy - merge runtime data with file status
        return f.mergeHealthyWorkloadData(containerInfo, result)
}

// handleRuntimeMismatch handles the case where file indicates running but runtime shows different status
func (f *fileStatusManager) handleRuntimeMismatch(
        ctx context.Context, workloadName string, result core.Workload, containerInfo rt.ContainerInfo,
) (core.Workload, error) {
        contextMsg := fmt.Sprintf("workload status mismatch: file indicates running, but runtime shows %s", containerInfo.State)
        if err := f.SetWorkloadStatus(ctx, workloadName, rt.WorkloadStatusUnhealthy, contextMsg); err != nil {
                slog.Warn("failed to update workload status to unhealthy", "workload", workloadName, "error", err)
        }

        // Convert to workload and return unhealthy status
        runtimeResult, err := types.WorkloadFromContainerInfo(&containerInfo, f.runConfigStore)
        if err != nil {
                return core.Workload{}, err
        }

        runtimeResult.Status = rt.WorkloadStatusUnhealthy
        runtimeResult.StatusContext = contextMsg
        runtimeResult.CreatedAt = result.CreatedAt // Keep the original file created time
        return runtimeResult, nil
}

// handleRuntimeMissing handles the case where the file indicates running or stopped but the runtime
// does not have the workload running. This can happen if using different versions of ToolHive, for example
// the CLI and UI have different versions.
func (f *fileStatusManager) handleRuntimeMissing(
        ctx context.Context, workloadName string, fileWorkload core.Workload,
) (core.Workload, error) {
        // Check if this is a remote workload using the Remote field
        if fileWorkload.Remote {
                // Remote workloads don't exist in the container runtime, so it's normal for them to be missing
                // Don't mark them as unhealthy
                return fileWorkload, nil
        }

        if fileWorkload.Status == rt.WorkloadStatusRunning || fileWorkload.Status == rt.WorkloadStatusStopped {
                // The workload cannot be running or stopped if the runtime container is not found
                contextMsg := fmt.Sprintf("workload %s not found in runtime, marking as unhealthy", workloadName)
                if err := f.setUnhealthyIfStillPresent(ctx, workloadName, contextMsg); err != nil {
                        return core.Workload{}, err
                }
                // The returned view reflects what this list observed regardless of
                // whether the guarded write persisted anything (the file may have been
                // removed or repurposed by a concurrent operation in the meantime).
                fileWorkload.Status = rt.WorkloadStatusUnhealthy
        }

        // If the workload has another status, like starting or stopping, we can keep it as is
        return fileWorkload, nil
}

// setUnhealthyIfStillPresent marks a workload unhealthy only if, re-checked
// under the file lock, its status file still exists and still claims running
// or stopped. handleRuntimeMissing decides to mark unhealthy from a snapshot
// that is stale by the time the write happens; an unconditional
// SetWorkloadStatus would recreate the file of a workload that `thv rm`
// deleted in the meantime (SetWorkloadStatus creates missing files), leaving
// a ghost "unhealthy" entry behind, or clobber a fresher status such as
// `removing`. Skipping in those cases is not an error: the persisted state is
// already newer than the observation that triggered the write.
func (f *fileStatusManager) setUnhealthyIfStillPresent(ctx context.Context, workloadName, contextMsg string) error {
        return f.withFileLock(ctx, workloadName, func(statusFilePath string) error {
                if _, err := os.Stat(statusFilePath); os.IsNotExist(err) {
                        slog.Debug("skipping unhealthy mark: status file no longer exists", "workload", workloadName)
                        return nil
                } else if err != nil {
                        return fmt.Errorf("failed to check status file for workload %s: %w", workloadName, err)
                }

                statusFile, err := f.readStatusFile(statusFilePath)
                if err != nil {
                        return fmt.Errorf("failed to read existing status for workload %s: %w", workloadName, err)
                }
                if statusFile.Status != rt.WorkloadStatusRunning && statusFile.Status != rt.WorkloadStatusStopped {
                        slog.Debug("skipping unhealthy mark: status changed since observation",
                                "workload", workloadName, "status", statusFile.Status)
                        return nil
                }

                statusFile.Status = rt.WorkloadStatusUnhealthy
                statusFile.StatusContext = contextMsg
                statusFile.UpdatedAt = time.Now()
                return f.writeStatusFile(statusFilePath, *statusFile)
        })
}

// isProxyUnhealthy checks if the proxy process is running for the workload.
// Returns (unhealthyWorkload, true) if proxy is not running, (emptyWorkload, false) if proxy is healthy or not applicable.
func (f *fileStatusManager) isProxyUnhealthy(
        ctx context.Context, workloadName string, result core.Workload, containerInfo rt.ContainerInfo, pid int,
) (core.Workload, bool) {
        // Use original container labels (before filtering) to get base name
        baseName := labels.GetContainerBaseName(containerInfo.Labels)
        if baseName == "" {
                return core.Workload{}, false // No proxy check needed
        }

        proxyRunning, err := process.FindProcess(pid)
        if err != nil {
                slog.Warn("unable to find process", "pid", pid, "error", err)
        } else if proxyRunning {
                return core.Workload{}, false // Proxy is healthy
        }

        // Proxy is not running, but workload should be running
        contextMsg := fmt.Sprintf("proxy process not running: workload shows running but proxy process for %s is not active",
                baseName)
        if err := f.SetWorkloadStatus(ctx, workloadName, rt.WorkloadStatusUnhealthy, contextMsg); err != nil {
                slog.Warn("failed to update workload status to unhealthy", "workload", workloadName, "error", err)
        }

        // Convert to workload and return unhealthy status
        runtimeResult, err := types.WorkloadFromContainerInfo(&containerInfo, f.runConfigStore)
        if err != nil {
                slog.Warn("failed to convert container info for unhealthy workload", "workload", workloadName, "error", err)
                return core.Workload{}, false // Return false to avoid double error handling
        }

        runtimeResult.Status = rt.WorkloadStatusUnhealthy
        runtimeResult.StatusContext = contextMsg
        runtimeResult.CreatedAt = result.CreatedAt // Keep the original file created time
        return runtimeResult, true
}

// mergeHealthyWorkloadData merges runtime container data with file-based status information
func (f *fileStatusManager) mergeHealthyWorkloadData(
        containerInfo rt.ContainerInfo, result core.Workload,
) (core.Workload, error) {
        // Runtime and proxy confirm workload is healthy - use runtime data but preserve file-based status info
        runtimeResult, err := types.WorkloadFromContainerInfo(&containerInfo, f.runConfigStore)
        if err != nil {
                return core.Workload{}, err
        }

        runtimeResult.Status = result.Status               // Keep the file status (running)
        runtimeResult.StatusContext = result.StatusContext // Keep the file status context
        runtimeResult.CreatedAt = result.CreatedAt         // Keep the file created time
        return runtimeResult, nil
}

// validateWorkloadInList validates a workload during list operations, similar to validateRunningWorkload
// but with different error handling to avoid disrupting the entire list operation.
func (f *fileStatusManager) validateWorkloadInList(
        ctx context.Context, workloadName string, fileWorkload core.Workload, containerInfo rt.ContainerInfo, pid int,
) (core.Workload, error) {
        // Only validate if file shows running status
        if fileWorkload.Status != rt.WorkloadStatusRunning {
                // For non-running workloads, just merge runtime data with file status
                runtimeWorkload, err := types.WorkloadFromContainerInfo(&containerInfo, f.runConfigStore)
                if err != nil {
                        return core.Workload{}, err
                }
                runtimeWorkload.Status = fileWorkload.Status
                runtimeWorkload.StatusContext = fileWorkload.StatusContext
                runtimeWorkload.CreatedAt = fileWorkload.CreatedAt
                return runtimeWorkload, nil
        }

        // For running workloads, apply full validation
        // Check if runtime status matches file status
        if containerInfo.State != rt.WorkloadStatusRunning {
                return f.handleRuntimeMismatch(ctx, workloadName, fileWorkload, containerInfo)
        }

        // Check if proxy process is running when workload is running
        if unhealthyWorkload, isUnhealthy := f.isProxyUnhealthy(ctx, workloadName, fileWorkload, containerInfo, pid); isUnhealthy {
                return unhealthyWorkload, nil
        }

        // Runtime and proxy confirm workload is healthy - merge runtime data with file status
        return f.mergeHealthyWorkloadData(containerInfo, fileWorkload)
}

// mergeRuntimeAndFileWorkloads returns a map of workloads that combines runtime containers and file-based workloads.
func (f *fileStatusManager) mergeRuntimeAndFileWorkloads(
        ctx context.Context,
        runtimeContainers []rt.ContainerInfo,
        fileWorkloadsWithPID map[string]workloadWithPID,
) map[string]core.Workload {
        runtimeWorkloadMap := make(map[string]rt.ContainerInfo)
        for _, container := range runtimeContainers {
                // Use base name from labels for matching, fall back to container name if not available
                baseName := labels.GetContainerBaseName(container.Labels)
                if baseName == "" {
                        baseName = container.Name // fallback for containers without base name label
                }
                runtimeWorkloadMap[baseName] = container
        }

        // Create result map to avoid duplicates and merge data
        workloadMap := make(map[string]core.Workload)

        // First, add all runtime workloads
        for _, container := range runtimeContainers {
                workload, err := types.WorkloadFromContainerInfo(&container, f.runConfigStore)
                if err != nil {
                        slog.Warn("failed to convert container info for workload", "workload", container.Name, "error", err)
                        continue
                }
                // Use base name for consistency with file workloads
                baseName := labels.GetContainerBaseName(container.Labels)
                if baseName == "" {
                        baseName = container.Name // fallback for containers without base name label
                }
                workloadMap[baseName] = workload
        }

        // Then, merge with file workloads, validating running workloads
        for name, fileWorkloadWithPID := range fileWorkloadsWithPID {
                fileWorkload := fileWorkloadWithPID.workload
                pid := fileWorkloadWithPID.pid

                if fileWorkload.Remote { // Remote workloads are not managed by the container runtime
                        continue // Skip remote workloads here, we add them in workload manager
                }
                if runtimeContainer, exists := runtimeWorkloadMap[name]; exists {
                        // Validate running workloads similar to GetWorkload
                        validatedWorkload, err := f.validateWorkloadInList(ctx, name, fileWorkload, runtimeContainer, pid)
                        if err != nil {
                                slog.Warn("failed to validate workload in list", "workload", name, "error", err)
                                // Fall back to basic merge without validation
                                if runtimeWorkload, exists := workloadMap[name]; exists {
                                        runtimeWorkload.Status = fileWorkload.Status
                                        runtimeWorkload.StatusContext = fileWorkload.StatusContext
                                        runtimeWorkload.CreatedAt = fileWorkload.CreatedAt
                                        workloadMap[name] = runtimeWorkload
                                } else {
                                        // Runtime workload not found, just use the file workload
                                        workloadMap[name] = fileWorkload
                                }
                        } else {
                                workloadMap[name] = validatedWorkload
                        }
                } else {
                        // File-only workload (runtime not available)
                        updatedWorkload, err := f.handleRuntimeMissing(ctx, name, fileWorkload)
                        if err != nil {
                                slog.Warn("failed to handle missing runtime for workload", "workload", name, "error", err)
                                workloadMap[name] = fileWorkload
                        } else {
                                workloadMap[name] = updatedWorkload
                        }
                }
        }
        return workloadMap
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package statuses

import (
        "context"

        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/core"
)

// NoopStatusManager is a no-op implementation of StatusManager that does nothing.
// All methods return zero values or empty results without performing any operations.
type NoopStatusManager struct{}

// NewNoopStatusManager creates a new NoopStatusManager instance.
func NewNoopStatusManager() StatusManager {
        return &NoopStatusManager{}
}

// GetWorkload returns an empty workload and nil error.
func (*NoopStatusManager) GetWorkload(_ context.Context, _ string) (core.Workload, error) {
        return core.Workload{}, nil
}

// ListWorkloads returns an empty slice of workloads.
func (*NoopStatusManager) ListWorkloads(_ context.Context, _ bool, _ []string) ([]core.Workload, error) {
        return []core.Workload{}, nil
}

// SetWorkloadStatus does nothing and returns nil.
func (*NoopStatusManager) SetWorkloadStatus(_ context.Context, _ string, _ rt.WorkloadStatus, _ string) error {
        return nil
}

// DeleteWorkloadStatus does nothing and returns nil.
func (*NoopStatusManager) DeleteWorkloadStatus(_ context.Context, _ string) error {
        return nil
}

// SetWorkloadPID does nothing and returns nil.
func (*NoopStatusManager) SetWorkloadPID(_ context.Context, _ string, _ int) error {
        return nil
}

// ResetWorkloadPID does nothing and returns nil.
func (*NoopStatusManager) ResetWorkloadPID(_ context.Context, _ string) error {
        return nil
}

// ResetWorkloadPIDIfMatch does nothing and returns nil.
func (*NoopStatusManager) ResetWorkloadPIDIfMatch(_ context.Context, _ string, _ int) error {
        return nil
}

// GetWorkloadPID returns 0 and nil error.
func (*NoopStatusManager) GetWorkloadPID(_ context.Context, _ string) (int, error) {
        return 0, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package statuses

import (
        "fmt"
        "os"
        "path/filepath"
        "strconv"
        "strings"

        "github.com/adrg/xdg"

        "github.com/stacklok/toolhive/pkg/container/runtime"
)

/*
 * NOTE: PID file functionality is deprecated. It should only be used by code which migrates PIDs to status files.
 */

// getOldPIDFilePath returns the legacy path to the PID file for a container (for backward compatibility)
// Note: containerBaseName is pre-sanitized by the caller
func getOldPIDFilePath(containerBaseName string) string {
        // Use the system temporary directory (old behavior)
        tmpDir := os.TempDir()
        // Clean the path to satisfy security scanners (containerBaseName is already sanitized)
        return filepath.Clean(filepath.Join(tmpDir, fmt.Sprintf("toolhive-%s.pid", containerBaseName)))
}

// getPIDFilePath returns the path to the PID file for a container
// It first tries the new XDG location, then falls back to the old temp directory location
func getPIDFilePath(containerBaseName string) (string, error) {
        // Return empty path in Kubernetes runtime since PID files are not used
        if runtime.IsKubernetesRuntime() {
                return "", fmt.Errorf("PID file operations are not supported in Kubernetes runtime")
        }

        // Get the new XDG-based path
        pidPath, err := xdg.DataFile(filepath.Join("toolhive", "pids", fmt.Sprintf("toolhive-%s.pid", containerBaseName)))
        if err != nil {
                return "", fmt.Errorf("failed to get PID file path: %w", err)
        }
        return pidPath, nil
}

// getPIDFilePathWithFallback returns the path to an existing PID file for a container
// It checks both the new XDG location and the old temp directory location
// Note: containerBaseName is pre-sanitized by the caller
func getPIDFilePathWithFallback(containerBaseName string) (string, error) {
        // Return empty path in Kubernetes runtime since PID files are not used
        if runtime.IsKubernetesRuntime() {
                return "", fmt.Errorf("PID file operations are not supported in Kubernetes runtime")
        }

        // First try the new XDG-based path
        newPath, err := getPIDFilePath(containerBaseName)
        if err != nil {
                return "", err
        }

        // Check if new file exists - prefer it if it does
        if _, err := os.Stat(newPath); err == nil {
                return newPath, nil
        }

        // Fall back to old location if new doesn't exist
        // Clean the path to satisfy security scanners (containerBaseName is already sanitized)
        oldPath := filepath.Clean(getOldPIDFilePath(containerBaseName))
        if _, err := os.Stat(oldPath); err == nil {
                return oldPath, nil
        }

        // If neither exists, return the new path (for new files)
        return newPath, nil
}

// readPIDFile reads the process ID from a file
// It checks both the new XDG location and the old temp directory location
// Note: containerBaseName is pre-sanitized by the caller
func readPIDFile(containerBaseName string) (int, error) {
        // Skip PID file operations in Kubernetes runtime
        if runtime.IsKubernetesRuntime() {
                return 0, fmt.Errorf("PID file operations are not supported in Kubernetes runtime")
        }

        // Get the PID file path with fallback
        pidFilePath, err := getPIDFilePathWithFallback(containerBaseName)
        if err != nil {
                return 0, fmt.Errorf("failed to get PID file path: %w", err)
        }

        // Read the PID from the file
        // Clean the path to satisfy security scanners (containerBaseName is already sanitized)
        cleanPidPath := filepath.Clean(pidFilePath)
        pidBytes, err := os.ReadFile(cleanPidPath)
        if err != nil {
                // If we can't read from the new location, try the old location explicitly
                oldPath := getOldPIDFilePath(containerBaseName)
                if oldPath != pidFilePath {
                        // Clean the path to satisfy security scanners (containerBaseName is already sanitized)
                        cleanOldPath := filepath.Clean(oldPath)
                        pidBytes, err = os.ReadFile(cleanOldPath)
                        if err != nil {
                                return 0, fmt.Errorf("failed to read PID file from both new and old locations: %w", err)
                        }
                } else {
                        return 0, fmt.Errorf("failed to read PID file: %w", err)
                }
        }

        // Parse the PID
        pidStr := strings.TrimSpace(string(pidBytes))
        pid, err := strconv.Atoi(pidStr)
        if err != nil {
                return 0, fmt.Errorf("failed to parse PID: %w", err)
        }

        return pid, nil
}

// removePIDFile removes the PID file
// It attempts to remove from both the new XDG location and the old temp directory location
func removePIDFile(containerBaseName string) error {
        // Skip PID file operations in Kubernetes runtime
        if runtime.IsKubernetesRuntime() {
                return nil
        }

        var lastErr error

        // Try to remove from the new location
        newPath, err := getPIDFilePath(containerBaseName)
        if err != nil {
                return fmt.Errorf("failed to get PID file path: %w", err)
        }

        if err := os.Remove(newPath); err != nil && !os.IsNotExist(err) {
                lastErr = err
        }

        // Also try to remove from the old location (cleanup legacy files)
        // Clean the path to satisfy security scanners (containerBaseName is already sanitized)
        oldPath := filepath.Clean(getOldPIDFilePath(containerBaseName))
        if err := os.Remove(oldPath); err != nil && !os.IsNotExist(err) {
                // If we couldn't remove either file and both had errors, return the error
                if lastErr != nil {
                        return fmt.Errorf("failed to remove PID files: new location: %v, old location: %w", lastErr, err)
                }
                lastErr = err
        }

        // If at least one was removed successfully (or didn't exist), consider it success
        if lastErr != nil && !os.IsNotExist(lastErr) {
                return lastErr
        }

        return nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package statuses provides an interface and implementation for managing workload statuses.
package statuses

import (
        "context"
        "fmt"
        "log/slog"

        "github.com/stacklok/toolhive-core/env"
        rt "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/core"
        "github.com/stacklok/toolhive/pkg/state"
        "github.com/stacklok/toolhive/pkg/workloads/types"
)

// StatusManager is an interface for fetching and retrieving workload statuses.
//
//go:generate mockgen -destination=mocks/mock_status_manager.go -package=mocks -source=status.go StatusManager
type StatusManager interface {
        // GetWorkload retrieves details of a workload by its name.
        GetWorkload(ctx context.Context, workloadName string) (core.Workload, error)
        // ListWorkloads returns details of all workloads.
        ListWorkloads(ctx context.Context, listAll bool, labelFilters []string) ([]core.Workload, error)
        // SetWorkloadStatus sets the status of a workload by its name.
        // Note that this does not return errors, but logs them instead.
        // This method will do nothing if the workload does not exist.
        // This method will preserve the PID of the workload if it was previously set.
        SetWorkloadStatus(ctx context.Context, workloadName string, status rt.WorkloadStatus, contextMsg string) error
        // DeleteWorkloadStatus removes the status of a workload by its name.
        DeleteWorkloadStatus(ctx context.Context, workloadName string) error
        // SetWorkloadPID sets the PID of a workload by its name.
        // This method will do nothing if the workload does not exist.
        SetWorkloadPID(ctx context.Context, workloadName string, pid int) error
        // ResetWorkloadPID resets the PID of a workload to 0.
        // This method will do nothing if the workload does not exist.
        ResetWorkloadPID(ctx context.Context, workloadName string) error
        // ResetWorkloadPIDIfMatch resets the PID of a workload to 0 only if the
        // current PID in the status file matches expectedPID. This prevents a
        // dying process from clobbering the PID written by a replacement process
        // that started in the meantime.
        ResetWorkloadPIDIfMatch(ctx context.Context, workloadName string, expectedPID int) error
        // GetWorkloadPID retrieves the PID of a workload by its name.
        // Returns 0 if the workload does not exist or if PID is not available.
        GetWorkloadPID(ctx context.Context, workloadName string) (int, error)
}

// NewStatusManagerFromRuntime creates a new instance of StatusManager from an existing runtime.
func NewStatusManagerFromRuntime(runtime rt.Runtime, runConfigStore state.Store) StatusManager {
        return &runtimeStatusManager{
                runtime:        runtime,
                runConfigStore: runConfigStore,
        }
}

// NewStatusManager creates a new status manager instance using the appropriate implementation
// based on the runtime environment. If running in Kubernetes, it returns the runtime-based
// implementation. Otherwise, it returns the file-based implementation.
func NewStatusManager(runtime rt.Runtime) (StatusManager, error) {
        return NewStatusManagerWithEnv(runtime, &env.OSReader{})
}

// NewStatusManagerWithEnv creates a new status manager instance using the provided environment reader.
// This allows for dependency injection of environment variable access for testing.
func NewStatusManagerWithEnv(runtime rt.Runtime, envReader env.Reader) (StatusManager, error) {
        if rt.IsKubernetesRuntimeWithEnv(envReader) {
                runConfigStore, err := state.NewRunConfigStore(state.DefaultAppName)
                if err != nil {
                        return nil, fmt.Errorf("failed to create run config store: %w", err)
                }
                return NewStatusManagerFromRuntime(runtime, runConfigStore), nil
        }
        return NewFileStatusManager(runtime)
}

// runtimeStatusManager is an implementation of StatusManager that uses the state
// returned by the underlying runtime. This reflects the existing behaviour of
// ToolHive at the time of writing.
type runtimeStatusManager struct {
        runtime        rt.Runtime
        runConfigStore state.Store
}

func (r *runtimeStatusManager) GetWorkload(ctx context.Context, workloadName string) (core.Workload, error) {
        if err := types.ValidateWorkloadName(workloadName); err != nil {
                return core.Workload{}, err
        }

        info, err := r.runtime.GetWorkloadInfo(ctx, workloadName)
        if err != nil {
                // The error from the runtime is already wrapped in context.
                return core.Workload{}, err
        }

        return types.WorkloadFromContainerInfo(&info, r.runConfigStore)
}

func (r *runtimeStatusManager) ListWorkloads(ctx context.Context, listAll bool, labelFilters []string) ([]core.Workload, error) {
        // List containers
        containers, err := r.runtime.ListWorkloads(ctx)
        if err != nil {
                return nil, fmt.Errorf("failed to list containers: %w", err)
        }

        // Parse the filters into a format we can use for matching.
        parsedFilters, err := types.ParseLabelFilters(labelFilters)
        if err != nil {
                return nil, fmt.Errorf("failed to parse label filters: %w", err)
        }

        // Filter containers to only show those managed by ToolHive
        var workloads []core.Workload
        for _, c := range containers {
                // If the caller did not set `listAll` to true, only include running containers.
                if c.IsRunning() || listAll {
                        workload, err := types.WorkloadFromContainerInfo(&c, r.runConfigStore)
                        if err != nil {
                                return nil, err
                        }
                        // If label filters are provided, check if the workload matches them.
                        if types.MatchesLabelFilters(workload.Labels, parsedFilters) {
                                workloads = append(workloads, workload)
                        }
                }
        }

        return workloads, nil
}

func (*runtimeStatusManager) SetWorkloadStatus(
        _ context.Context,
        workloadName string,
        status rt.WorkloadStatus,
        contextMsg string,
) error {
        // TODO: This will need to handle concurrent updates.
        slog.Debug("workload status set", "workload", workloadName, "status", status, "context", contextMsg)
        return nil
}

func (*runtimeStatusManager) DeleteWorkloadStatus(_ context.Context, _ string) error {
        // TODO: This will need to handle concurrent updates.
        // Noop
        return nil
}

func (*runtimeStatusManager) SetWorkloadPID(_ context.Context, workloadName string, pid int) error {
        // Noop for runtime status manager
        slog.Debug("workload PID set (noop for runtime status manager)", "workload", workloadName, "pid", pid)
        return nil
}

func (*runtimeStatusManager) ResetWorkloadPID(_ context.Context, workloadName string) error {
        // Noop for runtime status manager
        slog.Debug("workload PID reset (noop for runtime status manager)", "workload", workloadName)
        return nil
}

func (*runtimeStatusManager) ResetWorkloadPIDIfMatch(_ context.Context, workloadName string, expectedPID int) error {
        // Noop for runtime status manager
        slog.Debug("workload PID conditional reset (noop for runtime status manager)",
                "workload", workloadName, "expected_pid", expectedPID)
        return nil
}

func (*runtimeStatusManager) GetWorkloadPID(_ context.Context, workloadName string) (int, error) {
        // Noop for runtime status manager - always return 0
        slog.Debug("workload PID requested (noop for runtime status manager, returning 0)", "workload", workloadName)
        return 0, nil
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package types

import (
        "fmt"

        "github.com/stacklok/toolhive/pkg/labels"
)

// ParseLabelFilters parses label filters from a slice of strings and validates them.
func ParseLabelFilters(labelFilters []string) (map[string]string, error) {
        filters := make(map[string]string, len(labelFilters))
        for _, filter := range labelFilters {
                key, value, err := labels.ParseLabel(filter)
                if err != nil {
                        return nil, fmt.Errorf("invalid label filter '%s': %w", filter, err)
                }
                filters[key] = value
        }
        return filters, nil
}

// MatchesLabelFilters checks if workload labels match all the specified filters
func MatchesLabelFilters(workloadLabels, filters map[string]string) bool {
        for filterKey, filterValue := range filters {
                workloadValue, exists := workloadLabels[filterKey]
                if !exists || workloadValue != filterValue {
                        return false
                }
        }
        return true
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package types

import (
        "context"
        "encoding/json"
        "fmt"
        "io"
        "log/slog"
        "net/http"

        "github.com/stacklok/toolhive-core/httperr"
        "github.com/stacklok/toolhive/pkg/container/runtime"
        "github.com/stacklok/toolhive/pkg/core"
        "github.com/stacklok/toolhive/pkg/labels"
        "github.com/stacklok/toolhive/pkg/state"
        "github.com/stacklok/toolhive/pkg/transport"
        "github.com/stacklok/toolhive/pkg/transport/types"
)

// minimalRunConfig represents just the fields we need from a run configuration
type minimalRunConfig struct {
        Group     string `json:"group,omitempty" yaml:"group,omitempty"`
        ProxyMode string `json:"proxy_mode,omitempty" yaml:"proxy_mode,omitempty"`
}

// loadRunConfigFields attempts to load specific fields from the runconfig
// using the provided store. Returns empty struct if runconfig doesn't exist.
func loadRunConfigFields(ctx context.Context, store state.Store, name string) (*minimalRunConfig, error) {
        reader, err := store.GetReader(ctx, name)
        if err != nil {
                // If the run config doesn't exist, return empty config (not an error).
                // This also handles the race where a workload is deleted between listing
                // and reading its config.
                if httperr.Code(err) == http.StatusNotFound {
                        return &minimalRunConfig{}, nil
                }
                return nil, fmt.Errorf("failed to read run config for workload %q: %w", name, err)
        }
        defer func() {
                if err := reader.Close(); err != nil {
                        slog.Warn("failed to close run config reader", "workload", name, "error", err)
                }
        }()

        var config minimalRunConfig
        if err := json.NewDecoder(reader).Decode(&config); err != nil {
                // EOF from an empty reader (e.g. KubernetesStore) means no config exists
                if err == io.EOF {
                        return &minimalRunConfig{}, nil
                }
                return nil, fmt.Errorf("failed to decode run config for workload %q: %w", name, err)
        }
        return &config, nil
}

// WorkloadFromContainerInfo creates a Workload struct from the runtime container info.
// The runConfigStore is used to load run configuration fields (proxy mode, group)
// without hitting the real filesystem, enabling proper dependency injection for tests.
func WorkloadFromContainerInfo(container *runtime.ContainerInfo, runConfigStore state.Store) (core.Workload, error) {
        // Get workload name (base name) from labels for user-facing display
        name := labels.GetContainerBaseName(container.Labels)
        if name == "" {
                // Fallback to full container name if base name is not available
                containerName := labels.GetContainerName(container.Labels)
                if containerName == "" {
                        name = container.Name // Final fallback to container name
                } else {
                        name = containerName
                }
        }

        // Get port from labels
        port, err := labels.GetPort(container.Labels)
        if err != nil {
                port = 0
        }

        transportTypeLabel := labels.GetTransportType(container.Labels)

        tType, err := types.ParseTransportType(transportTypeLabel)
        if err != nil {
                // If we can't parse the transport type, default to SSE.
                tType = types.TransportTypeSSE
        }

        ctx := context.Background()
        runConfig, err := loadRunConfigFields(ctx, runConfigStore, name)
        if err != nil {
                return core.Workload{}, err
        }

        // Generate URL for the MCP server
        url := ""
        if port > 0 {
                url = transport.GenerateMCPServerURL(tType.String(), runConfig.ProxyMode, transport.LocalhostIPv4, port, name, "")
        }

        // Filter out standard ToolHive labels to show only user-defined labels
        userLabels := make(map[string]string)
        for key, value := range container.Labels {
                if !labels.IsStandardToolHiveLabel(key) {
                        userLabels[key] = value
                }
        }

        // Calculate the effective proxy mode that clients should use
        effectiveProxyMode := GetEffectiveProxyMode(tType, runConfig.ProxyMode)

        // Translate to the domain model.
        return core.Workload{
                Name:          name, // Use the calculated workload name (base name), not container name
                Package:       container.Image,
                URL:           url,
                TransportType: tType,
                ProxyMode:     effectiveProxyMode,
                Status:        container.State,
                StatusContext: container.Status,
                CreatedAt:     container.Created,
                Port:          port,
                Labels:        userLabels,
                Group:         runConfig.Group,
                StartedAt:     container.StartedAt,
        }, nil
}

// GetEffectiveProxyMode determines the effective proxy mode that clients should use.
// For stdio transports, this returns the proxy mode (sse or streamable-http).
// For direct transports (sse/streamable-http), this returns the transport type as the proxy mode.
//
// Prefer types.EffectiveProxyMode for new code operating on typed values.
func GetEffectiveProxyMode(transportType types.TransportType, proxyMode string) string {
        return types.EffectiveProxyMode(transportType, types.ProxyMode(proxyMode)).String()
}

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package types contains types and validation functions for workloads in ToolHive.
// This is separated to avoid circular dependencies with the core package.
package types

import (
        "errors"
        "fmt"
        "net/http"
        "path/filepath"
        "regexp"
        "strings"

        "github.com/stacklok/toolhive-core/httperr"
)

// ErrInvalidWorkloadName is returned when a workload name fails validation.
var ErrInvalidWorkloadName = httperr.WithCode(
        errors.New("invalid workload name"),
        http.StatusBadRequest,
)

// workloadNamePattern validates workload names to prevent path traversal attacks
// and other security issues. Workload names should only contain alphanumeric
// characters, hyphens, underscores, and dots.
var workloadNamePattern = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)

// commandInjectionPattern detects potentially dangerous command injection patterns
var commandInjectionPattern = regexp.MustCompile(`[$&;|]|\$\(|\` + "`")

// WorkloadNameIssue represents a specific issue found in a workload name
type WorkloadNameIssue struct {
        Type        string // "empty", "path_traversal", "absolute_path", "command_injection", "null_bytes", "invalid_chars", "too_long"
        Description string
        Position    int // For character-specific issues
}

// analyzeWorkloadName performs comprehensive analysis of a workload name and returns all issues found.
// This shared logic is used by both ValidateWorkloadName and SanitizeWorkloadName.
func analyzeWorkloadName(name string) []WorkloadNameIssue {
        var issues []WorkloadNameIssue

        if name == "" {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "empty",
                        Description: "workload name cannot be empty",
                })
                return issues
        }

        // Check for null bytes
        if strings.Contains(name, "\x00") {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "null_bytes",
                        Description: "workload name contains null bytes",
                })
        }

        // Use filepath.Clean to normalize the path and check for changes
        cleanName := filepath.Clean(name)
        if cleanName != name {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "path_normalization",
                        Description: "workload name requires path normalization",
                })
        }

        // Check if the cleaned path tries to escape current directory
        if rel, err := filepath.Rel(".", cleanName); err != nil || strings.HasPrefix(rel, "..") {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "path_traversal",
                        Description: "workload name contains path traversal",
                })
        }

        // Check for absolute paths
        if filepath.IsAbs(cleanName) {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "absolute_path",
                        Description: "workload name cannot be an absolute path",
                })
        }

        // Check for command injection patterns
        if commandInjectionPattern.MatchString(name) {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "command_injection",
                        Description: "workload name contains potentially dangerous characters",
                })
        }

        // Check against allowed pattern
        if !workloadNamePattern.MatchString(name) {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "invalid_chars",
                        Description: "workload name can only contain alphanumeric characters, dots, hyphens, and underscores",
                })
        }

        // Check length limit
        if len(name) > 100 {
                issues = append(issues, WorkloadNameIssue{
                        Type:        "too_long",
                        Description: "workload name too long (max 100 characters)",
                })
        }

        return issues
}

// ValidateWorkloadName checks if the provided workload name is valid.
// This function performs strict validation and rejects invalid names.
func ValidateWorkloadName(name string) error {
        issues := analyzeWorkloadName(name)

        if len(issues) == 0 {
                return nil
        }

        // Return the first critical issue found
        issue := issues[0]
        return fmt.Errorf("%w: %s", ErrInvalidWorkloadName, issue.Description)
}

// SanitizeWorkloadName sanitizes a user-provided workload name to ensure it's safe for file paths.
// It applies the same security analysis as ValidateWorkloadName but transforms invalid characters
// instead of rejecting them. This provides a more permissive approach for user-facing scenarios
// where we want to accept user input and make it safe rather than rejecting it.
// Returns the sanitized name and a boolean indicating whether the name was modified.
func SanitizeWorkloadName(name string) (string, bool) {
        if name == "" {
                return "", false
        }

        original := name
        result := name
        modified := false

        // Apply fixes based on the issues found
        issues := analyzeWorkloadName(name)

        for _, issue := range issues {
                var wasModified bool
                result, wasModified = applySanitizationFix(result, issue.Type)
                if wasModified {
                        modified = true
                }
        }

        // Ensure we don't return an empty string after sanitization
        if result == "" {
                result = "workload"
                modified = true
        }

        return result, modified || (result != original)
}

// applySanitizationFix applies a specific sanitization fix to the input string
func applySanitizationFix(input, issueType string) (string, bool) {
        switch issueType {
        case "null_bytes":
                return strings.ReplaceAll(input, "\x00", ""), true

        case "path_normalization":
                return filepath.Clean(input), true

        case "path_traversal":
                return strings.ReplaceAll(input, "..", "--"), true

        case "absolute_path":
                return strings.TrimLeft(input, "/\\"), true

        case "command_injection":
                return commandInjectionPattern.ReplaceAllString(input, "-"), true

        case "invalid_chars":
                return sanitizeInvalidChars(input)

        case "too_long":
                return truncateIfTooLong(input)

        default:
                return input, false
        }
}

// sanitizeInvalidChars sanitizes characters to only allow alphanumeric, dots, hyphens, and underscores
func sanitizeInvalidChars(input string) (string, bool) {
        var sanitized strings.Builder
        modified := false

        for _, c := range input {
                if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '.' || c == '-' || c == '_' {
                        sanitized.WriteRune(c)
                } else {
                        sanitized.WriteRune('-')
                        modified = true
                }
        }

        return sanitized.String(), modified
}

// truncateIfTooLong truncates the input if it's longer than 100 characters
func truncateIfTooLong(input string) (string, bool) {
        if len(input) > 100 {
                return input[:100], true
        }
        return input, false
}