From 21d50189af56932dd61c6659b5be8c6bb789363b Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Mon, 29 Jul 2024 00:52:17 -0700
Subject: [PATCH] tests: add JA4 pcap tests

---
 tests/pcap/src/client_hello.rs        | 10 ++++++++++
 tests/pcap/tests/s2n_client_hellos.rs | 28 +++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/tests/pcap/src/client_hello.rs b/tests/pcap/src/client_hello.rs
index 563d5e37104..c6728ed49a3 100644
--- a/tests/pcap/src/client_hello.rs
+++ b/tests/pcap/src/client_hello.rs
@@ -23,6 +23,16 @@ impl ClientHello {
         self.0.packet.metadata(Self::JA3_STR).map(str::to_owned)
     }
 
+    const JA4_HASH: &'static str = "tls.handshake.ja4";
+    pub fn ja4_hash(&self) -> Option<String> {
+        self.0.packet.metadata(Self::JA4_HASH).map(str::to_owned)
+    }
+
+    const JA4_STR: &'static str = "tls.handshake.ja4_r";
+    pub fn ja4_string(&self) -> Option<String> {
+        self.0.packet.metadata(Self::JA4_STR).map(str::to_owned)
+    }
+
     pub fn message(&self) -> &HandshakeMessage {
         &self.0
     }
diff --git a/tests/pcap/tests/s2n_client_hellos.rs b/tests/pcap/tests/s2n_client_hellos.rs
index 1f073aefc33..8e96c575c6b 100644
--- a/tests/pcap/tests/s2n_client_hellos.rs
+++ b/tests/pcap/tests/s2n_client_hellos.rs
@@ -6,6 +6,7 @@ use pcap::all_pcaps;
 use pcap::client_hello::ClientHello as PcapHello;
 use pcap::handshake_message::Builder;
 use s2n_tls::client_hello::{ClientHello as S2NHello, FingerprintType};
+use s2n_tls::fingerprint;
 
 fn get_s2n_hello(pcap_hello: &PcapHello) -> Result<Box<S2NHello>> {
     let bytes = pcap_hello.message().bytes();
@@ -14,9 +15,9 @@ fn get_s2n_hello(pcap_hello: &PcapHello) -> Result<Box<S2NHello>> {
     Ok(r?)
 }
 
-fn test_all_client_hellos<F>(test_fn: F) -> Result<()>
+fn test_all_client_hellos<F>(mut test_fn: F) -> Result<()>
 where
-    F: FnOnce(PcapHello, Box<S2NHello>) -> Result<()> + Copy,
+    F: FnMut(PcapHello, Box<S2NHello>) -> Result<()>,
 {
     let pcaps = all_pcaps();
     for pcap in pcaps {
@@ -62,3 +63,26 @@ fn ja3_fingerprints() -> Result<()> {
         Ok(())
     })
 }
+
+#[test]
+fn ja4_fingerprints() -> Result<()> {
+    let mut builder = fingerprint::Builder::new(FingerprintType::JA4)?;
+
+    test_all_client_hellos(|pcap_hello, s2n_hello| {
+        let mut fingerprint = builder.build(&s2n_hello)?;
+
+        let s2n_ja4_hash = fingerprint
+            .hash()
+            .context("s2n failed to calculate ja4 hash")?
+            .to_owned();
+
+        let s2n_ja4_str = fingerprint
+            .raw()
+            .context("s2n failed to calculate ja4 string")?
+            .to_owned();
+
+        assert_eq!(pcap_hello.ja4_hash(), Some(s2n_ja4_hash));
+        assert_eq!(pcap_hello.ja4_string(), Some(s2n_ja4_str));
+        Ok(())
+    })
+}