From 708b0a2ef740a8937ca02a9d7ad6ca1d553ac9ad Mon Sep 17 00:00:00 2001
From: Joseph Birr-Pixton <jpixton@gmail.com>
Date: Sat, 8 Jun 2019 21:39:29 +0100
Subject: [PATCH] Benchmark memory usage

---
 admin/bench-measure.mk     |  9 +++++
 examples/internal/bench.rs | 72 +++++++++++++++++++++++++++++++++++---
 2 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/admin/bench-measure.mk b/admin/bench-measure.mk
index a175a00c..d6fd3b7a 100644
--- a/admin/bench-measure.mk
+++ b/admin/bench-measure.mk
@@ -26,3 +26,12 @@ measure: ./target/release/examples/bench
 	$^ handshake TLS13_AES_256_GCM_SHA384
 	$^ handshake-resume TLS13_AES_256_GCM_SHA384
 	$^ handshake-ticket TLS13_AES_256_GCM_SHA384
+
+MEMUSAGE=/usr/bin/time -f %M
+memory: ./target/release/examples/bench
+	$(MEMUSAGE) $^ memory TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384 100
+	$(MEMUSAGE) $^ memory TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384 1000
+	$(MEMUSAGE) $^ memory TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384 5000
+	$(MEMUSAGE) $^ memory TLS13_AES_256_GCM_SHA384 100
+	$(MEMUSAGE) $^ memory TLS13_AES_256_GCM_SHA384 1000
+	$(MEMUSAGE) $^ memory TLS13_AES_256_GCM_SHA384 5000
diff --git a/examples/internal/bench.rs b/examples/internal/bench.rs
index cecbca71..4462d720 100644
--- a/examples/internal/bench.rs
+++ b/examples/internal/bench.rs
@@ -6,7 +6,7 @@
 use std::time::{Duration, Instant};
 use std::sync::Arc;
 use std::fs;
-use std::io::{self, Write};
+use std::io::{self, Write, Read};
 use std::env;
 
 use rustls;
@@ -300,15 +300,22 @@ fn bench_handshake(version: rustls::ProtocolVersion,
              (rounds as f64) / server_time);
 }
 
-fn do_handshake(client: &mut ClientSession, server: &mut ServerSession) {
-    while server.is_handshaking() || client.is_handshaking() {
+fn do_handshake_step(client: &mut ClientSession, server: &mut ServerSession) -> bool {
+    if server.is_handshaking() || client.is_handshaking() {
         transfer(client, server);
         server.process_new_packets().unwrap();
         transfer(server, client);
         client.process_new_packets().unwrap();
+        true
+    } else {
+        false
     }
 }
 
+fn do_handshake(client: &mut ClientSession, server: &mut ServerSession) {
+    while do_handshake_step(client, server) {}
+}
+
 fn bench_bulk(version: rustls::ProtocolVersion, suite: &'static rustls::SupportedCipherSuite,
               plaintext_size: u64) {
     let client_config =
@@ -364,6 +371,46 @@ fn bench_bulk(version: rustls::ProtocolVersion, suite: &'static rustls::Supporte
              total_mbs / time_recv);
 }
 
+fn bench_memory(version: rustls::ProtocolVersion,
+                suite: &'static rustls::SupportedCipherSuite,
+                session_count: u64) {
+    let client_config =
+        Arc::new(make_client_config(version, suite, ClientAuth::No, Resumption::No));
+    let server_config = Arc::new(make_server_config(version, suite, ClientAuth::No, Resumption::No));
+
+    if !suite.usable_for_version(version) {
+        return;
+    }
+
+    // The target here is to end up with session_count post-handshake
+    // server and client sessions.
+    let session_count = (session_count / 2) as usize;
+    let mut servers = Vec::with_capacity(session_count);
+    let mut clients = Vec::with_capacity(session_count);
+
+    for _i in 0..session_count {
+        servers.push(ServerSession::new(&server_config));
+        let dns_name = webpki::DNSNameRef::try_from_ascii_str("localhost").unwrap();
+        clients.push(ClientSession::new(&client_config, dns_name));
+    }
+
+    for _step in 0..5 {
+        for (mut client, mut server) in clients.iter_mut().zip(servers.iter_mut()) {
+            do_handshake_step(&mut client, &mut server);
+        }
+    }
+
+    for client in clients.iter_mut() {
+        client.write_all(&[0u8; 1024]).unwrap();
+    }
+
+    for (client, server) in clients.iter_mut().zip(servers.iter_mut()) {
+        transfer(client, server);
+        let mut buf = [0u8; 1024];
+        server.read(&mut buf).unwrap();
+    }
+}
+
 fn lookup_suite(name: &str) -> &'static rustls::SupportedCipherSuite {
     for suite in &rustls::ALL_CIPHERSUITES {
         if format!("{:?}", suite.suite).to_lowercase() == name.to_lowercase() {
@@ -384,7 +431,7 @@ fn selected_tests(mut args: env::Args) {
                 Some(suite) => {
                     let len = args.next()
                         .map(|arg| arg.parse::<u64>()
-                             .expect("3rd arg must be integer"))
+                             .expect("3rd arg must be plaintext size integer"))
                         .unwrap_or(1048576);
                     let suite = lookup_suite(&suite);
                     bench_bulk(rustls::ProtocolVersion::TLSv1_3, suite, len);
@@ -417,6 +464,23 @@ fn selected_tests(mut args: env::Args) {
             }
         }
 
+        "memory" => {
+            match args.next() {
+                Some(suite) => {
+                    let count = args.next()
+                        .map(|arg| arg.parse::<u64>()
+                             .expect("3rd arg must be session count integer"))
+                        .unwrap_or(1000000);
+                    let suite = lookup_suite(&suite);
+                    bench_memory(rustls::ProtocolVersion::TLSv1_3, suite, count);
+                    bench_memory(rustls::ProtocolVersion::TLSv1_2, suite, count);
+                }
+                None => {
+                    panic!("memory needs ciphersuite argument");
+                }
+            }
+        }
+
         _ => {
             panic!("unsupported mode {:?}", mode);
         }