diff --git a/shared/comm-lib/src/auth/types.rs b/shared/comm-lib/src/auth/types.rs --- a/shared/comm-lib/src/auth/types.rs +++ b/shared/comm-lib/src/auth/types.rs @@ -36,6 +36,10 @@ } } } + + pub fn is_services_token(&self) -> bool { + matches!(self, Self::ServicesToken(_)) + } } impl std::fmt::Display for AuthorizationCredential { diff --git a/shared/comm-lib/src/blob/client.rs b/shared/comm-lib/src/blob/client.rs --- a/shared/comm-lib/src/blob/client.rs +++ b/shared/comm-lib/src/blob/client.rs @@ -21,7 +21,9 @@ }, }; -use super::types::http::RemoveHoldersResponse; +use super::types::http::{ + BlobSizesRequest, BlobSizesResponse, RemoveHoldersResponse, +}; #[derive(From, Error, Debug, Display)] pub enum BlobServiceError { @@ -280,6 +282,51 @@ Ok(result) } + /// Fetches blob sizes for requested blob hashes. + /// This endpoint is callable only by other services. + /// + /// # Example + /// ```ignore + /// let client = + /// BlobServiceClient::new("http://localhost:50053".parse()?); + /// let blob_hashes = vec!["blob1".to_string(), "blob2".to_string()]; + /// + /// let response = client + /// .fetch_blob_sizes(BlobSizesRequest { blob_hashes }).await? + /// + /// let blob1_size = response.blob_sizes.get("blob1").unwrap(); + /// let all_blobs_size = response.total_size(); + /// ``` + pub async fn fetch_blob_sizes( + &self, + request: BlobSizesRequest, + ) -> BlobResult { + self.ensure_caller_is_service("fetch_blob_sizes")?; + + let url = self + .blob_service_url + .join("/metadata/get_blob_sizes") + .map_err(|err| BlobServiceError::URLError(err.to_string()))?; + + trace!("Request payload: {:?}", request); + let response = self + .request(Method::POST, url)? + .json(&request) + .send() + .await?; + + if !response.status().is_success() { + return error_response_result(response).await; + } + + let result: BlobSizesResponse = response.json().await?; + debug!( + "Request successful. Fetched sizes for {} blobs.", + result.blob_sizes.len() + ); + Ok(result) + } + /// Uploads a blob. Returns `BlobServiceError::AlreadyExists` if blob with given hash /// already exists. /// @@ -428,6 +475,26 @@ None => Ok(request), } } + + fn ensure_caller_is_service( + &self, + client_func_name: &'static str, + ) -> BlobResult<()> { + if self + .auth_credential + .as_ref() + .filter(|it| it.is_services_token()) + .is_some() + { + return Ok(()); + }; + + error!( + "Called service-only BlobServiceClient::{client_func_name}() {}", + "from outside a service or authenticated with client CSAT!" + ); + Err(BlobServiceError::UnexpectedError) + } } fn handle_http_error(status_code: StatusCode) -> BlobServiceError { diff --git a/shared/comm-lib/src/blob/types.rs b/shared/comm-lib/src/blob/types.rs --- a/shared/comm-lib/src/blob/types.rs +++ b/shared/comm-lib/src/blob/types.rs @@ -111,6 +111,13 @@ pub struct BlobSizesResponse { pub blob_sizes: HashMap, } + + impl BlobSizesResponse { + /// Returns total size of all requested blobs. + pub fn total_size(&self) -> u64 { + self.blob_sizes.values().sum() + } + } } /// Blob owning information - stores both blob_hash and holder