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

iota_sdk/
lib.rs
