gitdataai/libs/git/ssh/mod.rs

use crate::error::GitError;
use crate::hook::HookService;
use anyhow::Context;
use base64::Engine;
use config::AppConfig;
use db::cache::AppCache;
use db::database::AppDatabase;
use deadpool_redis::cluster::Pool as RedisPool;
use models::users::{user, user_token};
use russh::keys::PrivateKey;
use russh::server::Server;
use russh::{MethodKind, MethodSet, SshId, server::Config};
use sea_orm::prelude::*;
use sha2::{Digest, Sha256};
use slog::{Logger, error, info};
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;

pub mod authz;
pub mod handle;
pub mod server;

#[derive(Clone)]
pub struct SSHHandle {
    pub db: AppDatabase,
    pub app: AppConfig,
    pub cache: AppCache,
    pub redis_pool: RedisPool,
    pub logger: Logger,
}

impl SSHHandle {
    pub async fn run(&self) {
        let this = self.clone();
        tokio::spawn(async move {
            if let Err(e) = this.run_ssh().await {
                error!(this.logger, "SSH server error: {}", e);
            }
        });
    }
    pub fn new(
        db: AppDatabase,
        app: AppConfig,
        cache: AppCache,
        redis_pool: RedisPool,
        logger: Logger,
    ) -> Self {
        SSHHandle {
            db,
            app,
            cache,
            redis_pool,
            logger,
        }
    }
    pub async fn run_ssh(&self) -> anyhow::Result<()> {
        info!(self.logger, "SSH server starting");
        let private_key_content = self.app.ssh_server_private_key()?;
        if private_key_content.is_empty() {
            return Err(anyhow::anyhow!("SSH server private key is not configured"));
        }

        let preview = if private_key_content.len() > 100 {
            format!("{}...", &private_key_content[..100])
        } else {
            private_key_content.clone()
        };
        info!(
            self.logger,
            "Loading SSH private key (hex, {} bytes)",
            private_key_content.len()
        );

        let private_key_bytes = hex::decode(&private_key_content).with_context(|| {
            format!(
                "Failed to decode hex-encoded SSH private key. Preview: {}",
                preview
            )
        })?;

        info!(
            self.logger,
            "Hex decoded to {} bytes",
            private_key_bytes.len()
        );

        let private_key_pem = std::str::from_utf8(&private_key_bytes)
            .with_context(|| "Decoded SSH private key is not valid UTF-8")?;

        if let Some(first_line) = private_key_pem.lines().next() {
            info!(self.logger, "PEM format starts with: {}", first_line);
        }

        info!(
            self.logger,
            "Complete private key content:\n{}", private_key_pem
        );

        let private_key = {
            match ssh_key::PrivateKey::from_openssh(private_key_pem) {
                Ok(ssh_key) => {
                    info!(self.logger, "Successfully parsed with ssh-key crate");
                    let openssh_pem = ssh_key
                        .to_openssh(ssh_key::LineEnding::LF)
                        .with_context(|| "Failed to serialize to OpenSSH format")?;

                    PrivateKey::from_str(&openssh_pem)
                        .with_context(|| "Failed to parse with russh after ssh-key conversion")?
                }
                Err(e) => {
                    info!(
                        self.logger,
                        "ssh-key from_openssh failed: {}, trying direct russh parse", e
                    );
                    PrivateKey::from_str(private_key_pem).with_context(|| {
                        format!("Failed to parse SSH private key with both methods")
                    })?
                }
            }
        };
        info!(self.logger, "SSH private key loaded");
        let mut config = Config::default();
        config.keys = vec![private_key];
        let version = format!("SSH-2.0-GitdataAI {}", env!("CARGO_PKG_VERSION"));
        config.server_id = SshId::Standard(version);
        let mut method = MethodSet::empty();
        method.push(MethodKind::PublicKey);
        method.push(MethodKind::Password);
        config.methods = method;
        config.inactivity_timeout = Some(Duration::from_secs(300));
        config.keepalive_interval = Some(Duration::from_secs(60));
        config.keepalive_max = 3;

        info!(
            self.logger,
            "SSH server configured with methods: {:?}", config.methods
        );
        let token_service = SshTokenService::new(self.db.clone());
        let http = Arc::new(reqwest::Client::new());
        let hook = crate::hook::HookService::new(
            self.db.clone(),
            self.cache.clone(),
            self.redis_pool.clone(),
            self.logger.clone(),
            self.app.clone(),
            http,
        );
        let mut server = server::SSHServer::new(
            self.db.clone(),
            self.cache.clone(),
            self.redis_pool.clone(),
            self.logger.clone(),
            token_service,
            hook,
        );

        // Start the rate limiter cleanup background task so the HashMap
        // doesn't grow unbounded over time.
        let ssh_port = self.app.ssh_port()?;
        let bind_addr = format!("0.0.0.0:{}", ssh_port);
        let public_host = self.app.ssh_domain()?;
        let msg = if ssh_port == 22 {
            format!(
                "SSH server listening on port 22. Please use port {} for SSH connections.",
                ssh_port
            )
        } else {
            format!(
                "SSH server listening on port {} (public: {}). Please use port {} for SSH connections.",
                ssh_port, public_host, ssh_port
            )
        };
        info!(self.logger, "{}", msg);
        server.run_on_address(Arc::new(config), bind_addr).await?;
        Ok(())
    }
}

