From 6066acad46ad0fcc0536fe4719b4f26998519eea Mon Sep 17 00:00:00 2001
From: Al Liu <scygliu1@gmail.com>
Date: Tue, 11 Jun 2024 18:06:49 +0800
Subject: [PATCH] Fix two test cases (#1129)

---
 Cargo.lock                          |  2 +-
 crates/core/Cargo.toml              |  2 +-
 crates/core/src/config/secret.rs    |  8 +++++++-
 crates/core/src/node.rs             | 19 ++++++++++++-------
 crates/core/src/transport/crypto.rs | 20 +++++++++++---------
 5 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8693f6f7c..8a117d6cc 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1541,7 +1541,7 @@ dependencies = [
  "parking_lot",
  "pav_regression",
  "pico-args",
- "pkcs1",
+ "pkcs8",
  "rand",
  "redb",
  "reqwest",
diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml
index a4b1a6f3b..1e7dd4a07 100644
--- a/crates/core/Cargo.toml
+++ b/crates/core/Cargo.toml
@@ -65,7 +65,7 @@ wasmer = { features = ["sys"], workspace = true }
 xz2 = { version = "0.1" }
 reqwest = { version = "0.12", features = ["json"] }
 rsa = { version = "0.9", features = ["serde", "pem"] }
-pkcs1 = { version = "0.7", features = ["std", "pem"] }
+pkcs8 = { version = "0.10", features = ["std", "pem"] }
 
 # Tracing deps
 opentelemetry = "0.22"
diff --git a/crates/core/src/config/secret.rs b/crates/core/src/config/secret.rs
index 05253e70c..e07119ed8 100644
--- a/crates/core/src/config/secret.rs
+++ b/crates/core/src/config/secret.rs
@@ -245,7 +245,13 @@ mod tests {
 
         // write secrets to files
         transport_keypair_file
-            .write_all(transport_keypair.secret().to_bytes().unwrap().as_slice())
+            .write_all(
+                transport_keypair
+                    .secret()
+                    .to_pkcs8_pem()
+                    .unwrap()
+                    .as_slice(),
+            )
             .unwrap();
         nonce_file.write_all(&nonce).unwrap();
         cipher_file.write_all(&cipher).unwrap();
diff --git a/crates/core/src/node.rs b/crates/core/src/node.rs
index 37e275682..fb28dd006 100644
--- a/crates/core/src/node.rs
+++ b/crates/core/src/node.rs
@@ -1014,7 +1014,7 @@ impl Display for PeerId {
 
 #[cfg(test)]
 mod tests {
-    use std::net::Ipv4Addr;
+    use std::net::{Ipv4Addr, Ipv6Addr};
 
     use super::*;
 
@@ -1022,12 +1022,17 @@ mod tests {
     async fn test_hostname_resolution() {
         let addr = Address::Hostname("localhost".to_string());
         let socket_addr = NodeConfig::parse_socket_addr(&addr).await.unwrap();
-        assert_eq!(
-            socket_addr,
-            SocketAddr::new(
-                IpAddr::V4(Ipv4Addr::LOCALHOST),
-                crate::config::default_http_gateway_port()
-            )
+        assert!(
+            socket_addr
+                == SocketAddr::new(
+                    IpAddr::V4(Ipv4Addr::LOCALHOST),
+                    crate::config::default_http_gateway_port()
+                )
+                || socket_addr
+                    == SocketAddr::new(
+                        IpAddr::V6(Ipv6Addr::LOCALHOST),
+                        crate::config::default_http_gateway_port()
+                    )
         );
 
         let addr = Address::Hostname("google.com".to_string());
diff --git a/crates/core/src/transport/crypto.rs b/crates/core/src/transport/crypto.rs
index ed24ae328..70964b98a 100644
--- a/crates/core/src/transport/crypto.rs
+++ b/crates/core/src/transport/crypto.rs
@@ -1,6 +1,5 @@
-use pkcs1::{EncodeRsaPrivateKey, EncodeRsaPublicKey};
 use rand::rngs::OsRng;
-use rsa::{Pkcs1v15Encrypt, RsaPrivateKey, RsaPublicKey};
+use rsa::{pkcs8, Pkcs1v15Encrypt, RsaPrivateKey, RsaPublicKey};
 use serde::{Deserialize, Serialize};
 
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
@@ -60,7 +59,9 @@ impl TransportPublicKey {
 
 impl std::fmt::Display for TransportPublicKey {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let encoded = self.0.to_pkcs1_der().map_err(|_| std::fmt::Error)?;
+        use pkcs8::EncodePublicKey;
+
+        let encoded = self.0.to_public_key_der().map_err(|_| std::fmt::Error)?;
         write!(
             f,
             "{}",
@@ -85,19 +86,20 @@ pub(crate) struct TransportSecretKey(RsaPrivateKey);
 
 impl TransportSecretKey {
     pub fn decrypt(&self, data: &[u8]) -> rsa::Result<Vec<u8>> {
-        let padding = Pkcs1v15Encrypt;
-        self.0.decrypt(padding, data)
+        self.0.decrypt(Pkcs1v15Encrypt, data)
     }
 
-    pub fn to_bytes(&self) -> Result<Vec<u8>, pkcs1::Error> {
+    pub fn to_pkcs8_pem(&self) -> Result<Vec<u8>, pkcs8::Error> {
+        use pkcs8::EncodePrivateKey;
+
         #[cfg(unix)]
-        let line_endings = pkcs1::LineEnding::LF;
+        let line_endings = pkcs8::LineEnding::LF;
 
         #[cfg(windows)]
-        let line_endings = pkcs1::LineEnding::CRLF;
+        let line_endings = pkcs8::LineEnding::CRLF;
 
         self.0
-            .to_pkcs1_pem(line_endings)
+            .to_pkcs8_pem(line_endings)
             .map(|s| s.as_str().as_bytes().to_vec())
     }
 }