From a6501594a5c2bc695c7c481adfd02cadbd972c88 Mon Sep 17 00:00:00 2001 From: Devrandom Date: Thu, 5 May 2022 18:04:55 +0200 Subject: [PATCH] Improve ShutdownScript::new_witness_program --- lightning/src/ln/channel.rs | 4 ++-- lightning/src/ln/script.rs | 20 +++++--------------- lightning/src/ln/shutdown_tests.rs | 4 ++-- 3 files changed, 9 insertions(+), 19 deletions(-) diff --git a/lightning/src/ln/channel.rs b/lightning/src/ln/channel.rs index 8ebaa96b6..2e5cdd4a8 100644 --- a/lightning/src/ln/channel.rs +++ b/lightning/src/ln/channel.rs @@ -6363,8 +6363,8 @@ mod tests { use bitcoin::hashes::sha256::Hash as Sha256; use bitcoin::hashes::Hash; use bitcoin::hash_types::WPubkeyHash; - use core::num::NonZeroU8; use bitcoin::bech32::u5; + use bitcoin::util::address::WitnessVersion; use prelude::*; struct TestFeeEstimator { @@ -6428,7 +6428,7 @@ mod tests { fn upfront_shutdown_script_incompatibility() { let features = InitFeatures::known().clear_shutdown_anysegwit(); let non_v0_segwit_shutdown_script = - ShutdownScript::new_witness_program(NonZeroU8::new(16).unwrap(), &[0, 40]).unwrap(); + ShutdownScript::new_witness_program(WitnessVersion::V16, &[0, 40]).unwrap(); let seed = [42; 32]; let network = Network::Testnet; diff --git a/lightning/src/ln/script.rs b/lightning/src/ln/script.rs index 218b48276..7abe3060f 100644 --- a/lightning/src/ln/script.rs +++ b/lightning/src/ln/script.rs @@ -5,13 +5,13 @@ use bitcoin::blockdata::script::{Builder, Script}; use bitcoin::hashes::Hash; use bitcoin::hash_types::{WPubkeyHash, WScriptHash}; use bitcoin::secp256k1::PublicKey; +use bitcoin::util::address::WitnessVersion; use ln::features::InitFeatures; use ln::msgs::DecodeError; use util::ser::{Readable, Writeable, Writer}; use core::convert::TryFrom; -use core::num::NonZeroU8; use io; /// A script pubkey for shutting down a channel as defined by [BOLT #2]. @@ -84,9 +84,9 @@ impl ShutdownScript { /// # Errors /// /// This function may return an error if `program` is invalid for the segwit `version`. - pub fn new_witness_program(version: NonZeroU8, program: &[u8]) -> Result { + pub fn new_witness_program(version: WitnessVersion, program: &[u8]) -> Result { let script = Builder::new() - .push_int(version.get().into()) + .push_int(version as i64) .push_slice(&program) .into_script(); Self::try_from(script) @@ -180,7 +180,6 @@ mod shutdown_script_tests { use bitcoin::secp256k1::{PublicKey, SecretKey}; use ln::features::InitFeatures; use core::convert::TryFrom; - use core::num::NonZeroU8; use bitcoin::util::address::WitnessVersion; fn pubkey() -> bitcoin::util::key::PublicKey { @@ -239,9 +238,7 @@ mod shutdown_script_tests { #[test] fn generates_segwit_from_non_v0_witness_program() { let witness_program = Script::new_witness_program(WitnessVersion::V16, &[0; 40]); - - let version = NonZeroU8::new(WitnessVersion::V16 as u8).unwrap(); - let shutdown_script = ShutdownScript::new_witness_program(version, &[0; 40]).unwrap(); + let shutdown_script = ShutdownScript::new_witness_program(WitnessVersion::V16, &[0; 40]).unwrap(); assert!(shutdown_script.is_compatible(&InitFeatures::known())); assert!(!shutdown_script.is_compatible(&InitFeatures::known().clear_shutdown_anysegwit())); assert_eq!(shutdown_script.into_inner(), witness_program); @@ -253,12 +250,6 @@ mod shutdown_script_tests { assert!(ShutdownScript::try_from(op_return).is_err()); } - #[test] - fn fails_from_invalid_segwit_version() { - let version = NonZeroU8::new(17).unwrap(); - assert!(ShutdownScript::new_witness_program(version, &[0; 40]).is_err()); - } - #[test] fn fails_from_invalid_segwit_v0_witness_program() { let witness_program = Script::new_witness_program(WitnessVersion::V0, &[0; 2]); @@ -270,7 +261,6 @@ mod shutdown_script_tests { let witness_program = Script::new_witness_program(WitnessVersion::V16, &[0; 42]); assert!(ShutdownScript::try_from(witness_program).is_err()); - let version = NonZeroU8::new(WitnessVersion::V16 as u8).unwrap(); - assert!(ShutdownScript::new_witness_program(version, &[0; 42]).is_err()); + assert!(ShutdownScript::new_witness_program(WitnessVersion::V16, &[0; 42]).is_err()); } } diff --git a/lightning/src/ln/shutdown_tests.rs b/lightning/src/ln/shutdown_tests.rs index 42c0e3b2f..59004547b 100644 --- a/lightning/src/ln/shutdown_tests.rs +++ b/lightning/src/ln/shutdown_tests.rs @@ -26,11 +26,11 @@ use util::config::UserConfig; use bitcoin::blockdata::script::Builder; use bitcoin::blockdata::opcodes; use bitcoin::network::constants::Network; +use bitcoin::util::address::WitnessVersion; use regex; use core::default::Default; -use core::num::NonZeroU8; use ln::functional_test_utils::*; use ln::msgs::OptionalField::Present; @@ -654,7 +654,7 @@ fn test_unsupported_anysegwit_shutdown_script() { // Check that using an unsupported shutdown script fails and a supported one succeeds. let supported_shutdown_script = chanmon_cfgs[1].keys_manager.get_shutdown_scriptpubkey(); let unsupported_shutdown_script = - ShutdownScript::new_witness_program(NonZeroU8::new(16).unwrap(), &[0, 40]).unwrap(); + ShutdownScript::new_witness_program(WitnessVersion::V16, &[0, 40]).unwrap(); chanmon_cfgs[1].keys_manager .expect(OnGetShutdownScriptpubkey { returns: unsupported_shutdown_script.clone() }) .expect(OnGetShutdownScriptpubkey { returns: supported_shutdown_script }); -- 2.39.5