rust: bssl-tls-tokio: Rework the tests

Signed-off-by: Xiangfei Ding <xfding@google.com>
Change-Id: Ic2234ae5e6ffddb995721529b97db7406a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/100068
Reviewed-by: Adam Langley <agl@google.com>
Reviewed-by: Rudolf Polzer <rpolzer@google.com>
diff --git a/rust/bssl-tls-tokio/src/tests.rs b/rust/bssl-tls-tokio/src/tests.rs
index 408c67c..cb4dc27 100644
--- a/rust/bssl-tls-tokio/src/tests.rs
+++ b/rust/bssl-tls-tokio/src/tests.rs
@@ -20,47 +20,39 @@
 const RSA_SERVER_CERT: &[u8] = include_bytes!("../../test-data/BoringSSLServerTest-RSA.crt");
 const RSA_SERVER_KEY: &[u8] = include_bytes!("../../test-data/BoringSSLServerTest-RSA.key");
 
-use bssl_tls::{
-    context::TlsContextBuilder,
-    credentials::{
-        Certificate,
-        TlsCredentialBuilder, //
-    },
+use bssl_tls::credentials::{
+    Certificate,
+    TlsCredential,
+    TlsCredentialBuilder, //
 };
 use bssl_x509::{
     certificates::X509Certificate,
     keys::PrivateKey,
     params::Trust,
-    store::X509StoreBuilder, //
+    store::{X509Store, X509StoreBuilder}, //
 };
 
-/// Returns a `(server_ctx_builder, client_ctx_builder)` pair with certs/keys
-/// pre-configured for TLS testing.
-fn tls_ctx_builders() -> (TlsContextBuilder, TlsContextBuilder) {
+/// Returns a server credential with the test certificate chain and private key.
+fn server_credential() -> TlsCredential {
     let ca = Certificate::parse_one_from_pem(CA, None).unwrap();
     let server_cert = Certificate::parse_one_from_pem(RSA_SERVER_CERT, None).unwrap();
     let server_key = PrivateKey::from_pem(RSA_SERVER_KEY, || unreachable!()).unwrap();
+    let mut builder = TlsCredentialBuilder::new();
+    builder
+        .with_certificate_chain(&[server_cert, ca])
+        .unwrap()
+        .with_private_key(server_key)
+        .unwrap();
+    builder.build().unwrap()
+}
 
-    let mut server_builder = TlsContextBuilder::new_tls();
-    let server_cred = {
-        let mut builder = TlsCredentialBuilder::new();
-        builder
-            .with_certificate_chain(&[server_cert, ca])
-            .unwrap()
-            .with_private_key(server_key)
-            .unwrap();
-        builder.build().unwrap()
-    };
-    server_builder.with_credential(server_cred).unwrap();
-
-    let mut client_builder = TlsContextBuilder::new_tls();
-    let mut cert_store = X509StoreBuilder::new();
-    cert_store
+/// Returns a certificate store trusting the test CA.
+fn client_cert_store() -> X509Store {
+    let mut store = X509StoreBuilder::new();
+    store
         .set_trust(Trust::SslServer)
         .unwrap()
         .add_cert(X509Certificate::parse_one_from_pem(CA).unwrap())
         .unwrap();
-    client_builder.with_certificate_store(&cert_store.build());
-
-    (server_builder, client_builder)
+    store.build()
 }
diff --git a/rust/bssl-tls-tokio/src/tests/convenience.rs b/rust/bssl-tls-tokio/src/tests/convenience.rs
index 7d459b0..35ca349 100644
--- a/rust/bssl-tls-tokio/src/tests/convenience.rs
+++ b/rust/bssl-tls-tokio/src/tests/convenience.rs
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use bssl_tls::context::TlsContextBuilder;
+
 use tokio::io::{
     AsyncReadExt,
     AsyncWriteExt, //
@@ -21,10 +23,14 @@
 
 #[tokio::test]
 async fn high_level_tokio() -> Result<(), bssl_tls::errors::Error> {
-    let (server_builder, client_builder) = super::tls_ctx_builders();
-    let connector = client_builder.build_tokio_connector();
+    let mut server_builder = TlsContextBuilder::new_tls();
+    server_builder.with_credential(super::server_credential())?;
     let acceptor = server_builder.build_tokio_acceptor();
 
+    let mut client_builder = TlsContextBuilder::new_tls();
+    client_builder.with_certificate_store(&super::client_cert_store());
+    let connector = client_builder.build_tokio_connector();
+
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
     let addr = listener.local_addr().unwrap();
 
diff --git a/rust/bssl-tls-tokio/src/tests/datagram.rs b/rust/bssl-tls-tokio/src/tests/datagram.rs
index e88ea33..ed81644 100644
--- a/rust/bssl-tls-tokio/src/tests/datagram.rs
+++ b/rust/bssl-tls-tokio/src/tests/datagram.rs
@@ -24,24 +24,9 @@
         DtlsMode,
         TlsContextBuilder, //
     },
-    credentials::{
-        Certificate,
-        TlsCredentialBuilder, //
-    },
     errors::Error, //
 };
-use bssl_x509::{
-    certificates::X509Certificate,
-    keys::PrivateKey,
-    params::Trust,
-    store::X509StoreBuilder, //
-};
 
-use super::{
-    CA,
-    RSA_SERVER_CERT,
-    RSA_SERVER_KEY, //
-};
 use crate::{
     TokioDatagramIo,
     new_std_datagram_with_tokio, //
@@ -51,36 +36,15 @@
     TlsConnection<Server, DtlsMode>,
     TlsConnection<Client, DtlsMode>,
 ) {
-    let ca = Certificate::parse_one_from_pem(CA, None).unwrap();
-    let server_cert = Certificate::parse_one_from_pem(RSA_SERVER_CERT, None).unwrap();
-    let server_key = PrivateKey::from_pem(RSA_SERVER_KEY, || unreachable!()).unwrap();
-
     let mut server_ctx_builder = TlsContextBuilder::new_dtls();
-    let server_cred = {
-        let mut builder = TlsCredentialBuilder::new();
-        builder
-            .with_certificate_chain(&[server_cert, ca])
-            .unwrap()
-            .with_private_key(server_key)
-            .unwrap();
-        builder.build().unwrap()
-    };
-    server_ctx_builder.with_credential(server_cred).unwrap();
-    let server_ctx = server_ctx_builder.build();
-    let server_conn = server_ctx.new_server_connection().build();
+    server_ctx_builder
+        .with_credential(super::server_credential())
+        .unwrap();
+    let server_conn = server_ctx_builder.build().new_server_connection().build();
 
     let mut client_ctx_builder = TlsContextBuilder::new_dtls();
-    let ca = X509Certificate::parse_one_from_pem(CA).unwrap();
-    let mut cert_store = X509StoreBuilder::new();
-    cert_store
-        .set_trust(Trust::SslServer)
-        .unwrap()
-        .add_cert(ca)
-        .unwrap();
-    let cert_store = cert_store.build();
-    client_ctx_builder.with_certificate_store(&cert_store);
-    let client_ctx = client_ctx_builder.build();
-    let client_conn = client_ctx.new_client_connection().build();
+    client_ctx_builder.with_certificate_store(&super::client_cert_store());
+    let client_conn = client_ctx_builder.build().new_client_connection().build();
 
     (server_conn, client_conn)
 }
diff --git a/rust/bssl-tls-tokio/src/tests/transport.rs b/rust/bssl-tls-tokio/src/tests/transport.rs
index 5301825..519d810 100644
--- a/rust/bssl-tls-tokio/src/tests/transport.rs
+++ b/rust/bssl-tls-tokio/src/tests/transport.rs
@@ -32,9 +32,16 @@
 };
 
 fn dumb_server_client() -> (TlsConnection<Server>, TlsConnection<Client>) {
-    let (server_builder, client_builder) = super::tls_ctx_builders();
+    use bssl_tls::context::TlsContextBuilder;
+
+    let mut server_builder = TlsContextBuilder::new_tls();
+    server_builder.with_credential(super::server_credential()).unwrap();
     let server_conn = server_builder.build().new_server_connection().build();
+
+    let mut client_builder = TlsContextBuilder::new_tls();
+    client_builder.with_certificate_store(&super::client_cert_store());
     let client_conn = client_builder.build().new_client_connection().build();
+
     (server_conn, client_conn)
 }
 
diff --git a/rust/bssl-tls-tokio/tests/hyper.rs b/rust/bssl-tls-tokio/tests/hyper.rs
index 27b52a7..294026a 100644
--- a/rust/bssl-tls-tokio/tests/hyper.rs
+++ b/rust/bssl-tls-tokio/tests/hyper.rs
@@ -16,11 +16,14 @@
 
 use bssl_tls::alpn::H2;
 use bssl_tls::context::TlsContextBuilder;
-use bssl_tls::credentials::{Certificate, TlsCredentialBuilder};
-use bssl_tls_tokio::TokioTlsExt;
+use bssl_tls::credentials::{Certificate, TlsCredential, TlsCredentialBuilder};
 use bssl_tls_tokio::hyper::HyperBsslConnector;
+use bssl_tls_tokio::TokioTlsExt;
 use bssl_x509::{
-    certificates::X509Certificate, keys::PrivateKey, params::Trust, store::X509StoreBuilder,
+    certificates::X509Certificate,
+    keys::PrivateKey,
+    params::Trust,
+    store::{X509Store, X509StoreBuilder},
 };
 use hyper::body::{Body, Bytes, Frame};
 use hyper::service::service_fn;
@@ -36,6 +39,29 @@
 const RSA_SERVER_CERT: &[u8] = include_bytes!("../../test-data/BoringSSLServerTest-RSA.crt");
 const RSA_SERVER_KEY: &[u8] = include_bytes!("../../test-data/BoringSSLServerTest-RSA.key");
 
+fn server_credential() -> TlsCredential {
+    let ca = Certificate::parse_one_from_pem(CA, None).unwrap();
+    let server_cert = Certificate::parse_one_from_pem(RSA_SERVER_CERT, None).unwrap();
+    let server_key = PrivateKey::from_pem(RSA_SERVER_KEY, || unreachable!()).unwrap();
+    let mut builder = TlsCredentialBuilder::new();
+    builder
+        .with_certificate_chain(&[server_cert, ca])
+        .unwrap()
+        .with_private_key(server_key)
+        .unwrap();
+    builder.build().unwrap()
+}
+
+fn client_cert_store() -> X509Store {
+    let mut store = X509StoreBuilder::new();
+    store
+        .set_trust(Trust::SslServer)
+        .unwrap()
+        .add_cert(X509Certificate::parse_one_from_pem(CA).unwrap())
+        .unwrap();
+    store.build()
+}
+
 /// A body that yields a single data frame, or is empty.
 struct SimpleBody(Option<Bytes>);
 
@@ -94,20 +120,10 @@
     let addr = listener.local_addr().unwrap();
 
     // Set up the TLS server context.
-    let ca = Certificate::parse_one_from_pem(CA, None).unwrap();
-    let server_cert = Certificate::parse_one_from_pem(RSA_SERVER_CERT, None).unwrap();
-    let server_key = PrivateKey::from_pem(RSA_SERVER_KEY, || unreachable!()).unwrap();
     let mut server_ctx_builder = TlsContextBuilder::new_tls();
-    let server_cred = {
-        let mut builder = TlsCredentialBuilder::new();
-        builder
-            .with_certificate_chain(&[server_cert, ca])
-            .unwrap()
-            .with_private_key(server_key)
-            .unwrap();
-        builder.build().unwrap()
-    };
-    server_ctx_builder.with_credential(server_cred).unwrap();
+    server_ctx_builder
+        .with_credential(server_credential())
+        .unwrap();
     server_ctx_builder.set_alpn_protocols([H2]).unwrap();
     let acceptor = server_ctx_builder.build_tokio_acceptor();
 
@@ -134,13 +150,7 @@
         let mock_connector = MockTcpConnector { addr };
 
         let mut client_ctx_builder = TlsContextBuilder::new_tls();
-        let mut cert_store = X509StoreBuilder::new();
-        cert_store
-            .set_trust(Trust::SslServer)
-            .unwrap()
-            .add_cert(X509Certificate::parse_one_from_pem(CA).unwrap())
-            .unwrap();
-        client_ctx_builder.with_certificate_store(&cert_store.build());
+        client_ctx_builder.with_certificate_store(&client_cert_store());
         client_ctx_builder.set_alpn_protocols([H2]).unwrap();
         let connector = client_ctx_builder.build_tokio_connector();