//! russh server: host key, pubkey auth, exec and env handling, piping, timeouts.

use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::OpenOptions;
use std::io::{self, Write};
use std::net::SocketAddr;
use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
use std::os::unix::process::ExitStatusExt;
use std::path::{Path, PathBuf};
use std::process::ExitStatus;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use bytes::Bytes;
use russh::keys::ssh_key::LineEnding;
use russh::keys::ssh_key::private::{Ed25519Keypair, KeypairData};
use russh::keys::{Algorithm, EcdsaCurve, HashAlg, PrivateKey, PublicKey};
use russh::server::{Auth, Config, Handle, Handler, Msg, Session, run_stream};
use russh::{Channel, ChannelId, ChannelOpenFailure, Disconnect, MethodKind, MethodSet, Preferred};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::process::{Child, ChildStdin};
use tokio::sync::{Notify, mpsc};

use crate::backend::GitBackend;
use crate::error::{Error, Result};
use crate::parse::{LFS_AUTHENTICATE_PREFIX, parse_exec, parse_lfs_authenticate};

/// Stdout and stderr read size per message sent to the client.
const PUMP_CHUNK: usize = 32 * 1024;
/// Client input chunks buffered in front of the child's stdin. This bounds memory and
/// propagates backpressure to the SSH read loop.
const STDIN_QUEUE: usize = 32;
/// Constant delay applied to failed authentication attempts.
const AUTH_REJECTION_TIME: Duration = Duration::from_secs(1);
/// Failed publickey attempts before the connection is dropped.
const MAX_AUTH_ATTEMPTS: usize = 6;

#[derive(Debug, Clone)]
pub struct SshConfig {
    pub listen: SocketAddr,
    pub host_key_path: PathBuf,
    pub auth_timeout: Duration,
    pub inactivity_timeout: Duration,
}

/// Runs the SSH server until it fails. Creates the ed25519 host key on first boot.
pub async fn run<B: GitBackend>(config: SshConfig, backend: Arc<B>) -> Result<()> {
    let host_key = load_or_create_host_key(&config.host_key_path)?;
    let russh_config = Arc::new(Config {
        keys: vec![host_key],
        methods: publickey_only(),
        auth_rejection_time: AUTH_REJECTION_TIME,
        max_auth_attempts: MAX_AUTH_ATTEMPTS,
        inactivity_timeout: Some(config.inactivity_timeout),
        preferred: preferred_algorithms(),
        ..Config::default()
    });

    let listener = TcpListener::bind(config.listen).await?;
    tracing::info!(addr = %config.listen, "ssh server listening");

    loop {
        let (stream, peer) = match listener.accept().await {
            Ok(v) => v,
            Err(e) => {
                // Typically EMFILE. Back off briefly instead of tearing the server down.
                tracing::warn!(error = %e, "ssh accept failed");
                tokio::time::sleep(Duration::from_millis(100)).await;
                continue;
            }
        };
        let authed = Arc::new(AtomicBool::new(false));
        let client = Client {
            backend: Arc::clone(&backend),
            peer: Some(peer),
            authed: Arc::clone(&authed),
            user_id: None,
            git_protocol: None,
            channels: HashMap::new(),
        };
        let russh_config = Arc::clone(&russh_config);
        let auth_timeout = config.auth_timeout;
        tokio::spawn(async move {
            let running = match run_stream(russh_config, stream, client).await {
                Ok(running) => running,
                Err(e) => {
                    tracing::debug!(%peer, error = %e, "ssh handshake failed");
                    return;
                }
            };
            // Drop unauthenticated connections after the auth timeout.
            let handle = running.handle();
            tokio::spawn(async move {
                tokio::time::sleep(auth_timeout).await;
                if !authed.load(Ordering::Acquire) {
                    let _ = handle
                        .disconnect(
                            Disconnect::ByApplication,
                            "authentication timeout".into(),
                            String::new(),
                        )
                        .await;
                }
            });
            if let Err(e) = running.await {
                tracing::debug!(%peer, error = %e, "ssh session ended with error");
            }
        });
    }
}

/// Publickey is the only accepted method.
fn publickey_only() -> MethodSet {
    let mut methods = MethodSet::empty();
    methods.push(MethodKind::PublicKey);
    methods
}

