belt-kwp: add wrap_key_in_place and unwrap_key_in_place methods

belt-kwp/CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+### Added
+- `wrap_key_in_place` and `unwrap_key_in_place` methods ([#98])
+
+[#98]: https://github.com/RustCrypto/key-wraps/pull/98
+
 ## 0.2.0 (2026-05-27)
 ### Added
 - Optional `zeroize` support ([#88])

belt-kwp/src/lib.rs

@@ -62,14 +62,37 @@ impl BeltKwp {
                 expected: x.len() + IV_LEN,
             });
         }
-        let out = &mut out[..out_len];
 
-        let (l, r) = out.split_at_mut(x.len());
-        l.copy_from_slice(x);
-        r.copy_from_slice(iv);
+        out[..x.len()].copy_from_slice(x);
+        self.wrap_key_in_place(out, x.len(), iv)
+    }
+
+    /// Wrap the key held in `buf[..key_len]` with given `iv` in place and
+    /// return the wrapped key, i.e. `buf[..key_len + IV_LEN]`.
+    ///
+    /// `key_len` must be at least 16 bytes.
+    /// Size of `buf` must be bigger or equal to `key_len` + [`IV_LEN`].
+    #[inline]
+    pub fn wrap_key_in_place<'a>(
+        &self,
+        buf: &'a mut [u8],
+        key_len: usize,
+        iv: &[u8; IV_LEN],
+    ) -> Result<&'a [u8], Error> {
+        if key_len < 16 {
+            return Err(Error::InvalidDataSize);
+        }
+
+        let out_len = key_len.checked_add(IV_LEN).ok_or(Error::InvalidDataSize)?;
+        if buf.len() < out_len {
+            return Err(Error::InvalidOutputSize { expected: out_len });
+        }
+
+        let buf = &mut buf[..out_len];
+        buf[key_len..].copy_from_slice(iv);
 
-        belt_wblock_enc(out, &self.key).map_err(|_| Error::InvalidDataSize)?;
-        Ok(out)
+        belt_wblock_enc(buf, &self.key).map_err(|_| Error::InvalidDataSize)?;
+        Ok(buf)
     }
 
     /// Wrap fixed-size key `x` with given `iv` and return resulting array.
@@ -126,9 +149,28 @@ impl BeltKwp {
         let out = &mut out[..y.len()];
         out.copy_from_slice(y);
 
-        belt_wblock_dec(out, &self.key).map_err(|_| Error::InvalidDataSize)?;
+        self.unwrap_key_in_place(out, iv)
+    }
+
+    /// Unwrap the key held in `buf` with given `iv` in place and return the
+    /// unwrapped key, i.e. `buf[..buf.len() - IV_LEN]`.
+    ///
+    /// Size of `buf` must be bigger or equal to 32 bytes. `buf` is zeroized if
+    /// the integrity check does not pass.
+    #[inline]
+    pub fn unwrap_key_in_place<'a>(
+        &self,
+        buf: &'a mut [u8],
+        iv: &[u8; IV_LEN],
+    ) -> Result<&'a [u8], Error> {
+        if buf.len() < 32 {
+            return Err(Error::InvalidDataSize);
+        }
+
+        belt_wblock_dec(buf, &self.key).map_err(|_| Error::InvalidDataSize)?;
 
-        let (key, rem) = out.split_at_mut(y.len() - IV_LEN);
+        let key_len = buf.len() - IV_LEN;
+        let (key, rem) = buf.split_at_mut(key_len);
 
         let calc_iv = u128::from_ne_bytes(rem.try_into().unwrap());
         let expected_iv = u128::from_ne_bytes(*iv);

belt-kwp/tests/mod.rs

@@ -1,50 +1,57 @@
 //! Test vectors from STB 4.101.31-2020 (section A.10, tables A.21-A.22):
 //! https://apmi.bsu.by/assets/files/std/belt-spec371.pdf
-use belt_kwp::{BeltKwp, KeyInit, cipher::consts::U32};
+use belt_kwp::{BeltKwp, IV_LEN, KeyInit, cipher::consts::U32};
 use hex_literal::hex;
 
