cli: implement local download fallback

Implements an automatic local download fallback, similar to SSH
(cc @roblourens). If the initial download results in an error, either
in making the request or a 5xx, it'll try to fall back to making the
request locally and streaming it over the tunnel.

This abstracts the request client behing a "SimpleHttp" trait which
either uses to the native reqwest or uses the 'delegated' mode over the
socket.
This commit is contained in:
Connor Peet
2022-11-09 15:27:47 -08:00
parent 489d16dff3
commit d31573550f
8 changed files with 577 additions and 84 deletions

View File

@@ -2,7 +2,7 @@
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
use std::io;
use std::{io, task::Poll};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
@@ -57,3 +57,41 @@ where
Ok(bytes_so_far)
}
/// Helper used when converting Future interfaces to poll-based interfaces.
/// Stores excess data that can be reused on future polls.
#[derive(Default)]
pub(crate) struct ReadBuffer(Option<(Vec<u8>, usize)>);
impl ReadBuffer {
/// Removes any data stored in the read buffer
pub fn take_data(&mut self) -> Option<(Vec<u8>, usize)> {
self.0.take()
}
/// Writes as many bytes as possible to the readbuf, stashing any extra.
pub fn put_data(
&mut self,
target: &mut tokio::io::ReadBuf<'_>,
bytes: Vec<u8>,
start: usize,
) -> Poll<std::io::Result<()>> {
if bytes.is_empty() {
self.0 = None;
// should not return Ok(), since if nothing is written to the target
// it signals EOF. Instead wait for more data from the source.
return Poll::Pending;
}
if target.remaining() >= bytes.len() - start {
target.put_slice(&bytes[start..]);
self.0 = None;
} else {
let end = start + target.remaining();
target.put_slice(&bytes[start..end]);
self.0 = Some((bytes, end));
}
Poll::Ready(Ok(()))
}
}