in-progress work (does not buidl)

Author: adamspofford-dfinity

ic-agent/src/agent/agent_config.rs

@@ -5,9 +5,10 @@ use crate::{
 use reqwest::Client;
 use std::{sync::Arc, time::Duration};
 
-use super::route_provider::RouteProvider;
+use super::{route_provider::RouteProvider, IC_ROOT_KEY};
 
 /// A configuration for an agent.
+#[derive(Clone)]
 #[non_exhaustive]
 pub struct AgentConfig {
     /// See [`with_nonce_factory`](super::AgentBuilder::with_nonce_factory).
@@ -31,6 +32,8 @@ pub struct AgentConfig {
     /// See [`with_call_v3_endpoint`](super::AgentBuilder::with_call_v3_endpoint).
     #[cfg(feature = "experimental_sync_call")]
     pub use_call_v3_endpoint: bool,
+    /// See [`with_preset_root_key`](super::AgentBuilder::with_preset_root_key).
+    pub root_key: Option<Vec<u8>>,
 }
 
 impl Default for AgentConfig {
@@ -47,6 +50,7 @@ impl Default for AgentConfig {
             max_tcp_error_retries: 0,
             #[cfg(feature = "experimental_sync_call")]
             use_call_v3_endpoint: false,
+            root_key: None,
         }
     }
 }

ic-agent/src/agent/builder.rs

@@ -1,13 +1,22 @@
 use url::Url;
 
 use crate::{
-    agent::{agent_config::AgentConfig, Agent},
+    agent::{
+        agent_config::AgentConfig,
+        route_provider::{
+            dynamic_routing::{
+                dynamic_route_provider::{DynamicRouteProviderBuilder, IC0_SEED_DOMAIN},
+                node::Node,
+                snapshot::latency_based_routing::LatencyRoutingSnapshot,
+            },
+            RouteProvider,
+        },
+        Agent,
+    },
     AgentError, Identity, NonceFactory, NonceGenerator,
 };
 use std::sync::Arc;
 
-use super::route_provider::RouteProvider;
-
 /// A builder for an [`Agent`].
 #[derive(Default)]
 pub struct AgentBuilder {
@@ -20,19 +29,8 @@ impl AgentBuilder {
         Agent::new(self.config)
     }
 
-    #[cfg(all(feature = "reqwest", not(target_family = "wasm")))]
     /// Set the dynamic transport layer for the [`Agent`], performing continuos discovery of the API boundary nodes and routing traffic via them based on the latencies.
     pub async fn with_discovery_transport(self, client: reqwest::Client) -> Self {
-        use crate::agent::http_transport::{
-            dynamic_routing::{
-                dynamic_route_provider::{DynamicRouteProviderBuilder, IC0_SEED_DOMAIN},
-                node::Node,
-                snapshot::latency_based_routing::LatencyRoutingSnapshot,
-            },
-            route_provider::RouteProvider,
-            ReqwestTransport,
-        };
-
         // TODO: This is a temporary solution to get the seed node.
         let seed = Node::new(IC0_SEED_DOMAIN).unwrap();
 
@@ -46,10 +44,8 @@ impl AgentBuilder {
 
         let route_provider = Arc::new(route_provider) as Arc<dyn RouteProvider>;
 
-        let transport = ReqwestTransport::create_with_client_route(route_provider, client)
-            .expect("failed to create transport");
-
-        self.with_transport(transport)
+        self.with_http_client(client)
+            .with_arc_route_provider(route_provider)
     }
 
     /// Set the URL of the [Agent].

ic-agent/src/agent/mod.rs

@@ -30,6 +30,7 @@ use rangemap::{RangeInclusiveMap, RangeInclusiveSet, StepFns};
 use reqwest::{Body, Client, Request};
 use route_provider::RouteProvider;
 use time::OffsetDateTime;
+use url::Url;
 
 #[cfg(test)]
 mod agent_test;
@@ -144,6 +145,7 @@ type AgentFuture<'a, V> = Pin<Box<dyn Future<Output = Result<V, AgentError>> + '
 /// This agent does not understand Candid, and only acts on byte buffers.
 #[derive(Clone)]
 pub struct Agent {
+    // If adding any BN-specific fields, exclude them from clone_with_url.
     nonce_factory: Arc<dyn NonceGenerator>,
     identity: Arc<dyn Identity>,
     ingress_expiry: Duration,
@@ -176,11 +178,13 @@ impl Agent {
 
     /// Create an instance of an [`Agent`].
     pub fn new(config: agent_config::AgentConfig) -> Result<Agent, AgentError> {
-        Ok(Agent {
+        let agent = Agent {
             nonce_factory: config.nonce_factory,
             identity: config.identity,
             ingress_expiry: config.ingress_expiry.unwrap_or(DEFAULT_INGRESS_EXPIRY),
-            root_key: Arc::new(RwLock::new(IC_ROOT_KEY.to_vec())),
+            root_key: Arc::new(RwLock::new(
+                config.root_key.unwrap_or_else(|| IC_ROOT_KEY.to_vec()),
+            )),
             client: config.client.unwrap_or_else(|| {
                 #[cfg(not(target_family = "wasm"))]
                 {
@@ -213,7 +217,9 @@ impl Agent {
                     false
                 }
             },
-        })
+        };
+        agent.route_provider.notify_start(agent.clone());
+        Ok(agent)
     }
 
     /// Set the identity provider for signing messages.
@@ -1244,6 +1250,12 @@ impl Agent {
             Ok((status, body))
         }
     }
+
+    fn clone_with_url(&self, url: Url) -> Self {
+        let mut clone = self.clone();
+        clone.route_provider = Arc::new(url);
+        clone
+    }
 }
 
 const DEFAULT_INGRESS_EXPIRY: Duration = Duration::from_secs(240);

ic-agent/src/agent/route_provider.rs

@@ -5,7 +5,7 @@ use std::{
 };
 use url::Url;
 
-use crate::agent::AgentError;
+use crate::{agent::AgentError, Agent};
 
 pub mod dynamic_routing;
 
@@ -33,6 +33,8 @@ pub trait RouteProvider: std::fmt::Debug + Send + Sync {
     /// appearing first. The returned vector can contain fewer than `n` URLs if
     /// fewer are available.
     fn n_ordered_routes(&self, n: usize) -> Result<Vec<Url>, AgentError>;
+
+    fn notify_start(&self, agent: Agent) {}
 }
 
 /// A simple implementation of the [`RouteProvider`] which produces an even distribution of the urls from the input ones.