quick_flash/
storage.rs

1use crate::credentials::{Credentials, StorageType};
2use anyhow::{self, Context};
3use chrono::{DateTime, Utc};
4use s3::{self, serde_types::Object};
5use serde::{Deserialize, Serialize};
6use std::path::{Path, PathBuf};
7
8#[derive(Serialize, Deserialize, Clone, PartialEq, Debug)]
9pub struct FirmwareMetadata {
10    pub name: String,
11    pub version: String,
12    pub last_modified: i64,
13}
14
15#[derive(Serialize, Deserialize, Debug)]
16pub struct Firmware {
17    pub name: String,
18    pub version: String,
19    pub chip: String,
20    pub path: PathBuf,
21}
22
23#[derive(Serialize, Deserialize)]
24struct Manifest {
25    chip: String,
26}
27
28pub struct Storage {
29    bucket: Box<s3::Bucket>,
30}
31
32impl Storage {
33    pub fn new(creds: &Credentials) -> anyhow::Result<Self> {
34        let region = match creds.storage_type {
35            StorageType::R2 => s3::Region::R2 {
36                account_id: creds.storage_account_id.clone(),
37            },
38        };
39
40        let bucket = s3::Bucket::new(
41            &creds.storage_name,
42            region,
43            s3::creds::Credentials {
44                access_key: Some(creds.storage_access_key.clone()),
45                secret_key: Some(creds.storage_secret_key.clone()),
46                security_token: None,
47                session_token: None,
48                expiration: None,
49            },
50        )?;
51        Ok(Storage { bucket })
52    }
53
54    pub fn is_available(&self) -> anyhow::Result<()> {
55        match self.bucket.exists() {
56            Ok(true) => Ok(()),
57            Ok(false) => anyhow::bail!("Bucket does not exist"),
58            Err(e) => Err(e.into()),
59        }
60    }
61
62    fn list_common_prefixes(&self, prefix: String) -> anyhow::Result<Vec<String>> {
63        let response = self
64            .bucket
65            .list(prefix, Some("/".to_string()))?
66            .first()
67            .cloned()
68            .context("No response data received")?;
69
70        Ok(response
71            .common_prefixes
72            .context("No common prefixes received")?
73            .iter()
74            .map(|p| p.prefix.strip_suffix("/").unwrap_or(&p.prefix).to_owned())
75            .collect())
76    }
77
78    fn list_object_metadata(&self, prefix: String) -> anyhow::Result<Vec<FirmwareMetadata>> {
79        let response = self
80            .bucket
81            .list(prefix, None)?
82            .first()
83            .cloned()
84            .context("No response data received")?;
85
86        let filter = |o: &Object| {
87            if let Some(key) = o.key.strip_suffix("/manifest.json") {
88                let parts = key.split("/").collect::<Vec<&str>>();
89                Some(FirmwareMetadata {
90                    name: parts[0].to_owned(),
91                    version: parts[1].to_owned(),
92                    last_modified: DateTime::parse_from_rfc3339(&o.last_modified)
93                        .ok()?
94                        .with_timezone(&Utc)
95                        .timestamp(),
96                })
97            } else {
98                None
99            }
100        };
101
102        Ok(response.contents.iter().filter_map(filter).collect())
103    }
104
105    pub fn list_firmwares(&self) -> anyhow::Result<Vec<FirmwareMetadata>> {
106        let prefixes = self.list_common_prefixes("".to_string())?;
107        let mut ret = Vec::<FirmwareMetadata>::new();
108
109        for prefix in prefixes {
110            if let Some(f) = self
111                .list_object_metadata(prefix)?
112                .iter()
113                .max_by_key(|f| f.last_modified)
114            {
115                ret.push(f.clone())
116            }
117        }
118
119        Ok(ret)
120    }
121
122    pub fn list_firmware_versions(
123        &self,
124        firmware_name: &str,
125    ) -> anyhow::Result<Vec<FirmwareMetadata>> {
126        let mut firmware_name = firmware_name.to_owned();
127        firmware_name.push('/');
128        self.list_object_metadata(firmware_name.clone())
129    }
130
131    pub fn download_firmware(
132        &self,
133        name: &str,
134        version: &str,
135        cache_base: &Path,
136    ) -> anyhow::Result<Firmware> {
137        let cache_base = cache_base.to_path_buf().join(name).join(version);
138        let cache_firmware = cache_base.join("firmware.elf");
139        let cache_manifest = cache_base.join("manifest.json");
140
141        if !cache_base.exists() {
142            std::fs::create_dir_all(&cache_base)?;
143            let bucket_base = format!("{}/{}", name, version);
144            let bucket_firmware = format!("{}/firmware.elf", bucket_base);
145            let bucket_manifest = format!("{}/manifest.json", bucket_base);
146
147            eprintln!("Downloading firmware to {}...", cache_base.display());
148            let firmware = self.bucket.get_object(&bucket_firmware)?;
149            std::fs::write(&cache_firmware, firmware.bytes())?;
150            let manifest = self.bucket.get_object(&bucket_manifest)?;
151            std::fs::write(&cache_manifest, manifest.bytes())?;
152        }
153
154        let manifest: Manifest = serde_json::from_str(&std::fs::read_to_string(cache_manifest)?)?;
155
156        Ok(Firmware {
157            name: name.to_owned(),
158            version: version.to_owned(),
159            chip: manifest.chip,
160            path: cache_firmware,
161        })
162    }
163}