From fdc2e90ef2e64224f4f778f2a26bbfeb02afde2b Mon Sep 17 00:00:00 2001
From: edef <edef@edef.eu>
Date: Fri, 27 Oct 2023 07:36:06 +0000
Subject: feat(tvix/nix-compat): add fixed-length base32 decoding

Change-Id: Iea9d51157c1e38e1734dc1ba752eba9b2372c61b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9850
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
---
 tvix/nix-compat/src/nixbase32.rs | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

(limited to 'tvix')

diff --git a/tvix/nix-compat/src/nixbase32.rs b/tvix/nix-compat/src/nixbase32.rs
index fae80cd4a042..d9821bc495bb 100644
--- a/tvix/nix-compat/src/nixbase32.rs
+++ b/tvix/nix-compat/src/nixbase32.rs
@@ -20,6 +20,8 @@ pub enum Nixbase32DecodeError {
     CharacterNotInAlphabet(u8),
     #[error("nonzero carry")]
     NonzeroCarry(),
+    #[error("invalid length")]
+    InvalidLength,
 }
 
 /// Returns encoded input
@@ -74,6 +76,25 @@ pub fn decode(input: impl AsRef<[u8]>) -> Result<Vec<u8>, Nixbase32DecodeError>
     let output_len = decode_len(input.len());
     let mut output: Vec<u8> = vec![0x00; output_len];
 
+    decode_inner(input, &mut output)?;
+    Ok(output)
+}
+
+pub fn decode_fixed<const K: usize>(
+    input: impl AsRef<[u8]>,
+) -> Result<[u8; K], Nixbase32DecodeError> {
+    let input = input.as_ref();
+
+    if input.len() != encode_len(K) {
+        return Err(Nixbase32DecodeError::InvalidLength);
+    }
+
+    let mut output = [0; K];
+    decode_inner(input, &mut output)?;
+    Ok(output)
+}
+
+fn decode_inner(input: &[u8], output: &mut [u8]) -> Result<(), Nixbase32DecodeError> {
     // loop over all characters in reverse, and keep the iteration count in n.
     let mut carry = 0;
     let mut mask = 0;
@@ -100,7 +121,7 @@ pub fn decode(input: impl AsRef<[u8]>) -> Result<Vec<u8>, Nixbase32DecodeError>
         return Err(Nixbase32DecodeError::NonzeroCarry());
     }
 
-    Ok(output)
+    Ok(())
 }
 
 #[cold]
@@ -160,6 +181,18 @@ mod tests {
         }
     }
 
+    #[test]
+    fn decode_fixed() {
+        assert_eq!(
+            super::decode_fixed("00bgd045z0d4icpbc2yyz4gx48ak44la").unwrap(),
+            hex!("8a12321522fd91efbd60ebb2481af88580f61600")
+        );
+        assert_eq!(
+            super::decode_fixed::<32>("00").unwrap_err(),
+            super::Nixbase32DecodeError::InvalidLength
+        );
+    }
+
     #[test]
     fn encode_len() {
         assert_eq!(super::encode_len(0), 0);
-- 
cgit 1.4.1