From 36bd00a046757dd726d99e49d0e695ab95a46aff Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Fri, 29 Mar 2024 17:52:12 +0800
Subject: [PATCH] Add p521 to magiskboot

---
 native/src/Cargo.lock      | 16 ++++++++++++++++
 native/src/Cargo.toml      |  2 ++
 native/src/boot/Cargo.toml |  2 ++
 native/src/boot/sign.rs    | 25 ++++++++++++++++++++++++-
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/native/src/Cargo.lock b/native/src/Cargo.lock
index 4472ba2d3..4efbb9ead 100644
--- a/native/src/Cargo.lock
+++ b/native/src/Cargo.lock
@@ -488,10 +488,12 @@ dependencies = [
  "cxx-gen",
  "der",
  "digest",
+ "ecdsa",
  "fdt",
  "num-traits",
  "p256",
  "p384",
+ "p521",
  "pb-rs",
  "quick-protobuf",
  "rsa",
@@ -624,6 +626,20 @@ dependencies = [
  "sha2",
 ]
 
+[[package]]
+name = "p521"
+version = "0.13.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fc9e2161f1f215afdfce23677034ae137bbd45016a880c2eb3ba8eb95f085b2"
+dependencies = [
+ "base16ct",
+ "ecdsa",
+ "elliptic-curve",
+ "primeorder",
+ "rand_core",
+ "sha2",
+]
+
 [[package]]
 name = "pb-rs"
 version = "0.10.0"
diff --git a/native/src/Cargo.toml b/native/src/Cargo.toml
index 02ba1efd3..f6654269c 100644
--- a/native/src/Cargo.toml
+++ b/native/src/Cargo.toml
@@ -18,6 +18,8 @@ sha2 = "0.10"
 digest = "0.10"
 p256 = "0.13"
 p384 = "0.13"
+p521 = "0.13"
+ecdsa = "0.16"
 rsa = "0.9"
 x509-cert = "0.2"
 der = "0.7"
diff --git a/native/src/boot/Cargo.toml b/native/src/boot/Cargo.toml
index 79807f067..f36f97402 100644
--- a/native/src/boot/Cargo.toml
+++ b/native/src/boot/Cargo.toml
@@ -23,6 +23,8 @@ sha2 = { workspace = true }
 digest = { workspace = true }
 p256 = { workspace = true }
 p384 = { workspace = true }
+p521 = { workspace = true }
+ecdsa = { workspace = true }
 rsa = { workspace = true, features = ["sha2"] }
 x509-cert = { workspace = true }
 der = { workspace = true, features = ["derive"] }
diff --git a/native/src/boot/sign.rs b/native/src/boot/sign.rs
index 3b578af63..2b9386957 100644
--- a/native/src/boot/sign.rs
+++ b/native/src/boot/sign.rs
@@ -1,6 +1,7 @@
 use der::referenced::OwnedToRef;
 use der::{Decode, DecodePem, Encode, Sequence, SliceReader};
 use digest::DynDigest;
+use ecdsa;
 use p256::ecdsa::{
     Signature as P256Signature, SigningKey as P256SigningKey, VerifyingKey as P256VerifyingKey,
 };
@@ -8,6 +9,10 @@ use p256::pkcs8::DecodePrivateKey;
 use p384::ecdsa::{
     Signature as P384Signature, SigningKey as P384SigningKey, VerifyingKey as P384VerifyingKey,
 };
+use p521::{
+    ecdsa::{Signature as P521Signature, SigningKey as P521SigningKey},
+    NistP521,
+};
 use rsa::pkcs1v15::{
     Signature as RsaSignature, SigningKey as RsaSigningKey, VerifyingKey as RsaVerifyingKey,
 };
@@ -16,7 +21,7 @@ use rsa::signature::hazmat::{PrehashSigner, PrehashVerifier};
 use rsa::signature::SignatureEncoding;
 use rsa::{RsaPrivateKey, RsaPublicKey};
 use sha1::Sha1;
-use sha2::{Sha256, Sha384};
+use sha2::{Sha256, Sha384, Sha512};
 use x509_cert::der::asn1::{OctetString, PrintableString};
 use x509_cert::der::Any;
 use x509_cert::spki::AlgorithmIdentifier;
@@ -27,6 +32,8 @@ use base::{log_err, LoggedResult, MappedFile, ResultExt, StrErr, Utf8CStr};
 
 use crate::ffi::BootImage;
 
+type P521VerifyingKey = ecdsa::VerifyingKey<NistP521>;
+
 #[allow(clippy::upper_case_acronyms)]
 pub enum SHA {
     SHA1(Sha1),
@@ -82,6 +89,7 @@ enum SigningKey {
     SHA256withRSA(RsaSigningKey<Sha256>),
     SHA256withECDSA(P256SigningKey),
     SHA384withECDSA(P384SigningKey),
+    SHA521withECDSA(P521SigningKey),
 }
 
 #[allow(clippy::large_enum_variant)]
@@ -89,6 +97,7 @@ enum VerifyingKey {
     SHA256withRSA(RsaVerifyingKey<Sha256>),
     SHA256withECDSA(P256VerifyingKey),
     SHA384withECDSA(P384VerifyingKey),
+    SHA521withECDSA(P521VerifyingKey),
 }
 
 struct Verifier {
@@ -108,6 +117,9 @@ impl Verifier {
         } else if let Ok(ec) = P384VerifyingKey::try_from(key.clone()) {
             digest = Box::<Sha384>::default();
             VerifyingKey::SHA384withECDSA(ec)
+        } else if let Ok(ec) = P521VerifyingKey::try_from(key.clone()) {
+            digest = Box::<Sha512>::default();
+            VerifyingKey::SHA521withECDSA(ec)
         } else {
             return Err(log_err!("Unsupported private key"));
         };
@@ -133,6 +145,10 @@ impl Verifier {
                 let sig = P384Signature::from_slice(signature)?;
                 key.verify_prehash(hash.as_ref(), &sig).log()
             }
+            VerifyingKey::SHA521withECDSA(key) => {
+                let sig = P521Signature::from_slice(signature)?;
+                key.verify_prehash(hash.as_ref(), &sig).log()
+            }
         };
     }
 }
@@ -154,6 +170,9 @@ impl Signer {
         } else if let Ok(ec) = P384SigningKey::from_pkcs8_der(key) {
             digest = Box::<Sha384>::default();
             SigningKey::SHA384withECDSA(ec)
+        } else if let Ok(ec) = ecdsa::SigningKey::<NistP521>::from_pkcs8_der(key) {
+            digest = Box::<Sha512>::default();
+            SigningKey::SHA521withECDSA(P521SigningKey::from(ec))
         } else {
             return Err(log_err!("Unsupported private key"));
         };
@@ -179,6 +198,10 @@ impl Signer {
                 let sig: P384Signature = key.sign_prehash(hash.as_ref())?;
                 sig.to_vec()
             }
+            SigningKey::SHA521withECDSA(key) => {
+                let sig: P521Signature = key.sign_prehash(hash.as_ref())?;
+                sig.to_vec()
+            }
         };
         Ok(v)
     }