From 1dffba9de6a0e30e1cb63c9462c88c2f6587d2f0 Mon Sep 17 00:00:00 2001
From: Dmitry Ivanov <ivadmi5@gmail.com>
Date: Fri, 23 Sep 2022 18:30:44 +0300
Subject: [PATCH] Write more tests for the proxy... (#1918)

And change a few more things in the process.
---
 proxy/src/auth/backend/console.rs | 12 ++++++++
 proxy/src/auth/credentials.rs     | 11 +++-----
 proxy/src/cancellation.rs         | 46 +++++++++++++++++++++++++++++++
 proxy/src/parse.rs                | 28 ++++++++++++++++++-
 4 files changed, 89 insertions(+), 8 deletions(-)

diff --git a/proxy/src/auth/backend/console.rs b/proxy/src/auth/backend/console.rs
index e5ee07813c..a351b82c6a 100644
--- a/proxy/src/auth/backend/console.rs
+++ b/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);
+    }
+}
diff --git a/proxy/src/auth/credentials.rs b/proxy/src/auth/credentials.rs
index ea71eba010..e43bcf8791 100644
--- a/proxy/src/auth/credentials.rs
+++ b/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.
diff --git a/proxy/src/cancellation.rs b/proxy/src/cancellation.rs
index b7412b6f5b..92f8e35dab 100644
--- a/proxy/src/cancellation.rs
+++ b/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(())
+    }
+}
diff --git a/proxy/src/parse.rs b/proxy/src/parse.rs
index 8a05ff9c82..cbd48d91e9 100644
--- a/proxy/src/parse.rs
+++ b/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"))));
+    }
+}