Implement downloading Arweave transaction data

Author: mikonieminen

src/arweave.rs

@@ -1,30 +1,30 @@
+use bytes::{Buf, BufMut};
+use chrono::Utc;
+use data_encoding::BASE64URL_NOPAD;
 use derive_more::{Display, Error};
-use diesel::expression::UncheckedBind;
 use futures::future::BoxFuture;
-use futures::{pin_mut, stream, FutureExt, StreamExt, TryFutureExt};
-use http::header::CONTENT_LENGTH;
-use http::{StatusCode, Uri};
-use log::{debug, error, info, warn};
+use futures::{pin_mut, FutureExt, StreamExt};
+use http::header::{ToStrError, CONTENT_LENGTH};
+use http::Uri;
+use log::{debug, error, info};
 use serde::{Deserialize, Serialize};
-use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::convert::Infallible;
 use std::fmt::{self, Debug};
-use std::pin::Pin;
-use std::sync::atomic::{AtomicU16, AtomicU8, Ordering};
+use std::sync::atomic::{AtomicU16, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
-use tokio::io::{AsyncWrite, AsyncWriteExt};
+use tokio::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, BufWriter};
 
 use std::num::ParseIntError;
 use std::str::FromStr;
 use url::Url;
 
 use crate::dynamic_source::DynamicSource;
 use crate::http::reqwest::execute_with_retry;
-use crate::http::{Client, ClientAccess};
+use crate::http::{Client, ClientAccess, RetryAfter};
 use crate::key_manager::public_key_to_address;
-use crate::pool::Pool;
+use crate::retry::{retry, RetryControl};
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
 #[serde(transparent)]
@@ -307,6 +307,12 @@ impl FromStr for TransactionSize {
     }
 }
 
+impl Into<u64> for TransactionSize {
+    fn into(self) -> u64 {
+        self.0 as u64
+    }
+}
+
 impl TryFrom<String> for TransactionSize {
     type Error = ParseIntError;
 
@@ -386,6 +392,11 @@ pub struct Offset {
     size: u64,
 }
 
+#[derive(Clone, Debug, Deserialize)]
+struct Chunk<'a> {
+    chunk: &'a [u8],
+}
+
 #[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
 pub struct Node(String);
 
@@ -539,66 +550,6 @@ where
     .boxed()
 }
 
