refactor(tvix/*store): helper for channel creation from url r/6802

This moves the repetitive code to parse a URL and create a channel
connected to it into `tvix_castore::channel::from_url`.

Part of b/308

Change-Id: Idd342cd71cad5e78a9b258b38c1b227993e75310
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9707
Tested-by: BuildkiteCI
Reviewed-by: Connor Brewster <cbrewster@hey.com>

1 files changed, 68 insertions, 0 deletions

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)
+    }
+}