From a0912994f3758392784a14d2db63ee59d6e1165c Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Tue, 30 Sep 2025 16:14:06 +0800
Subject: [PATCH] internal/zuc: support fast forward

---
 internal/zuc/eea.go      | 20 ++++++++++++++++++--
 internal/zuc/eea_test.go | 27 +++++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/internal/zuc/eea.go b/internal/zuc/eea.go
index 1715288..423cdb9 100644
--- a/internal/zuc/eea.go
+++ b/internal/zuc/eea.go
@@ -28,8 +28,8 @@ type eea struct {
 }
 
 const (
-	magic         = "zuceea"
-	stateSize     = (16 + 6) * 4 // zucState32 size in bytes
+	magic            = "zuceea"
+	stateSize        = (16 + 6) * 4 // zucState32 size in bytes
 	minMarshaledSize = len(magic) + stateSize + 8 + 4*3
 )
 
@@ -276,6 +276,21 @@ func (c *eea) reset(offset uint64) {
 	c.used = n * uint64(c.bucketSize)
 }
 
+func (c *eea) fastForward(offset uint64) {
+	// fast forward, check and adjust state if needed
+	var n uint64
+	if c.bucketSize > 0 {
+		n = offset / uint64(c.bucketSize)
+		expectedStateIndex := int(n)
+		if expectedStateIndex > c.stateIndex && expectedStateIndex < len(c.states) {
+			c.stateIndex = int(n)
+			c.zucState32 = *c.states[n]
+			c.xLen = 0
+			c.used = n * uint64(c.bucketSize)
+		}
+	}
+}
+
 // seek sets the offset for the next XORKeyStream operation.
 //
 // If the offset is less than the current offset, the state will be reset to the initial state.
@@ -283,6 +298,7 @@ func (c *eea) reset(offset uint64) {
 // If the offset is greater than the current offset, the function will forward the state to the offset.
 // Note: This method is not thread-safe.
 func (c *eea) seek(offset uint64) {
+	c.fastForward(offset)
 	if offset < c.used {
 		c.reset(offset)
 	}
diff --git a/internal/zuc/eea_test.go b/internal/zuc/eea_test.go
index 1921c0b..05486c6 100644
--- a/internal/zuc/eea_test.go
+++ b/internal/zuc/eea_test.go
@@ -275,7 +275,7 @@ func TestEEAXORKeyStreamAtWithBucketSize(t *testing.T) {
 		}
 		clear(dst)
 		bucketCipher.XORKeyStreamAt(dst[513:768], src[513:768], 513)
-		if bucketCipher.stateIndex != 0 {
+		if bucketCipher.stateIndex != 4 {
 			t.Fatalf("expected=%d, result=%d\n", 0, bucketCipher.stateIndex)
 		}
 		if len(bucketCipher.states) != 7 {
@@ -296,6 +296,29 @@ func TestEEAXORKeyStreamAtWithBucketSize(t *testing.T) {
 			t.Fatalf("expected=%x, result=%x\n", expected[512:768], dst[512:768])
 		}
 	})
+
+	t.Run("Rotate end to start, end to start", func(t *testing.T) {
+		bucketCipher, err := NewEEACipherWithBucketSize(key, zucEEATests[0].count, zucEEATests[0].bearer, zucEEATests[0].direction, 128)
+		if err != nil {
+			t.Error(err)
+		}
+		clear(dst)
+		for i := len(src) - RoundBytes; i >= 0; i -= RoundBytes {
+			offset := i
+			bucketCipher.XORKeyStreamAt(dst[offset:offset+RoundBytes], src[offset:offset+RoundBytes], uint64(offset))
+			if !bytes.Equal(expected[offset:offset+RoundBytes], dst[offset:offset+RoundBytes]) {
+				t.Fatalf("at %d, expected=%x, result=%x\n", offset, expected[offset:offset+RoundBytes], dst[offset:offset+RoundBytes])
+			}
+		}
+		clear(dst)
+		for i := len(src) - RoundBytes; i >= 0; i -= RoundBytes {
+			offset := i
+			bucketCipher.XORKeyStreamAt(dst[offset:offset+RoundBytes], src[offset:offset+RoundBytes], uint64(offset))
+			if !bytes.Equal(expected[offset:offset+RoundBytes], dst[offset:offset+RoundBytes]) {
+				t.Fatalf("at %d, expected=%x, result=%x\n", offset, expected[offset:offset+RoundBytes], dst[offset:offset+RoundBytes])
+			}
+		}
+	})
 }
 
 func TestMarshalUnmarshalBinary(t *testing.T) {
@@ -394,7 +417,7 @@ func TestUnmarshalBinary_InvalidRemainingBytes(t *testing.T) {
 	data[xLenOffset+3] = 8 // xLen = 8
 
 	// Truncate data so remaining bytes < xLen
-	truncated := data[:minMarshaledSize + 4]
+	truncated := data[:minMarshaledSize+4]
 
 	c2 := NewEmptyCipher()
 	err = c2.UnmarshalBinary(truncated)