diff options
| -rw-r--r-- | tvix/castore/src/blobservice/grpc.rs | 49 | ||||
| -rw-r--r-- | tvix/castore/src/channel.rs | 68 | ||||
| -rw-r--r-- | tvix/castore/src/directoryservice/grpc.rs | 48 | ||||
| -rw-r--r-- | tvix/castore/src/errors.rs | 6 | ||||
| -rw-r--r-- | tvix/castore/src/lib.rs | 1 | ||||
| -rw-r--r-- | tvix/store/src/pathinfoservice/grpc.rs | 45 |
6 files changed, 88 insertions, 129 deletions
diff --git a/tvix/castore/src/blobservice/grpc.rs b/tvix/castore/src/blobservice/grpc.rs index 627f011a5041..d0f619fcebcb 100644 --- a/tvix/castore/src/blobservice/grpc.rs +++ b/tvix/castore/src/blobservice/grpc.rs @@ -9,7 +9,7 @@ use std::{ task::Poll, }; use tokio::io::AsyncWriteExt; -use tokio::{net::UnixStream, task::JoinHandle}; +use tokio::task::JoinHandle; use tokio_stream::{wrappers::ReceiverStream, StreamExt}; use tokio_util::{ io::{CopyToBytes, SinkWriter}, @@ -44,49 +44,10 @@ impl BlobService for GRPCBlobService { /// - In the case of unix sockets, there must be a path, but may not be a host. /// - In the case of non-unix sockets, there must be a host, but no path. fn from_url(url: &url::Url) -> Result<Self, crate::Error> { - // Start checking for the scheme to start with grpc+. - match url.scheme().strip_prefix("grpc+") { - None => Err(crate::Error::StorageError("invalid scheme".to_string())), - Some(rest) => { - let channel = if rest == "unix" { - if url.host_str().is_some() { - return Err(crate::Error::StorageError( - "host may not be set".to_string(), - )); - } - let path = url.path().to_string(); - tonic::transport::Endpoint::try_from("http://[::]:50051") // doesn't matter - .unwrap() - .connect_with_connector_lazy(tower::service_fn( - move |_: tonic::transport::Uri| UnixStream::connect(path.clone()), - )) - } else { - // ensure path is empty, not supported with gRPC. - if !url.path().is_empty() { - return Err(crate::Error::StorageError( - "path may not be set".to_string(), - )); - } - - // clone the uri, and drop the grpc+ from the scheme. - // Recreate a new uri with the `grpc+` prefix dropped from the scheme. - // We can't use `url.set_scheme(rest)`, as it disallows - // setting something http(s) that previously wasn't. - let url = { - let url_str = url.to_string(); - let s_stripped = url_str.strip_prefix("grpc+").unwrap(); - url::Url::parse(s_stripped).unwrap() - }; - tonic::transport::Endpoint::try_from(url.to_string()) - .unwrap() - .connect_lazy() - }; - - Ok(Self::from_client( - proto::blob_service_client::BlobServiceClient::new(channel), - )) - } - } + let channel = crate::channel::from_url(url)?; + Ok(Self::from_client( + proto::blob_service_client::BlobServiceClient::new(channel), + )) } #[instrument(skip(self, digest), fields(blob.digest=%digest))] diff --git a/tvix/castore/src/channel.rs b/tvix/castore/src/channel.rs new file mode 100644 index 000000000000..53f9bd33442f --- /dev/null +++ b/tvix/castore/src/channel.rs @@ -0,0 +1,68 @@ +use tokio::net::UnixStream; +use tonic::transport::Channel; + +/// Turn a [url::Url] to a [Channel] if it can be parsed successfully. +/// It supports `grpc+unix:/path/to/socket`, +/// as well as the regular schemes supported by tonic, prefixed with grpc+, +/// for example `grpc+http://[::1]:8000`. +pub fn from_url(url: &url::Url) -> Result<Channel, self::Error> { + // Start checking for the scheme to start with grpc+. + // If it doesn't start with that, bail out. + match url.scheme().strip_prefix("grpc+") { + None => Err(Error::MissingGRPCPrefix()), + Some(rest) => { + if rest == "unix" { + if url.host_str().is_some() { + return Err(Error::HostSetForUnixSocket()); + } + + let url = url.clone(); + Ok( + tonic::transport::Endpoint::from_static("http://[::]:50051") // doesn't matter + .connect_with_connector_lazy(tower::service_fn( + move |_: tonic::transport::Uri| { + UnixStream::connect(url.path().to_string().clone()) + }, + )), + ) + } else { + // ensure path is empty, not supported with gRPC. + if !url.path().is_empty() { + return Err(Error::PathMayNotBeSet()); + } + + // Stringify the URL and remove the grpc+ prefix. + // We can't use `url.set_scheme(rest)`, as it disallows + // setting something http(s) that previously wasn't. + let url = url.to_string().strip_prefix("grpc+").unwrap().to_owned(); + + // Use the regular tonic transport::Endpoint logic to + Ok(tonic::transport::Endpoint::try_from(url) + .unwrap() + .connect_lazy()) + } + } + } +} + +/// Errors occuring when trying to connect to a backend +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("grpc+ prefix is missing from Url")] + MissingGRPCPrefix(), + + #[error("host may not be set for unix domain sockets")] + HostSetForUnixSocket(), + + #[error("path may not be set")] + PathMayNotBeSet(), + + #[error("transport error: {0}")] + TransportError(tonic::transport::Error), +} + +impl From<tonic::transport::Error> for Error { + fn from(value: tonic::transport::Error) -> Self { + Self::TransportError(value) + } +} diff --git a/tvix/castore/src/directoryservice/grpc.rs b/tvix/castore/src/directoryservice/grpc.rs index c3207ad064c3..43d460fae629 100644 --- a/tvix/castore/src/directoryservice/grpc.rs +++ b/tvix/castore/src/directoryservice/grpc.rs @@ -6,7 +6,6 @@ use crate::proto::{self, get_directory_request::ByWhat}; use crate::{B3Digest, Error}; use async_stream::try_stream; use futures::Stream; -use tokio::net::UnixStream; use tokio::spawn; use tokio::sync::mpsc::UnboundedSender; use tokio::task::JoinHandle; @@ -42,49 +41,10 @@ impl DirectoryService for GRPCDirectoryService { /// - In the case of unix sockets, there must be a path, but may not be a host. /// - In the case of non-unix sockets, there must be a host, but no path. fn from_url(url: &url::Url) -> Result<Self, crate::Error> { - // Start checking for the scheme to start with grpc+. - match url.scheme().strip_prefix("grpc+") { - None => Err(crate::Error::StorageError("invalid scheme".to_string())), - Some(rest) => { - let channel = if rest == "unix" { - if url.host_str().is_some() { - return Err(crate::Error::StorageError( - "host may not be set".to_string(), - )); - } - let path = url.path().to_string(); - tonic::transport::Endpoint::try_from("http://[::]:50051") // doesn't matter - .unwrap() - .connect_with_connector_lazy(tower::service_fn( - move |_: tonic::transport::Uri| UnixStream::connect(path.clone()), - )) - } else { - // ensure path is empty, not supported with gRPC. - if !url.path().is_empty() { - return Err(crate::Error::StorageError( - "path may not be set".to_string(), - )); - } - - // clone the uri, and drop the grpc+ from the scheme. - // Recreate a new uri with the `grpc+` prefix dropped from the scheme. - // We can't use `url.set_scheme(rest)`, as it disallows - // setting something http(s) that previously wasn't. - let url = { - let url_str = url.to_string(); - let s_stripped = url_str.strip_prefix("grpc+").unwrap(); - url::Url::parse(s_stripped).unwrap() - }; - tonic::transport::Endpoint::try_from(url.to_string()) - .unwrap() - .connect_lazy() - }; - - Ok(Self::from_client( - proto::directory_service_client::DirectoryServiceClient::new(channel), - )) - } - } + let channel = crate::channel::from_url(url)?; + Ok(Self::from_client( + proto::directory_service_client::DirectoryServiceClient::new(channel), + )) } async fn get( diff --git a/tvix/castore/src/errors.rs b/tvix/castore/src/errors.rs index 3b23f972b045..9f596afd4e26 100644 --- a/tvix/castore/src/errors.rs +++ b/tvix/castore/src/errors.rs @@ -34,6 +34,12 @@ impl From<Error> for Status { } } +impl From<crate::channel::Error> for Error { + fn from(value: crate::channel::Error) -> Self { + Self::StorageError(value.to_string()) + } +} + // TODO: this should probably go somewhere else? impl From<Error> for std::io::Error { fn from(value: Error) -> Self { diff --git a/tvix/castore/src/lib.rs b/tvix/castore/src/lib.rs index d51022428e5b..5a031aec7686 100644 --- a/tvix/castore/src/lib.rs +++ b/tvix/castore/src/lib.rs @@ -2,6 +2,7 @@ mod digests; mod errors; pub mod blobservice; +pub mod channel; pub mod directoryservice; pub mod fixtures; pub mod import; diff --git a/tvix/store/src/pathinfoservice/grpc.rs b/tvix/store/src/pathinfoservice/grpc.rs index 9bb09e86137f..5363109efc28 100644 --- a/tvix/store/src/pathinfoservice/grpc.rs +++ b/tvix/store/src/pathinfoservice/grpc.rs @@ -3,7 +3,6 @@ use crate::proto::{self, ListPathInfoRequest, PathInfo}; use async_stream::try_stream; use futures::Stream; use std::{pin::Pin, sync::Arc}; -use tokio::net::UnixStream; use tonic::{async_trait, transport::Channel, Code}; use tvix_castore::{ blobservice::BlobService, directoryservice::DirectoryService, proto as castorepb, Error, @@ -40,46 +39,10 @@ impl PathInfoService for GRPCPathInfoService { _blob_service: Arc<dyn BlobService>, _directory_service: Arc<dyn DirectoryService>, ) -> Result<Self, tvix_castore::Error> { - // Start checking for the scheme to start with grpc+. - match url.scheme().strip_prefix("grpc+") { - None => Err(Error::StorageError("invalid scheme".to_string())), - Some(rest) => { - let channel = if rest == "unix" { - if url.host_str().is_some() { - return Err(Error::StorageError("host may not be set".to_string())); - } - let path = url.path().to_string(); - tonic::transport::Endpoint::try_from("http://[::]:50051") // doesn't matter - .unwrap() - .connect_with_connector_lazy(tower::service_fn( - move |_: tonic::transport::Uri| UnixStream::connect(path.clone()), - )) - } else { - // ensure path is empty, not supported with gRPC. - if !url.path().is_empty() { - return Err(tvix_castore::Error::StorageError( - "path may not be set".to_string(), - )); - } - - // clone the uri, and drop the grpc+ from the scheme. - // Recreate a new uri with the `grpc+` prefix dropped from the scheme. - // We can't use `url.set_scheme(rest)`, as it disallows - // setting something http(s) that previously wasn't. - let url = { - let url_str = url.to_string(); - let s_stripped = url_str.strip_prefix("grpc+").unwrap(); - url::Url::parse(s_stripped).unwrap() - }; - tonic::transport::Endpoint::try_from(url.to_string()) - .unwrap() - .connect_lazy() - }; - Ok(Self::from_client( - proto::path_info_service_client::PathInfoServiceClient::new(channel), - )) - } - } + let channel = tvix_castore::channel::from_url(url)?; + Ok(Self::from_client( + proto::path_info_service_client::PathInfoServiceClient::new(channel), + )) } async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, Error> { |