Sign and verify P2WPKH (#26)

SchnobiTobi · web-flow · commit 9abefabc04c0 · 2024-08-22T16:22:26.000+02:00
diff --git a/src/error.rs b/src/error.rs
@@ -10,7 +10,7 @@ pub enum Error {
   },
   #[snafu(display("Failed to parse private key"))]
   PrivateKeyParse { source: bitcoin::key::Error },
-  #[snafu(display("Unsuported address `{address}`, only P2TR allowed"))]
+  #[snafu(display("Unsuported address `{address}`, only P2TR or P2WPKH allowed"))]
   UnsupportedAddress { address: String },
   #[snafu(display("Decode error for signature `{signature}`"))]
   SignatureDecode {
@@ -60,4 +60,10 @@ pub enum Error {
   SigHashTypeUnsupported { sighash_type: String },
   #[snafu(display("Not key path spend"))]
   NotKeyPathSpend,
+  #[snafu(display("Invalid public key"))]
+  InvalidPublicKey,
+  #[snafu(display("Invalid witness"))]
+  InvalidWitness,
+  #[snafu(display("Public key does not match"))]
+  PublicKeyMismatch,
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -2,7 +2,7 @@ use {
   base64::{engine::general_purpose, Engine},
   bitcoin::{
     absolute::LockTime,
-    address::AddressType,
+    address::Payload,
     blockdata::script,
     consensus::Decodable,
     consensus::Encodable,
@@ -13,7 +13,8 @@ use {
     secp256k1::{self, schnorr::Signature, Message, Secp256k1, XOnlyPublicKey},
     sighash::{self, SighashCache, TapSighashType},
     transaction::Version,
-    Address, Amount, OutPoint, PrivateKey, ScriptBuf, Sequence, Transaction, TxIn, TxOut, Witness,
+    Address, Amount, EcdsaSighashType, OutPoint, PrivateKey, PublicKey, ScriptBuf, Sequence,
+    Transaction, TxIn, TxOut, Witness,
   },
   bitcoin_hashes::{sha256, Hash},
   error::Error,
@@ -237,7 +238,7 @@ mod tests {
       "3B5fQsEXEaV8v6U3ejYc8XaKXAkyQj2MjV",
       "",
       "AkcwRAIgM2gBAQqvZX15ZiysmKmQpDrG83avLIT492QBzLnQIxYCIBaTpOaD20qRlEylyxFSeEA2ba9YOixpX8z46TSDtS40ASECx/EgAxlkQpQ9hYjgGu6EBCPMVPwVIVJqO4XCsMvViHI=").unwrap_err().to_string(),
-      "Unsuported address `3B5fQsEXEaV8v6U3ejYc8XaKXAkyQj2MjV`, only P2TR allowed"
+      "Unsuported address `3B5fQsEXEaV8v6U3ejYc8XaKXAkyQj2MjV`, only P2TR or P2WPKH allowed"
     )
   }
 
@@ -263,4 +264,88 @@ mod tests {
       "Decode error for signature `AkcwRAIgM2gBAQqvZX15ZiysmKmQpDrG83avLIT492QBzLnQIxYCIBaTpOaD20qRlEylyxFSeEA2ba9YOixpX8z46TSDtS40ASECx/EgAxlkQpQ9hYjgGu6EBCPMVPwVIVJqO4XCsMvViH`"
     )
   }
+
+  #[test]
+  fn simple_verify_and_falsify_p2wpkh() {
+    assert!(
+      verify_simple_encoded(
+        SEGWIT_ADDRESS,
+        "Hello World",
+        "AkcwRAIgZRfIY3p7/DoVTty6YZbWS71bc5Vct9p9Fia83eRmw2QCICK/ENGfwLtptFluMGs2KsqoNSk89pO7F29zJLUx9a/sASECx/EgAxlkQpQ9hYjgGu6EBCPMVPwVIVJqO4XCsMvViHI="
+      ).is_ok()
+    );
+
+    assert!(
+      verify_simple_encoded(
+        SEGWIT_ADDRESS,
+        "Hello World - this should fail",
+        "AkcwRAIgZRfIY3p7/DoVTty6YZbWS71bc5Vct9p9Fia83eRmw2QCICK/ENGfwLtptFluMGs2KsqoNSk89pO7F29zJLUx9a/sASECx/EgAxlkQpQ9hYjgGu6EBCPMVPwVIVJqO4XCsMvViHI="
+      ).is_err()
+    );
+
+    assert!(
+      verify_simple_encoded(
+        SEGWIT_ADDRESS,
+        "Hello World",
+        "AkgwRQIhAOzyynlqt93lOKJr+wmmxIens//zPzl9tqIOua93wO6MAiBi5n5EyAcPScOjf1lAqIUIQtr3zKNeavYabHyR8eGhowEhAsfxIAMZZEKUPYWI4BruhAQjzFT8FSFSajuFwrDL1Yhy"
+      ).is_ok()
+    );
+
+    assert!(
+      verify_simple_encoded(
+        SEGWIT_ADDRESS,
+        "",
+        "AkcwRAIgM2gBAQqvZX15ZiysmKmQpDrG83avLIT492QBzLnQIxYCIBaTpOaD20qRlEylyxFSeEA2ba9YOixpX8z46TSDtS40ASECx/EgAxlkQpQ9hYjgGu6EBCPMVPwVIVJqO4XCsMvViHI="
+      ).is_ok()
+    );
+
+    assert!(
+      verify_simple_encoded(
+        SEGWIT_ADDRESS,
+        "fail",
+        "AkcwRAIgM2gBAQqvZX15ZiysmKmQpDrG83avLIT492QBzLnQIxYCIBaTpOaD20qRlEylyxFSeEA2ba9YOixpX8z46TSDtS40ASECx/EgAxlkQpQ9hYjgGu6EBCPMVPwVIVJqO4XCsMvViHI="
+      ).is_err()
+    );
+
+    assert!(
+      verify_simple_encoded(
+        SEGWIT_ADDRESS,
+        "",
+        "AkgwRQIhAPkJ1Q4oYS0htvyuSFHLxRQpFAY56b70UvE7Dxazen0ZAiAtZfFz1S6T6I23MWI2lK/pcNTWncuyL8UL+oMdydVgzAEhAsfxIAMZZEKUPYWI4BruhAQjzFT8FSFSajuFwrDL1Yhy"
+      ).is_ok()
+    );
+  }
+
+  #[test]
+  fn simple_sign_p2wpkh() {
+    assert_eq!(
+      sign_simple_encoded(SEGWIT_ADDRESS, "Hello World", WIF_PRIVATE_KEY).unwrap(),
+      "AkgwRQIhAOzyynlqt93lOKJr+wmmxIens//zPzl9tqIOua93wO6MAiBi5n5EyAcPScOjf1lAqIUIQtr3zKNeavYabHyR8eGhowEhAsfxIAMZZEKUPYWI4BruhAQjzFT8FSFSajuFwrDL1Yhy"
+    );
+
+    assert_eq!(
+      sign_simple_encoded(SEGWIT_ADDRESS, "", WIF_PRIVATE_KEY).unwrap(),
+      "AkgwRQIhAPkJ1Q4oYS0htvyuSFHLxRQpFAY56b70UvE7Dxazen0ZAiAtZfFz1S6T6I23MWI2lK/pcNTWncuyL8UL+oMdydVgzAEhAsfxIAMZZEKUPYWI4BruhAQjzFT8FSFSajuFwrDL1Yhy"
+    );
+  }
+
+  #[test]
+  fn roundtrip_p2wpkh_simple() {
+    assert!(verify_simple_encoded(
+      SEGWIT_ADDRESS,
+      "Hello World",
+      &sign_simple_encoded(SEGWIT_ADDRESS, "Hello World", WIF_PRIVATE_KEY).unwrap()
+    )
+    .is_ok());
+  }
+
+  #[test]
+  fn roundtrip_p2wpkh_full() {
+    assert!(verify_full_encoded(
+      SEGWIT_ADDRESS,
+      "Hello World",
+      &sign_full_encoded(SEGWIT_ADDRESS, "Hello World", WIF_PRIVATE_KEY).unwrap()
+    )
+    .is_ok());
+  }
 }
diff --git a/src/sign.rs b/src/sign.rs
@@ -1,7 +1,6 @@
 use super::*;
 
-/// Signs the BIP-322 simple from encoded values, i.e. address encoding, message string and
-/// WIF private key string. Returns a base64 encoded witness stack.
+/// Signs the BIP-322 simple from spec-compliant string encodings.
 pub fn sign_simple_encoded(address: &str, message: &str, wif_private_key: &str) -> Result<String> {
   let address = Address::from_str(address)
     .context(error::AddressParse { address })?
@@ -20,8 +19,7 @@ pub fn sign_simple_encoded(address: &str, message: &str, wif_private_key: &str)
   Ok(general_purpose::STANDARD.encode(buffer))
 }
 
-/// Signs the BIP-322 full from encoded values, i.e. address encoding, message string and
-/// WIF private key string. Returns a base64 encoded transaction.
+/// Signs the BIP-322 full from spec-compliant string encodings.
 pub fn sign_full_encoded(address: &str, message: &str, wif_private_key: &str) -> Result<String> {
   let address = Address::from_str(address)
     .context(error::AddressParse { address })?
@@ -54,32 +52,86 @@ pub fn sign_full(
   message: &[u8],
   private_key: PrivateKey,
 ) -> Result<Transaction> {
-  if let bitcoin::address::Payload::WitnessProgram(witness_program) = address.payload() {
-    if witness_program.version().to_num() != 1 {
+  let to_spend = create_to_spend(address, message)?;
+  let mut to_sign = create_to_sign(&to_spend, None)?;
+
+  let witness =
+    if let bitcoin::address::Payload::WitnessProgram(witness_program) = address.payload() {
+      let version = witness_program.version().to_num();
+      let program_len = witness_program.program().len();
+
+      match version {
+        0 => {
+          if program_len != 20 {
+            return Err(Error::NotKeyPathSpend);
+          }
+          create_message_signature_p2wpkh(&to_spend, &to_sign, private_key)
+        }
+        1 => {
+          if program_len != 32 {
+            return Err(Error::NotKeyPathSpend);
+          }
+          create_message_signature_taproot(&to_spend, &to_sign, private_key)
+        }
+        _ => {
+          return Err(Error::UnsupportedAddress {
+            address: address.to_string(),
+          })
+        }
+      }
+    } else {
       return Err(Error::UnsupportedAddress {
         address: address.to_string(),
       });
-    }
-
-    if witness_program.program().len() != 32 {
-      return Err(Error::NotKeyPathSpend);
-    }
-  } else {
-    return Err(Error::UnsupportedAddress {
-      address: address.to_string(),
-    });
-  };
+    };
 
-  let to_spend = create_to_spend(address, message)?;
-  let mut to_sign = create_to_sign(&to_spend, None)?;
-
-  let witness = create_message_signature(&to_spend, &to_sign, private_key);
   to_sign.inputs[0].final_script_witness = Some(witness);
 
   to_sign.extract_tx().context(error::TransactionExtract)
 }
 
-fn create_message_signature(
+fn create_message_signature_p2wpkh(
+  to_spend_tx: &Transaction,
+  to_sign: &Psbt,
+  private_key: PrivateKey,
+) -> Witness {
+  let secp = Secp256k1::new();
+  let sighash_type = EcdsaSighashType::All;
+  let mut sighash_cache = SighashCache::new(to_sign.unsigned_tx.clone());
+
+  let sighash = sighash_cache
+    .p2wpkh_signature_hash(
+      0,
+      &to_spend_tx.output[0].script_pubkey,
+      to_spend_tx.output[0].value,
+      sighash_type,
+    )
+    .expect("signature hash should compute");
+
+  let sig = secp.sign_ecdsa(
+    &secp256k1::Message::from_digest_slice(sighash.as_ref())
+      .expect("should be cryptographically secure hash"),
+    &private_key.inner,
+  );
+
+  let witness = sighash_cache
+    .witness_mut(0)
+    .expect("getting mutable witness reference should work");
+
+  witness.push(
+    bitcoin::ecdsa::Signature {
+      sig,
+      hash_ty: sighash_type,
+    }
+    .to_vec(),
+  );
+
+  witness.push(private_key.public_key(&secp).to_bytes());
+
+  witness.to_owned()
+}
+
+fn create_message_signature_taproot(
   to_spend_tx: &Transaction,
   to_sign: &Psbt,
   private_key: PrivateKey,
@@ -88,6 +140,7 @@ fn create_message_signature(
 
   let secp = Secp256k1::new();
   let key_pair = Keypair::from_secret_key(&secp, &private_key.inner);
+
   let (x_only_public_key, _parity) = XOnlyPublicKey::from_keypair(&key_pair);
   to_sign.inputs[0].tap_internal_key = Some(x_only_public_key);
 
diff --git a/src/verify.rs b/src/verify.rs
