use std::collections::HashMap;
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use anyhow::{bail, Context, Result};
use bytes::Bytes;
use russh::keys::{Algorithm, PrivateKey};
use russh::server::{self, Auth, Msg, Server as _, Session};
use russh::{Channel, ChannelId};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::{mpsc, Mutex};
use crate::httpgit::AppState;
use crate::service::Service;
struct Inner {
state: AppState,
}
#[derive(Clone)]
struct GitHandler {
inner: Arc<Inner>,
stdins: Arc<Mutex<HashMap<ChannelId, mpsc::UnboundedSender<Bytes>>>>,
}
impl server::Server for GitHandler {
type Handler = Self;
fn new_client(&mut self, _addr: Option<SocketAddr>) -> Self {
GitHandler { inner: self.inner.clone(), stdins: Arc::new(Mutex::new(HashMap::new())) }
}
}
impl server::Handler for GitHandler {
type Error = russh::Error;
async fn auth_publickey(&mut self, _user: &str, _key: &russh::keys::PublicKey) -> Result<Auth, Self::Error> {
Ok(Auth::Accept)
}
async fn channel_open_session(
&mut self,
_channel: Channel<Msg>,
reply: server::ChannelOpenHandle,
_session: &mut Session,
) -> Result<(), Self::Error> {
reply.accept().await;
Ok(())
}
async fn exec_request(
&mut self,
channel: ChannelId,
data: &[u8],
session: &mut Session,
) -> Result<(), Self::Error> {
let cmd_line = String::from_utf8_lossy(data).into_owned();
match self.start(channel, &cmd_line, session.handle()).await {
Ok(()) => {
session.channel_success(channel)?;
}
Err(e) => {
session.data(channel, Bytes::from(format!("nugitea: {e}\n")))?;
session.channel_failure(channel)?;
session.exit_status_request(channel, 1)?;
session.close(channel)?;
}
}
Ok(())
}
async fn data(&mut self, channel: ChannelId, data: &[u8], _session: &mut Session) -> Result<(), Self::Error> {
if let Some(tx) = self.stdins.lock().await.get(&channel) {
let _ = tx.send(Bytes::copy_from_slice(data));
}
Ok(())
}
async fn channel_eof(&mut self, channel: ChannelId, _session: &mut Session) -> Result<(), Self::Error> {
self.stdins.lock().await.remove(&channel);
Ok(())
}
}
impl GitHandler {
async fn start(&self, channel: ChannelId, cmd_line: &str, handle: server::Handle) -> Result<()> {
let fields = shlex::split(cmd_line).context("unsupported command")?;
let [verb, repo_ref] = fields.as_slice() else {
bail!("unsupported command");
};
let service = Service::from_ssh_verb(verb).with_context(|| format!("unsupported command {verb:?}"))?;
let state = &self.inner.state;
let repo_name = state.resolve(repo_ref).await.with_context(|| format!("repository {repo_ref:?} not found"))?;
if !service.allow(state.auth.as_ref(), &repo_name) {
bail!("forbidden");
}
let stream = state.storage.connect_git_exec(&repo_name, service).await?;
let (mut read_half, write_half) = stream.into_split();
let (tx, mut rx) = mpsc::unbounded_channel::<Bytes>();
self.stdins.lock().await.insert(channel, tx);
let stdins = self.stdins.clone();
tokio::spawn(async move {
let write_task = tokio::spawn(async move {
let mut write_half = write_half;
while let Some(chunk) = rx.recv().await {
if write_half.write_all(&chunk).await.is_err() {
break;
}
}
let _ = write_half.shutdown().await;
});
let mut buf = [0u8; 32 * 1024];
let mut ok = true;
loop {
match read_half.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if handle.data(channel, Bytes::copy_from_slice(&buf[..n])).await.is_err() {
ok = false;
break;
}
}
Err(_) => {
ok = false;
break;
}
}
}
write_task.abort();
stdins.lock().await.remove(&channel);
let _ = handle.exit_status_request(channel, if ok { 0 } else { 1 }).await;
let _ = handle.eof(channel).await;
let _ = handle.close(channel).await;
});
Ok(())
}
}
pub async fn serve(state: AppState, addr: SocketAddr, host_key_path: PathBuf) -> Result<()> {
let key = load_or_create_host_key(&host_key_path)?;
let config = Arc::new(server::Config { keys: vec![key], ..Default::default() });
let mut handler = GitHandler { inner: Arc::new(Inner { state }), stdins: Arc::new(Mutex::new(HashMap::new())) };
println!("ssh: listening on {addr}");
handler.run_on_address(config, addr).await?;
Ok(())
}
fn load_or_create_host_key(path: &Path) -> Result<PrivateKey> {
if let Ok(data) = std::fs::read(path) {
return PrivateKey::from_openssh(&data).context("parse ssh host key");
}
let key = PrivateKey::random(&mut rand::rng(), Algorithm::Ed25519).context("generate ssh host key")?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let pem = key.to_openssh(russh::keys::ssh_key::LineEnding::LF)?;
std::fs::write(path, pem.as_bytes())?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?;
}
Ok(key)
}