From 49f3208e7f2ed45aabb730fb78a1cb06814d7556 Mon Sep 17 00:00:00 2001
From: Jimmy Cuadra <jimmy@jimmycuadra.com>
Date: Tue, 9 Jul 2019 01:52:10 -0700
Subject: [PATCH] Add get and get_mut methods to SignatureMap.

---
 src/signatures.rs | 51 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 43 insertions(+), 8 deletions(-)

diff --git a/src/signatures.rs b/src/signatures.rs
index 5bcf581f..efacc138 100644
--- a/src/signatures.rs
+++ b/src/signatures.rs
@@ -131,18 +131,41 @@ impl SignatureMap {
         server_name: &str,
         signature_set: SignatureSet,
     ) -> Result<Option<SignatureSet>, Error> {
-        let url_string = format!("https://{}", server_name);
-        let url = Url::parse(&url_string)
-            .map_err(|_| Error::new(format!("invalid server name: {}", server_name)))?;
-
-        let host = match url.host() {
-            Some(host) => host.to_owned(),
-            None => return Err(Error::new(format!("invalid server name: {}", server_name))),
-        };
+        let host = server_name_to_host(server_name)?;
 
         Ok(self.map.insert(host, signature_set))
     }
 
+    /// Gets a reference to the signature set for the given server, if any.
+    ///
+    /// # Parameters
+    ///
+    /// * server_name: The hostname or IP of the homeserver, e.g. `example.com`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the given server name cannot be parsed as a valid host.
+    pub fn get(&self, server_name: &str) -> Result<Option<&SignatureSet>, Error> {
+        let host = server_name_to_host(server_name)?;
+
+        Ok(self.map.get(&host))
+    }
+
+    /// Gets a mutable reference to the signature set for the given server, if any.
+    ///
+    /// # Parameters
+    ///
+    /// * server_name: The hostname or IP of the homeserver, e.g. `example.com`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the given server name cannot be parsed as a valid host.
+    pub fn get_mut(&mut self, server_name: &str) -> Result<Option<&mut SignatureSet>, Error> {
+        let host = server_name_to_host(server_name)?;
+
+        Ok(self.map.get_mut(&host))
+    }
+
     /// The number of servers in the collection.
     pub fn len(&self) -> usize {
         self.map.len()
@@ -363,3 +386,15 @@ fn split_id(id: &str) -> Result<(Algorithm, String), SplitError<'_>> {
 
     Ok((algorithm, signature_id[1].to_string()))
 }
+
+/// Attempts to convert a server name as a string into a `url::Host`.
+fn server_name_to_host(server_name: &str) -> Result<Host, Error> {
+    let url_string = format!("https://{}", server_name);
+    let url = Url::parse(&url_string)
+        .map_err(|_| Error::new(format!("invalid server name: {}", server_name)))?;
+
+    match url.host() {
+        Some(host) => Ok(host.to_owned()),
+        None => Err(Error::new(format!("invalid server name: {}", server_name))),
+    }
+}