feat: change default behavior for endpoints (#505)

* feat: change default behavior for endpoints

* made default use cache_endpoint

* fixed default access

* fixing lint errors

* removed override function

* minor fixes

* fix: renaming variables

* fix: changed error messages

---------

Co-authored-by: Aya Lubbad <aya@momentohq.com>

Author: alubbad

src/credential_provider.rs

@@ -31,6 +31,7 @@ pub struct CredentialProvider {
     pub(crate) token_endpoint: String,
     pub(crate) endpoint_security: EndpointSecurity,
     pub(crate) use_private_endpoints: bool,
+    pub(crate) use_endpoints_http_api: bool,
 }
 
 impl Display for CredentialProvider {
@@ -53,6 +54,7 @@ impl Debug for CredentialProvider {
             .field("token_endpoint", &self.token_endpoint)
             .field("endpoint_security", &self.endpoint_security)
             .field("use_private_endpoints", &self.use_private_endpoints)
+            .field("use_endpoints_http_api", &self.use_endpoints_http_api)
             .finish()
     }
 }
@@ -196,6 +198,14 @@ impl CredentialProvider {
     /// addresses to connect to.
     pub fn with_private_endpoints(mut self) -> CredentialProvider {
         self.use_private_endpoints = true;
+        self.use_endpoints_http_api = true;
+        self
+    }
+    /// Directs the ProtosocketCacheClient to look up public endpoints when discovering
+    /// addresses to connect to.
+    pub fn with_endpoints(mut self) -> CredentialProvider {
+        self.use_private_endpoints = false;
+        self.use_endpoints_http_api = true;
         self
     }
 
@@ -236,6 +246,7 @@ fn process_v1_token(auth_token_bytes: Vec<u8>) -> MomentoResult<CredentialProvid
         token_endpoint: https_endpoint(get_token_endpoint(&json.endpoint)),
         endpoint_security: EndpointSecurity::Tls,
         use_private_endpoints: false,
+        use_endpoints_http_api: false,
     })
 }
 

src/protosocket/cache/connection_manager.rs

@@ -26,6 +26,7 @@ use std::{
 };
 
 use crate::{MomentoError, MomentoResult};
+use std::net::ToSocketAddrs;
 
 #[derive(Debug)]
 struct BackgroundAddressLoader {
@@ -122,39 +123,75 @@ impl ClientConnector for ProtosocketConnectionManager {
         let address = match self.credential_provider.endpoint_security {
             EndpointSecurity::Tls => {
                 log::debug!("selecting address from address provider for TLS endpoint");
-                if self
-                    .address_provider
-                    .get_addresses()
-                    .for_az(self.az_id.as_deref())
-                    .is_empty()
-                {
-                    if let Err(e) = self.address_provider.try_refresh_addresses().await {
-                        log::warn!("error refreshing address list: {e:?}");
+                if self.credential_provider.use_endpoints_http_api {
+                    if self
+                        .address_provider
+                        .get_addresses()
+                        .for_az(self.az_id.as_deref())
+                        .is_empty()
+                    {
+                        if let Err(e) = self.address_provider.try_refresh_addresses().await {
+                            log::warn!("error refreshing address list: {e:?}");
+                        }
                     }
+                    let addresses = self
+                        .address_provider
+                        .get_addresses()
+                        .for_az(self.az_id.as_deref());
+                    if addresses.is_empty() {
+                        return Err(protosocket_rpc::Error::IoFailure(
+                            std::io::Error::new(
+                                std::io::ErrorKind::AddrNotAvailable,
+                                "No addresses available from address provider",
+                            )
+                            .into(),
+                        ));
+                    }
+                    addresses[self
+                        .connection_sequence
+                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+                        % addresses.len()]
+                } else {
+                    // Use the modified cache_endpoint with :9004 appended and https:// prefix removed
+                    let mut cache_endpoint = self
+                        .credential_provider
+                        .cache_endpoint
+                        .strip_prefix("https://")
+                        .unwrap_or(&self.credential_provider.cache_endpoint)
+                        .to_string();
+                    cache_endpoint.push_str(":9004");
+
+                    cache_endpoint
+                        .to_socket_addrs()
+                        .map_err(|e| {
+                            protosocket_rpc::Error::IoFailure(
+                                std::io::Error::other(format!(
+                                    "could not parse address from endpoint: {}: {:?}",
+                                    &self.credential_provider.cache_endpoint, e
+                                ))
+                                .into(),
+                            )
+                        })?
+                        .next()
+                        .ok_or_else(|| {
+                            protosocket_rpc::Error::IoFailure(
+                                std::io::Error::new(
+                                    std::io::ErrorKind::AddrNotAvailable,
+                                    format!(
+                                        "Unable to connect: endpoint '{}' did not resolve.",
+                                        cache_endpoint
+                                    ),
+                                )
+                                .into(),
+                            )
+                        })?
                 }
-                let addresses = self
-                    .address_provider
-                    .get_addresses()
-                    .for_az(self.az_id.as_deref());
-                if addresses.is_empty() {
-                    return Err(protosocket_rpc::Error::IoFailure(
-                        std::io::Error::new(
-                            std::io::ErrorKind::AddrNotAvailable,
-                            "No addresses available from address provider",
-                        )
-                        .into(),
-                    ));
-                }
-                addresses[self
-                    .connection_sequence
-                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
-                    % addresses.len()]
             }
             _ => {
                 log::debug!("using endpoint address directly for endpoint override");
                 self.credential_provider
                     .cache_endpoint
-                    .parse()
+                    .to_socket_addrs()
                     .map_err(|e| {
                         protosocket_rpc::Error::IoFailure(
                             std::io::Error::other(format!(
@@ -164,10 +201,19 @@ impl ClientConnector for ProtosocketConnectionManager {
                             .into(),
                         )
                     })?
+                    .next()
+                    .ok_or_else(|| {
+                        protosocket_rpc::Error::IoFailure(
+                            std::io::Error::new(
+                                std::io::ErrorKind::AddrNotAvailable,
+                                "Failed to resolve endpoint hostname into a valid address",
+                            )
+                            .into(),
+                        )
+                    })?
             }
         };
         log::debug!("connecting over protosocket to {address}");
-
         let unauthenticated_client = create_protosocket_connection(
             self.credential_provider.clone(),
             self.runtime.clone(),