/// Host and user key algorithms this server negotiates. `ssh-rsa` (SHA-1) is not listed.
fn preferred_algorithms() -> Preferred {
    Preferred {
        key: Cow::Owned(vec![
            Algorithm::Ed25519,
            Algorithm::SkEd25519,
            Algorithm::SkEcdsaSha2NistP256,
            Algorithm::Ecdsa {
                curve: EcdsaCurve::NistP256,
            },
            Algorithm::Ecdsa {
                curve: EcdsaCurve::NistP384,
            },
            Algorithm::Ecdsa {
                curve: EcdsaCurve::NistP521,
            },
            Algorithm::Rsa {
                hash: Some(HashAlg::Sha512),
            },
            Algorithm::Rsa {
                hash: Some(HashAlg::Sha256),
            },
        ]),
        ..Preferred::default()
    }
}

/// Key types accepted for authentication. RSA keys appear as `ssh-rsa` key blobs whatever the
/// signature algorithm is, so this check cannot tell SHA-1 signatures apart. See `run`'s docs.
fn key_type_allowed(algorithm: &Algorithm) -> bool {
    matches!(
        algorithm,
        Algorithm::Ed25519
            | Algorithm::SkEd25519
            | Algorithm::SkEcdsaSha2NistP256
            | Algorithm::Ecdsa { .. }
            | Algorithm::Rsa { .. }
    )
}

/// Accepts exactly `version=<digit>`, e.g. `version=2`.
fn valid_git_protocol(value: &str) -> bool {
    let bytes = value.as_bytes();
    bytes.len() == 9 && value.starts_with("version=") && bytes[8].is_ascii_digit()
}

fn load_or_create_host_key(path: &Path) -> Result<PrivateKey> {
    if path.exists() {
        let key = PrivateKey::read_openssh_file(path)
            .map_err(|e| Error::Ssh(format!("host key: {e}")))?;
        if key.algorithm() != Algorithm::Ed25519 {
            return Err(Error::Ssh("host key must be ed25519".into()));
        }
        return Ok(key);
    }

    let mut seed = [0u8; 32];
    getrandom::fill(&mut seed).map_err(|e| Error::Ssh(format!("host key entropy: {e}")))?;
    let key = PrivateKey::new(KeypairData::Ed25519(Ed25519Keypair::from_seed(&seed)), "")
        .map_err(|e| Error::Ssh(format!("host key generation: {e}")))?;
    let pem = key
        .to_openssh(LineEnding::LF)
        .map_err(|e| Error::Ssh(format!("host key encoding: {e}")))?;

    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    // create_new so an existing key is never overwritten. The mode applies at creation.
    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .mode(0o600)
        .open(path)?;
    file.write_all(pem.as_bytes())?;
    file.sync_all()?;
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?;
    tracing::info!(path = %path.display(), "generated ssh host key");
    Ok(key)
}

/// Per-connection SSH handler.
struct Client<B> {
    backend: Arc<B>,
    peer: Option<SocketAddr>,
    authed: Arc<AtomicBool>,
    user_id: Option<String>,
    git_protocol: Option<String>,
    /// Open exec channels. Entries are removed on channel EOF or close, which drops the
    /// sender and closes the child's stdin.
    channels: HashMap<ChannelId, mpsc::Sender<Bytes>>,
}

impl<B: GitBackend> Client<B> {
    async fn lookup(&self, key: &PublicKey) -> Option<String> {
        if !key_type_allowed(&key.algorithm()) {
            return None;
        }
        let blob = key.to_bytes().ok()?;
        self.backend
            .authenticate_key(key.algorithm().as_str(), &blob)
            .await
    }

    /// Answers `git-lfs-authenticate` with the JSON reply on stdout and exit status 0. Nothing is
    /// left open afterwards: the channel is closed once the reply is queued.
    async fn exec_lfs_authenticate(
        &self,
        session: &mut Session,
        channel: ChannelId,
        user_id: &str,
        command: &str,
    ) -> std::result::Result<(), russh::Error> {
        let (path, operation) = match parse_lfs_authenticate(command) {
            Ok(parsed) => parsed,
            Err(e) => return reject_exec(session, channel, &e.to_string()),
        };
        let reply = match self
            .backend
            .lfs_authenticate(user_id, &path, operation)
            .await
        {
            Ok(reply) => reply,
            Err(msg) => return reject_exec(session, channel, &msg),
        };
        session.channel_success(channel)?;
        session.data(channel, reply)?;
        session.exit_status_request(channel, 0)?;
        session.eof(channel)?;
        session.close(channel)?;
        Ok(())
    }
}

impl<B: GitBackend> Handler for Client<B> {
    type Error = russh::Error;

    async fn auth_publickey_offered(
        &mut self,
        _user: &str,
        key: &PublicKey,
    ) -> std::result::Result<Auth, Self::Error> {
        // Probe before the client signs. Reject early so unknown keys never get a signature.
        if self.lookup(key).await.is_some() {
            Ok(Auth::Accept)
        } else {
            Ok(Auth::reject())
        }
    }

