Fix getNumberOfKeys minimal length of 4 bytes

massalabs · Jan 8, 2024 · 0fde237 · 0fde237
1 parent 4cc03c2
commit 0fde237
Show file tree

Hide file tree

Showing 4 changed files with 250 additions and 152 deletions.
diff --git a/assembly/std/__tests__/index.spec.ts b/assembly/std/__tests__/index.spec.ts
@@ -46,6 +46,7 @@ const keySerOk3: StaticArray<u8> = [
 ];
 
 const keysSerKo8: StaticArray<u8> = []; // edge case
+const keysSerKo9: StaticArray<u8> = [0, 0]; // length < 4 (Minimal length)
 
 describe('index tests', () => {
   it('derOpKeys ok 1', () => {
@@ -73,6 +74,11 @@ describe('index tests', () => {
     expect(res8.length).toBe(0);
   });
 
+  it('derOpKeys ko 8', () => {
+    let res9 = derKeys(keysSerKo9);
+    expect(res9.length).toBe(0);
+  });
+
   it('derOpKeys ok 3', () => {
     let res = derKeys(keySerOk3);
     expect(res.length).toBe(1);

diff --git a/assembly/std/op-datastore/util.ts b/assembly/std/op-datastore/util.ts
@@ -7,6 +7,8 @@
  *
  */
 function getNumberOfKeys(keysSer: StaticArray<u8>): u32 {
+  // check if keysSer is more than 4 bytes
+  assert(keysSer.length >= 4, 'Invalid keysSer length');
   // The first 4 bytes of the input array represent the number of keys
   let arr = new Uint8Array(4);
   arr[0] = keysSer[0];
@@ -46,7 +48,8 @@ function getNumberOfKeys(keysSer: StaticArray<u8>): u32 {
  *
  */
 export function derKeys(keysSer: StaticArray<u8>): Array<StaticArray<u8>> {
-  if (keysSer.length == 0) return [];
+  // check if keysSer is more than 4 bytes (for the number of keys)
+  if (keysSer.length < 4) return [];
 
   const keyCount: u32 = getNumberOfKeys(keysSer);
   const keys = new Array<StaticArray<u8>>(keyCount);