crank/crates/crank-core/src/ext/auth.rs

use std::sync::Arc;

use async_trait::async_trait;
use time::OffsetDateTime;

use crate::{
    IssueAgentTokenRequest, IssueOneTimeAgentTokenRequest, IssuedAgentTokenResponse,
    MachineAccessMode, Membership, MembershipRole, OperationSecurityLevel, PlatformApiKeyScope,
    User, UserId, WorkspaceId,
};

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VerifiedMachineCredential {
    pub machine_access_mode: MachineAccessMode,
    pub max_security_level: OperationSecurityLevel,
    pub scopes: Vec<PlatformApiKeyScope>,
}

#[derive(Debug, thiserror::Error)]
#[error("machine credential verifier failed")]
pub struct MachineCredentialVerifierError;

#[async_trait]
pub trait MachineCredentialVerifier: Send + Sync {
    async fn verify_bearer_token(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
        token: &str,
    ) -> Result<Option<VerifiedMachineCredential>, MachineCredentialVerifierError>;
}

pub type SharedMachineCredentialVerifier = Arc<dyn MachineCredentialVerifier>;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TokenIssuerActor {
    pub user_id: UserId,
    pub workspace_id: WorkspaceId,
    pub role: MembershipRole,
}

#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum TokenIssuerError {
    #[error("machine tokens are not available in this edition")]
    NotSupportedInEdition,
    #[error("invalid grant: {0}")]
    InvalidGrant(String),
    #[error("agent key is unknown")]
    AgentKeyUnknown,
    #[error("operation is not strict")]
    OperationNotStrict,
    #[error("operation is not published for agent")]
    OperationNotPublishedForAgent,
    #[error("registry failure: {0}")]
    RegistryFailure(String),
    #[error("replay guard failure: {0}")]
    ReplayGuardFailure(String),
}

#[async_trait]
pub trait MachineTokenIssuer: Send + Sync {
    async fn issue_short_lived(
        &self,
        request: IssueAgentTokenRequest,
        actor: &TokenIssuerActor,
    ) -> Result<IssuedAgentTokenResponse, TokenIssuerError>;

    async fn issue_one_time(
        &self,
        request: IssueOneTimeAgentTokenRequest,
        actor: &TokenIssuerActor,
    ) -> Result<IssuedAgentTokenResponse, TokenIssuerError>;
}

pub type SharedMachineTokenIssuer = Arc<dyn MachineTokenIssuer>;

#[derive(Clone, Copy, Debug, Default)]
pub struct NoMachineTokenIssuer;

#[async_trait]
impl MachineTokenIssuer for NoMachineTokenIssuer {
    async fn issue_short_lived(
        &self,
        _request: IssueAgentTokenRequest,
        _actor: &TokenIssuerActor,
    ) -> Result<IssuedAgentTokenResponse, TokenIssuerError> {
        Err(TokenIssuerError::NotSupportedInEdition)
    }

