chore: pad and unpad BB

AztecProtocol · May 8, 2024 · 04a3bd0 · 04a3bd0
1 parent 8c19449
commit 04a3bd0
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 8 deletions.
diff --git a/yarn-project/circuits.js/src/barretenberg/crypto/aes128/index.test.ts b/yarn-project/circuits.js/src/barretenberg/crypto/aes128/index.test.ts
@@ -9,14 +9,31 @@ describe('aes128', () => {
     aes128 = new Aes128();
   });
 
+  // PKCS#7 padding
+  const pad = (data: Buffer): Buffer => {
+    const rawLength = data.length;
+    const numPaddingBytes = 16 - (rawLength % 16);
+    const paddingBuffer = Buffer.alloc(numPaddingBytes);
+    paddingBuffer.fill(numPaddingBytes);
+    return Buffer.concat([data, paddingBuffer]);
+  };
+
+  // PKCS#7 padding removal
+  const removePadding = (paddedBuffer: Buffer): Buffer => {
+    const paddingToRemove = paddedBuffer[paddedBuffer.length - 1];
+    return paddedBuffer.subarray(0, paddedBuffer.length - paddingToRemove);
+  };
+
   it('should correctly encrypt input', () => {
     const data = randomBytes(32);
     const key = randomBytes(16);
     const iv = randomBytes(16);
 
+    const paddedData = pad(data);
+
     const cipher = createCipheriv('aes-128-cbc', key, iv);
     cipher.setAutoPadding(false);
-    const expected = Buffer.concat([cipher.update(data), cipher.final()]);
+    const expected = Buffer.concat([cipher.update(paddedData), cipher.final()]);
 
     const result: Buffer = aes128.encryptBufferCBC(data, iv, key);
 
@@ -28,13 +45,15 @@ describe('aes128', () => {
     const key = randomBytes(16);
     const iv = randomBytes(16);
 
+    const paddedData = pad(data);
+
     const cipher = createCipheriv('aes-128-cbc', key, iv);
     cipher.setAutoPadding(false);
-    const ciphertext = Buffer.concat([cipher.update(data), cipher.final()]);
+    const ciphertext = Buffer.concat([cipher.update(paddedData), cipher.final()]);
 
     const decipher = createDecipheriv('aes-128-cbc', key, iv);
     decipher.setAutoPadding(false);
-    const expected = Buffer.concat([decipher.update(ciphertext), decipher.final()]);
+    const expected = removePadding(Buffer.concat([decipher.update(ciphertext), decipher.final()]));
 
     const result: Buffer = aes128.decryptBufferCBC(ciphertext, iv, key);
 

diff --git a/yarn-project/circuits.js/src/barretenberg/crypto/aes128/index.ts b/yarn-project/circuits.js/src/barretenberg/crypto/aes128/index.ts
@@ -17,11 +17,9 @@ export class Aes128 {
     const rawLength = data.length;
     const numPaddingBytes = 16 - (rawLength % 16);
     const paddingBuffer = Buffer.alloc(numPaddingBytes);
-    // input num bytes needs to be a multiple of 16
+    // input num bytes needs to be a multiple of 16 and at least 1 byte
     // node uses PKCS#7-Padding scheme, where padding byte value = the number of padding bytes
-    if (numPaddingBytes != 0) {
-      paddingBuffer.fill(numPaddingBytes);
-    }
+    paddingBuffer.fill(numPaddingBytes);
     const input = Buffer.concat([data, paddingBuffer]);
 
     const api = BarretenbergSync.getSingleton();
@@ -39,8 +37,10 @@ export class Aes128 {
    */
   public decryptBufferCBC(data: Uint8Array, iv: Uint8Array, key: Uint8Array) {
     const api = BarretenbergSync.getSingleton();
-    return Buffer.from(
+    const paddedBuffer = Buffer.from(
       api.aesDecryptBufferCbc(new RawBuffer(data), new RawBuffer(iv), new RawBuffer(key), data.length),
     );
+    const paddingToRemove = paddedBuffer[paddedBuffer.length - 1];
+    return paddedBuffer.subarray(0, paddedBuffer.length - paddingToRemove);
   }
 }