    async fn auth_publickey(
        &mut self,
        _user: &str,
        key: &PublicKey,
    ) -> std::result::Result<Auth, Self::Error> {
        match self.lookup(key).await {
            Some(user_id) => {
                tracing::info!(peer = ?self.peer, %user_id, "ssh publickey auth accepted");
                self.user_id = Some(user_id);
                self.authed.store(true, Ordering::Release);
                Ok(Auth::Accept)
            }
            None => Ok(Auth::reject()),
        }
    }

    async fn channel_open_session(
        &mut self,
        _channel: Channel<Msg>,
        reply: russh::server::ChannelOpenHandle,
        _session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        if self.user_id.is_none() {
            reply
                .reject(ChannelOpenFailure::AdministrativelyProhibited)
                .await;
        } else {
            reply.accept().await;
        }
        Ok(())
    }

    async fn env_request(
        &mut self,
        _channel: ChannelId,
        variable_name: &str,
        variable_value: &str,
        _session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        // Only GIT_PROTOCOL is kept. Everything else is ignored, never forwarded.
        if variable_name == "GIT_PROTOCOL" && valid_git_protocol(variable_value) {
            self.git_protocol = Some(variable_value.to_owned());
        }
        Ok(())
    }

    async fn exec_request(
        &mut self,
        channel: ChannelId,
        data: &[u8],
        session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        if self.channels.contains_key(&channel) {
            session.channel_failure(channel)?;
            return Ok(());
        }
        let Some(user_id) = self.user_id.clone() else {
            session.channel_failure(channel)?;
            return Ok(());
        };
        let command = match std::str::from_utf8(data) {
            Ok(command) => command,
            Err(_) => return reject_exec(session, channel, "invalid command"),
        };
        if command.starts_with(LFS_AUTHENTICATE_PREFIX) {
            return self
                .exec_lfs_authenticate(session, channel, &user_id, command)
                .await;
        }
        let (service, path) = match parse_exec(command) {
            Ok(parsed) => parsed,
            Err(e) => return reject_exec(session, channel, &e.to_string()),
        };

        let mut cmd = match self
            .backend
            .prepare(&user_id, service, &path, self.git_protocol.as_deref())
            .await
        {
            Ok(cmd) => cmd,
            Err(msg) => return reject_exec(session, channel, &msg),
        };
        cmd.stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .kill_on_drop(true);

        let child = match cmd.spawn() {
            Ok(child) => child,
            Err(e) => {
                tracing::error!(error = %e, "failed to spawn git");
                return reject_exec(session, channel, "failed to start git");
            }
        };

        session.channel_success(channel)?;
        let (tx, rx) = mpsc::channel(STDIN_QUEUE);
        self.channels.insert(channel, tx);
        tokio::spawn(run_child(session.handle(), channel, child, rx));
        Ok(())
    }

    async fn data(
        &mut self,
        channel: ChannelId,
        data: &[u8],
        _session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        if let Some(tx) = self.channels.get(&channel) {
            // Err means the child has exited. Its output is still flushed, so this is not fatal.
            let _ = tx.send(Bytes::copy_from_slice(data)).await;
        }
        Ok(())
    }

    async fn channel_eof(
        &mut self,
        channel: ChannelId,
        _session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        self.channels.remove(&channel);
        Ok(())
    }

    async fn channel_close(
        &mut self,
        channel: ChannelId,
        _session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        self.channels.remove(&channel);
        Ok(())
    }

    async fn pty_request(
        &mut self,
        channel: ChannelId,
        _term: &str,
        _col_width: u32,
        _row_height: u32,
        _pix_width: u32,
        _pix_height: u32,
        _modes: &[(russh::Pty, u32)],
        session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        session.channel_failure(channel)?;
        Ok(())
    }

    async fn shell_request(
        &mut self,
        channel: ChannelId,
        session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        session.channel_failure(channel)?;
        Ok(())
    }

    async fn subsystem_request(
        &mut self,
        channel: ChannelId,
        _name: &str,
        session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        session.channel_failure(channel)?;
        Ok(())
    }

    async fn x11_request(
        &mut self,
        channel: ChannelId,
        _single_connection: bool,
        _x11_auth_protocol: &str,
        _x11_auth_cookie: &str,
        _x11_screen_number: u32,
        session: &mut Session,
    ) -> std::result::Result<(), Self::Error> {
        session.channel_failure(channel)?;
        Ok(())
    }

    async fn tcpip_forward(
        &mut self,
        _address: &str,
        _port: &mut u32,
        _session: &mut Session,
    ) -> std::result::Result<bool, Self::Error> {
        Ok(false)
    }