    async fn issue_one_time(
        &self,
        _request: IssueOneTimeAgentTokenRequest,
        _actor: &TokenIssuerActor,
    ) -> Result<IssuedAgentTokenResponse, TokenIssuerError> {
        Err(TokenIssuerError::NotSupportedInEdition)
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum IdentityProviderKind {
    Password,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AuthenticatedIdentity {
    pub user: User,
    pub memberships: Vec<Membership>,
    pub current_workspace_id: Option<WorkspaceId>,
}

#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum IdentityError {
    #[error("not supported for this provider")]
    NotSupportedForProvider,
    #[error("bad credentials")]
    BadCredentials,
    #[error("account disabled")]
    AccountDisabled,
    #[error("internal: {0}")]
    Internal(String),
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LoginPayload {
    pub email: String,
    pub password: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SsoAuthorizeRequest {
    pub workspace_id: WorkspaceId,
    pub provider_id: String,
    pub base_origin: String,
    pub return_to: Option<String>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SsoAuthorizeRedirect {
    pub authorization_url: String,
    pub state: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SsoCallbackRequest {
    pub workspace_id: WorkspaceId,
    pub provider_id: String,
    pub code: String,
    pub base_origin: String,
    pub return_to: Option<String>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorPendingIdentity {
    pub user_id: UserId,
    pub workspace_id: Option<WorkspaceId>,
    pub return_to: Option<String>,
    pub expires_at: OffsetDateTime,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorStatus {
    pub enabled: bool,
    pub pending_setup: bool,
    pub recovery_codes_remaining: usize,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorSetup {
    pub secret: String,
    pub otpauth_url: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorActivation {
    pub recovery_codes: Vec<String>,
}

#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct TwoFactorChallenge {
    pub code: Option<String>,
    pub recovery_code: Option<String>,
}

pub enum LoginOutcome {
    Authenticated(AuthenticatedIdentity),
    TwoFactorRequired(TwoFactorPendingIdentity),
}

#[async_trait]
pub trait IdentityProvider: Send + Sync {
    fn id(&self) -> &str;

    fn kind(&self) -> IdentityProviderKind;

    async fn login_password(&self, payload: LoginPayload) -> Result<LoginOutcome, IdentityError>;

    async fn begin_sso(
        &self,
        _request: SsoAuthorizeRequest,
    ) -> Result<SsoAuthorizeRedirect, IdentityError> {
        Err(IdentityError::NotSupportedForProvider)
    }

    async fn complete_sso(
        &self,
        _request: SsoCallbackRequest,
    ) -> Result<LoginOutcome, IdentityError> {
        Err(IdentityError::NotSupportedForProvider)
    }

    async fn get_two_factor_status(
        &self,
        _user_id: &UserId,
    ) -> Result<TwoFactorStatus, IdentityError> {
        Err(IdentityError::NotSupportedForProvider)
    }

    async fn begin_two_factor_setup(
        &self,
        _user_id: &UserId,
        _email: &str,
    ) -> Result<TwoFactorSetup, IdentityError> {
        Err(IdentityError::NotSupportedForProvider)
    }

    async fn activate_two_factor(
        &self,
        _user_id: &UserId,
        _code: &str,
    ) -> Result<TwoFactorActivation, IdentityError> {
        Err(IdentityError::NotSupportedForProvider)
    }

    async fn disable_two_factor(
        &self,
        _user_id: &UserId,
        _challenge: TwoFactorChallenge,
    ) -> Result<(), IdentityError> {
        Err(IdentityError::NotSupportedForProvider)
    }

    async fn verify_two_factor_login(
        &self,
        _pending: &TwoFactorPendingIdentity,
        _challenge: TwoFactorChallenge,
    ) -> Result<AuthenticatedIdentity, IdentityError> {
        Err(IdentityError::NotSupportedForProvider)
    }
}

pub type SharedIdentityProvider = Arc<dyn IdentityProvider>;

#[cfg(test)]
mod tests {
    use super::{MachineTokenIssuer, NoMachineTokenIssuer, TokenIssuerActor, TokenIssuerError};
    use crate::{
        AgentTokenGrantType, IssueAgentTokenRequest, IssueOneTimeAgentTokenRequest, MembershipRole,
        UserId, WorkspaceId,
    };

    fn actor() -> TokenIssuerActor {
        TokenIssuerActor {
            user_id: UserId::new("user_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            role: MembershipRole::Owner,
        }
    }

    #[tokio::test]
    async fn no_machine_token_issuer_rejects_short_lived_issue() {
        let result = NoMachineTokenIssuer
            .issue_short_lived(
                IssueAgentTokenRequest {
                    grant_type: AgentTokenGrantType::AgentKey,
                    agent_key: Some("crk_agent_secret".to_owned()),
                    refresh_token: None,
                    scope: vec![],
                },
                &actor(),
            )
            .await;

        assert_eq!(result, Err(TokenIssuerError::NotSupportedInEdition));
    }

    #[tokio::test]
    async fn no_machine_token_issuer_rejects_one_time_issue() {
        let result = NoMachineTokenIssuer
            .issue_one_time(
                IssueOneTimeAgentTokenRequest {
                    agent_key: "crk_agent_secret".to_owned(),
                    operation_id: "op_01".into(),
                    scope: vec![],
                },
                &actor(),
            )
            .await;

        assert_eq!(result, Err(TokenIssuerError::NotSupportedInEdition));
    }
}