sugar_cli/upload/methods/
aws.rs

1use std::{fs, sync::Arc};
2
3use async_trait::async_trait;
4use ini::ini;
5use s3::{bucket::Bucket, creds::Credentials, region::Region};
6use tokio::task::JoinHandle;
7
8use crate::{
9    common::*,
10    config::*,
11    upload::{
12        assets::{AssetPair, DataType},
13        uploader::{AssetInfo, ParallelUploader, Prepare},
14    },
15};
16
17// Maximum number of times to retry each individual upload.
18const MAX_RETRY: u8 = 3;
19
20pub struct AWSMethod {
21    pub bucket: Arc<Bucket>,
22    pub directory: String,
23    pub domain: String,
24}
25
26impl AWSMethod {
27    pub async fn new(config_data: &ConfigData) -> Result<Self> {
28        let profile = &config_data
29            .aws_config
30            .as_ref()
31            .ok_or_else(|| anyhow!("AWS values not specified in config file!"))?
32            .profile;
33
34        let credentials = Credentials::from_profile(Some(profile))?;
35        let region = AWSMethod::load_region(config_data)?;
36
37        if let Some(config) = &config_data.aws_config {
38            let domain = if let Some(domain) = &config.domain {
39                match url::Url::parse(domain) {
40                    Ok(url) => url.to_string(),
41                    Err(error) => {
42                        return Err(anyhow!("Malformed domain URL ({})", error.to_string()))
43                    }
44                }
45            } else {
46                format!("https://{}.s3.amazonaws.com", &config.bucket)
47            };
48
49            Ok(Self {
50                bucket: Arc::new(Bucket::new(&config.bucket, region, credentials)?),
51                directory: config.directory.clone(),
52                domain,
53            })
54        } else {
55            Err(anyhow!("Missing AwsConfig 'bucket' value in config file."))
56        }
57    }
58
59    fn load_region(config_data: &ConfigData) -> Result<Region> {
60        let home_dir = dirs::home_dir().expect("Couldn't find home dir.");
61        let credentials = home_dir.join(Path::new(".aws/credentials"));
62        let configuration = ini!(credentials
63            .to_str()
64            .ok_or_else(|| anyhow!("Failed to load AWS credentials"))?);
65
66        let profile = &config_data
67            .aws_config
68            .as_ref()
69            .ok_or_else(|| anyhow!("AWS values not specified in config file!"))?
70            .profile;
71
72        let region = &configuration
73            .get(profile)
74            .ok_or_else(|| anyhow!("Profile not found in AWS credentials file!"))?
75            .get("region")
76            .ok_or_else(|| anyhow!("Region not found in AWS credentials file!"))?
77            .as_ref()
78            .ok_or_else(|| anyhow!("Region not found in AWS credentials file!"))?
79            .to_string();
80
81        Ok(region.parse()?)
82    }
83
84    async fn send(
85        bucket: Arc<Bucket>,
86        directory: String,
87        domain: String,
88        asset_info: AssetInfo,
89    ) -> Result<(String, String)> {
90        let data = match asset_info.data_type {
91            DataType::Image => fs::read(&asset_info.content)?,
92            DataType::Metadata => asset_info.content.into_bytes(),
93            DataType::Animation => fs::read(&asset_info.content)?,
94        };
95
96        // Take care of any spaces in the directory path.
97        let directory = directory.replace(' ', "_");
98
99        let path = Path::new(&directory).join(&asset_info.name);
100        let path_str = path
101            .to_str()
102            .ok_or_else(|| anyhow!("Failed to convert S3 bucket directory path to string."))?;
103
104        let mut retry = MAX_RETRY;
105        // send data to AWS S3 with a simple retry logic (mitigates dns lookup errors)
106        loop {
107            match bucket
108                .put_object_with_content_type(path_str, &data, &asset_info.content_type)
109                .await
110            {
111                Ok((_, code)) => match code {
112                    200 => {
113                        break;
114                    }
115                    _ => {
116                        return Err(anyhow!(
117                            "Failed to upload {} to S3 with Http Code: {code}",
118                            asset_info.name
119                        ));
120                    }
121                },
122                Err(error) => {
123                    if retry == 0 {
124                        return Err(error.into());
125                    }
126                    // we try one more time before reporting the error
127                    retry -= 1;
128                }
129            }
130        }
131
132        let link = url::Url::parse(&domain)?.join(path_str)?;
133
134        Ok((asset_info.asset_id, link.to_string()))
135    }
136}
137
138#[async_trait]
139impl Prepare for AWSMethod {
140    async fn prepare(
141        &self,
142        _sugar_config: &SugarConfig,
143        _asset_pairs: &HashMap<isize, AssetPair>,
144        _asset_indices: Vec<(DataType, &[isize])>,
145    ) -> Result<()> {
146        // nothing to do here
147        Ok(())
148    }
149}
150
151#[async_trait]
152impl ParallelUploader for AWSMethod {
153    fn upload_asset(&self, asset_info: AssetInfo) -> JoinHandle<Result<(String, String)>> {
154        let bucket = self.bucket.clone();
155        let directory = self.directory.clone();
156        let domain = self.domain.clone();
157
158        tokio::spawn(async move { AWSMethod::send(bucket, directory, domain, asset_info).await })
159    }
160}