Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

First phase/step of the auth flow implementation. Also collapse the r… #343

Merged
merged 4 commits into from
Jan 25, 2024
Merged

First phase/step of the auth flow implementation. Also collapse the r… #343

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ede11ed
First phase/step of the auth flow implementation. Also collapse the r…
ianoc Jan 24, 2024
e4569f8
Add simple helper for common case of not having a body or custom headers
ianoc Jan 25, 2024
6930723
Use ?
ianoc Jan 25, 2024
e60b080
Review comments
ianoc Jan 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 67 additions & 9 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,12 @@ anyhow = "1.0.79"
async-recursion = "1.0.5"
async-stream = "0.3.5"
async-trait = "0.1.77"
base64 = "0.21.7"
bytes = "1.4.0"
clap = { version = "4.4.18", features = ["derive", "env"] }
console = "0.15.8"
flate2 = "1.0.28"
futures = "0.3.30"
futures-core = "0.3.30"
http = "0.2.11"
hyper = { version = "0.14.27", features = ["full"] }
Expand Down
121 changes: 70 additions & 51 deletions src/registry/http/blob.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,8 @@ use crate::hash::sha256_value::Sha256Value;
use crate::registry::ops::BYTES_IN_MB;
use crate::registry::BlobStore;
use anyhow::{bail, Context, Error};
use http::StatusCode;
use http::Uri;
use http::{Response, StatusCode};

use hyper::Body;

use indicatif::ProgressBar;
use sha2::Digest;
Expand All @@ -18,19 +16,17 @@ use tokio::sync::Mutex;
use tokio_stream::StreamExt;
use tokio_util::io::ReaderStream;

use super::util::{dump_body_to_string, redirect_uri_fetch};
use super::util::dump_body_to_string;

#[async_trait::async_trait]
impl BlobStore for super::HttpRegistry {
async fn blob_exists(&self, digest: &str) -> Result<bool, Error> {
let uri = self.repository_uri_from_path(format!("/blobs/{}", digest))?;

let mut r = redirect_uri_fetch(
&self.http_client,
|req| req.method(http::Method::HEAD),
&uri,
)
.await?;
let mut r = self
.http_client
.request_simple(&uri, http::Method::HEAD, 3)
.await?;

if r.status() == StatusCode::NOT_FOUND {
Ok(false)
Expand All @@ -56,9 +52,10 @@ impl BlobStore for super::HttpRegistry {
let target_file = target_file.to_path_buf();

let uri = self.repository_uri_from_path(format!("/blobs/{}", digest))?;
let mut response =
redirect_uri_fetch(&self.http_client, |req| req.method(http::Method::GET), &uri)
.await?;
let mut response = self
.http_client
.request_simple(&uri, http::Method::GET, 3)
.await?;

if response.status() != StatusCode::OK {
bail!(
Expand Down Expand Up @@ -123,11 +120,12 @@ impl BlobStore for super::HttpRegistry {
) -> Result<(), Error> {
let post_target_uri = self.repository_uri_from_path("/blobs/uploads/")?;
// We expect our POST request to get a location header of where to perform the real upload to.
let req_builder = http::request::Builder::default()
.method(http::Method::POST)
.uri(post_target_uri.clone());
let request = req_builder.body(Body::from(""))?;
let mut r: Response<Body> = self.http_client.request(request).await?;

let mut r = self
.http_client
.request_simple(&post_target_uri, http::Method::POST, 0)
.await?;

if r.status() != StatusCode::ACCEPTED {
bail!(
"Expected to get a ACCEPTED/202 for upload post request to {:?}, but got {:?}",
Expand All @@ -145,41 +143,62 @@ impl BlobStore for super::HttpRegistry {
bail!("Was a redirection response code, but missing Location header, invalid response from server, body:\n{:#?}", body);
};

let f = tokio::fs::File::open(local_path).await?;

let mut file_reader_stream = ReaderStream::new(f);

let total_uploaded_bytes = Arc::new(Mutex::new(0));
let stream_byte_ref = Arc::clone(&total_uploaded_bytes);

let stream = async_stream::try_stream! {

while let Some(chunk) = file_reader_stream.next().await {
let chunk = chunk?;
let mut cntr = stream_byte_ref.lock().await;
*cntr += chunk.len();

if let Some(progress_bar) = &progress_bar {
progress_bar.set_position(*cntr as u64 / BYTES_IN_MB);
}
yield chunk
}
};

let body =
hyper::Body::wrap_stream::<_, bytes::Bytes, Box<dyn std::error::Error + Send + Sync>>(
stream,
);

let req_builder = http::request::Builder::default()
.method(http::Method::PUT)
.uri(location_uri.clone())
.header("Content-Length", length.to_string())
.header("Content-Type", "application/octet-stream");

let request = req_builder.body(body)?;

let mut r: Response<Body> = self.http_client.request(request).await?;
struct Context {
progress_bar: Option<ProgressBar>,
length: u64,
local_path: std::path::PathBuf,
}
let mut r = self
.http_client
.request(
&location_uri,
Arc::new(Context {
progress_bar,
length,
local_path: local_path.to_path_buf(),
}),
|context, builder| async move {
let f = tokio::fs::File::open(context.local_path.clone()).await?;

let stream = futures::stream::unfold(
(context.progress_bar.clone(), ReaderStream::new(f), 0),
|(progress_bar_cp, mut reader_stream, read_bytes)| async move {
let nxt_chunk = reader_stream.next().await?;

match nxt_chunk {
Ok(chunk) => {
let read_bytes: usize = read_bytes + chunk.len();
if let Some(progress_bar) = &progress_bar_cp {
progress_bar.set_position(read_bytes as u64 / BYTES_IN_MB);
}
Some((Ok(chunk), (progress_bar_cp, reader_stream, read_bytes)))
}
Err(ex) => {
let e: Box<dyn std::error::Error + Send + Sync> = Box::new(ex);
Some((Err(e), (progress_bar_cp, reader_stream, read_bytes)))
}
}
},
);

let body = hyper::Body::wrap_stream::<
_,
bytes::Bytes,
Box<dyn std::error::Error + Send + Sync>,
>(stream);

builder
.method(http::Method::PUT)
.header("Content-Length", context.length)
.header("Content-Type", "application/octet-stream")
.body(body)
.map_err(|e| e.into())
},
0,
)
.await?;

let total_uploaded_bytes: usize = {
let m = total_uploaded_bytes.lock().await;
Expand Down
21 changes: 13 additions & 8 deletions src/registry/http/copy_operations.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,6 @@ use anyhow::{bail, Error};

use http::StatusCode;

use super::util::redirect_uri_fetch;

#[async_trait::async_trait]
impl CopyOperations for super::HttpRegistry {
fn registry_name(&self) -> RegistryName {
Expand All @@ -21,12 +19,19 @@ impl CopyOperations for super::HttpRegistry {
digest, source_registry_name
))?;

let r = redirect_uri_fetch(
&self.http_client,
|req| req.method(http::Method::POST),
&uri,
)
.await?;
let r = self
.http_client
.request(
&uri,
(),
|_, c| async {
c.method(http::Method::POST)
.body(hyper::Body::from(""))
.map_err(|e| e.into())
},
3,
)
.await?;

if r.status() == StatusCode::CREATED {
Ok(())
Expand Down
Loading
Loading