Download raw images for flasher type device types

Change-type: minor Signed-off-by: Rahul Thakoor <[email protected]>
balena-os · May 3, 2024 · ff1a530 · ff1a530
1 parent 68880e3
commit ff1a530
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 30 deletions.
diff --git a/src/stage1/api_calls.rs b/src/stage1/api_calls.rs
@@ -6,7 +6,22 @@ use reqwest::{blocking::Client, header};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 
-use crate::common::{Error, ErrorKind, Result, ToError};
+use crate::{
+    common::{Error, ErrorKind, Result, ToError},
+    stage1::defs::{
+        DEV_TYPE_BBB, DEV_TYPE_BBG, DEV_TYPE_GEN_AMD64, DEV_TYPE_GEN_X86_64, DEV_TYPE_INTEL_NUC,
+        DEV_TYPE_JETSON_XAVIER,
+    },
+};
+
+pub const FLASHER_DEVICES: [&str; 6] = [
+    DEV_TYPE_INTEL_NUC,
+    DEV_TYPE_GEN_AMD64,
+    DEV_TYPE_GEN_X86_64,
+    DEV_TYPE_BBG,
+    DEV_TYPE_BBB,
+    DEV_TYPE_JETSON_XAVIER,
+];
 
 const OS_VERSION_URL_ENDPOINT: &str = "/v6/release";
 
@@ -23,6 +38,8 @@ struct ImageRequestData {
     version: String,
     #[serde(rename = "fileType")]
     file_type: String,
+    #[serde(rename = "imageType")]
+    image_type: Option<String>,
 }
 /// Structs corresponding to API response for endpoint /v6/releases
 #[derive(Serialize, Deserialize, Debug)]
@@ -121,10 +138,21 @@ pub(crate) fn get_os_image(
 
     let request_url = format!("{}{}", api_endpoint, OS_IMG_URL);
 
-    let post_data = ImageRequestData {
-        device_type: String::from(device),
-        version: String::from(version),
-        file_type: String::from(".gz"),
+    let post_data = if FLASHER_DEVICES.contains(&device) {
+        debug!("Downloading raw image for device type {device}");
+        ImageRequestData {
+            device_type: String::from(device),
+            version: String::from(version),
+            file_type: String::from(".gz"),
+            image_type: Some(String::from("raw")),
+        }
+    } else {
+        ImageRequestData {
+            device_type: String::from(device),
+            version: String::from(version),
+            file_type: String::from(".gz"),
+            image_type: None,
+        }
     };
 
     debug!("get_os_image: request_url: '{}'", request_url);

diff --git a/src/stage1/image_retrieval.rs b/src/stage1/image_retrieval.rs
@@ -10,7 +10,6 @@ use crate::{
     common::{
         defs::NIX_NONE,
         disk_util::{Disk, PartitionIterator, PartitionReader},
-        is_admin,
         loop_device::LoopDevice,
         path_append,
         stream_progress::StreamProgress,
@@ -30,6 +29,7 @@ use crate::{
 use flate2::{Compression, GzBuilder};
 use nix::mount::{mount, umount, MsFlags};
 
+#[allow(dead_code)]
 pub const FLASHER_DEVICES: [&str; 5] = [
     DEV_TYPE_INTEL_NUC,
     DEV_TYPE_GEN_X86_64,
@@ -148,6 +148,7 @@ fn determine_version(ver_str: &str, versions: &Versions) -> Result<Version> {
     }
 }
 
+#[allow(dead_code)]
 pub(crate) fn extract_image<P1: AsRef<Path>, P2: AsRef<Path>>(
     stream: Box<dyn Read>,
     image_file_name: P1,
@@ -351,30 +352,22 @@ pub(crate) fn download_image(
         format!("balena-cloud-{}-{}.img.gz", device_type, version),
     );
 
-    if FLASHER_DEVICES.contains(&device_type) {
-        if !is_admin()? {
-            error!("please run this program as root");
-            return Err(Error::displayed());
-        }
-        extract_image(stream, &img_file_name, device_type, work_dir)?;
-    } else {
-        debug!("Downloading file '{}'", img_file_name.display());
-        let mut file = File::create(&img_file_name).upstream_with_context(&format!(
-            "Failed to create file: '{}'",
-            img_file_name.display()
-        ))?;
-
-        // TODO: show progress
-        let mut progress = StreamProgress::new(stream, 10, Level::Info, None);
-        copy(&mut progress, &mut file).upstream_with_context(&format!(
-            "Failed to write downloaded data to '{}'",
-            img_file_name.display()
-        ))?;
-        info!(
-            "The balena OS image was successfully written to '{}'",
-            img_file_name.display()
-        );
-    }
+    debug!("Downloading file '{}'", img_file_name.display());
+    let mut file = File::create(&img_file_name).upstream_with_context(&format!(
+        "Failed to create file: '{}'",
+        img_file_name.display()
+    ))?;
+
+    // TODO: show progress
+    let mut progress = StreamProgress::new(stream, 10, Level::Info, None);
+    copy(&mut progress, &mut file).upstream_with_context(&format!(
+        "Failed to write downloaded data to '{}'",
+        img_file_name.display()
+    ))?;
+    info!(
+        "The balena OS image was successfully written to '{}'",
+        img_file_name.display()
+    );
 
     Ok(img_file_name)
 }