From d1cf43af46a4630f77fe48aee2adbf739a71b97b Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Wed, 1 Dec 2021 14:15:08 +0100
Subject: [PATCH] Move client certificate setup into epp_connect()

---
 epp-client/src/config.rs              | 38 ---------------------
 epp-client/src/connection/registry.rs | 48 +++++++++++++++++++++++----
 2 files changed, 42 insertions(+), 44 deletions(-)

diff --git a/epp-client/src/config.rs b/epp-client/src/config.rs
index 7a34dd5..8559371 100644
--- a/epp-client/src/config.rs
+++ b/epp-client/src/config.rs
@@ -38,15 +38,9 @@
 //! ```
 
 use std::collections::HashMap;
-use std::fs::File;
-use std::io::{BufReader, Seek, SeekFrom};
 
-use rustls::{Certificate, PrivateKey};
-use rustls_pemfile;
 use serde::{Deserialize, Serialize};
 
-use crate::error::Error;
-
 /// Paths to the client certificate and client key PEM files
 #[derive(Serialize, Deserialize, Debug)]
 pub struct EppClientTlsFiles {
@@ -91,36 +85,4 @@ impl EppClientConnection {
     pub fn ext_uris(&self) -> Option<&Vec<String>> {
         self.ext_uris.as_ref()
     }
-    /// Returns the parsed client certificate and private key for client TLS auth
-    pub fn tls_files(&self) -> Result<Option<(Vec<Certificate>, PrivateKey)>, Error> {
-        let (certs_file, key_file) = match &self.tls_files {
-            Some(files) => (&files.cert_chain, &files.key),
-            None => return Ok(None),
-        };
-
-        let certs = rustls_pemfile::certs(&mut BufReader::new(File::open(certs_file)?))?
-            .into_iter()
-            .map(Certificate)
-            .collect::<Vec<_>>();
-
-        let mut r = BufReader::new(File::open(key_file).unwrap());
-
-        let mut rsa_keys = rustls_pemfile::rsa_private_keys(&mut r).unwrap();
-        if rsa_keys.len() > 1 {
-            warn!("Multiple RSA keys found in PEM file {}", key_file);
-        } else if let Some(key) = rsa_keys.pop() {
-            return Ok(Some((certs, rustls::PrivateKey(key))));
-        }
-
-        r.seek(SeekFrom::Start(0))?;
-
-        let mut pkcs8_keys = rustls_pemfile::pkcs8_private_keys(&mut r).unwrap();
-        if pkcs8_keys.len() > 1 {
-            warn!("Multiple PKCS8 keys found in PEM file {}", key_file);
-        } else if let Some(key) = pkcs8_keys.pop() {
-            return Ok(Some((certs, rustls::PrivateKey(key))));
-        }
-
-        Err(Error::Other("No private key found in PEM file".to_owned()))
-    }
 }
diff --git a/epp-client/src/connection/registry.rs b/epp-client/src/connection/registry.rs
index e2e18eb..6991ba8 100644
--- a/epp-client/src/connection/registry.rs
+++ b/epp-client/src/connection/registry.rs
@@ -1,10 +1,15 @@
 //! Manages registry connections and reading/writing to them
 
-use rustls::{OwnedTrustAnchor, RootCertStore};
 use std::convert::TryInto;
+use std::fs::File;
+use std::io::{BufReader, Seek, SeekFrom};
 use std::sync::Arc;
 use std::{error::Error, io as stdio, net::ToSocketAddrs};
 use std::{io, str, u32};
+
+use rustls::{Certificate, PrivateKey};
+use rustls::{OwnedTrustAnchor, RootCertStore};
+use rustls_pemfile;
 use tokio::{io::AsyncReadExt, io::AsyncWriteExt, net::TcpStream};
 use tokio_rustls::{client::TlsStream, rustls::ClientConfig, TlsConnector};
 
@@ -137,11 +142,42 @@ pub async fn epp_connect(
         .with_safe_defaults()
         .with_root_certificates(roots);
 
-    let config = match registry_creds.tls_files()? {
-        Some((cert_chain, key)) => match builder.with_single_cert(cert_chain, key) {
-            Ok(config) => config,
-            Err(e) => return Err(format!("Failed to set client TLS credentials: {}", e).into()),
-        },
+    let config = match &registry_creds.tls_files {
+        Some(files) => {
+            let (certs_file, key_file) = (&files.cert_chain, &files.key);
+            let certs = rustls_pemfile::certs(&mut BufReader::new(File::open(certs_file)?))?
+                .into_iter()
+                .map(Certificate)
+                .collect::<Vec<_>>();
+
+            let mut key;
+            let mut r = BufReader::new(File::open(key_file).unwrap());
+            let mut rsa_keys = rustls_pemfile::rsa_private_keys(&mut r).unwrap();
+            if rsa_keys.len() > 1 {
+                warn!("Multiple RSA keys found in PEM file {}", key_file);
+            }
+            key = rsa_keys.pop();
+
+            if key.is_none() {
+                r.seek(SeekFrom::Start(0))?;
+                let mut pkcs8_keys = rustls_pemfile::pkcs8_private_keys(&mut r).unwrap();
+                if pkcs8_keys.len() > 1 {
+                    warn!("Multiple PKCS8 keys found in PEM file {}", key_file);
+                }
+                key = pkcs8_keys.pop();
+            }
+
+            match key {
+                Some(key) => builder
+                    .with_single_cert(certs, PrivateKey(key))
+                    .map_err(|e| error::Error::Other(e.to_string()))?,
+                None => {
+                    return Err(error::Error::Other(
+                        "No private key found in PEM file".to_owned(),
+                    ))
+                }
+            }
+        }
         None => builder.with_no_client_auth(),
     };