iota_config/
object_storage_config.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{env, fs, path::PathBuf, sync::Arc, time::Duration};
6
7use anyhow::{Context, Result, anyhow, bail};
8use clap::*;
9use object_store::{ClientOptions, DynObjectStore, aws::AmazonS3Builder};
10use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
11use serde::{Deserialize, Serialize};
12use tracing::info;
13
14/// Object-store type.
15#[derive(Debug, Copy, Clone, PartialEq, Eq, Deserialize, Serialize, ValueEnum)]
16pub enum ObjectStoreType {
17    /// Local file system
18    File,
19    /// AWS S3
20    S3,
21    /// Google Cloud Store
22    GCS,
23    /// Azure Blob Store
24    Azure,
25}
26
27#[derive(Default, Debug, Clone, Deserialize, Serialize, Args)]
28#[serde(rename_all = "kebab-case")]
29pub struct ObjectStoreConfig {
30    /// Which object storage to use. If not specified, defaults to local file
31    /// system.
32    #[serde(skip_serializing_if = "Option::is_none")]
33    #[arg(value_enum)]
34    pub object_store: Option<ObjectStoreType>,
35    /// Path of the local directory. Only relevant is `--object-store` is File
36    #[serde(skip_serializing_if = "Option::is_none")]
37    #[arg(long)]
38    pub directory: Option<PathBuf>,
39    /// Name of the bucket to use for the object store. Must also set
40    /// `--object-store` to a cloud object storage to have any effect.
41    #[serde(skip_serializing_if = "Option::is_none")]
42    #[arg(long)]
43    pub bucket: Option<String>,
44    /// When using Amazon S3 as the object store, set this to an access key that
45    /// has permission to read from and write to the specified S3 bucket.
46    #[serde(skip_serializing_if = "Option::is_none")]
47    #[arg(long)]
48    pub aws_access_key_id: Option<String>,
49    /// When using Amazon S3 as the object store, set this to the secret access
50    /// key that goes with the specified access key ID.
51    #[serde(skip_serializing_if = "Option::is_none")]
52    #[arg(long)]
53    pub aws_secret_access_key: Option<String>,
54    /// When using Amazon S3 as the object store, set this to bucket endpoint
55    #[serde(skip_serializing_if = "Option::is_none")]
56    #[arg(long)]
57    pub aws_endpoint: Option<String>,
58    /// When using Amazon S3 as the object store, set this to the region
59    /// that goes with the specified bucket
60    #[serde(skip_serializing_if = "Option::is_none")]
61    #[arg(long)]
62    pub aws_region: Option<String>,
63    #[serde(skip_serializing_if = "Option::is_none")]
64    #[arg(long)]
65    pub aws_profile: Option<String>,
66    /// Enable virtual hosted style requests
67    #[serde(default)]
68    #[arg(long, default_value_t = true)]
69    pub aws_virtual_hosted_style_request: bool,
70    /// Allow unencrypted HTTP connection to AWS.
71    #[serde(default)]
72    #[arg(long, default_value_t = true)]
73    pub aws_allow_http: bool,
74    /// When using Google Cloud Storage as the object store, set this to the
75    /// path to the JSON file that contains the Google credentials.
76    #[serde(skip_serializing_if = "Option::is_none")]
77    #[arg(long)]
78    pub google_service_account: Option<String>,
79    /// When using Google Cloud Storage as the object store and writing to a
80    /// bucket with Requester Pays enabled, set this to the project_id
81    /// you want to associate the write cost with.
82    #[serde(skip_serializing_if = "Option::is_none")]
83    #[arg(long)]
84    pub google_project_id: Option<String>,
85    /// When using Microsoft Azure as the object store, set this to the
86    /// azure account name
87    #[serde(skip_serializing_if = "Option::is_none")]
88    #[arg(long)]
89    pub azure_storage_account: Option<String>,
90    /// When using Microsoft Azure as the object store, set this to one of the
91    /// keys in storage account settings
92    #[serde(skip_serializing_if = "Option::is_none")]
93    #[arg(long)]
94    pub azure_storage_access_key: Option<String>,
95    #[serde(default = "default_object_store_connection_limit")]
96    #[arg(long, default_value_t = 20)]
97    pub object_store_connection_limit: usize,
98    #[serde(default)]
99    #[arg(long, default_value_t = false)]
100    pub no_sign_request: bool,
101}
102
103fn default_object_store_connection_limit() -> usize {
104    20
105}
106
107/// How long establishing a connection may take before it is abandoned.
108pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
109
110/// How long a transfer may go without receiving any bytes before it is
111/// abandoned.
112pub const TRANSFER_STALL_TIMEOUT: Duration = Duration::from_secs(60);
113
114/// Client options with no overall request timeout, since a transfer takes as
115/// long as the object is large, but with the connect phase bounded.
116fn transfer_client_options() -> ClientOptions {
117    ClientOptions::new()
118        .with_timeout_disabled()
119        .with_connect_timeout(CONNECT_TIMEOUT)
120        .with_pool_idle_timeout(Duration::from_secs(300))
121}
122
123impl ObjectStoreConfig {
124    fn new_local_fs(&self) -> Result<Arc<DynObjectStore>, anyhow::Error> {
125        info!(directory=?self.directory, object_store_type="File", "Object Store");
126        if let Some(path) = &self.directory {
127            fs::create_dir_all(path).context(anyhow!(
128                "failed to create local directory: {}",
129                path.display()
130            ))?;
131            let store = object_store::local::LocalFileSystem::new_with_prefix(path)
132                .context(anyhow!("failed to create local object store"))?;
133            Ok(Arc::new(store))
134        } else {
135            bail!("no directory provided for local fs storage");
136        }
137    }
138    fn new_s3(&self) -> Result<Arc<DynObjectStore>, anyhow::Error> {
139        use object_store::limit::LimitStore;
140
141        info!(bucket=?self.bucket, object_store_type="S3", "Object Store");
142
143        let mut builder = AmazonS3Builder::new()
144            .with_client_options(transfer_client_options())
145            .with_imdsv1_fallback();
146
147        if self.aws_virtual_hosted_style_request {
148            builder = builder.with_virtual_hosted_style_request(true);
149        }
150        if self.aws_allow_http {
151            builder = builder.with_allow_http(true);
152        }
153        if let Some(region) = &self.aws_region {
154            builder = builder.with_region(region);
155        }
156        if let Some(bucket) = &self.bucket {
157            builder = builder.with_bucket_name(bucket);
158        }
159
160        if let Some(key_id) = &self.aws_access_key_id {
161            builder = builder.with_access_key_id(key_id);
162        } else if let Ok(secret) = env::var("ARCHIVE_READ_AWS_ACCESS_KEY_ID") {
163            builder = builder.with_access_key_id(secret);
164        } else if let Ok(secret) = env::var("FORMAL_SNAPSHOT_WRITE_AWS_ACCESS_KEY_ID") {
165            builder = builder.with_access_key_id(secret);
166        } else if let Ok(secret) = env::var("DB_SNAPSHOT_READ_AWS_ACCESS_KEY_ID") {
167            builder = builder.with_access_key_id(secret);
168        }
169
170        if let Some(secret) = &self.aws_secret_access_key {
171            builder = builder.with_secret_access_key(secret);
172        } else if let Ok(secret) = env::var("ARCHIVE_READ_AWS_SECRET_ACCESS_KEY") {
173            builder = builder.with_secret_access_key(secret);
174        } else if let Ok(secret) = env::var("FORMAL_SNAPSHOT_WRITE_AWS_SECRET_ACCESS_KEY") {
175            builder = builder.with_secret_access_key(secret);
176        } else if let Ok(secret) = env::var("DB_SNAPSHOT_READ_AWS_SECRET_ACCESS_KEY") {
177            builder = builder.with_secret_access_key(secret);
178        }
179
180        if let Some(endpoint) = &self.aws_endpoint {
181            builder = builder.with_endpoint(endpoint);
182        }
183        Ok(Arc::new(LimitStore::new(
184            builder.build().context("invalid s3 config")?,
185            self.object_store_connection_limit,
186        )))
187    }
188    fn new_gcs(&self) -> Result<Arc<DynObjectStore>, anyhow::Error> {
189        use object_store::{gcp::GoogleCloudStorageBuilder, limit::LimitStore};
190
191        info!(bucket=?self.bucket, object_store_type="GCS", "Object Store");
192
193        let mut builder = GoogleCloudStorageBuilder::new();
194
195        if let Some(bucket) = &self.bucket {
196            builder = builder.with_bucket_name(bucket);
197        }
198        if let Some(account) = &self.google_service_account {
199            builder = builder.with_service_account_path(account);
200        }
201
202        let mut client_options = transfer_client_options();
203        if let Some(google_project_id) = &self.google_project_id {
204            let x_project_header = HeaderName::from_static("x-goog-user-project");
205            let iam_req_header = HeaderName::from_static("userproject");
206
207            let mut headers = HeaderMap::new();
208            headers.insert(x_project_header, HeaderValue::from_str(google_project_id)?);
209            headers.insert(iam_req_header, HeaderValue::from_str(google_project_id)?);
210            client_options = client_options.with_default_headers(headers);
211        }
212        builder = builder.with_client_options(client_options);
213
214        Ok(Arc::new(LimitStore::new(
215            builder.build().context("invalid gcs config")?,
216            self.object_store_connection_limit,
217        )))
218    }
219    fn new_azure(&self) -> Result<Arc<DynObjectStore>, anyhow::Error> {
220        use object_store::{azure::MicrosoftAzureBuilder, limit::LimitStore};
221
222        info!(bucket=?self.bucket, account=?self.azure_storage_account,
223          object_store_type="Azure", "Object Store");
224
225        let mut builder =
226            MicrosoftAzureBuilder::new().with_client_options(transfer_client_options());
227
228        if let Some(bucket) = &self.bucket {
229            builder = builder.with_container_name(bucket);
230        }
231        if let Some(account) = &self.azure_storage_account {
232            builder = builder.with_account(account)
233        }
234        if let Some(key) = &self.azure_storage_access_key {
235            builder = builder.with_access_key(key)
236        }
237
238        Ok(Arc::new(LimitStore::new(
239            builder.build().context("invalid azure config")?,
240            self.object_store_connection_limit,
241        )))
242    }
243    pub fn make(&self) -> Result<Arc<DynObjectStore>, anyhow::Error> {
244        match &self.object_store {
245            Some(ObjectStoreType::File) => self.new_local_fs(),
246            Some(ObjectStoreType::S3) => self.new_s3(),
247            Some(ObjectStoreType::GCS) => self.new_gcs(),
248            Some(ObjectStoreType::Azure) => self.new_azure(),
249            _ => bail!("at least one storage backend should be provided"),
250        }
251    }
252}
iota_config/object_storage_config.rs

iota_config/
object_storage_config.rs
