134 lines
4.4 KiB
Rust
134 lines
4.4 KiB
Rust
use std::os::unix::fs::PermissionsExt;
|
|
use std::path::PathBuf;
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
use tokio::io::AsyncWriteExt;
|
|
use tracing::info;
|
|
|
|
use crate::commands::run::to_file_name;
|
|
|
|
pub const INTERNAL_CONFIG_FOLDER_NAME: &str = "YANPM";
|
|
const FILE_SIZE_LIMIT: usize = 10 * 1024 * 1024; // 10MB
|
|
|
|
pub struct WriteConfigCommand {
|
|
nginx_config_dir: PathBuf,
|
|
}
|
|
|
|
impl WriteConfigCommand {
|
|
pub fn new(nginx_config_dir: PathBuf) -> Self {
|
|
Self { nginx_config_dir }
|
|
}
|
|
pub async fn write_config(
|
|
&self,
|
|
config_name: &str,
|
|
timestamp: u64,
|
|
content: &str,
|
|
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
|
let filename = to_file_name(config_name, timestamp)?;
|
|
let path = self.nginx_config_dir.clone();
|
|
// ensure main config dir exists
|
|
tokio::fs::create_dir_all(&path).await?;
|
|
info!("Writing config to {:?}", path.join(&filename));
|
|
|
|
// create YANPM subdir where fragment files live
|
|
let yanpm_dir = path.join(INTERNAL_CONFIG_FOLDER_NAME);
|
|
tokio::fs::create_dir_all(&yanpm_dir).await?;
|
|
let final_path = yanpm_dir.join(&filename);
|
|
|
|
// limit size to 10MB
|
|
if content.len() > FILE_SIZE_LIMIT {
|
|
return Err(format!(
|
|
"content exceeds {}MB size limit",
|
|
FILE_SIZE_LIMIT / (1024 * 1024)
|
|
)
|
|
.into());
|
|
}
|
|
|
|
// create a temporary filename in the same directory for atomic replace
|
|
let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos();
|
|
let tmp_filename = format!("{}.tmp.{}.{}", filename, std::process::id(), now);
|
|
// create tmp file in the same directory as final file to ensure atomic rename
|
|
let tmp_path = yanpm_dir.join(tmp_filename);
|
|
|
|
let mut file = tokio::fs::OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.truncate(true)
|
|
.open(&tmp_path)
|
|
.await?;
|
|
file.write_all(content.as_bytes()).await?;
|
|
// ensure data is flushed to disk; propagate errors
|
|
file.sync_all().await?;
|
|
|
|
// atomically move the tmp file into the YANPM dir
|
|
tokio::fs::rename(&tmp_path, &final_path).await?;
|
|
|
|
// set explicit permissions (rw-r-----)
|
|
tokio::fs::set_permissions(&final_path, std::fs::Permissions::from_mode(0o640)).await?;
|
|
info!("Config written and permissions set for {:?}", final_path);
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{INTERNAL_CONFIG_FOLDER_NAME, WriteConfigCommand};
|
|
use std::time::SystemTime;
|
|
use std::time::UNIX_EPOCH;
|
|
|
|
#[tokio::test]
|
|
async fn write_config_success_and_cleanup() {
|
|
let base = std::env::temp_dir().join(format!(
|
|
"yanpm_test_{}_{}",
|
|
std::process::id(),
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
// ensure clean
|
|
let _ = tokio::fs::remove_dir_all(&base).await;
|
|
let cmd = WriteConfigCommand::new(base.clone());
|
|
|
|
let config_name = "unittest";
|
|
let timestamp = 42u64;
|
|
let content = "hello world";
|
|
|
|
cmd.write_config(config_name, timestamp, content)
|
|
.await
|
|
.expect("write should succeed");
|
|
|
|
let filename = super::to_file_name(config_name, timestamp).unwrap();
|
|
let final_path = base.join(INTERNAL_CONFIG_FOLDER_NAME).join(&filename);
|
|
let data = tokio::fs::read_to_string(&final_path)
|
|
.await
|
|
.expect("file should exist");
|
|
assert_eq!(data, content);
|
|
|
|
// cleanup
|
|
tokio::fs::remove_dir_all(&base).await.expect("cleanup");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn write_config_size_limit() {
|
|
let base = std::env::temp_dir().join(format!(
|
|
"yanpm_test_{}_{}",
|
|
std::process::id(),
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
let _ = tokio::fs::remove_dir_all(&base).await;
|
|
let cmd = WriteConfigCommand::new(base.clone());
|
|
|
|
// exceed 10MB limit
|
|
let large = vec![b'a'; 10 * 1024 * 1024 + 1];
|
|
let large_str = String::from_utf8_lossy(&large).to_string();
|
|
|
|
let res = cmd.write_config("big", 1, &large_str).await;
|
|
assert!(res.is_err());
|
|
|
|
let _ = tokio::fs::remove_dir_all(&base).await;
|
|
}
|
|
}
|