Skip to content

Commit

Permalink
refactor (to squash)
Browse files Browse the repository at this point in the history
  • Loading branch information
@aetimmes
aetimmes committed Dec 31, 2024
1 parent b165389 commit 6b0507a
Showing 1 changed file with 40 additions and 112 deletions.
152 changes: 40 additions & 112 deletions sources/api/corndog/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,51 +110,28 @@ where
}

/// Validates sysctl key/value pairs for systemd-sysctl compatibility
fn validate_sysctl_entry<K>(key: K, value: &str) -> Result<()>
fn validate_sysctl_entry<K>(key: K, value: &str) -> Result<()>
where
K: AsRef<str>,
{
let key = key.as_ref();

// Basic validation of key format
if !key.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '*') {
return error::InvalidSysctlKeySnafu { key: key.to_string() }.fail();
}

// Ensure value doesn't contain newlines which could break the config file format
if value.contains('\n') {
return error::InvalidSysctlValueSnafu {
return error::InvalidSysctlValueSnafu {
key: key.to_string(),
value: value.to_string()
value: value.to_string()
}.fail();
}

Ok(())
}

/// Creates a backup of the existing sysctl config file
fn backup_sysctl_config() -> Result<Option<NamedTempFile>> {
const SYSCTL_CONFIG: &str = "/etc/sysctl.d/95-corndog.conf";

if !Path::new(SYSCTL_CONFIG).exists() {
return Ok(None);
}

let mut temp_file = NamedTempFile::new().context(error::TempFileSnafu)?;
let contents = fs::read_to_string(SYSCTL_CONFIG).context(error::ReadSysctlConfigSnafu)?;
temp_file.write_all(contents.as_bytes()).context(error::WriteTempFileSnafu)?;

Ok(Some(temp_file))
}

/// Restores sysctl config from backup
fn restore_sysctl_config(backup: NamedTempFile) -> Result<()> {
const SYSCTL_CONFIG: &str = "/etc/sysctl.d/95-corndog.conf";

fs::copy(backup.path(), SYSCTL_CONFIG).context(error::RestoreBackupSnafu)?;
Ok(())
}

/// Restarts systemd-sysctl service
fn restart_sysctl() -> Result<()> {
let status = Command::new("systemctl")
Expand All @@ -165,7 +142,7 @@ fn restart_sysctl() -> Result<()> {
if !status.success() {
return error::SystemctlFailedSnafu.fail();
}

Ok(())
}

Expand All @@ -182,28 +159,30 @@ where
validate_sysctl_entry(key, value)?;
}

// Create backup of existing config
let backup = backup_sysctl_config()?;

// Write new config
// Create config string
let mut config = String::new();
for (key, value) in sysctls {
config.push_str(&format!("{}={}\n", key.as_ref(), value));
}

fs::write(SYSCTL_CONFIG, config).context(error::WriteSysctlConfigSnafu)?;

