gitdataai/lib/git/ssh/authz.rs

use base64::{Engine as _, engine::general_purpose};
use db::{database::AppDatabase, sqlx};
use model::{
    repos::{RepoHistoryNameModel, RepoModel},
    users::{UserModel, UserSshKeyModel},
    workspace::{WkHistoryNameModel, WkMemberModel, WorkspaceModel},
};
use sha2::{Digest, Sha256};

use crate::errors::GitError;

pub struct SshAuthService {
    db: AppDatabase,
}

pub struct SshKeyUser {
    pub user: UserModel,
    pub key_id: i64,
    pub key_title: String,
}

impl SshAuthService {
    pub fn new(db: AppDatabase) -> Self {
        Self { db }
    }

    pub async fn find_repo(
        &self,
        namespace: &str,
        repo_name: &str,
    ) -> Result<RepoModel, GitError> {
        let namespace = self.find_namespace(namespace).await?;
        self.find_repository_by_name_and_wk(repo_name, namespace.id)
            .await
    }

    async fn find_namespace(
        &self,
        namespace: &str,
    ) -> Result<WorkspaceModel, GitError> {
        let workspace = sqlx::query_as::<_, WorkspaceModel>(
            "SELECT id, name, description, avatar_url, created_at FROM workspace WHERE name = $1",
        )
        .bind(namespace)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| GitError::Internal(e.to_string()))?;

        if let Some(ws) = workspace {
            return Ok(ws);
        }