-#[test]
-fn belt_kwp() {
-    // Table A.21
-    let x1 = hex!("B194BAC8 0A08F53B 366D008E 584A5DE4 8504FA9D 1BB6C7AC 252E72C2 02FDCE0D");
-    let i1 = hex!("5BE3D612 17B96181 FE6786AD 716B890B");
-    let k1 = hex!("E9DEE72C 8F0C0FA6 2DDB49F4 6F739647 06075316 ED247A37 39CBA383 03A98BF6");
-    let y1 = hex!(
+/// Table A.21: key `x`, header `i`, wrapping key `k`, and wrapped key `y`.
+const A21: ([u8; 32], [u8; IV_LEN], [u8; 32], [u8; 48]) = (
+    hex!("B194BAC8 0A08F53B 366D008E 584A5DE4 8504FA9D 1BB6C7AC 252E72C2 02FDCE0D"),
+    hex!("5BE3D612 17B96181 FE6786AD 716B890B"),
+    hex!("E9DEE72C 8F0C0FA6 2DDB49F4 6F739647 06075316 ED247A37 39CBA383 03A98BF6"),
+    hex!(
         "49A38EE1 08D6C742 E52B774F 00A6EF98 B106CBD1 3EA4FB06 80323051 BC04DF76"
         "E487B055 C69BCF54 1176169F 1DC9F6C8"
-    );
-
-    // Table A.22
-    let x2 = hex!("92632EE0 C21AD9E0 9A39343E 5C07DAA4 889B03F2 E6847EB1 52EC99F7 A4D9F154");
-    let i2 = hex!("B5EF68D8 E4A39E56 7153DE13 D72254EE");
-    let k2 = hex!("92BD9B1C E5D14101 5445FBC9 5E4D0EF2 682080AA 227D642F 2687F934 90405511");
-    let y2 = hex!(
+    ),
+);
+
+/// Table A.22, laid out like [`A21`].
+const A22: ([u8; 32], [u8; IV_LEN], [u8; 32], [u8; 48]) = (
+    hex!("92632EE0 C21AD9E0 9A39343E 5C07DAA4 889B03F2 E6847EB1 52EC99F7 A4D9F154"),
+    hex!("B5EF68D8 E4A39E56 7153DE13 D72254EE"),
+    hex!("92BD9B1C E5D14101 5445FBC9 5E4D0EF2 682080AA 227D642F 2687F934 90405511"),
+    hex!(
         "E12BDC1A E28257EC 703FCCF0 95EE8DF1 C1AB7638 9FE678CA F7C6F860 D5BB9C4F"
         "F33C657B 637C306A DD4EA779 9EB23D31"
-    );
+    ),
+);
 
+#[test]
+fn belt_kwp() {
     let mut buf = [0u8; 48];
-    let kw = BeltKwp::new(&k1.into());
-
-    let res = kw.wrap_key(&x1, &i1, &mut buf).unwrap();
-    assert_eq!(y1, res);
-    let res = kw.unwrap_key(&y1, &i1, &mut buf).unwrap();
-    assert_eq!(x1, res);
 
-    let res = kw.wrap_fixed_key::<U32>(&x1.into(), &i1);
-    assert_eq!(y1, res.0);
-    let res = kw.unwrap_fixed_key::<U32>(&res, &i1).unwrap();
-    assert_eq!(x1, res.0);
+    for (x, i, k, y) in [A21, A22] {
+        let kw = BeltKwp::new(&k.into());
 
-    let kw = BeltKwp::new(&k2.into());
+        assert_eq!(kw.wrap_key(&x, &i, &mut buf).unwrap(), y);
+        assert_eq!(kw.unwrap_key(&y, &i, &mut buf).unwrap(), x);
 
-    let res = kw.wrap_key(&x2, &i2, &mut buf).unwrap();
-    assert_eq!(y2, res);
-    let res = kw.unwrap_key(&y2, &i2, &mut buf).unwrap();
-    assert_eq!(x2, res);
+        let wrapped = kw.wrap_fixed_key::<U32>(&x.into(), &i);
+        assert_eq!(wrapped.0, y);
+        assert_eq!(kw.unwrap_fixed_key::<U32>(&wrapped, &i).unwrap().0, x);
+    }
+}
 
-    let res = kw.wrap_fixed_key::<U32>(&x2.into(), &i2);
-    assert_eq!(y2, res.0);
-    let res = kw.unwrap_fixed_key::<U32>(&res, &i2).unwrap();
-    assert_eq!(x2, res.0);
+#[test]
+fn belt_kwp_in_place() {
+    for (x, i, k, y) in [A21, A22] {
+        let kw = BeltKwp::new(&k.into());
+
+        // `buf` may be longer than the wrapped key; the tail is left alone.
+        let mut buf = [0xFFu8; 64];
+        buf[..x.len()].copy_from_slice(&x);
+
+        assert_eq!(kw.wrap_key_in_place(&mut buf, x.len(), &i).unwrap(), y);
+        assert_eq!(buf[y.len()..], [0xFF; 16]);
+        assert_eq!(kw.unwrap_key_in_place(&mut buf[..y.len()], &i).unwrap(), x);
+    }
 }