shadow_drive_cli/
utils.rs

1use anyhow::anyhow;
2use byte_unit::Byte;
3use chrono::DateTime;
4use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
5use reqwest::Response;
6use shadow_drive_sdk::constants::SHDW_DRIVE_OBJECT_PREFIX;
7use shadow_drive_sdk::error::{Error, FileError};
8use shadow_drive_sdk::models::{ShadowDriveResult, ShadowFile};
9use shadow_drive_sdk::ShadowDriveClient;
10use shadow_rpc_auth::HttpSenderWithHeaders;
11use solana_client::nonblocking;
12use solana_client::rpc_client::RpcClient;
13use solana_sdk::pubkey::Pubkey;
14use solana_sdk::signature::{Signature, Signer, SignerError};
15use std::io::stdin;
16use std::path::PathBuf;
17use std::str::FromStr;
18
19/// Maximum amount of files to batch into a single [store_files] request.
20pub const FILE_UPLOAD_BATCH_SIZE: usize = 5;
21
22/// Clap value parser for base58 string representations of [Pubkey].
23pub fn pubkey_arg(pubkey: &str) -> anyhow::Result<Pubkey> {
24    Pubkey::from_str(pubkey).map_err(|e| anyhow!("invalid pubkey: {}", e.to_string()))
25}
26
27/// To get around using a [Box<dyn Signer>] with [ShadowDriveClient].
28///
29/// TODO: cleanup if not necessary
30pub struct WrappedSigner(Box<dyn Signer>);
31
32impl WrappedSigner {
33    pub fn new(signer: Box<dyn Signer>) -> Self {
34        Self(signer)
35    }
36}
37
38impl Signer for WrappedSigner {
39    fn try_pubkey(&self) -> Result<Pubkey, SignerError> {
40        Ok(self.0.pubkey())
41    }
42
43    fn try_sign_message(&self, message: &[u8]) -> Result<Signature, SignerError> {
44        self.0.try_sign_message(message)
45    }
46
47    fn is_interactive(&self) -> bool {
48        self.0.is_interactive()
49    }
50}
51
52/// Further diagnostic printing wherever possible.
53pub fn process_shadow_api_response<T>(response: ShadowDriveResult<T>) -> anyhow::Result<T> {
54    match response {
55        Ok(response) => Ok(response),
56        Err(err) => match err {
57            Error::ShadowDriveServerError { status, message } => {
58                let err = format!(
59                    "Shadow Drive Server Error {}: {:#?}",
60                    status,
61                    message.to_string()
62                );
63                println!("{}", err);
64                Err(anyhow!("{}", err))
65            }
66            Error::FileSystemError(err) => {
67                let err = format!("Filesystem Error: {:#?}", err.to_string());
68                println!("{}", err);
69                Err(anyhow!("{}", err))
70            }
71            Error::FileValidationError(errs) => {
72                let mut err_vec = vec![];
73                for err in errs {
74                    let FileError { file, error } = err;
75                    let err = format!("File Validation Error for {}: {}", file, error);
76                    err_vec.push(err);
77                }
78                println!("{:#?}", err_vec);
79                Err(anyhow!("{:#?}", err_vec))
80            }
81            e => {
82                println!("{:#?}", e);
83                Err(anyhow!("{:#?}", e))
84            }
85        },
86    }
87}
88
89/// Generate a Shadow Drive file URL from storage account and filename.
90pub fn storage_object_url(storage_account: &Pubkey, file: &str) -> String {
91    format!(
92        "{}/{}/{}",
93        SHDW_DRIVE_OBJECT_PREFIX,
94        storage_account.to_string(),
95        file
96    )
97}
98
99/// Returns false when "Content-Type" header is not "text/plain".
100fn is_text_response(headers: &HeaderMap) -> anyhow::Result<bool> {
101    let content_type = headers
102        .get("content-type")
103        .and_then(|s| Some(s.to_str()))
104        .transpose()?;
105    Ok(content_type == Some("text/plain"))
106}
107
108/// Check with a HEAD that the URL exists and is a "text/plain" file.
109/// If so, return the response of a GET request.
110pub async fn get_text(url: &String) -> anyhow::Result<Response> {
111    let http_client = reqwest::Client::new();
112    let head_resp = http_client.head(url).send().await?;
113    if !is_text_response(head_resp.headers())? {
114        return Err(anyhow!("Not a text file at url {}", url));
115    }
116    Ok(http_client.get(url).send().await?)
117}
118
119#[derive(Debug)]
120pub struct FileMetadata {
121    pub timestamp: i64,
122    pub content_type: String,
123    pub last_modified: i64,
124    pub etag: String,
125    pub storage_account: String,
126    pub storage_owner: String,
127}
128
129impl FileMetadata {
130    pub fn from_headers(h: &HeaderMap) -> anyhow::Result<Self> {
131        let getter = |key| {
132            Ok::<_, anyhow::Error>(
133                h.get(key)
134                    .ok_or(anyhow!("Missing file metadata header: {}", key))?
135                    .to_str()?
136                    .to_string(),
137            )
138        };
139        let parse_timestamp = |key| {
140            let timestamp = getter(key)?;
141            let timestamp = DateTime::parse_from_rfc2822(&timestamp)?;
142            Ok::<_, anyhow::Error>(timestamp.timestamp())
143        };
144        let timestamp = parse_timestamp("date")?;
145        let last_modified = parse_timestamp("last-modified")?;
146        Ok(Self {
147            timestamp,
148            content_type: getter("content-type")?,
149            last_modified,
150            etag: getter("etag")?,
151            storage_account: getter("x-amz-meta-owner-account-pubkey")?,
152            storage_owner: getter("x-amz-meta-storage-account-pubkey")?,
153        })
154    }
155}
156
157/// Pulls "last-modified" from [HeaderMap], unaltered.
158pub fn last_modified(headers: &HeaderMap) -> anyhow::Result<String> {
159    Ok(headers
160        .get("last-modified")
161        .ok_or(anyhow!("'last modified' header not found"))?
162        .to_str()?
163        .to_string())
164}
165
166/// Convert a file size string to [Byte] object with the denoted size.
167pub fn parse_filesize(size: &str) -> anyhow::Result<Byte> {
168    Byte::from_str(size).map_err(|e| {
169        anyhow!(
170            "invalid filesize, \
171        expected a number followed by KB, MB, GB:\n{}",
172            e.to_string()
173        )
174    })
175}
176
177/// Confirm from the user that they definitely want some irreversible
178/// operation to occur.
179pub fn wait_for_user_confirmation(skip: bool) -> anyhow::Result<()> {
180    if skip {
181        return Ok(());
182    }
183    println!("Press ENTER to continue, or CTRL+C to abort");
184    let mut proceed = String::new();
185    stdin().read_line(&mut proceed)?;
186    Ok(())
187}
188
189/// We either create an authenticated client with default auth headers,
190/// or else we simply use the [RpcClient] provided by the normal
191/// [ShadowDriveClient] constructor.
192pub fn shadow_client_factory<T: Signer>(
193    signer: T,
194    url: &str,
195    auth: Option<String>,
196) -> ShadowDriveClient<T> {
197    if let Some(auth) = auth {
198        let mut headers = HeaderMap::new();
199        headers.append(
200            HeaderName::from_str("Authorization").unwrap(),
201            HeaderValue::from_str(&format!("Bearer {}", auth)).unwrap(),
202        );
203        let rpc_client = nonblocking::rpc_client::RpcClient::new_sender(
204            HttpSenderWithHeaders::new(url, Some(headers.clone())),
205            Default::default(),
206        );
207        let client = RpcClient::new_sender(
208            HttpSenderWithHeaders::new(url, Some(headers)),
209            Default::default(),
210        );
211        let balance = client.get_balance(&signer.pubkey());
212        match balance {
213            Ok(balance) => {
214                println!("{}: {} lamports", signer.pubkey().to_string(), balance);
215            }
216            Err(e) => {
217                println!("Failed to fetch balance: {:?}", e);
218            }
219        }
220        ShadowDriveClient::new_with_rpc(signer, rpc_client)
221    } else {
222        ShadowDriveClient::new(signer, url)
223    }
224}
225
226// TODO Maybe make this a result type.
227/// Factory function for a [ShadowFile], where we just use the path's
228/// basename. Panics if `path.file_name()` returns None.
229pub fn shadow_file_with_basename(path: &PathBuf) -> ShadowFile {
230    let basename = {
231        path.file_name()
232            .and_then(|s| s.to_str())
233            .unwrap()
234            .to_string()
235    };
236    ShadowFile::file(basename, path.clone())
237}