/// Direct sync service — calls HookService::sync_repo inline.
/// K8s StatefulSet HA scheduling ensures exclusive access per repo shard.
#[derive(Clone)]
pub struct ReceiveSyncService {
    hook: HookService,
}

impl ReceiveSyncService {
    pub fn new(hook: HookService) -> Self {
        Self { hook }
    }

    /// Execute a full repo sync synchronously.
    /// Returns Ok on success, Err on failure.
    pub async fn send(&self, task: RepoReceiveSyncTask) -> Result<(), crate::GitError> {
        let repo_id = task.repo_uid.to_string();
        slog::info!(self.hook.logger, "starting sync repo_id={}", repo_id);
        let res = self.hook.sync_repo(&repo_id).await;
        match &res {
            Ok(()) => {
                slog::info!(self.hook.logger, "sync completed repo_id={}", repo_id);
            }
            Err(e) => {
                slog::error!(self.hook.logger, "sync failed repo_id={} error={}", repo_id, e);
            }
        }
        res
    }
}

#[derive(Clone)]
pub struct RepoReceiveSyncTask {
    pub repo_uid: uuid::Uuid,
}

/// SSH token authentication service.
/// Uses the same token hash algorithm as user access keys (SHA256 + base64).
#[derive(Clone)]
pub struct SshTokenService {
    db: AppDatabase,
}

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

    fn hash_token(token: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(token.as_bytes());
        base64::prelude::BASE64_STANDARD.encode(hasher.finalize())
    }

    pub async fn find_user_by_token(&self, token: &str) -> Result<Option<user::Model>, GitError> {
        let token_hash = Self::hash_token(token);

        let token_model = user_token::Entity::find()
            .filter(user_token::Column::TokenHash.eq(&token_hash))
            .filter(user_token::Column::IsRevoked.eq(false))
            .one(self.db.reader())
            .await
            .map_err(|e| GitError::Internal(e.to_string()))?;

        let token_model = match token_model {
            Some(t) => t,
            None => return Ok(None),
        };

        // Check expiry
        if let Some(expires_at) = token_model.expires_at {
            if expires_at < chrono::Utc::now() {
                return Ok(None);
            }
        }

        let user_model = user::Entity::find()
            .filter(user::Column::Uid.eq(token_model.user))
            .one(self.db.reader())
            .await
            .map_err(|e| GitError::Internal(e.to_string()))?;

        Ok(user_model)
    }
}

pub async fn run_ssh(config: AppConfig, logger: Logger) -> anyhow::Result<()> {
    info!(logger, "SSH server initializing");
    let db = AppDatabase::init(&config).await?;
    let cache = AppCache::init(&config).await?;
    let redis_pool = cache.redis_pool().clone();
    SSHHandle::new(db, config.clone(), cache, redis_pool, logger)
        .run_ssh()
        .await?;
    Ok(())
}