e478cc4
use std::net::SocketAddr;
use std::str::FromStr;
use std::sync::Arc;
use russh::keys::{Algorithm, PrivateKey};
use russh::server::Server as _;
use russh::{MethodKind, MethodSet, SshId, server};
use tokio::sync::Semaphore;
use crate::error::{SshError, SshResult};
use crate::handler::GitSshHandler;
use crate::service::GitSshService;
use config::git::GitSshConfig;
pub struct GitSshServer {
service: Arc<GitSshService>,
session_slots: Arc<Semaphore>,
}
impl GitSshServer {
pub fn new(service: Arc<GitSshService>) -> Self {
let max = service.config.max_sessions.max(1);
Self {
service,
session_slots: Arc::new(Semaphore::new(max)),
}
}
}
impl server::Server for GitSshServer {
type Handler = GitSshHandler;
fn new_client(&mut self, _peer_addr: Option<SocketAddr>) -> Self::Handler {
let permit = self.session_slots.clone().try_acquire_owned().ok();
crate::metrics::record_session(permit.is_some());
GitSshHandler::new(self.service.clone(), _peer_addr, permit)
}
fn handle_session_error(&mut self, error: russh::Error) {
tracing::debug!("ssh session error: {error}");
}
}
pub fn load_host_key(path: &str) -> SshResult<PrivateKey> {
let data = std::fs::read_to_string(path)
.map_err(|e| SshError::KeyLoadFailed(format!("cannot read {}: {e}", path)))?;
if let Ok(key) = PrivateKey::from_openssh(&data) {
return Ok(key);
}
if let Ok(key) = PrivateKey::from_str(&data) {
return Ok(key);
}
Err(SshError::KeyLoadFailed(format!(
"unsupported host key format in {}",
path
)))
}
pub fn generate_host_key() -> SshResult<PrivateKey> {
let mut rng = rand::rng();
PrivateKey::random(&mut rng, Algorithm::Ed25519)
.map_err(|e| SshError::KeyLoadFailed(format!("failed to generate host key: {e}")))
}
pub fn build_server_config(config: &GitSshConfig) -> SshResult<server::Config> {
let host_key = load_host_key(&config.host_key_path)?;
let mut methods = MethodSet::empty();
methods.push(MethodKind::PublicKey);
Ok(server::Config {
keys: vec![host_key],
server_id: SshId::Standard(format!("SSH-2.0-gitssh-{}", env!("CARGO_PKG_VERSION")).into()),
methods,
max_auth_attempts: 10,
auth_rejection_time: config.auth_rejection_time,
inactivity_timeout: config.inactivity_timeout,
keepalive_interval: Some(config.keepalive_interval),
keepalive_max: config.keepalive_max as usize,
..Default::default()
})
}
pub async fn run_ssh(service: GitSshService) -> SshResult<()> {
let bind_addr = service.config.bind_addr.clone();
let port = service.config.port;
let config = Arc::new(build_server_config(&service.config)?);
let mut server = GitSshServer::new(Arc::new(service));
server
.run_on_address(config, format!("{bind_addr}:{port}"))
.await
.map_err(SshError::Io)?;
Ok(())
}
pub async fn serve_on_listener(
service: GitSshService,
listener: tokio::net::TcpListener,
) -> SshResult<()> {
let config = Arc::new(build_server_config(&service.config)?);
let mut server = GitSshServer::new(Arc::new(service));
let running = server.run_on_socket(config, &listener);
running.await.map_err(SshError::Io)?;
Ok(())
}