sgx-daemon/src/utils.rs
Noor 974906804e
Refactor host port logic to port range
HostConfig to use public_port_range and update port mapping logic
2025-03-28 01:47:04 +05:30

102 lines
3.1 KiB
Rust

use anyhow::{anyhow, Result};
use flate2::read::GzDecoder;
use reqwest::Client;
use std::io::BufReader;
use std::path::Path;
use tar::Archive;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::{fs, fs::File};
use crate::global::{PACKAGE_ARCHIVE_DIR_PATH, PACKAGE_ARCHIVE_POSTFIX, PACKAGE_DIR_PATH};
pub async fn handle_package(
package_url: String,
container_uuid: String,
delete_archive: bool,
) -> Result<String> {
let dir_path = Path::new(PACKAGE_ARCHIVE_DIR_PATH);
fs::create_dir_all(dir_path).await?;
let file_name = format!("{container_uuid}{PACKAGE_ARCHIVE_POSTFIX}",);
let file_path = dir_path.join(file_name);
if let Err(e) = download_file(&package_url, &file_path).await {
println!("Error downloading file: {:?}", e);
return Err(anyhow!("Error downloading file"));
}
let downloaded_file = std::fs::File::open(&file_path)?;
let mut reader = BufReader::new(downloaded_file);
let mut archive = Archive::new(GzDecoder::new(&mut reader));
if let Err(er) = archive.entries() {
dbg!(&er);
return Err(anyhow!("Error: file not an archive: {er:?}"));
};
let unarchive_dir = format!("{PACKAGE_DIR_PATH}/{container_uuid}");
fs::create_dir_all(Path::new(&unarchive_dir)).await?;
let top_level_directory = get_top_level_dir(file_path.to_string_lossy().to_string())
.ok_or(anyhow!("Error: failed get toplevel directory"))?;
archive.unpack(&unarchive_dir)?;
if delete_archive {
let _ = fs::remove_file(file_path).await;
}
Ok(format!("{unarchive_dir}/{top_level_directory}"))
}
fn get_top_level_dir(file_path: String) -> Option<String> {
let file = std::fs::File::open(file_path).ok()?;
let reader = BufReader::new(file);
let mut archive = Archive::new(GzDecoder::new(reader));
archive.entries().ok()?.flatten().find_map(|entry| {
entry
.path()
.ok()?
.components()
.next()?
.as_os_str()
.to_str()
.map(String::from)
})
}
pub async fn download_file(url: &str, file_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
let client = Client::new();
let response = client.get(url).send().await?;
let data = response.bytes().await?;
let mut file = File::create(file_path).await?;
file.write_all(&data).await?;
Ok(())
}
pub async fn is_port_available(port: u16) -> bool {
TcpListener::bind(&format!("127.0.0.1:{}", port))
.await
.is_ok()
}
pub async fn cleanup_enclave_disk_and_package(container_uuid: String) -> Result<()> {
let enclave_disk_dir_str = format!("{PACKAGE_DIR_PATH}/{container_uuid}");
let enclave_disk_path = Path::new(&enclave_disk_dir_str);
if enclave_disk_path.exists() {
std::fs::remove_dir_all(enclave_disk_path)?;
}
let enclave_archive_dir_str =
format!("{PACKAGE_ARCHIVE_DIR_PATH}/{container_uuid}{PACKAGE_ARCHIVE_POSTFIX}");
let enclave_archive_path = Path::new(&enclave_archive_dir_str);
if enclave_archive_path.exists() {
let _ = std::fs::remove_file(enclave_archive_path);
}
Ok(())
}