1 files changed, 22 insertions, 25 deletions

diff --git a/btc-wire/src/segwit.rs b/btc-wire/src/segwit.rs
index e754d71..956594a 100644
--- a/btc-wire/src/segwit.rs
+++ b/btc-wire/src/segwit.rs
@@ -13,21 +13,18 @@
   You should have received a copy of the GNU Affero General Public License along with
   TALER; see the file COPYING.  If not, see <http://www.gnu.org/licenses/>
 */
-use bech32::{u5, FromBase32, ToBase32, Variant};
+use bech32::Hrp;
 use common::{rand::rngs::OsRng, rand_slice};
 use std::cmp::Ordering;
 
 /// Encode metadata into a segwit address
-pub fn encode_segwit_addr(hrp: &str, metada: &[u8; 20]) -> String {
-    // We use the version 0 with bech32 encoding
-    let mut buf = vec![u5::try_from_u8(0).unwrap()];
-    buf.extend_from_slice(&metada.to_base32());
-    bech32::encode(hrp, buf, Variant::Bech32).unwrap()
+pub fn encode_segwit_addr(hrp: Hrp, metada: &[u8; 20]) -> String {
+    bech32::segwit::encode_v0(hrp, metada).unwrap()
 }
 
 /// Encode half of a 32B key into a segwit address
 fn encode_segwit_key_half(
-    hrp: &str,
+    hrp: Hrp,
     is_first: bool,
     prefix: &[u8; 4],
     key_half: &[u8; 16],
@@ -47,7 +44,7 @@ fn encode_segwit_key_half(
 }
 
 /// Encode a 32B key into two segwit adresses
-pub fn encode_segwit_key(hrp: &str, msg: &[u8; 32]) -> [String; 2] {
+pub fn encode_segwit_key(hrp: Hrp, msg: &[u8; 32]) -> [String; 2] {
     // Generate a random prefix
     let prefix = rand_slice();
     // Split key in half;
@@ -75,20 +72,20 @@ pub fn decode_segwit_msg(segwit_addrs: &[impl AsRef<str>]) -> Result<[u8; 32], D
     let decoded: Vec<(bool, [u8; 4], [u8; 16])> = segwit_addrs
         .iter()
         .filter_map(|addr| {
-            bech32::decode(addr.as_ref()).ok().and_then(|(_, wp, _)| {
-                // Skip version
-                let pg: Vec<u8> = Vec::from_base32(&wp[1..]).unwrap();
-                if pg.len() == 20 {
-                    let mut prefix: [u8; 4] = pg[..4].try_into().unwrap();
-                    let key_half: [u8; 16] = pg[4..].try_into().unwrap();
-                    let is_first = !pg[0] & 0b1000_0000 == 0;
-                    // Clear first bit
-                    prefix[0] &= 0b0111_1111;
-                    Some((is_first, prefix, key_half))
-                } else {
-                    None
-                }
-            })
+            bech32::segwit::decode(addr.as_ref())
+                .ok()
+                .and_then(|(_, _, pg)| {
+                    if pg.len() == 20 {
+                        let mut prefix: [u8; 4] = pg[..4].try_into().unwrap();
+                        let key_half: [u8; 16] = pg[4..].try_into().unwrap();
+                        let is_first = !pg[0] & 0b1000_0000 == 0;
+                        // Clear first bit
+                        prefix[0] &= 0b0111_1111;
+                        Some((is_first, prefix, key_half))
+                    } else {
+                        None
+                    }
+                })
         })
         .collect();
 
@@ -122,7 +119,7 @@ pub fn decode_segwit_msg(segwit_addrs: &[impl AsRef<str>]) -> Result<[u8; 32], D
 }
 
 // TODO find a way to hide that function while using it in test and benchmark
-pub fn rand_addresses(hrp: &str, key: &[u8; 32]) -> Vec<String> {
+pub fn rand_addresses(hrp: Hrp, key: &[u8; 32]) -> Vec<String> {
     use common::rand::prelude::SliceRandom;
 
     let mut rng_address: Vec<String> =
@@ -149,7 +146,7 @@ mod test {
     fn test_shuffle() {
         for _ in 0..1000 {
             let key = rand_slice();
-            let mut addresses = encode_segwit_key("test", &key);
+            let mut addresses = encode_segwit_key(bech32::hrp::TB, &key);
             addresses.shuffle(&mut OsRng);
             let decoded =
                 decode_segwit_msg(&addresses.iter().map(|s| s.as_str()).collect::<Vec<&str>>())
@@ -162,7 +159,7 @@ mod test {
     fn test_shuffle_many() {
         for _ in 0..1000 {
             let key = rand_slice();
-            let addresses = rand_addresses("test", &key);
+            let addresses = rand_addresses(bech32::hrp::TB, &key);
             let decoded =
                 decode_segwit_msg(&addresses.iter().map(|s| s.as_str()).collect::<Vec<&str>>())
                     .unwrap();