iota_sdk/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5//! The IOTA Rust SDK
6//!
7//! It aims at providing a similar SDK functionality like the one existing for
8//! [TypeScript](https://github.com/iotaledger/iota/tree/main/sdk/typescript/).
9//! IOTA Rust SDK builds on top of the [JSON RPC API](https://docs.iota.org/iota-api-ref)
10//! and therefore many of the return types are the ones specified in
11//! [iota_types].
12//!
13//! The API is split in several parts corresponding to different functionalities
14//! as following:
15//! * [CoinReadApi] - provides read-only functions to work with the coins
16//! * [EventApi] - provides event related functions functions to
17//! * [GovernanceApi] - provides functionality related to staking
18//! * [QuorumDriverApi] - provides functionality to execute a transaction block
19//!   and submit it to the fullnode(s)
20//! * [ReadApi] - provides functions for retrieving data about different objects
21//!   and transactions
22//! * <a href="../iota_transaction_builder/struct.TransactionBuilder.html"
23//!   title="struct
24//!   iota_transaction_builder::TransactionBuilder">TransactionBuilder</a> -
25//!   provides functions for building transactions
26//!
27//! # Usage
28//! The main way to interact with the API is through the [IotaClientBuilder],
29//! which returns an [IotaClient] object from which the user can access the
30//! various APIs.
31//!
32//! ## Getting Started
33//! Add the Rust SDK to the project by running `cargo add iota-sdk` in the root
34//! folder of your Rust project.
35//!
36//! The main building block for the IOTA Rust SDK is the [IotaClientBuilder],
37//! which provides a simple and straightforward way of connecting to an IOTA
38//! network and having access to the different available APIs.
39//!
40//! Below is a simple example which connects to a running IOTA local network,
41//! devnet, and testnet.
42//! To successfully run this program, make sure to spin up a local
43//! network with a local validator, a fullnode, and a faucet server
44//! (see [the README](https://github.com/iotaledger/iota/tree/develop/crates/iota-sdk/README.md#prerequisites) for more information).
45//!
46//! ```rust,no_run
47//! use iota_sdk::IotaClientBuilder;
48//!
49//! #[tokio::main]
50//! async fn main() -> Result<(), anyhow::Error> {
51//!     let iota = IotaClientBuilder::default()
52//!         .build("http://127.0.0.1:9000") // provide the IOTA network URL
53//!         .await?;
54//!     println!("IOTA local network version: {:?}", iota.api_version());
55//!
56//!     // local IOTA network, same result as above except using the dedicated function
57//!     let iota_local = IotaClientBuilder::default().build_localnet().await?;
58//!     println!("IOTA local network version: {:?}", iota_local.api_version());
59//!
60//!     // IOTA devnet running at `https://api.devnet.iota.cafe`
61//!     let iota_devnet = IotaClientBuilder::default().build_devnet().await?;
62//!     println!("IOTA devnet version: {:?}", iota_devnet.api_version());
63//!
64//!     // IOTA testnet running at `https://api.testnet.iota.cafe`
65//!     let iota_testnet = IotaClientBuilder::default().build_testnet().await?;
66//!     println!("IOTA testnet version: {:?}", iota_testnet.api_version());
67//!
68//!     // IOTA mainnet running at `https://api.mainnet.iota.cafe`
69//!     let iota_mainnet = IotaClientBuilder::default().build_mainnet().await?;
70//!     println!("IOTA mainnet version: {:?}", iota_mainnet.api_version());
71//!
72//!     Ok(())
73//! }
74//! ```
75//!
76//! ## Examples
77//!
78//! For detailed examples, please check the APIs docs and the examples folder
79//! in the [repository](https://github.com/iotaledger/iota/tree/main/crates/iota-sdk/examples).
80
81pub mod apis;
82pub mod error;
83pub mod iota_client_config;
84pub mod json_rpc_error;
85pub mod wallet_context;
86
87use std::{
88    collections::{HashMap, VecDeque},
89    fmt::{Debug, Formatter},
90    marker::PhantomData,
91    pin::Pin,
92    str::FromStr,
93    sync::Arc,
94    task::Poll,
95    time::Duration,
96};
97
98use async_trait::async_trait;
99use base64::Engine;
100use futures::TryStreamExt;
101pub use iota_json as json;
102use iota_json_rpc_api::{
103    CLIENT_SDK_TYPE_HEADER, CLIENT_SDK_VERSION_HEADER, CLIENT_TARGET_API_VERSION_HEADER,
104};
105pub use iota_json_rpc_types as rpc_types;
106use iota_json_rpc_types::{
107    IotaObjectDataFilter, IotaObjectDataOptions, IotaObjectResponse, IotaObjectResponseQuery, Page,
108};
109use iota_sdk_types::{Address, ObjectId, StructTag};
110use iota_transaction_builder::{DataReader, TransactionBuilder};
111pub use iota_types as types;
112use jsonrpsee::{
113    core::client::ClientT,
114    http_client::{HeaderMap, HeaderValue, HttpClient, HttpClientBuilder},
115    rpc_params,
116    ws_client::{PingConfig, WsClient, WsClientBuilder},
117};
118use reqwest::header::HeaderName;
119use rustls::crypto::{CryptoProvider, ring};
120use serde_json::Value;
121
122use crate::{
123    apis::{CoinReadApi, EventApi, GovernanceApi, QuorumDriverApi, ReadApi},
124    error::{Error, IotaRpcResult},
125};
126
127pub const IOTA_COIN_TYPE: &str = "0x2::iota::IOTA";
128pub const IOTA_LOCAL_NETWORK_URL: &str = "http://127.0.0.1:9000";
129pub const IOTA_LOCAL_NETWORK_URL_0: &str = "http://0.0.0.0:9000";
130pub const IOTA_LOCAL_NETWORK_GRAPHQL_URL: &str = "http://127.0.0.1:9125";
131pub const IOTA_LOCAL_NETWORK_GRPC_URL: &str = "http://127.0.0.1:50051";
132pub const IOTA_LOCAL_NETWORK_GAS_URL: &str = "http://127.0.0.1:9123/v1/gas";
133pub const IOTA_DEVNET_URL: &str = "https://api.devnet.iota.cafe";
134pub const IOTA_DEVNET_GRAPHQL_URL: &str = "https://graphql.devnet.iota.cafe";
135pub const IOTA_DEVNET_GRPC_URL: &str = "https://grpc.devnet.iota.cafe:443";
136pub const IOTA_DEVNET_GAS_URL: &str = "https://faucet.devnet.iota.cafe/v1/gas";
137pub const IOTA_TESTNET_URL: &str = "https://api.testnet.iota.cafe";
138pub const IOTA_TESTNET_GRAPHQL_URL: &str = "https://graphql.testnet.iota.cafe";
139pub const IOTA_TESTNET_GRPC_URL: &str = "https://grpc.testnet.iota.cafe:443";
140pub const IOTA_TESTNET_GAS_URL: &str = "https://faucet.testnet.iota.cafe/v1/gas";
141pub const IOTA_MAINNET_URL: &str = "https://api.mainnet.iota.cafe";
142pub const IOTA_MAINNET_GRAPHQL_URL: &str = "https://graphql.mainnet.iota.cafe";
143pub const IOTA_MAINNET_GRPC_URL: &str = "https://grpc.mainnet.iota.cafe:443";
144
145/// Builder for creating an [IotaClient] for connecting to the IOTA network.
146///
147/// By default `maximum concurrent requests` is set to 256 and `request timeout`
148/// is set to 60 seconds. These can be adjusted using
149/// [`Self::max_concurrent_requests()`], and the [`Self::request_timeout()`].
150/// If you use the WebSocket, consider setting `ws_ping_interval` appropriately
151/// to prevent an inactive WS subscription being disconnected due to proxy
152/// timeout.
153///
154/// # Examples
155///
156/// ```rust,no_run
157/// use iota_sdk::IotaClientBuilder;
158///
159/// #[tokio::main]
160/// async fn main() -> Result<(), anyhow::Error> {
161///     let iota = IotaClientBuilder::default()
162///         .build("http://127.0.0.1:9000")
163///         .await?;
164///
165///     println!("IOTA local network version: {:?}", iota.api_version());
166///     Ok(())
167/// }
168/// ```
169pub struct IotaClientBuilder {
170    request_timeout: Duration,
171    max_concurrent_requests: Option<usize>,
172    ws_url: Option<String>,
173    ws_ping_interval: Option<Duration>,
174    basic_auth: Option<(String, String)>,
175    tls_config: Option<rustls::ClientConfig>,
176    headers: Option<HashMap<String, String>>,
177}
178
179impl Default for IotaClientBuilder {
180    fn default() -> Self {
181        Self {
182            request_timeout: Duration::from_secs(60),
183            max_concurrent_requests: None,
184            ws_url: None,
185            ws_ping_interval: None,
186            basic_auth: None,
187            tls_config: None,
188            headers: None,
189        }
190    }
191}
192
193impl IotaClientBuilder {
194    /// Set the request timeout to the specified duration.
195    pub fn request_timeout(mut self, request_timeout: Duration) -> Self {
196        self.request_timeout = request_timeout;
197        self
198    }
199
200    /// Set the max concurrent requests allowed.
201    pub fn max_concurrent_requests(mut self, max_concurrent_requests: usize) -> Self {
202        self.max_concurrent_requests = Some(max_concurrent_requests);
203        self
204    }
205
206    /// Set the WebSocket URL for the IOTA network.
207    pub fn ws_url(mut self, url: impl AsRef<str>) -> Self {
208        self.ws_url = Some(url.as_ref().to_string());
209        self
210    }
211
212    /// Set the WebSocket ping interval.
213    pub fn ws_ping_interval(mut self, duration: Duration) -> Self {
214        self.ws_ping_interval = Some(duration);
215        self
216    }
217
218    /// Set the basic auth credentials for the HTTP client.
219    pub fn basic_auth(mut self, username: impl AsRef<str>, password: impl AsRef<str>) -> Self {
220        self.basic_auth = Some((username.as_ref().to_string(), password.as_ref().to_string()));
221        self
222    }
223
224    /// Set custom headers for the HTTP client
225    pub fn custom_headers(mut self, headers: HashMap<String, String>) -> Self {
226        self.headers = Some(headers);
227        self
228    }
229
230    /// Set a TLS configuration for the HTTP client.
231    pub fn tls_config(mut self, config: rustls::ClientConfig) -> Self {
232        self.tls_config = Some(config);
233        self
234    }
235
236    /// Return an [IotaClient] object connected to the IOTA network accessible
237    /// via the provided URI.
238    ///
239    /// # Examples
240    ///
241    /// ```rust,no_run
242    /// use iota_sdk::IotaClientBuilder;
243    ///
244    /// #[tokio::main]
245    /// async fn main() -> Result<(), anyhow::Error> {
246    ///     let iota = IotaClientBuilder::default()
247    ///         .build("http://127.0.0.1:9000")
248    ///         .await?;
249    ///
250    ///     println!("IOTA local version: {:?}", iota.api_version());
251    ///     Ok(())
252    /// }
253    /// ```
254    pub async fn build(self, http: impl AsRef<str>) -> IotaRpcResult<IotaClient> {
255        if CryptoProvider::get_default().is_none() {
256            ring::default_provider().install_default().ok();
257        }
258
259        let client_version = env!("CARGO_PKG_VERSION");
260        let mut headers = HeaderMap::new();
261        headers.insert(
262            CLIENT_TARGET_API_VERSION_HEADER,
263            // in rust, the client version is the same as the target api version
264            HeaderValue::from_static(client_version),
265        );
266        headers.insert(
267            CLIENT_SDK_VERSION_HEADER,
268            HeaderValue::from_static(client_version),
269        );
270        headers.insert(CLIENT_SDK_TYPE_HEADER, HeaderValue::from_static("rust"));
271
272        if let Some((username, password)) = self.basic_auth {
273            let auth =
274                base64::engine::general_purpose::STANDARD.encode(format!("{username}:{password}"));
275            headers.insert(
276                "authorization",
277                // reqwest::header::AUTHORIZATION,
278                HeaderValue::from_str(&format!("Basic {auth}")).unwrap(),
279            );
280        }
281
282        if let Some(custom_headers) = self.headers {
283            for (key, value) in custom_headers {
284                let header_name =
285                    HeaderName::from_str(&key).map_err(|e| Error::CustomHeaders(e.to_string()))?;
286                let header_value = HeaderValue::from_str(&value)
287                    .map_err(|e| Error::CustomHeaders(e.to_string()))?;
288                headers.insert(header_name, header_value);
289            }
290        }
291
292        let ws = if let Some(url) = self.ws_url {
293            let mut builder = WsClientBuilder::default()
294                .max_request_size(2 << 30)
295                .set_headers(headers.clone())
296                .request_timeout(self.request_timeout);
297
298            if let Some(duration) = self.ws_ping_interval {
299                builder = builder.enable_ws_ping(PingConfig::new().ping_interval(duration))
300            }
301
302            if let Some(max_concurrent_requests) = self.max_concurrent_requests {
303                builder = builder.max_concurrent_requests(max_concurrent_requests);
304            }
305
306            builder.build(url).await.ok()
307        } else {
308            None
309        };
310
311        let mut http_builder = HttpClientBuilder::default()
312            .max_request_size(2 << 30)
313            .set_headers(headers)
314            .request_timeout(self.request_timeout);
315
316        if let Some(max_concurrent_requests) = self.max_concurrent_requests {
317            http_builder = http_builder.max_concurrent_requests(max_concurrent_requests);
318        }
319
320        if let Some(tls_config) = self.tls_config {
321            http_builder = http_builder.with_custom_cert_store(tls_config);
322        }
323
324        let http = http_builder.build(http)?;
325
326        let info = Self::get_server_info(&http, &ws).await?;
327
328        let rpc = RpcClient { http, ws, info };
329        let api = Arc::new(rpc);
330        let read_api = Arc::new(ReadApi::new(api.clone()));
331        let quorum_driver_api = QuorumDriverApi::new(api.clone());
332        let event_api = EventApi::new(api.clone());
333        let transaction_builder = TransactionBuilder::new(read_api.clone());
334        let coin_read_api = CoinReadApi::new(api.clone());
335        let governance_api = GovernanceApi::new(api.clone());
336
337        Ok(IotaClient {
338            api,
339            transaction_builder,
340            read_api,
341            coin_read_api,
342            event_api,
343            quorum_driver_api,
344            governance_api,
345        })
346    }
347
348    /// Return an [IotaClient] object that is ready to interact with the local
349    /// development network (by default it expects the IOTA network to be up
350    /// and running at `127.0.0.1:9000`).
351    ///
352    /// For connecting to a custom URI, use the `build` function instead.
353    ///
354    /// # Examples
355    ///
356    /// ```rust,no_run
357    /// use iota_sdk::IotaClientBuilder;
358    ///
359    /// #[tokio::main]
360    /// async fn main() -> Result<(), anyhow::Error> {
361    ///     let iota = IotaClientBuilder::default().build_localnet().await?;
362    ///
363    ///     println!("IOTA local version: {:?}", iota.api_version());
364    ///     Ok(())
365    /// }
366    /// ```
367    pub async fn build_localnet(self) -> IotaRpcResult<IotaClient> {
368        self.build(IOTA_LOCAL_NETWORK_URL).await
369    }
370
371    /// Return an [IotaClient] object that is ready to interact with the IOTA
372    /// devnet.
373    ///
374    /// For connecting to a custom URI, use the `build` function instead.
375    ///
376    /// # Examples
377    ///
378    /// ```rust,no_run
379    /// use iota_sdk::IotaClientBuilder;
380    ///
381    /// #[tokio::main]
382    /// async fn main() -> Result<(), anyhow::Error> {
383    ///     let iota = IotaClientBuilder::default().build_devnet().await?;
384    ///
385    ///     println!("{:?}", iota.api_version());
386    ///     Ok(())
387    /// }
388    /// ```
389    pub async fn build_devnet(self) -> IotaRpcResult<IotaClient> {
390        self.build(IOTA_DEVNET_URL).await
391    }
392
393    /// Return an [IotaClient] object that is ready to interact with the IOTA
394    /// testnet.
395    ///
396    /// For connecting to a custom URI, use the `build` function instead.
397    ///
398    /// # Examples
399    ///
400    /// ```rust,no_run
401    /// use iota_sdk::IotaClientBuilder;
402    ///
403    /// #[tokio::main]
404    /// async fn main() -> Result<(), anyhow::Error> {
405    ///     let iota = IotaClientBuilder::default().build_testnet().await?;
406    ///
407    ///     println!("{:?}", iota.api_version());
408    ///     Ok(())
409    /// }
410    /// ```
411    pub async fn build_testnet(self) -> IotaRpcResult<IotaClient> {
412        self.build(IOTA_TESTNET_URL).await
413    }
414
415    /// Returns an [IotaClient] object that is ready to interact with the IOTA
416    /// mainnet.
417    ///
418    /// For connecting to a custom URI, use the `build` function instead.
419    ///
420    /// # Examples
421    ///
422    /// ```rust,no_run
423    /// use iota_sdk::IotaClientBuilder;
424    ///
425    /// #[tokio::main]
426    /// async fn main() -> Result<(), anyhow::Error> {
427    ///     let iota = IotaClientBuilder::default().build_mainnet().await?;
428    ///
429    ///     println!("{:?}", iota.api_version());
430    ///     Ok(())
431    /// }
432    /// ```
433    pub async fn build_mainnet(self) -> IotaRpcResult<IotaClient> {
434        self.build(IOTA_MAINNET_URL).await
435    }
436
437    /// Return the server information as a `ServerInfo` structure.
438    ///
439    /// Fails with an error if it cannot call the RPC discover.
440    async fn get_server_info(
441        http: &HttpClient,
442        ws: &Option<WsClient>,
443    ) -> Result<ServerInfo, Error> {
444        let rpc_spec: Value = http.request("rpc.discover", rpc_params![]).await?;
445        let version = rpc_spec
446            .pointer("/info/version")
447            .and_then(|v| v.as_str())
448            .ok_or_else(|| {
449                Error::Data("Fail parsing server version from rpc.discover endpoint.".into())
450            })?;
451        let rpc_methods = Self::parse_methods(&rpc_spec)?;
452
453        let subscriptions = if let Some(ws) = ws {
454            match ws.request("rpc.discover", rpc_params![]).await {
455                Ok(rpc_spec) => Self::parse_methods(&rpc_spec)?,
456                Err(_) => Vec::new(),
457            }
458        } else {
459            Vec::new()
460        };
461        let iota_system_state_v2_support =
462            rpc_methods.contains(&"iotax_getLatestIotaSystemStateV2".to_string());
463        Ok(ServerInfo {
464            rpc_methods,
465            subscriptions,
466            version: version.to_string(),
467            iota_system_state_v2_support,
468        })
469    }
470
471    fn parse_methods(server_spec: &Value) -> Result<Vec<String>, Error> {
472        let methods = server_spec
473            .pointer("/methods")
474            .and_then(|methods| methods.as_array())
475            .ok_or_else(|| {
476                Error::Data("Fail parsing server information from rpc.discover endpoint.".into())
477            })?;
478
479        Ok(methods
480            .iter()
481            .flat_map(|method| method["name"].as_str())
482            .map(|s| s.into())
483            .collect())
484    }
485}
486
487/// Provides all the necessary abstractions for interacting with the IOTA
488/// network.
489///
490/// # Usage
491///
492/// Use [IotaClientBuilder] to build an [IotaClient].
493///
494/// # Examples
495///
496/// ```rust,no_run
497/// use std::str::FromStr;
498///
499/// use iota_sdk::IotaClientBuilder;
500/// use iota_sdk_types::Address;
501///
502/// #[tokio::main]
503/// async fn main() -> Result<(), anyhow::Error> {
504///     let iota = IotaClientBuilder::default()
505///         .build("http://127.0.0.1:9000")
506///         .await?;
507///
508///     println!("{:?}", iota.available_rpc_methods());
509///     println!("{:?}", iota.available_subscriptions());
510///     println!("{:?}", iota.api_version());
511///
512///     let address = Address::from_str("0x0000....0000")?;
513///     let owned_objects = iota
514///         .read_api()
515///         .get_owned_objects(address, None, None, None)
516///         .await?;
517///
518///     println!("{:?}", owned_objects);
519///
520///     Ok(())
521/// }
522/// ```
523#[derive(Clone)]
524pub struct IotaClient {
525    api: Arc<RpcClient>,
526    transaction_builder: TransactionBuilder,
527    read_api: Arc<ReadApi>,
528    coin_read_api: CoinReadApi,
529    event_api: EventApi,
530    quorum_driver_api: QuorumDriverApi,
531    governance_api: GovernanceApi,
532}
533
534pub(crate) struct RpcClient {
535    http: HttpClient,
536    ws: Option<WsClient>,
537    info: ServerInfo,
538}
539
540impl Debug for RpcClient {
541    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
542        write!(
543            f,
544            "RPC client. Http: {:?}, Websocket: {:?}",
545            self.http, self.ws
546        )
547    }
548}
549
550/// Contains all the useful information regarding the API version, the available
551/// RPC calls, and subscriptions.
552struct ServerInfo {
553    rpc_methods: Vec<String>,
554    subscriptions: Vec<String>,
555    version: String,
556    iota_system_state_v2_support: bool,
557}
558
559impl IotaClient {
560    /// Return a list of RPC methods supported by the node the client is
561    /// connected to.
562    pub fn available_rpc_methods(&self) -> &Vec<String> {
563        &self.api.info.rpc_methods
564    }
565
566    /// Return a list of streaming/subscription APIs supported by the node the
567    /// client is connected to.
568    pub fn available_subscriptions(&self) -> &Vec<String> {
569        &self.api.info.subscriptions
570    }
571
572    /// Return the API version information as a string.
573    ///
574    /// The format of this string is `<major>.<minor>.<patch>`, e.g., `1.6.0`,
575    /// and it is retrieved from the OpenRPC specification via the discover
576    /// service method.
577    pub fn api_version(&self) -> &str {
578        &self.api.info.version
579    }
580
581    /// Verify if the API version matches the server version and returns an
582    /// error if they do not match.
583    pub fn check_api_version(&self) -> IotaRpcResult<()> {
584        let server_version = self.api_version();
585        let client_version = env!("CARGO_PKG_VERSION");
586        if server_version != client_version {
587            return Err(Error::ServerVersionMismatch {
588                client_version: client_version.to_string(),
589                server_version: server_version.to_string(),
590            });
591        };
592        Ok(())
593    }
594
595    /// Return a reference to the coin read API.
596    pub fn coin_read_api(&self) -> &CoinReadApi {
597        &self.coin_read_api
598    }
599
600    /// Return a reference to the event API.
601    pub fn event_api(&self) -> &EventApi {
602        &self.event_api
603    }
604
605    /// Return a reference to the governance API.
606    pub fn governance_api(&self) -> &GovernanceApi {
607        &self.governance_api
608    }
609
610    /// Return a reference to the quorum driver API.
611    pub fn quorum_driver_api(&self) -> &QuorumDriverApi {
612        &self.quorum_driver_api
613    }
614
615    /// Return a reference to the read API.
616    pub fn read_api(&self) -> &ReadApi {
617        &self.read_api
618    }
619
620    /// Return a reference to the transaction builder API.
621    pub fn transaction_builder(&self) -> &TransactionBuilder {
622        &self.transaction_builder
623    }
624
625    /// Return a reference to the underlying http client.
626    pub fn http(&self) -> &HttpClient {
627        &self.api.http
628    }
629
630    /// Return a reference to the underlying WebSocket client, if any.
631    pub fn ws(&self) -> Option<&WsClient> {
632        self.api.ws.as_ref()
633    }
634}
635
636#[async_trait]
637impl DataReader for ReadApi {
638    async fn get_owned_objects(
639        &self,
640        address: Address,
641        object_type: StructTag,
642        cursor: Option<ObjectId>,
643        limit: Option<usize>,
644        options: IotaObjectDataOptions,
645    ) -> Result<iota_json_rpc_types::ObjectsPage, anyhow::Error> {
646        let query = Some(IotaObjectResponseQuery {
647            filter: Some(IotaObjectDataFilter::StructType(object_type)),
648            options: Some(options),
649        });
650
651        Ok(self
652            .get_owned_objects(address, query, cursor, limit)
653            .await?)
654    }
655
656    async fn get_object_with_options(
657        &self,
658        object_id: ObjectId,
659        options: IotaObjectDataOptions,
660    ) -> Result<IotaObjectResponse, anyhow::Error> {
661        Ok(self.get_object_with_options(object_id, options).await?)
662    }
663
664    /// Return the reference gas price as a u64 or an error otherwise
665    async fn get_reference_gas_price(&self) -> Result<u64, anyhow::Error> {
666        Ok(self.get_reference_gas_price().await?)
667    }
668}
669
670/// A helper trait for repeatedly calling an async function which returns pages
671/// of data.
672pub trait PagedFn<O, C, F, E>: Sized + Fn(Option<C>) -> F
673where
674    O: Send,
675    C: Send,
676    F: futures::Future<Output = Result<Page<O, C>, E>> + Send,
677{
678    /// Get all items from the source and collect them into a vector.
679    fn collect<T>(self) -> impl futures::Future<Output = Result<T, E>>
680    where
681        T: Default + Extend<O>,
682    {
683        self.stream().try_collect::<T>()
684    }
685
686    /// Get a stream which will return all items from the source.
687    fn stream(self) -> PagedStream<O, C, F, E, Self> {
688        PagedStream::new(self)
689    }
690}
691
692impl<O, C, F, E, Fun> PagedFn<O, C, F, E> for Fun
693where
694    Fun: Fn(Option<C>) -> F,
695    O: Send,
696    C: Send,
697    F: futures::Future<Output = Result<Page<O, C>, E>> + Send,
698{
699}
700
701/// A stream which repeatedly calls an async function which returns a page of
702/// data.
703pub struct PagedStream<O, C, F, E, Fun> {
704    fun: Fun,
705    fut: Pin<Box<F>>,
706    next: VecDeque<O>,
707    has_next_page: bool,
708    _data: PhantomData<(E, C)>,
709}
710
711impl<O, C, F, E, Fun> PagedStream<O, C, F, E, Fun>
712where
713    Fun: Fn(Option<C>) -> F,
714{
715    pub fn new(fun: Fun) -> Self {
716        let fut = fun(None);
717        Self {
718            fun,
719            fut: Box::pin(fut),
720            next: Default::default(),
721            has_next_page: true,
722            _data: PhantomData,
723        }
724    }
725}
726
727impl<O, C, F, E, Fun> futures::Stream for PagedStream<O, C, F, E, Fun>
728where
729    O: Send,
730    C: Send,
731    F: futures::Future<Output = Result<Page<O, C>, E>> + Send,
732    Fun: Fn(Option<C>) -> F,
733{
734    type Item = Result<O, E>;
735
736    fn poll_next(
737        self: std::pin::Pin<&mut Self>,
738        cx: &mut std::task::Context<'_>,
739    ) -> Poll<Option<Self::Item>> {
740        let this = unsafe { self.get_unchecked_mut() };
741        if this.next.is_empty() && this.has_next_page {
742            match this.fut.as_mut().poll(cx) {
743                Poll::Ready(res) => match res {
744                    Ok(mut page) => {
745                        this.next.extend(page.data);
746                        this.has_next_page = page.has_next_page;
747                        if this.has_next_page {
748                            this.fut.set((this.fun)(page.next_cursor.take()));
749                        }
750                    }
751                    Err(e) => {
752                        this.has_next_page = false;
753                        return Poll::Ready(Some(Err(e)));
754                    }
755                },
756                Poll::Pending => return Poll::Pending,
757            }
758        }
759        Poll::Ready(this.next.pop_front().map(Ok))
760    }
761}
762
763#[cfg(test)]
764mod test {
765    use futures::StreamExt;
766    use iota_json_rpc_types::Page;
767
768    use super::*;
769
770    #[tokio::test]
771    async fn test_get_all_pages() {
772        let data = (0..10000).collect::<Vec<_>>();
773        struct Endpoint {
774            data: Vec<i32>,
775        }
776
777        impl Endpoint {
778            async fn get_page(&self, cursor: Option<usize>) -> anyhow::Result<Page<i32, usize>> {
779                const PAGE_SIZE: usize = 100;
780                anyhow::ensure!(cursor.is_none_or(|v| v < self.data.len()), "invalid cursor");
781                let index = cursor.unwrap_or_default();
782                let data = self.data[index..]
783                    .iter()
784                    .copied()
785                    .take(PAGE_SIZE)
786                    .collect::<Vec<_>>();
787                let has_next_page = self.data.len() > index + PAGE_SIZE;
788                Ok(Page {
789                    data,
790                    next_cursor: has_next_page.then_some(index + PAGE_SIZE),
791                    has_next_page,
792                })
793            }
794        }
795
796        let endpoint = Endpoint { data };
797
798        let mut stream = PagedFn::stream(async |cursor| endpoint.get_page(cursor).await);
799
800        assert_eq!(
801            stream
802                .by_ref()
803                .take(9999)
804                .try_collect::<Vec<_>>()
805                .await
806                .unwrap(),
807            endpoint.data[..9999]
808        );
809        assert_eq!(stream.by_ref().try_next().await.unwrap(), Some(9999));
810        assert!(stream.try_next().await.unwrap().is_none());
811
812        let mut bad_stream = PagedFn::stream(async |_| endpoint.get_page(Some(99999)).await);
813
814        assert!(bad_stream.try_next().await.is_err());
815    }
816}
iota_sdk/lib.rs

iota_sdk/
lib.rs
