refactor: use single method to perform length check

refactor: use single method to perform length check refactor: use single method to perform length check
AztecProtocol · Apr 19, 2023 · c3315d3 · c3315d3
1 parent 55e7ace
commit c3315d3
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git a/yarn-project/circuits.js/src/abis/abis.test.ts b/yarn-project/circuits.js/src/abis/abis.test.ts
@@ -45,6 +45,22 @@ describe('abis wasm bindings', () => {
     expect(res).toMatchSnapshot();
   });
 
+  it('compute function leaf should revert if buffer is over 4 bytes', () => {
+    expect(() => {
+      new FunctionLeafPreimage(Buffer.from([0, 0, 0, 0, 123]), true, Fr.ZERO, Fr.ZERO);
+    }).toThrow('Function selector must be 4 bytes long, got 5 bytes.');
+  });
+
+  it('function leaf toBuffer should revert if buffer is over 4 bytes ', () => {
+    const initBuffer = Buffer.from([0, 0, 0, 123]);
+    const largerBuffer = Buffer.from([0, 0, 0, 0, 123]);
+    expect(() => {
+      const leaf = new FunctionLeafPreimage(initBuffer, true, Fr.ZERO, Fr.ZERO);
+      leaf.functionSelector = largerBuffer;
+      leaf.toBuffer();
+    }).toThrow('Function selector must be 4 bytes long, got 5 bytes.');
+  });
+
   it('computes function tree root', async () => {
     const res = await computeFunctionTreeRoot(wasm, [new Fr(0n), new Fr(0n), new Fr(0n), new Fr(0n)]);
     expect(res).toMatchSnapshot();

diff --git a/yarn-project/circuits.js/src/structs/function_leaf_preimage.ts b/yarn-project/circuits.js/src/structs/function_leaf_preimage.ts
@@ -9,6 +9,13 @@ export class FunctionLeafPreimage {
   readonly FUNCTION_SELECTOR_LENGTH = 4;
 
   constructor(public functionSelector: Buffer, public isPrivate: boolean, public vkHash: Fr, public acirHash: Fr) {
+    this.assertFunctionSelectorLength(functionSelector);
+  }
+
+  /**
+   * Assert the function selector buffer length matches `FUNCTION_SELECTOR_LENGTH`
+   */
+  private assertFunctionSelectorLength(functionSelector: Buffer) {
     if (functionSelector.byteLength !== this.FUNCTION_SELECTOR_LENGTH) {
       throw new Error(
         `Function selector must be ${this.FUNCTION_SELECTOR_LENGTH} bytes long, got ${functionSelector.byteLength} bytes.`,
@@ -21,11 +28,7 @@ export class FunctionLeafPreimage {
    * @returns The buffer.
    */
   toBuffer(): Buffer {
-    if (this.functionSelector.byteLength !== this.FUNCTION_SELECTOR_LENGTH) {
-      throw new Error(
-        `Function selector must be ${this.FUNCTION_SELECTOR_LENGTH} bytes long, got ${this.functionSelector.byteLength} bytes.`,
-      );
-    }
+    this.assertFunctionSelectorLength(this.functionSelector);
     return serializeToBuffer(this.functionSelector, this.isPrivate, this.vkHash, this.acirHash);
   }