x.crypto.chacha20: mark State as public to make it reusable (#25449)

blackshirt · web-flow · commit 9eef7ecd7bed · 2025-10-07T11:39:25.000+03:00
diff --git a/vlib/x/crypto/chacha20/chacha_test.v b/vlib/x/crypto/chacha20/chacha_test.v
@@ -2,6 +2,39 @@ import rand
 import encoding.hex
 import x.crypto.chacha20
 
+// test for chacha20.State exposed api. The test material was taken from
+// Test Vector for the ChaCha20 Block Function ch 2.3.2. from RFC 8439 doc.
+// See https://datatracker.ietf.org/doc/html/rfc8439#section-2.3.2
+fn test_chacha20_state() ! {
+	// ChaCha state with the key setup.
+	s := chacha20.State([u32(0x61707865), 0x3320646e, 0x79622d32, 0x6b206574, 0x03020100, 0x07060504,
+		0x0b0a0908, 0x0f0e0d0c, 0x13121110, 0x17161514, 0x1b1a1918, 0x1f1e1d1c, 0x00000001,
+		0x09000000, 0x4a000000, 0x00000000]!)
+
+	// test for .clone
+	mut ws := s.clone()
+	for i := 0; i < 16; i++ {
+		assert s[i] == ws[i]
+	}
+	// After running 20 rounds (10 column rounds interleaved with 10 "diagonal rounds"),
+	// the ChaCha state looks like this:
+	expected := chacha20.State([u32(0x837778ab), 0xe238d763, 0xa67ae21e, 0x5950bb2f, 0xc4f2d0c7,
+		0xfc62bb2f, 0x8fa018fc, 0x3f5ec7b7, 0x335271c2, 0xf29489f3, 0xeabda8fc, 0x82e46ebd,
+		0xd19c12b4, 0xb04e16de, 0x9e83d0cb, 0x4e3c50a2]!)
+
+	// test .qround call
+	ws.qround(10)
+	for i := 0; i < 16; i++ {
+		assert ws[i] == expected[i]
+	}
+
+	// test for .reset call
+	ws.reset()
+	for i := 0; i < 16; i++ {
+		assert ws[i] == 0
+	}
+}
+
 // test for extended chacha20 construct with 64-bit counter support
 fn test_xchacha20_cipher_with_64_counter() ! {
 	key := rand.bytes(32)!
diff --git a/vlib/x/crypto/chacha20/stream.v b/vlib/x/crypto/chacha20/stream.v
@@ -197,7 +197,7 @@ fn (mut s Stream) keystream_full(mut dst []u8, src []u8) ! {
 fn (mut s Stream) keystream() ![]u8 {
 	// initializes current state and working state
 	mut awal := s.new_curr_state()
-	mut ws := clone_state(awal)
+	mut ws := awal.clone()
 
 	// precomputes cache counter-independent  values
 	if s.mode == .standard && !s.precomp {
@@ -222,27 +222,8 @@ fn (mut s Stream) keystream() ![]u8 {
 	//
 	// perform chacha20 quarter round n-times
 	n := if s.mode == .standard { 9 } else { default_qround_nr }
-	for i := 0; i < n; i++ {
-		// Column-round
-		//  0 |  1 |  2 |  3
-		//  4 |  5 |  6 |  7
-		//  8 |  9 | 10 | 11
-		// 12 | 13 | 14 | 15
-		qround_on_state(mut ws, 0, 4, 8, 12) // 0
-		qround_on_state(mut ws, 1, 5, 9, 13) // 1
-		qround_on_state(mut ws, 2, 6, 10, 14) // 2
-		qround_on_state(mut ws, 3, 7, 11, 15) // 3
+	ws.qround(n)
 
-		// Diagonal round.
-		//   0 \  1 \  2 \  3
-		//   5 \  6 \  7 \  4
-		//  10 \ 11 \  8 \  9
-		//  15 \ 12 \ 13 \ 14
-		qround_on_state(mut ws, 0, 5, 10, 15)
-		qround_on_state(mut ws, 1, 6, 11, 12)
-		qround_on_state(mut ws, 2, 7, 8, 13)
-		qround_on_state(mut ws, 3, 4, 9, 14)
-	}
 	// Adding the working state values with inital state values.
 	// We dont performs xor-ing here, its done on xor_key_stream and or keystream_full.
 	for i, _ in ws {
@@ -379,17 +360,53 @@ fn (b Stream) max_ctr() u64 {
 }
 
 // State represents the running 64-bytes of chacha20 stream,
-type State = [16]u32
+pub type State = [16]u32
 
+// clone returns a new copy of this state.
 @[direct_array_access; inline]
-fn clone_state(s State) State {
+pub fn (s State) clone() State {
 	mut sc := State{}
 	for i, v in s {
 		sc[i] = v
 	}
 	return sc
 }
 
+// reset resets internal values of this state.
+@[direct_array_access; inline]
+pub fn (mut s State) reset() {
+	for i, _ in s {
+		s[i] = u32(0)
+	}
+}
+
+// qround performs quarter round on the working state ws with round number specified in nr.
+// Its responsibility the user to provide the correct round number.
+@[direct_array_access]
+pub fn (mut ws State) qround(nr int) {
+	for i := 0; i < nr; i++ {
+		// Column-round
+		//  0 |  1 |  2 |  3
+		//  4 |  5 |  6 |  7
+		//  8 |  9 | 10 | 11
+		// 12 | 13 | 14 | 15
+		qround_on_state(mut ws, 0, 4, 8, 12) // 0
+		qround_on_state(mut ws, 1, 5, 9, 13) // 1
+		qround_on_state(mut ws, 2, 6, 10, 14) // 2
+		qround_on_state(mut ws, 3, 7, 11, 15) // 3
+
+		// Diagonal round.
+		//   0 \  1 \  2 \  3
+		//   5 \  6 \  7 \  4
+		//  10 \ 11 \  8 \  9
+		//  15 \ 12 \ 13 \ 14
+		qround_on_state(mut ws, 0, 5, 10, 15)
+		qround_on_state(mut ws, 1, 6, 11, 12)
+		qround_on_state(mut ws, 2, 7, 8, 13)
+		qround_on_state(mut ws, 3, 4, 9, 14)
+	}
+}
+
 // qround_on_state_with_quartet run qround_on_state by previously set up state values in offset
 // (a,b,c,d) with values from quartet (q0, q1, q2, q3)
 @[direct_array_access]
diff --git a/vlib/x/crypto/chacha20/xchacha.v b/vlib/x/crypto/chacha20/xchacha.v
@@ -52,23 +52,7 @@ fn hchacha20(key []u8, nonce []u8) ![]u8 {
 	x[15] = binary.little_endian_u32(nonce[12..16])
 
 	// After initialization, proceed through the ChaCha20 rounds as usual.
-	for i := 0; i < 10; i++ {
-		// Column round.
-		qround_on_state(mut x, 0, 4, 8, 12) // 0
-		qround_on_state(mut x, 1, 5, 9, 13) // 1
-		qround_on_state(mut x, 2, 6, 10, 14) // 2
-		qround_on_state(mut x, 3, 7, 11, 15) // 3
-
-		// Diagonal round.
-		//   0 \  1 \  2 \  3
-		//   5 \  6 \  7 \  4
-		//  10 \ 11 \  8 \  9
-		//  15 \ 12 \ 13 \ 14
-		qround_on_state(mut x, 0, 5, 10, 15)
-		qround_on_state(mut x, 1, 6, 11, 12)
-		qround_on_state(mut x, 2, 7, 8, 13)
-		qround_on_state(mut x, 3, 4, 9, 14)
-	}
+	x.qround(10)
 
 	// Once the 20 ChaCh20 rounds have been completed, the first 128 bits (16 bytes) and
 	// last 128 bits (16 bytes) of the ChaCha state (both little-endian) are
