ptth/crates/ptth_multi_call_server/src/download.rs

use std::{
	ffi::OsString,
	io::{
		self,
		Write,
	},
	time::Duration,
};

use anyhow::{
	anyhow,
	bail,
};

use futures_util::StreamExt;

use reqwest::{
	StatusCode,
};

use sha2::{
	Digest,
	Sha512,
};

use tokio::{
	sync::mpsc,
	task::{
		spawn,
		spawn_blocking,
	},
};

pub async fn main (args: &[OsString]) -> anyhow::Result <()> {
	let mut url = None;
	let mut expected_sha512 = None;
	
	let mut args = args [1..].into_iter ();
	
	loop {
		let arg = match args.next () {
			None => break,
			Some (x) => x,
		};
		
		match arg.to_str ().ok_or_else (|| anyhow! ("All arguments must be valid UTF-8"))? 
		{
			"--help" => println! ("For now, just look at the source code"),
			"--expect-sha512" => {
				let expected = args.next ().ok_or_else (|| anyhow! ("--expect-sha512 needs an argument"))?;
				expected_sha512 = Some (expected.to_str ().ok_or_else (|| anyhow! ("--expect-sha512's argument must be valid Unicode"))?);
			}
			arg => {
				url = Some (arg);
				break;
			},
		}
	}
	
	let url = match url {
		None => bail! ("URL argument is required"),
		Some (x) => x,
	};
	
	// Cookie 01FYZ3W64SM6KYNP48J6EWSCEF
	// Try to keep the Clients similar here
	
	let client = reqwest::Client::builder ()
	.connect_timeout (Duration::from_secs (30))
	.build ()?;
	
	let resp = client.get (url)
	.send ().await?;
	
	if resp.status () != StatusCode::OK {
		bail! ("Expected 200 OK, got {}", resp.status ());
	}
	
	let mut resp_stream = resp.bytes_stream ();
	
	// The hasher is owned by a task because it makes ownership simpler
	let (mut hash_tx, mut hash_rx) = mpsc::channel (1);
	
	let hasher_task = spawn_blocking (move || {
		let mut hasher = Sha512::new ();
		
		while let Some (chunk) = tokio::runtime::Handle::current ().block_on (hash_rx.recv ()) {
			hasher.update (&chunk);
		}
		
		anyhow::Result::<_>::Ok (hasher.finalize ())
	});
	
	while let Some (chunk) = resp_stream.next ().await {
		let chunk = chunk?;
		
		hash_tx.send (chunk.clone ()).await?;
		
		{
			let chunk = chunk.clone ();
			spawn_blocking (move || {
				io::stdout ().write_all (&chunk)?;
				anyhow::Result::<_>::Ok (())
			}).await??;
		}
	}
	
	drop (hash_tx);
	
	let hash = hasher_task.await??;
	
	let actual_sha512 = hex::encode (&hash);
	
	match expected_sha512 {
		None => eprintln! ("Actual SHA512 = {}", actual_sha512),
		Some (expected) => if ! actual_sha512.starts_with (&expected) {
			bail! ("Expected SHA512 prefix {}, actual SHA512 {}", expected, actual_sha512);
		},
	}
	
	Ok (())
}