feat: port high-level DApp crypto utilities from wallet plugin

src/crypto_utils.ts

@@ -1,5 +1,5 @@
 import forge from 'node-forge'
-import { BaseWallet, getBytes, SigningKey, solidityPackedKeccak256 } from "ethers"
+import { BaseWallet, getBytes, SigningKey, solidityPackedKeccak256, hexlify, Wallet } from "ethers"
 import { ctString, ctUint, ctUint256, itString, itUint, itUint256 } from './types';
 
 const BLOCK_SIZE = 16 // AES block size in bytes
@@ -170,16 +170,26 @@ export function sign(message: string, privateKey: string) {
     return new Uint8Array([...getBytes(sig.r), ...getBytes(sig.s), ...getBytes(`0x0${sig.v - 27}`)])
 }
 
+// Computes the COTI IT message hash shared by signInputText and buildItSignature.
+function buildItMessageHash(
+    signerAddress: string,
+    contractAddress: string,
+    functionSelector: string,
+    ct: bigint
+): string {
+    return solidityPackedKeccak256(
+        ["address", "address", "bytes4", "uint256"],
+        [signerAddress, contractAddress, functionSelector, ct]
+    )
+}
+
 export function signInputText(
     sender: { wallet: BaseWallet; userKey: string },
     contractAddress: string,
     functionSelector: string,
     ct: bigint
 ) {
-    const message = solidityPackedKeccak256(
-        ["address", "address", "bytes4", "uint256"],
-        [sender.wallet.address, contractAddress, functionSelector, ct]
-    )
+    const message = buildItMessageHash(sender.wallet.address, contractAddress, functionSelector, ct)
 
     return sign(message, sender.wallet.privateKey);
 }
@@ -253,7 +263,30 @@ export function buildStringInputText(
     return inputText
 }
 
