pg-api 0.3.8

//! Server-side PostgreSQL custom-format backup operations.
//!
//! The Nexus process is the only component that invokes `pg_dump` and
//! `pg_restore`.  Callers receive bytes over the authenticated HTTP API and
//! never provide a shell command or a filesystem path.

use crate::models::{DatabaseAccess, PostgresInstance};
use std::path::{Path, PathBuf};
use tokio::process::Command;
use uuid::Uuid;

pub const DEFAULT_MAX_BACKUP_BYTES: u64 = 1024 * 1024 * 1024;

fn max_backup_bytes() -> u64 {
    std::env::var("PG_API_MAX_BACKUP_BYTES")
        .ok()
        .and_then(|value| value.parse().ok())
        .filter(|value: &u64| *value > 0)
        .unwrap_or(DEFAULT_MAX_BACKUP_BYTES)
}

fn temporary_path() -> PathBuf {
    std::env::temp_dir().join(format!("pg-api-backup-{}.dump", Uuid::new_v4()))
}

async fn create_private_file(path: &Path) -> Result<(), String> {
    let file = tokio::fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(path)
        .await
        .map_err(|error| format!("could not create temporary backup: {error}"))?;
    drop(file);
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        tokio::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
            .await
            .map_err(|error| format!("could not secure temporary backup: {error}"))?;
    }
    Ok(())
}

fn connection_options(instance: &PostgresInstance, access: &DatabaseAccess) -> Vec<String> {
    vec![
        "--host".into(),
        instance.host.clone(),
        "--port".into(),
        instance.port.to_string(),
        "--username".into(),
        access.username.clone(),
    ]
}

fn command_error(command: &str, output: &std::process::Output) -> String {
    let stderr = String::from_utf8_lossy(&output.stderr);
    tracing::error!(command, stderr = %stderr, "backup command failed");
    format!("{command} failed")
}

/// Create and validate a custom-format dump for an authorized database.
pub async fn export_custom(
    instance: &PostgresInstance,
    access: &DatabaseAccess,
    exclude_tables: &[String],
) -> Result<Vec<u8>, String> {
    for table in exclude_tables {
        validate_exclude_table(table)?;
    }
    let path = temporary_path();
    create_private_file(&path).await?;
    let result = async {
        let mut command = Command::new("pg_dump");
        command
            .args(["--format=custom", "--no-owner", "--no-privileges", "--file"])
            .arg(&path)
            .args(
                exclude_tables
                    .iter()
                    .flat_map(|table| ["--exclude-table".to_owned(), table.clone()]),
            )
            .args(connection_options(instance, access))
            .arg(&access.database)
            .env("PGPASSWORD", &access.password);
        let output = command
            .output()
            .await
            .map_err(|error| format!("could not execute pg_dump: {error}"))?;
        if !output.status.success() {
            return Err(command_error("pg_dump", &output));
        }
        let metadata = tokio::fs::metadata(&path)
            .await
            .map_err(|error| format!("could not inspect backup: {error}"))?;
        if metadata.len() > max_backup_bytes() {
            return Err("backup exceeds configured size limit".into());
        }
        let mut list = Command::new("pg_restore");
        list.arg("--list").arg(&path);
        let list_output = list
            .output()
            .await
            .map_err(|error| format!("could not execute pg_restore: {error}"))?;
        if !list_output.status.success() {
            return Err(command_error("pg_restore --list", &list_output));
        }
        tokio::fs::read(&path)
            .await
            .map_err(|error| format!("could not read backup: {error}"))
    }
    .await;
    let _ = tokio::fs::remove_file(&path).await;
    result
}

fn validate_exclude_table(table: &str) -> Result<(), String> {
    if table.is_empty()
        || !table.starts_with("public.")
        || !table
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'.')
    {
        return Err("exclude_table must be a simple public schema identifier".into());
    }
    Ok(())
}

/// Restore a custom-format dump into an authorized database.
pub async fn import_custom(
    instance: &PostgresInstance,
    access: &DatabaseAccess,
    bytes: &[u8],
    clean: bool,
) -> Result<(), String> {
    if bytes.is_empty() {
        return Err("backup payload is empty".into());
    }
    if bytes.len() as u64 > max_backup_bytes() {
        return Err("backup exceeds configured size limit".into());
    }
    let path = temporary_path();
    create_private_file(&path).await?;
    let result = async {
        tokio::fs::write(&path, bytes)
            .await
            .map_err(|error| format!("could not stage backup: {error}"))?;
        let mut list = Command::new("pg_restore");
        let list_output = list
            .arg("--list")
            .arg(&path)
            .output()
            .await
            .map_err(|error| format!("could not execute pg_restore: {error}"))?;
        if !list_output.status.success() {
            return Err(command_error("pg_restore --list", &list_output));
        }
        let mut command = Command::new("pg_restore");
        command
            .args(["--exit-on-error", "--no-owner", "--no-privileges"])
            .args(if clean {
                vec!["--if-exists", "--clean"]
            } else {
                vec![]
            })
            .args(["--dbname"])
            .arg(&access.database)
            .args(connection_options(instance, access))
            .arg(&path)
            .env("PGPASSWORD", &access.password);
        let output = command
            .output()
            .await
            .map_err(|error| format!("could not execute pg_restore: {error}"))?;
        if !output.status.success() {
            return Err(command_error("pg_restore", &output));
        }
        Ok(())
    }
    .await;
    let _ = tokio::fs::remove_file(&path).await;
    result
}

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

    #[test]
    fn temporary_paths_are_unique_and_private_in_shape() {
        let first = temporary_path();
        let second = temporary_path();
        assert_ne!(first, second);
        assert!(
            first
                .file_name()
                .unwrap()
                .to_string_lossy()
                .starts_with("pg-api-backup-")
        );
    }

    #[test]
    fn exclude_tables_are_limited_to_simple_public_identifiers() {
        assert!(validate_exclude_table("public.bak_old").is_ok());
        assert!(validate_exclude_table("public.bak_old;DROP TABLE users").is_err());
        assert!(validate_exclude_table("pg_catalog.pg_authid").is_err());
    }
}