change counter to uint32

src/chacha20.test.ts

@@ -13,7 +13,7 @@ describe('ChaCha', () => {
     it('should calculate quarter round correctly', async () => {
         const key: UInt32[] = Array(8).fill(new UInt32(0));
         const nonce: UInt32[] = Array(3).fill(new UInt32(0));
-        const counter = 0;
+        const counter = UInt32.from(0);
 
         const chacha = new ChaChaState(key, nonce, counter);
 
@@ -40,9 +40,10 @@ describe('ChaCha', () => {
     it('should add two ChaChaState instances correctly', async () => {
         const zeroKey: UInt32[] = Array(8).fill(new UInt32(0));
         const zeroNonce: UInt32[] = Array(3).fill(new UInt32(0));
+        const zeroCounter = UInt32.from(0)
 
-        const state1 = new ChaChaState(zeroKey, zeroNonce, 0);
-        const state2 = new ChaChaState(zeroKey, zeroNonce, 0);
+        const state1 = new ChaChaState(zeroKey, zeroNonce, zeroCounter);
+        const state2 = new ChaChaState(zeroKey, zeroNonce, zeroCounter);
 
         state1.state[0] = UInt32.from(0x11111111);
         state1.state[1] = UInt32.from(0x01020304);
@@ -88,7 +89,7 @@ describe('ChaCha', () => {
             UInt32.from(0x00000000)
         ];
 
-        const counter = 1;
+        const counter = UInt32.from(1);
         const chachaState = new ChaChaState(key, nonce, counter);
 
         const expectedState: UInt32[] = [
@@ -124,7 +125,7 @@ describe('ChaCha', () => {
             UInt32.from(0x00000000)
         ];
 
-        const counter = 1;
+        const counter = UInt32.from(1);
 
         const expectedState: UInt32[] = [
             UInt32.from(0x10f1e7e4), UInt32.from(0xd13b5915), UInt32.from(0x500fdd1f), UInt32.from(0xa32071c4),
@@ -158,7 +159,7 @@ describe('ChaCha', () => {
             UInt32.from(0x00000000)
         ];
 
-        const counter = 1;
+        const counter = UInt32.from(1);
 
         const plaintext: UInt32[] = [
             UInt32.from(0x4c616469), UInt32.from(0x65732061), UInt32.from(0x6e642047), UInt32.from(0x656e746c),

src/chacha20.ts

@@ -7,38 +7,48 @@ export { ChaChaState, chacha20 };
  * 
  * @param {UInt32[]} key - The key used for encryption (256-bit).
  * @param {UInt32[]} nonce - The nonce used for encryption (96-bit).
- * @param {number} counter - The initial block counter.
+ * @param {UInt32} counter - The initial block counter.
  * @param {UInt32[]} plaintext - The plaintext to be encrypted or decrypted.
  * @returns {UInt32[]} - The resulting ciphertext or decrypted text as an array of UInt32.
  */
-function chacha20(key: UInt32[], nonce: UInt32[], counter: number, plaintext: UInt32[]): UInt32[] {
+function chacha20(key: UInt32[], nonce: UInt32[], counter: UInt32, plaintext: UInt32[]): UInt32[] {
     // Initialize the result array with the same length as the plaintext, filled with zeros.
     const res: UInt32[] = Array(plaintext.length).fill(UInt32.from(0));
 
     /**
      * Processes a block of 16 UInt32 words, encrypting or decrypting it using the ChaCha20 block function.
      * 
-     * @param {number} offset - The block offset in the plaintext.
-     * @param {number} length - The number of words to process (should be 16 for full blocks).
+     * @param {UInt32} offset - The block offset in the plaintext, representing which block of 16 words to process.
+     * @param {UInt32} length - The number of words to process in this block (typically 16, but could be less for the last block).
      */
-    function processBlock(offset: number, length: number) {
-        const keyStream = ChaChaState.chacha20Block(key, nonce, counter + offset);
-        for (let t = 0; t < length; t++) {
+    function processBlock(offset: UInt32, length: UInt32) {
+        // Generate the keystream block using the ChaCha20 block function.
+        const keyStream = ChaChaState.chacha20Block(key, nonce, counter.add(offset));
+
+        // Precompute the base index for this block in the plaintext array.
+        const baseIndex = Number(offset.toBigint() * 16n);
+
+        // Process each word in the block.
+        for (let t = 0; t < length.toBigint(); t++) {
+            const index = baseIndex + Number(t);  // Calculate the index in the plaintext array.
+
             // XOR the plaintext with the keystream to produce the ciphertext.
-            res[offset * 16 + t] = UInt32.from(plaintext[offset * 16 + t].toBigint() ^ keyStream[t].toBigint());
+            res[index] = plaintext[index].xor(keyStream[Number(t)]);
         }
     }
 
     // Determine the number of full 16-word blocks in the plaintext.
     const numFullBlocks = Math.floor(plaintext.length / 16);
+
+    // Process each full block of 16 words.
     for (let j = 0; j < numFullBlocks; j++) {
-        processBlock(j, 16);
+        processBlock(UInt32.from(j), UInt32.from(16));
     }
 
     // Process any remaining words in the plaintext that do not fill a full block.
     const remaining = plaintext.length % 16;
     if (remaining > 0) {
-        processBlock(numFullBlocks, remaining);
+        processBlock(UInt32.from(numFullBlocks), UInt32.from(remaining));
     }
 
     return res;
@@ -61,14 +71,14 @@ class ChaChaState {
      * 
      * @param {UInt32[]} key - The key used in encryption.
      * @param {UInt32[]} nonce - The nonce value.
-     * @param {number} counter - The block counter.
+     * @param {UInt32} counter - The block counter.
      */
-    constructor(key: UInt32[], nonce: UInt32[], counter: number) {
+    constructor(key: UInt32[], nonce: UInt32[], counter: UInt32) {
         // Initialize the state array with ChaCha constants, key, counter, and nonce.
         this.state = [
             UInt32.from(0x61707865), UInt32.from(0x3320646e), UInt32.from(0x79622d32), UInt32.from(0x6b206574), // ChaCha constants
             ...key,
-            UInt32.from(counter),
+            counter,
             ...nonce,
         ];
     }
@@ -164,10 +174,10 @@ class ChaChaState {
      * 
      * @param {UInt32[]} key - The encryption key (256-bit).
      * @param {UInt32[]} nonce - The nonce (96-bit).
-     * @param {number} counter - The block counter.
+     * @param {UInt32} counter - The block counter.
      * @returns {UInt32[]} - The resulting keystream block as an array of UInt32.
      */
-    static chacha20Block(key: UInt32[], nonce: UInt32[], counter: number): UInt32[] {
+    static chacha20Block(key: UInt32[], nonce: UInt32[], counter: UInt32): UInt32[] {
         // Initialize the state and a working copy of the state.
         const state = new ChaChaState(key, nonce, counter);
         const workingState = new ChaChaState(key, nonce, counter);