    async fn agent_request(
        &mut self,
        _channel: ChannelId,
        _session: &mut Session,
    ) -> std::result::Result<bool, Self::Error> {
        Ok(false)
    }
}

/// Sends a rejection as the spec requires: the message on stderr, exit status 1, then EOF and close.
/// The request itself is acknowledged so git reads the message instead of failing at the protocol level.
fn reject_exec(
    session: &mut Session,
    channel: ChannelId,
    msg: &str,
) -> std::result::Result<(), russh::Error> {
    session.channel_success(channel)?;
    session.extended_data(channel, 1, format!("{msg}\n"))?;
    session.exit_status_request(channel, 1)?;
    session.eof(channel)?;
    session.close(channel)?;
    Ok(())
}

/// Runs one git process for an exec channel. Stdin is fed from `stdin_rx`, stdout goes to
/// channel data and stderr to extended data 1. Exit status, EOF and close are sent last.
async fn run_child(
    handle: Handle,
    channel: ChannelId,
    mut child: Child,
    stdin_rx: mpsc::Receiver<Bytes>,
) {
    // Raised by a pump when the client has gone away, so the child is killed.
    let abort = Arc::new(Notify::new());

    let feeder = tokio::spawn(feed_stdin(child.stdin.take(), stdin_rx));
    let stdout_task = tokio::spawn(pump(
        child.stdout.take(),
        handle.clone(),
        channel,
        None,
        Arc::clone(&abort),
    ));
    let stderr_task = tokio::spawn(pump(
        child.stderr.take(),
        handle.clone(),
        channel,
        Some(1),
        Arc::clone(&abort),
    ));

    let status = tokio::select! {
        status = child.wait() => status,
        () = abort.notified() => {
            let _ = child.start_kill();
            child.wait().await
        }
    };

    // Wait for the pumps so every byte of output is queued before exit status and EOF.
    let _ = stdout_task.await;
    let _ = stderr_task.await;
    feeder.abort();

    let _ = handle.exit_status_request(channel, exit_code(status)).await;
    let _ = handle.eof(channel).await;
    let _ = handle.close(channel).await;
}

async fn feed_stdin(stdin: Option<ChildStdin>, mut rx: mpsc::Receiver<Bytes>) {
    let Some(mut stdin) = stdin else {
        return;
    };
    while let Some(chunk) = rx.recv().await {
        if stdin.write_all(&chunk).await.is_err() {
            break;
        }
    }
    // Dropping stdin here closes the child's input.
    let _ = stdin.shutdown().await;
}

/// Forwards one child output stream to the client. `ext` = Some(n) sends extended data
/// type n. Sets `abort` when the client has gone away.
async fn pump<R>(
    reader: Option<R>,
    handle: Handle,
    channel: ChannelId,
    ext: Option<u32>,
    abort: Arc<Notify>,
) where
    R: AsyncRead + Unpin + Send + 'static,
{
    let Some(mut reader) = reader else {
        return;
    };
    let mut buf = vec![0u8; PUMP_CHUNK];
    loop {
        let n = match reader.read(&mut buf).await {
            Ok(0) | Err(_) => break,
            Ok(n) => n,
        };
        let chunk = Bytes::copy_from_slice(buf.get(..n).unwrap_or_default());
        let sent = match ext {
            None => handle.data(channel, chunk).await.is_ok(),
            Some(kind) => handle.extended_data(channel, kind, chunk).await.is_ok(),
        };
        if !sent {
            abort.notify_one();
            break;
        }
    }
}

fn exit_code(status: io::Result<ExitStatus>) -> u32 {
    match status {
        Ok(status) => status
            .code()
            .and_then(|c| u32::try_from(c).ok())
            .or_else(|| {
                status
                    .signal()
                    .map(|sig| 128 + u32::try_from(sig).unwrap_or(0))
            })
            .unwrap_or(1),
        Err(_) => 1,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn git_protocol_is_strict() {
        assert!(valid_git_protocol("version=2"));
        assert!(valid_git_protocol("version=1"));
        assert!(!valid_git_protocol("version=10"));
        assert!(!valid_git_protocol("version=x"));
        assert!(!valid_git_protocol("version=2\n"));
        assert!(!valid_git_protocol("version=2;x"));
        assert!(!valid_git_protocol(""));
    }

    #[test]
    fn key_types_exclude_dsa() {
        assert!(key_type_allowed(&Algorithm::Ed25519));
        assert!(key_type_allowed(&Algorithm::Rsa { hash: None }));
        assert!(!key_type_allowed(&Algorithm::Dsa));
    }
}