-// fn get_peers<'a, HttpClient>(
-//     client: &'a HttpClient,
-//     node: &'a Node,
-//     timeout: Option<Duration>,
-// ) -> BoxFuture<'a, Result<Vec<Node>, FetchPeersError<HttpClient::Error>>>
-// where
-//     HttpClient:
-//         Client<Request = reqwest::Request, Response = reqwest::Response> + Send + Sync + 'static,
-//     HttpClient::Error: From<reqwest::Error>,
-// {
-//     async move {
-//         debug!("Get peers for {}", node);
-//         let url =
-//             match Url::from_str(&format!("http://{}", node)).and_then(|base| base.join("/peers")) {
-//                 Ok(url) => url,
-//                 Err(err) => {
-//                     debug!(
-//                         "Failed to build URL request peers for node: {}, {:?}",
-//                         node, err
-//                     );
-//                     return Err(FetchPeersError::UnsupportedPeerAddress(node.clone()));
-//                 }
-//             };
-
-//         let res = match get(client, url, timeout).await {
-//             Ok(res) => match res.error_for_status() {
-//                 Ok(res) => res,
-//                 Err(err) => {
-//                     debug!(
-//                         "Request for fetching peers failed, peer: {}, err: {:?}",
-//                         node, err
-//                     );
-//                     return Err(FetchPeersError::HttpClientError(err.into()));
-//                 }
-//             },
-//             Err(err) => {
-//                 error!(
-//                     "Request for fetching peers failed, peer: {}, err: {:?}",
-//                     node, err
-//                 );
-//                 return Err(FetchPeersError::ArweaveError(err));
-//             }
-//         };
-
-//         let peers: Vec<Node> = match res.json().await {
-//             Ok(peers) => peers,
-//             Err(err) => {
-//                 error!(
-//                     "Failed to deserialize peers, peer: {}, error: {:?}",
-//                     node, err
-//                 );
-//                 return Err(FetchPeersError::ResponseDeserializationError);
-//             }
-//         };
-
-//         Ok(peers)
-//     }
-//     .boxed()
-// }
-
 #[derive(Debug, Display, Error, Clone, PartialEq)]
 pub enum ArweaveError {
     MalformedRequest,
@@ -834,7 +785,7 @@ impl Arweave {
         Context: ClientAccess<HttpClient>,
         HttpClient: Client<Request = reqwest::Request, Response = reqwest::Response>,
         HttpClient::Error: From<reqwest::Error>,
-        Output: AsyncWrite + Unpin,
+        Output: AsyncWrite + AsyncSeek + Unpin,
     {
         let client = ctx.get_http_client();
         let base_url = if let Some(ref peer) = peer {
@@ -854,11 +805,12 @@ impl Arweave {
             ArweaveError::UnknownErr
         })?;
 
-        let mut download_size = 0;
-
-        while download_size < size {
-            let chunk_offset = offset - download_size;
+        info!("Transaction offset={}, size={}", offset, size);
 
+        let end_offset = offset;
+        let start_offset = offset - size + 1;
+        let mut chunk_offset = start_offset;
+        while chunk_offset < end_offset + 1 {
             let url = self
                 .base_url
                 .join(&format!("/chunk/{}", &chunk_offset))
@@ -869,24 +821,58 @@ impl Arweave {
 
             let mut res = get(client, url, None).await?;
 
-            let chunk_size = res.headers().get(CONTENT_LENGTH).ok_or_else(|| {
-                error!("Could not read chunk size, missing Content-Length header");
-                ArweaveError::RequestFailed
-            })?;
+            let content_length: u64 = res
+                .headers()
+                .get(CONTENT_LENGTH)
+                .and_then(|h| h.to_str().ok())
+                .and_then(|s| s.parse::<u64>().ok())
+                .ok_or_else(|| {
+                    error!("Could not read chunk size, missing Content-Length header");
+                    ArweaveError::RequestFailed
+                })?;
 
-            let chunk_size: u64 = chunk_size.to_str().unwrap().parse().unwrap();
-            info!("Chunk: offset={}, size={}", chunk_offset, chunk_size);
+            let file_offset = chunk_offset - start_offset;
 
+            // This is not strictly needed at the moment, but keeping as
+            // an example, because this will be needed once we add retry
+            // and concurrency
+            output
+                .seek(std::io::SeekFrom::Start(file_offset))
+                .await
+                .map_err(|err| {
+                    error!(
+                        "Failed to seek into the right position in the output: {:?}",
+                        err
+                    );
+                    ArweaveError::RequestFailed
+                })?;
+
+            let mut buf = Vec::with_capacity(content_length as usize);
             while let Some(chunk) = res.chunk().await.map_err(|err| {
                 error!("Failed to read chunk data: {:?}", err);
                 ArweaveError::RequestFailed
             })? {
-                download_size += chunk.len() as u64;
-                output.write_all(&chunk).await.map_err(|err| {
+                buf.write_all(&chunk).await.map_err(|err| {
                     error!("Failed to write chunk data to output: {:?}", err);
                     ArweaveError::RequestFailed
                 })?;
             }
+            let chunk: Chunk = serde_json::from_slice(buf.as_slice()).unwrap();
+            let data = BASE64URL_NOPAD.decode(chunk.chunk).unwrap();
+
+            output.write_all(data.as_slice()).await.map_err(|err| {
+                error!("Failed to write chunk data: {:?}", err);
+                ArweaveError::UnknownErr
+            })?;
+
+            let chunk_size = data.len() as u64;
+
+            info!(
+                "Got chunk: offset={}, chunk_size={}",
+                chunk_offset, chunk_size
+            );
+
+            chunk_offset += chunk_size;
         }
 
         Ok(())

src/bin/download_transaction.rs

@@ -1,4 +1,7 @@
-use std::path::PathBuf;
+use std::{
+    fs::{self, DirBuilder},
+    path::PathBuf,
+};
 
 use clap::Parser;
 
@@ -7,7 +10,7 @@ use log::info;
 use tokio::fs::File;
 use url::Url;
 use validator::{
-    arweave::{Arweave, ArweaveContext, ArweaveError, TransactionId},
+    arweave::{Arweave, ArweaveContext, Transaction, TransactionId},
     http::{reqwest::ReqwestClient, ClientAccess},
 };
 
@@ -46,7 +49,7 @@ impl ArweaveContext<ReqwestClient> for Context {
 }
 
 #[tokio::main]
-async fn main() -> Result<(), ArweaveError> {
+async fn main() {
     dotenv::dotenv().ok();
     env_logger::init_from_env(Env::default().default_filter_or("info"));
 
@@ -55,18 +58,40 @@ async fn main() -> Result<(), ArweaveError> {
     let arweave = Arweave::new(args.arweave_gateway.clone());
     let ctx = Context::new(reqwest::Client::default());
 
-    let mut tx_data_file = args.tx_cache.clone();
+    if !args
+        .tx_cache
+        .try_exists()
+        .expect("Failed to check if tx cache folder exists")
+    {
+        DirBuilder::new()
+            .recursive(true)
+            .create(args.tx_cache.clone())
+            .expect("Failed to create tx cache dir");
+    }
+
+    let mut tx_data_file = args.tx_cache;
     tx_data_file.push(format!("{}", args.tx));
-    // let tx_data_file = fs::canonicalize(tx_data_file)
-    //     .expect("Failed to create absolute file path from provided information");
 
-    info!("Write tx data to {:?}", tx_data_file);
+    let Transaction { data_size, .. } = arweave
+        .get_transaction_info(&ctx, &args.tx)
+        .await
+        .expect("Failed to fetch transaction size info");
 
-    let mut output = File::create(tx_data_file)
+    let mut output = File::create(tx_data_file.clone())
         .await
         .expect("Failed to open file for writing");
 
+    output
+        .set_len(data_size.into())
+        .await
+        .expect("Failed to set file size to match transaction data size");
+
     arweave
         .download_transaction_data(&ctx, &args.tx, &mut output, None)
         .await
+        .expect("Failed to download transaction data");
+
+    let file_path = fs::canonicalize(tx_data_file)
+        .expect("Failed to create absolute file path from provided information");
+    info!("Wrote tx data to {}", file_path.to_str().unwrap());
 }