feat(provider): add ProviderBuilder

Makes it easier to configure the `Provider`
n0-computer · Jan 25, 2023 · bf01702 · bf01702
1 parent dbbf510
commit bf01702
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 17 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -29,7 +29,7 @@ mod tests {
         let db = provider::create_db(vec![provider::DataSource::File(path.clone())]).await?;
         let hash = *db.iter().next().unwrap().0;
         let addr = "127.0.0.1:4443".parse().unwrap();
-        let mut provider = provider::Provider::new(db);
+        let mut provider = provider::Provider::builder().database(db).build()?;
         let peer_id = provider.peer_id();
         let token = provider.auth_token();
 
@@ -88,7 +88,7 @@ mod tests {
 
             let db = provider::create_db(vec![provider::DataSource::File(path)]).await?;
             let hash = *db.iter().next().unwrap().0;
-            let mut provider = provider::Provider::new(db);
+            let mut provider = provider::Provider::builder().database(db).build()?;
             let peer_id = provider.peer_id();
             let token = provider.auth_token();
 
@@ -133,7 +133,7 @@ mod tests {
         tokio::fs::write(&path, content).await?;
         let db = provider::create_db(vec![provider::DataSource::File(path)]).await?;
         let hash = *db.iter().next().unwrap().0;
-        let mut provider = provider::Provider::new(db);
+        let mut provider = provider::Provider::builder().database(db).build()?;
         let peer_id = provider.peer_id();
         let token = provider.auth_token();
 

diff --git a/src/main.rs b/src/main.rs
@@ -166,11 +166,12 @@ async fn main() -> Result<()> {
             if let Some(addr) = addr {
                 opts.addr = addr;
             }
-            let mut provider = provider::Provider::new(db);
+            let mut provider_builder = provider::Provider::builder().database(db);
             if let Some(ref hex) = auth_token {
                 let auth_token = AuthToken::from_str(hex)?;
-                provider.set_auth_token(auth_token);
+                provider_builder = provider_builder.auth_token(auth_token);
             }
+            let mut provider = provider_builder.build()?;
 
             println!("PeerID: {}", provider.peer_id());
             println!("Auth token: {}", provider.auth_token());

diff --git a/src/provider.rs b/src/provider.rs
@@ -31,21 +31,56 @@ const MAX_STREAMS: u64 = 10;
 
 pub type Database = Arc<HashMap<bao::Hash, Data>>;
 
+#[derive(Debug)]
 pub struct Provider {
     keypair: Keypair,
     auth_token: AuthToken,
     db: Database,
 }
 
+/// Builder to configure a `Provider`.
+#[derive(Debug, Default)]
+pub struct ProviderBuilder {
+    auth_token: Option<AuthToken>,
+    keypair: Option<Keypair>,
+    db: Option<Database>,
+}
+
+impl ProviderBuilder {
+    /// Set the authentication token, if none is provided a new one is generated.
+    pub fn auth_token(mut self, auth_token: AuthToken) -> Self {
+        self.auth_token = Some(auth_token);
+        self
+    }
+
+    /// Set the keypair, if none is provided a new one is generated.
+    pub fn keypair(mut self, keypair: Keypair) -> Self {
+        self.keypair = Some(keypair);
+        self
+    }
+
+    /// Set the database.
+    pub fn database(mut self, db: Database) -> Self {
+        self.db = Some(db);
+        self
+    }
+
+    /// Consumes the builder and constructs a `Provider`.
+    pub fn build(self) -> Result<Provider> {
+        ensure!(self.db.is_some(), "missing database");
+
+        Ok(Provider {
+            auth_token: self.auth_token.unwrap_or_else(AuthToken::generate),
+            keypair: self.keypair.unwrap_or_else(Keypair::generate),
+            db: self.db.unwrap(),
+        })
+    }
+}
+
 impl Provider {
-    pub fn new(db: Database) -> Self {
-        let keypair = Keypair::generate();
-        let auth_token = AuthToken::generate();
-        Provider {
-            keypair,
-            db,
-            auth_token,
-        }
+    /// Returns a new `ProviderBuilder`.
+    pub fn builder() -> ProviderBuilder {
+        ProviderBuilder::default()
     }
 
     pub fn peer_id(&self) -> PeerId {
@@ -56,10 +91,6 @@ impl Provider {
         self.auth_token
     }
 
-    pub fn set_auth_token(&mut self, auth_token: AuthToken) {
-        self.auth_token = auth_token;
-    }
-
     pub async fn run(&mut self, opts: Options) -> Result<()> {
         let server_config = tls::make_server_config(&self.keypair)?;
         let tls = s2n_quic::provider::tls::rustls::Server::from(server_config);