// Attempt to restart systemd-sysctl
if let Err(e) = restart_sysctl() {
// On failure, restore backup if we had one
if let Some(backup) = backup {
if let Err(restore_err) = restore_sysctl_config(backup) {
error!("Failed to restore sysctl config backup: {}", restore_err);
}
// Try to restart again with old config
restart_sysctl()?;
}
return Err(e);

// Append to file, creating it if it doesn't exist
let mut file = fs::OpenOptions::new()
.create(true)
.append(true)
.open(SYSCTL_CONFIG)
.context(error::OpenSysctlConfigSnafu)?;

file.write_all(config.as_bytes())
.context(error::WriteSysctlConfigSnafu)?;

// Restart systemd-sysctl to apply changes
let status = Command::new("systemctl")
.args(["restart", "systemd-sysctl"])
.status()
.context(error::SystemctlSnafu)?;

if !status.success() {
return error::SystemctlFailedSnafu.fail();
}

Ok(())
Expand Down Expand Up @@ -432,38 +411,23 @@ mod error {
value: String,
},

#[snafu(display("Failed to create temporary file: {}", source))]
TempFile {
#[snafu(display("Failed to execute systemctl: {}", source))]
Systemctl {
source: io::Error,
},

#[snafu(display("Failed to read existing sysctl config: {}", source))]
ReadSysctlConfig {
source: io::Error,
},
#[snafu(display("systemd-sysctl restart failed"))]
SystemctlFailed,

#[snafu(display("Failed to write to temporary file: {}", source))]
WriteTempFile {
#[snafu(display("Failed to open sysctl config for writing: {}", source))]
OpenSysctlConfig {
source: io::Error,
},

#[snafu(display("Failed to write sysctl config: {}", source))]
#[snafu(display("Failed to write to sysctl config: {}", source))]
WriteSysctlConfig {
source: io::Error,
},

#[snafu(display("Failed to restore sysctl config backup: {}", source))]
RestoreBackup {
source: io::Error,
},

#[snafu(display("Failed to execute systemctl: {}", source))]
Systemctl {
source: io::Error,
},

#[snafu(display("systemd-sysctl restart failed"))]
SystemctlFailed,
}
}
type Result<T> = std::result::Result<T, error::Error>;
Expand Down Expand Up @@ -538,55 +502,19 @@ mod test {
assert!(validate_sysctl_entry("kernel.shmmax", "value\nwith\nnewlines").is_err());
}

#[test]
fn test_sysctl_config_backup_restore() -> Result<()> {
let temp_dir = TempDir::new().unwrap();
let config_path = temp_dir.path().join("95-corndog.conf");
let test_content = "net.ipv4.ip_forward=1\nkernel.shmmax=68719476736\n";

// Create initial config
fs::write(&config_path, test_content).unwrap();

// Test backup creation
let backup = {
let old_path = env::var("SYSCTL_CONFIG").ok();
env::set_var("SYSCTL_CONFIG", config_path.to_str().unwrap());
let result = backup_sysctl_config();
if let Some(path) = old_path {
env::set_var("SYSCTL_CONFIG", path);
}
result?
};

assert!(backup.is_some());

// Verify backup contents
let backup_content = fs::read_to_string(backup.as_ref().unwrap().path()).unwrap();
assert_eq!(backup_content, test_content);

// Test restore
fs::write(&config_path, "changed content").unwrap();
restore_sysctl_config(backup.unwrap())?;

let restored_content = fs::read_to_string(&config_path).unwrap();
assert_eq!(restored_content, test_content);

Ok(())
}

#[test]
fn test_set_sysctls() -> Result<()> {
let temp_dir = TempDir::new().unwrap();
let config_path = temp_dir.path().join("95-corndog.conf");

let mut sysctls = HashMap::new();
sysctls.insert("net.ipv4.ip_forward", "1".to_string());
sysctls.insert("kernel.shmmax", "68719476736".to_string());

// Override config path for testing
let old_path = env::var("SYSCTL_CONFIG").ok();
env::set_var("SYSCTL_CONFIG", config_path.to_str().unwrap());

// Mock systemctl command for testing
let result = {
struct MockCommand;
Expand All @@ -597,27 +525,27 @@ mod test {
Ok(std::process::ExitStatus::from_raw(0))
}
}

let original_command = Command::new;
Command::new = |_| MockCommand;
let result = set_sysctls(sysctls);
Command::new = original_command;
result
};

// Restore original config path
if let Some(path) = old_path {
env::set_var("SYSCTL_CONFIG", path);
}

// Check result
assert!(result.is_ok());

// Verify written config
let content = fs::read_to_string(&config_path).unwrap();
assert!(content.contains("net.ipv4.ip_forward=1\n"));
assert!(content.contains("kernel.shmmax=68719476736\n"));

Ok(())
}
}

0 comments on commit 6b0507a

Please sign in to comment.