Write more tests for the proxy... (#1918)

And change a few more things in the process.

proxy/src/auth/backend/console.rs

@@ -259,3 +259,15 @@ fn parse_host_port(input: &str) -> Option<(&str, u16)> {
     let (host, port) = input.split_once(':')?;
     Some((host, port.parse().ok()?))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_host_port() {
+        let (host, port) = parse_host_port("127.0.0.1:5432").expect("failed to parse");
+        assert_eq!(host, "127.0.0.1");
+        assert_eq!(port, 5432);
+    }
+}

proxy/src/auth/credentials.rs

@@ -54,13 +54,10 @@ impl<'a> ClientCredentials<'a> {
         let dbname = get_param("database")?;
 
         // Project name might be passed via PG's command-line options.
-        let project_a = params.options_raw().and_then(|options| {
-            for opt in options {
-                if let Some(value) = opt.strip_prefix("project=") {
-                    return Some(Cow::Borrowed(value));
-                }
-            }
-            None
+        let project_a = params.options_raw().and_then(|mut options| {
+            options
+                .find_map(|opt| opt.strip_prefix("project="))
+                .map(Cow::Borrowed)
         });
 
         // Alternative project name is in fact a subdomain from SNI.

proxy/src/cancellation.rs

@@ -52,6 +52,16 @@ impl CancelMap {
         let session = Session::new(key, self);
         f(session).await
     }
+
+    #[cfg(test)]
+    fn contains(&self, session: &Session) -> bool {
+        self.0.lock().contains_key(&session.key)
+    }
+
+    #[cfg(test)]
+    fn is_empty(&self) -> bool {
+        self.0.lock().is_empty()
+    }
 }
 
 /// This should've been a [`std::future::Future`], but
@@ -104,3 +114,39 @@ impl<'a> Session<'a> {
         self.key
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use once_cell::sync::Lazy;
+
+    #[tokio::test]
+    async fn check_session_drop() -> anyhow::Result<()> {
+        static CANCEL_MAP: Lazy<CancelMap> = Lazy::new(Default::default);
+
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        let task = tokio::spawn(CANCEL_MAP.with_session(|session| async move {
+            assert!(CANCEL_MAP.contains(&session));
+
+            tx.send(()).expect("failed to send");
+            let () = futures::future::pending().await; // sleep forever
+
+            Ok(())
+        }));
+
+        // Wait until the task has been spawned.
+        let () = rx.await.context("failed to hear from the task")?;
+
+        // Drop the session's entry by cancelling the task.
+        task.abort();
+        let error = task.await.expect_err("task should have failed");
+        if !error.is_cancelled() {
+            anyhow::bail!(error);
+        }
+
+        // Check that the session has been dropped.
+        assert!(CANCEL_MAP.is_empty());
+
+        Ok(())
+    }
+}

proxy/src/parse.rs

@@ -1,6 +1,5 @@
 //! Small parsing helpers.
 
-use std::convert::TryInto;
 use std::ffi::CStr;
 
 pub fn split_cstr(bytes: &[u8]) -> Option<(&CStr, &[u8])> {
@@ -10,9 +9,36 @@ pub fn split_cstr(bytes: &[u8]) -> Option<(&CStr, &[u8])> {
     Some((unsafe { CStr::from_bytes_with_nul_unchecked(cstr) }, other))
 }
 
+/// See <https://doc.rust-lang.org/std/primitive.slice.html#method.split_array_ref>.
 pub fn split_at_const<const N: usize>(bytes: &[u8]) -> Option<(&[u8; N], &[u8])> {
     (bytes.len() >= N).then(|| {
         let (head, tail) = bytes.split_at(N);
         (head.try_into().unwrap(), tail)
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_split_cstr() {
+        assert!(split_cstr(b"").is_none());
+        assert!(split_cstr(b"foo").is_none());
+
+        let (cstr, rest) = split_cstr(b"\0").expect("uh-oh");
+        assert_eq!(cstr.to_bytes(), b"");
+        assert_eq!(rest, b"");
+
+        let (cstr, rest) = split_cstr(b"foo\0bar").expect("uh-oh");
+        assert_eq!(cstr.to_bytes(), b"foo");
+        assert_eq!(rest, b"bar");
+    }
+
+    #[test]
+    fn test_split_at_const() {
+        assert!(split_at_const::<0>(b"").is_some());
+        assert!(split_at_const::<1>(b"").is_none());
+        assert!(matches!(split_at_const::<1>(b"ok"), Some((b"o", b"k"))));
+    }
+}