+/**
+ * Decrypts a 64-bit ctUint ciphertext using the user's AES key.
+ *
+ * - A zero ciphertext is short-circuited to `0n` without validating the key,
+ *   since it represents uninitialized/empty on-chain storage. This allows DApps
+ *   to read empty balances before the user has configured their AES key.
+ *   Note: this means `decryptUint(0n, invalidKey)` returns `0n` without throwing.
+ * - For non-zero ciphertexts, key validation is performed by `encodeKey`
+ *   (strips "0x", lowercases, enforces 128-bit hex). Invalid keys throw
+ *   rather than producing garbage.
+ *
+ * @param ciphertext - The encrypted 64-bit value (bigint).
+ * @param userKey - The AES key (32 hex chars, optionally "0x"-prefixed).
+ * @returns The decrypted plaintext as a bigint.
+ * @throws Error if the key is invalid (null, wrong length, non-hex) and ciphertext is non-zero.
+ */
 export function decryptUint(ciphertext: ctUint, userKey: string): bigint {
+    // A zero ciphertext represents uninitialized/empty storage, which decrypts
+    // to plaintext 0. Short-circuit before touching the key so callers can read
+    // empty values without a valid key (and to avoid returning AES garbage).
+    if (ciphertext === 0n) {
+        return 0n
+    }
+
     // Convert ciphertext to Uint8Array
     let ctArray = new Uint8Array()
 
@@ -269,6 +302,7 @@ export function decryptUint(ciphertext: ctUint, userKey: string): bigint {
     const cipher = ctArray.subarray(0, BLOCK_SIZE)
     const r = ctArray.subarray(BLOCK_SIZE)
 
+    // encodeKey validates and normalizes the key (strips "0x", enforces 128-bit)
     const userKeyBytes = encodeKey(userKey)
 
     // Decrypt the cipher
@@ -360,10 +394,14 @@ export function encodeString(str: string): Uint8Array {
 }
 
 export function encodeKey(userKey: string): Uint8Array {
+    // Validate and normalize the key (strips "0x", lowercases, enforces 128-bit)
+    // so that every encrypt/decrypt path rejects malformed keys consistently
+    // instead of silently producing NaN/garbage bytes.
+    const normalizedKey = normalizeAesKey(userKey)
     const keyBytes = new Uint8Array(16)
 
     for (let i = 0; i < 32; i += 2) {
-        keyBytes[i / 2] = Number.parseInt(userKey.slice(i, i + 2), HEX_BASE)
+        keyBytes[i / 2] = Number.parseInt(normalizedKey.slice(i, i + 2), HEX_BASE)
     }
 
     return keyBytes
@@ -544,4 +582,59 @@ export function prepareIT256(
         },
         signature
     }
-}
+}
+// ------------- Wallet Plugin Additions -------------
+
+/**
+ * Validates and normalizes an AES key: ensures it is present, strips the "0x"
+ * prefix, and lowercases it. COTI uses a 128-bit AES key, so only 32-character
+ * hex strings are accepted.
+ *
+ * @param aesKey - The AES key, optionally prefixed with "0x".
+ * @returns The normalized lowercase hex string.
+ * @throws Error if the key is empty/null/undefined, contains non-hex characters, or is not 32 hex characters.
+ */
+export function normalizeAesKey(aesKey: string | null | undefined): string {
+    if (!aesKey) {
+        throw new Error("AES key is required")
+    }
+
+    const trimmed = aesKey.startsWith("0x") ? aesKey.slice(2) : aesKey
+    const lowered = trimmed.toLowerCase()
+
+    if (!/^[0-9a-f]+$/.test(lowered)) {
+        throw new Error("Invalid AES key: contains non-hexadecimal characters")
+    }
+
+    if (lowered.length !== 32) {
+        throw new Error(`Invalid AES key: expected 32 hex characters (128-bit), got ${lowered.length}`)
+    }
+
+    return lowered
+}
+
+/**
+ * Builds a COTI input-text (IT) signature over (signer, contract, selector, ciphertext).
+ *
+ * @param signerAddress - Address of the signer; must match the address derived from privateKey.
+ * @param contractAddress - Target contract address.
+ * @param functionSelector - 4-byte function selector (e.g. "0x11223344").
+ * @param ciphertext - The encrypted value being signed.
+ * @param privateKey - Signer's private key.
+ * @returns The 65-byte signature as a hex string.
+ * @throws Error if signerAddress does not match the address derived from privateKey.
+ */
+export function buildItSignature(
+    signerAddress: string,
+    contractAddress: string,
+    functionSelector: string,
+    ciphertext: bigint,
+    privateKey: string
+): string {
+    const wallet = new Wallet(privateKey);
+    if ( wallet.address.toLowerCase() !== signerAddress.toLowerCase()) {
+        throw new Error("Invalid signer: signerAddress does not match the address derived from privateKey");
+    }
+    let signature = signInputText({ wallet, userKey: ''}, contractAddress, functionSelector, ciphertext);
+    return hexlify(signature);
+}

tests/integration/validation.test.ts

