gitdataai/lib/git/ssh/mod.rs

use std::{sync::Arc, time::Duration};

use anyhow::Context;
use argon2::{
    Argon2,
    password_hash::{PasswordHash, PasswordVerifier},
};
use cache::AppCache;
use config::AppConfig;
use db::{database::AppDatabase, sqlx};
use deadpool_redis::cluster::Pool as RedisPool;
use model::users::{UserModel, UserTokenModel};
use russh::{
    MethodKind, MethodSet, SshId,
    server::{Config, Server},
};

use crate::errors::GitError;

pub mod authz;
pub mod branch_protect;
pub mod forward;
pub mod git_service;
pub mod handler;
pub mod rate_limit;
pub mod ref_update;
pub mod server;

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

impl SSHHandle {
    pub fn new(
        db: AppDatabase,
        app: AppConfig,
        cache: AppCache,
        redis_pool: RedisPool,
    ) -> Self {
        SSHHandle {
            db,
            app,
            cache,
            redis_pool,
        }
    }

    pub async fn run_ssh(&self) -> anyhow::Result<()> {
        tracing::info!("SSH server starting");
        let key_file = self.app.ssh_server_private_key_file()?;
        if key_file.is_empty() {
            return Err(anyhow::anyhow!(
                "SSH server private key file is not configured (APP_SSH_SERVER_PRIVATE_KEY_FILE)"
            ));
        }

        tracing::info!("Loading SSH private key from file: {}", key_file);

        let private_key_pem = std::fs::read_to_string(&key_file)
            .with_context(|| format!("Failed to read SSH private key file: {}", key_file))?;

        let private_key = russh::keys::decode_secret_key(&private_key_pem, None)
            .or_else(|e| {
                tracing::info!("decode_secret_key failed: {}, trying from_openssh", e);
                russh::keys::ssh_key::PrivateKey::from_openssh(&private_key_pem)
                    .map_err(|e2| anyhow::anyhow!(
                        "Failed to parse SSH private key from {}: decode_secret_key={}, from_openssh={}",
                        key_file, e, e2
                    ))
            })?;

        tracing::info!("SSH private key loaded");
        let mut config = Config::default();
        config.keys = vec![private_key];
        let version = format!("SSH-2.0-Work {}", env!("CARGO_PKG_VERSION"));
        config.server_id = SshId::Standard(version.into());
        config.methods = MethodSet::empty();
        config.methods.push(MethodKind::PublicKey);
        config.methods.push(MethodKind::Password);
        config.auth_rejection_time = Duration::from_secs(5);
        config.inactivity_timeout = Some(Duration::from_secs(300));
        config.keepalive_interval = Some(Duration::from_secs(60));
        config.keepalive_max = 3;

        tracing::info!(
            "SSH server configured with methods: {:?}",
            config.methods
        );
        let token_service = SshTokenService::new(self.db.clone());
        let mut ssh_server = server::SSHServer::new(
            self.db.clone(),
            self.cache.clone(),
            self.redis_pool.clone(),
            token_service,
        );

        let _cleanup = ssh_server.rate_limiter.clone().start_cleanup();

        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
            )
        };
        tracing::info!("{}", msg);
        ssh_server
            .run_on_address(Arc::new(config), bind_addr)
            .await?;
        Ok(())
    }
}
#[derive(Clone)]
pub struct SshTokenService {
    db: AppDatabase,
}

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

    pub async fn find_user_by_token(
        &self,
        token: &str,
    ) -> Result<Option<UserModel>, GitError> {
        let token_models = sqlx::query_as::<_, UserTokenModel>(
            "SELECT id, user, name, token_hash, scopes, expires_at, is_revoked, created_at, updated_at FROM user_token WHERE is_revoked = false",
        )
            .fetch_all(self.db.reader())
            .await
            .map_err(|e| GitError::Internal(e.to_string()))?;

        for token_model in token_models {
            if token_model
                .expires_at
                .map(|expires_at| expires_at < chrono::Utc::now())
                .unwrap_or(false)
            {
                continue;
            }

            let Ok(hash) = PasswordHash::new(&token_model.token_hash) else {
                tracing::warn!(
                    token_id = token_model.id,
                    "invalid stored SSH token hash"
                );
                continue;
            };

            if Argon2::default()
                .verify_password(token.as_bytes(), &hash)
                .is_err()
            {
                continue;
            }

            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(token_model.user)
                .fetch_optional(self.db.reader())
                .await
                .map_err(|e| GitError::Internal(e.to_string()))?;

            return Ok(user_model);
        }

        Ok(None)
    }
}

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