        let history = sqlx::query_as::<_, WkHistoryNameModel>(
            "SELECT id, wk, name, changed_by, created_at FROM wk_history_name WHERE name = $1",
        )
        .bind(namespace)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| GitError::Internal(e.to_string()))?;

        if let Some(history) = history {
            let ws = sqlx::query_as::<_, WorkspaceModel>(
                "SELECT id, name, description, avatar_url, created_at FROM workspace WHERE id = $1",
            )
            .bind(history.wk)
            .fetch_optional(self.db.reader())
            .await
            .map_err(|e| GitError::Internal(e.to_string()))?;

            if let Some(ws) = ws {
                return Ok(ws);
            }
        }

        Err(GitError::NotFound("Workspace not found".to_string()))
    }

    async fn find_repository_by_name_and_wk(
        &self,
        repo_name: &str,
        wk_id: uuid::Uuid,
    ) -> Result<RepoModel, GitError> {
        let repo = sqlx::query_as::<_, RepoModel>(
            "SELECT id, wk, name, description, default_branch, visibility, size_bytes, is_archived, is_template, is_mirror, created_by, storage_path, created_at, updated_at, deleted_at FROM repo WHERE name = $1 AND wk = $2",
        )
            .bind(repo_name)
            .bind(wk_id)
            .fetch_optional(self.db.reader())
            .await
            .map_err(|e| GitError::Internal(e.to_string()))?;

        if let Some(repo) = repo {
            return Ok(repo);
        }

        let history = sqlx::query_as::<_, RepoHistoryNameModel>(
            "SELECT id, repo, name, changed_by, created_at FROM repo_history_name WHERE name = $1",
        )
        .bind(repo_name)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| GitError::Internal(e.to_string()))?;

        if let Some(history) = history {
            let repo = sqlx::query_as::<_, RepoModel>(
                "SELECT id, wk, name, description, default_branch, visibility, size_bytes, is_archived, is_template, is_mirror, created_by, storage_path, created_at, updated_at, deleted_at FROM repo WHERE id = $1 AND wk = $2",
            )
                .bind(history.repo)
                .bind(wk_id)
                .fetch_optional(self.db.reader())
                .await
                .map_err(|e| GitError::Internal(e.to_string()))?;

            if let Some(repo) = repo {
                return Ok(repo);
            }
        }

        Err(GitError::NotFound("Repository not found".to_string()))
    }

    pub async fn find_user_by_public_key(
        &self,
        public_key_str: &str,
    ) -> Result<Option<SshKeyUser>, sqlx::Error> {
        let fingerprint =
            match self.generate_fingerprint_from_public_key(public_key_str) {
                Ok(fp) => fp,
                Err(e) => {
                    tracing::error!(
                        "failed to generate SSH key fingerprint error={}",
                        e
                    );
                    return Ok(None);
                }
            };

        let fingerprint_preview = if fingerprint.len() > 16 {
            format!("{}...", &fingerprint[..16])
        } else {
            fingerprint.clone()
        };
        tracing::info!(
            "looking up user with SSH key fingerprint={}",
            fingerprint_preview
        );

        let ssh_key = sqlx::query_as::<_, UserSshKeyModel>(
            "SELECT id, \"user\", title, public_key, fingerprint, key_type, key_bits, is_verified, last_used_at, expires_at, is_revoked, created_at, updated_at FROM user_ssh_key WHERE fingerprint = $1 AND is_revoked = false",
        )
            .bind(&fingerprint)
            .fetch_optional(self.db.reader())
            .await?;

        let ssh_key = match ssh_key {
            Some(key) => key,
            None => {
                tracing::warn!("no SSH key found fingerprint={}", fingerprint);
                return Ok(None);
            }
        };

        if self.is_key_expired(&ssh_key) {
            tracing::warn!(
                "SSH key expired key_id={} expires_at={:?}",
                ssh_key.id,
                ssh_key.expires_at
            );
            return Ok(None);
        }

        let user_model = sqlx::query_as::<_, UserModel>(
            "SELECT id, username, display_name, avatar_url, website_url, allow_use, can_search, last_sign_in_at, created_at, updated_at FROM \"user\" WHERE id = $1",
        )
            .bind(ssh_key.user)
            .fetch_optional(self.db.reader())
            .await?;

        if let Some(user) = user_model {
            tracing::info!(
                "SSH key matched user={} key={}",
                user.username,
                ssh_key.title
            );
            return Ok(Some(SshKeyUser {
                user,
                key_id: ssh_key.id,
                key_title: ssh_key.title,
            }));
        }

        Ok(None)
    }

    fn is_key_expired(&self, ssh_key: &UserSshKeyModel) -> bool {
        if let Some(expires_at) = ssh_key.expires_at {
            let now = chrono::Utc::now();
            now >= expires_at
        } else {
            false
        }
    }

    pub fn update_key_last_used_async(&self, key_id: i64) {
        let db_clone = self.db.clone();
        tokio::spawn(async move {
            if let Err(e) =
                Self::update_key_last_used_sync(db_clone, key_id).await
            {
                tracing::warn!(
                    "failed to update key last_used key_id={} error={}",
                    key_id,
                    e
                );
            }
        });
    }

    async fn update_key_last_used_sync(
        db: AppDatabase,
        key_id: i64,
    ) -> Result<(), sqlx::Error> {
        let now = chrono::Utc::now();
        sqlx::query("UPDATE user_ssh_key SET last_used_at = $1, updated_at = $2 WHERE id = $3")
            .bind(now)
            .bind(now)
            .bind(key_id)
            .execute(db.writer())
            .await?;

        tracing::info!("updated key last_used key_id={}", key_id);
        Ok(())
    }

    pub async fn check_repo_permission(
        &self,
        user: &UserModel,
        repo: &RepoModel,
        is_write: bool,
    ) -> bool {
        if repo.created_by == user.id {
            tracing::info!(
                "user is repo owner user={} repo={}",
                user.username,
                repo.name
            );
            return true;
        }

        if !is_write && repo.visibility == "public" {
            tracing::info!("public repo allows read access repo={}", repo.name);
            return true;
        }

        let wk_id = repo.wk;
        if self
            .check_wk_member_permission(user, wk_id, is_write)
            .await
            .unwrap_or(false)
        {
            tracing::info!(
                "user has workspace member access user={} repo={}",
                user.username,
                repo.name
            );
            return true;
        }

        tracing::warn!(
            "access denied user={} repo={} is_write={}",
            user.username,
            repo.name,
            is_write
        );
        false
    }

    async fn check_wk_member_permission(
        &self,
        user: &UserModel,
        wk_id: uuid::Uuid,
        is_write: bool,
    ) -> Result<bool, sqlx::Error> {
        let member = sqlx::query_as::<_, WkMemberModel>(
            "SELECT wk, \"user\", owner, admin, join_at, leave_at FROM wk_member WHERE wk = $1 AND \"user\" = $2 AND leave_at IS NULL",
        )
            .bind(wk_id)
            .bind(user.id)
            .fetch_optional(self.db.reader())
            .await?;

        if let Some(member) = member {
            if member.owner || member.admin {
                return Ok(true);
            }
            Ok(!is_write)
        } else {
            Ok(false)
        }
    }

    fn generate_fingerprint_from_public_key(
        &self,
        public_key_str: &str,
    ) -> Result<String, String> {
        let key_data_base64 = public_key_str
            .split_whitespace()
            .nth(1)
            .ok_or("Invalid SSH key format")?;

        let key_data = general_purpose::STANDARD
            .decode(key_data_base64)
            .map_err(|e| format!("Base64 decode error: {}", e))?;

        let mut hasher = Sha256::new();
        hasher.update(&key_data);
        let hash = hasher.finalize();

        let mut fingerprint = String::with_capacity(51);
        fingerprint.push_str("SHA256:");
        fingerprint.push_str(&general_purpose::STANDARD_NO_PAD.encode(&hash));

        Ok(fingerprint)
    }
}