@@ -30,37 +30,32 @@ function createTestSender() {
 const describeWithEnv = HAS_ENV ? describe : describe.skip
 describeWithEnv('Integration: Input Validation', () => {
     describe('Invalid user key format', () => {
-        test('prepareIT produces incorrect result with user key that is too short', () => {
+        test('prepareIT throws with user key that is too short', () => {
             const sender = {
                 wallet: new Wallet(TEST_CONSTANTS.PRIVATE_KEY),
                 userKey: '1234567890123456789012345678901' // 31 chars instead of 32
             }
-            // Function doesn't throw, but produces incorrect encryption
-            const result = prepareIT(
+            // Encryption now validates the key and rejects non-128-bit keys
+            expect(() => prepareIT(
                 12345n,
                 sender,
                 TEST_CONSTANTS.CONTRACT_ADDRESS,
                 TEST_CONSTANTS.FUNCTION_SELECTOR
-            )
-            expect(result).toHaveProperty('ciphertext')
-            // Decryption with correct key will fail
-            const decrypted = decryptUint(result.ciphertext, TEST_CONSTANTS.USER_KEY)
-            expect(decrypted).not.toBe(12345n)
+            )).toThrow('expected 32 hex characters')
         })
 
-        test('prepareIT produces incorrect result with user key that is too long', () => {
+        test('prepareIT throws with user key that is too long', () => {
             const sender = {
                 wallet: new Wallet(TEST_CONSTANTS.PRIVATE_KEY),
                 userKey: '12345678901234567890123456789012a' // 33 chars instead of 32
             }
-            // Function processes first 32 chars, doesn't throw
-            const result = prepareIT(
+            // Encryption now validates the key and rejects non-128-bit keys
+            expect(() => prepareIT(
                 12345n,
                 sender,
                 TEST_CONSTANTS.CONTRACT_ADDRESS,
                 TEST_CONSTANTS.FUNCTION_SELECTOR
-            )
-            expect(result).toHaveProperty('ciphertext')
+            )).toThrow('expected 32 hex characters')
         })
 
         test('prepareIT produces incorrect result with user key containing invalid hex', () => {

tests/unit/error.handling.test.ts

@@ -105,31 +105,18 @@ describe('Unit: Error Handling', () => {
     })
 
     describe('encodeKey error cases', () => {
-        test('produces incorrect result when user key length is not 32 hex characters', () => {
+        test('throws when user key length is not 32 hex characters', () => {
             const shortKey = '1234567890123456789012345678901' // 31 chars
-            const result = encodeKey(shortKey)
-            // Function doesn't throw, but produces incorrect result (last byte is incomplete)
-            expect(result).toBeInstanceOf(Uint8Array)
-            expect(result.length).toBe(16)
+            expect(() => encodeKey(shortKey)).toThrow('expected 32 hex characters')
         })
 
-        test('produces incorrect result when user key has invalid hex characters', () => {
+        test('throws when user key has invalid hex characters', () => {
             const invalidKey = '1234567890123456789012345678901g' // 'g' is invalid hex
-            const result = encodeKey(invalidKey)
-            // Function doesn't throw, but produces incorrect values
-            // parseInt('1g', 16) returns 1 (stops at invalid char), so result[15] = 1
-            expect(result).toBeInstanceOf(Uint8Array)
-            expect(result.length).toBe(16)
-            // The function processes what it can, invalid hex is handled gracefully
-            expect(result[15]).toBe(1) // '1g' parses as '1'
+            expect(() => encodeKey(invalidKey)).toThrow('non-hexadecimal')
         })
 
-        test('produces zero-filled array when user key is empty', () => {
-            const result = encodeKey('')
-            // Function doesn't throw, but produces zeros
-            expect(result).toBeInstanceOf(Uint8Array)
-            expect(result.length).toBe(16)
-            expect(result.every(b => b === 0)).toBe(true)
+        test('throws when user key is empty', () => {
+            expect(() => encodeKey('')).toThrow('AES key is required')
         })
     })
 
@@ -228,17 +215,16 @@ describe('Unit: Error Handling', () => {
     })
 
     describe('decryptUint error cases', () => {
-        test('produces incorrect result when user key has wrong length', () => {
+        test('throws when user key has wrong length', () => {
             const { ciphertext } = prepareIT(
                 12345n,
                 createTestSender(),
                 TEST_CONSTANTS.CONTRACT_ADDRESS,
                 TEST_CONSTANTS.FUNCTION_SELECTOR
             )
             const wrongKey = '1234567890123456789012345678901' // 31 chars
-            // Function doesn't throw, but produces incorrect decryption
-            const decrypted = decryptUint(ciphertext, wrongKey)
-            expect(decrypted).not.toBe(12345n)
+            // decryptUint now validates the key and rejects non-128-bit keys
+            expect(() => decryptUint(ciphertext, wrongKey)).toThrow('expected 32 hex characters')
         })
 
         test('produces incorrect result when user key has invalid hex', () => {
@@ -261,17 +247,16 @@ describe('Unit: Error Handling', () => {
     })
 
     describe('decryptUint256 error cases', () => {
-        test('produces incorrect result when user key has wrong length', () => {
+        test('throws when user key has wrong length', () => {
             const { ciphertext } = prepareIT256(
                 2n ** 200n,
                 createTestSender(),
                 TEST_CONSTANTS.CONTRACT_ADDRESS,
                 TEST_CONSTANTS.FUNCTION_SELECTOR
             )
             const wrongKey = '1234567890123456789012345678901' // 31 chars
-            // Function doesn't throw, but produces incorrect decryption
-            const decrypted = decryptUint256(ciphertext, wrongKey)
-            expect(decrypted).not.toBe(2n ** 200n)
+            // decryptUint256 delegates to encodeKey, which now validates the key
+            expect(() => decryptUint256(ciphertext, wrongKey)).toThrow('expected 32 hex characters')
         })
 
         test('produces incorrect result when ciphertext structure is invalid', () => {
@@ -285,17 +270,16 @@ describe('Unit: Error Handling', () => {
     })
 
     describe('decryptString error cases', () => {
-        test('produces incorrect result when user key has wrong length', () => {
+        test('throws when user key has wrong length', () => {
             const { ciphertext } = buildStringInputText(
                 'Hello',
                 createTestSender(),
                 TEST_CONSTANTS.CONTRACT_ADDRESS,
                 TEST_CONSTANTS.FUNCTION_SELECTOR
             )
             const wrongKey = '1234567890123456789012345678901' // 31 chars
-            // Function doesn't throw, but produces incorrect decryption
-            const decrypted = decryptString(ciphertext, wrongKey)
-            expect(decrypted).not.toBe('Hello')
+            // decryptString delegates to decryptUint, which now validates the key
+            expect(() => decryptString(ciphertext, wrongKey)).toThrow('expected 32 hex characters')
         })
 
         test('handles empty ciphertext array', () => {

tests/unit/stringCrypt.test.ts

@@ -0,0 +1,58 @@
+import { ethers } from "ethers";
+import { 
+    buildStringInputText, 
+    decryptString, 
+    generateRandomAesKeySizeNumber 
+} from "../../src/crypto_utils";
+
+describe("String Encryption and Decryption (decryptString)", () => {
+    let wallet: ethers.BaseWallet;
+    let userKey: string;
+    let sender: { wallet: ethers.BaseWallet; userKey: string };
+    const contractAddress = "0x0000000000000000000000000000000000000000";
+    const functionSelector = "0x12345678";
+
+    beforeAll(() => {
+        wallet = ethers.Wallet.createRandom();
+        // generateRandomAesKeySizeNumber() returns 16 raw bytes (a binary string).
+        // The SDK's documented AES key format is 32 hex chars, so convert it.
+        userKey = Buffer.from(generateRandomAesKeySizeNumber(), "binary").toString("hex");
+        sender = { wallet, userKey };
+    });
+
+    it("should correctly encrypt and decrypt a string perfectly matching the 8-byte chunk size", () => {
+        const targetString = "12345678"; // Exact 8 bytes
+        const it = buildStringInputText(targetString, sender, contractAddress, functionSelector);
+
+        const decrypted = decryptString(it.ciphertext, userKey);
+        expect(decrypted).toBe(targetString);
+    });
+
+    it("should correctly encrypt and decrypt a string with padding (less than 8 bytes)", () => {
+        const targetString = "Hello"; // 5 bytes => pads 3 bytes
+        const it = buildStringInputText(targetString, sender, contractAddress, functionSelector);
+
+        const decrypted = decryptString(it.ciphertext, userKey);
+        expect(decrypted).toBe(targetString);
+        expect(decrypted.length).toBe(targetString.length); // Verifies trailing null bytes are gone
+    });
+
+    it("should correctly encrypt and decrypt a multi-chunk string with padding", () => {
+        const targetString = "Hello World!"; // 12 bytes => 1 chunk (8) + 1 chunk (4 + pads 4)
+        const it = buildStringInputText(targetString, sender, contractAddress, functionSelector);
+
+        // Assert there are 2 chunks under the hood
+        expect(it.ciphertext.value.length).toBe(2); 
+
+        const decrypted = decryptString(it.ciphertext, userKey);
+        expect(decrypted).toBe(targetString);
+    });
+
+    it("should preserve strings with null-like characters strictly intentionally inside the string", () => {
+        const targetString = "a\0b\0c"; // Internal null characters
+        const it = buildStringInputText(targetString, sender, contractAddress, functionSelector);
+
+        const decrypted = decryptString(it.ciphertext, userKey);
+        expect(decrypted).toBe(targetString);
+    });
+});