From 90fa2233a8ae64ad4f072bdbbd53a1951da5ff2a Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Fri, 6 Dec 2024 16:36:24 +0800
Subject: [PATCH] zuc: eea seek benchmark test and refactor

---
 zuc/eea.go      | 62 +++++++++++++++++++++++++------------------------
 zuc/eea_test.go | 34 +++++++++++++++++++++++++--
 2 files changed, 64 insertions(+), 32 deletions(-)

diff --git a/zuc/eea.go b/zuc/eea.go
index 0d32110..2fe8204 100644
--- a/zuc/eea.go
+++ b/zuc/eea.go
@@ -8,9 +8,12 @@ import (
 )
 
 const (
+	// number of words in a round
 	RoundWords = 32
-	WordSize   = 4
-	WordMask   = WordSize - 1
+	// number of bytes in a word
+	WordSize = 4
+	WordMask = WordSize - 1
+	// number of bytes in a round
 	RoundBytes = RoundWords * WordSize
 )
 
@@ -18,8 +21,8 @@ type eea struct {
 	zucState32
 	x         [WordSize]byte // remaining bytes buffer
 	xLen      int            // number of remaining bytes
-	initState zucState32
-	used      uint64
+	initState zucState32     // initial state for reset
+	used      uint64         // number of key bytes processed, current offset
 }
 
 // NewCipher create a stream cipher based on key and iv aguments.
@@ -105,61 +108,60 @@ func (c *eea) reset() {
 	c.used = 0
 }
 
-func (c *eea) XORKeyStreamAt(dst, src []byte, offset uint64) {
-	if len(dst) < len(src) {
-		panic("zuc: output smaller than input")
-	}
-	if alias.InexactOverlap(dst[:len(src)], src) {
-		panic("zuc: invalid buffer overlap")
-	}
+// 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.
+// If the offset is equal to the current offset, the function behaves the same as XORKeyStream.
+// 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) {
 	if offset < c.used {
-		// reset the state to the initial state
 		c.reset()
 	}
-
 	if offset == c.used {
-		c.XORKeyStream(dst, src)
 		return
 	}
-
-	offsetDiff := offset - c.used
-	if offsetDiff <= uint64(c.xLen) {
-		c.xLen -= int(offsetDiff)
-		c.used += offsetDiff
+	gap := offset - c.used
+	if gap <= uint64(c.xLen) {
+		// offset is within the remaining key bytes
+		c.xLen -= int(gap)
+		c.used += gap
 		if c.xLen > 0 {
-			copy(c.x[:], c.x[offsetDiff:])
+			// adjust remaining key bytes
+			copy(c.x[:], c.x[gap:])
 		}
-		c.XORKeyStream(dst, src)
 		return
 	}
-
 	// consumed all remaining key bytes first
 	if c.xLen > 0 {
 		c.used += uint64(c.xLen)
-		offsetDiff -= uint64(c.xLen)
+		gap -= uint64(c.xLen)
 		c.xLen = 0
 	}
 
 	// forward the state to the offset
+	c.used += gap
 	stepLen := uint64(RoundBytes)
 	var keyStream [RoundWords]uint32
-	for ; offsetDiff >= uint64(stepLen); offsetDiff -= stepLen {
+	for gap >= stepLen {
 		genKeyStream(keyStream[:], &c.zucState32)
-		c.used += stepLen
+		gap -= stepLen
 	}
 
-	if offsetDiff > 0 {
-		numWords := (offsetDiff + WordMask) / WordSize
+	if gap > 0 {
+		numWords := (gap + WordMask) / WordSize
 		genKeyStream(keyStream[:numWords], &c.zucState32)
-		partiallyUsed := int(offsetDiff & WordMask)
-		c.used += numWords * WordSize
+		partiallyUsed := int(gap & WordMask)
 		if partiallyUsed > 0 {
 			// save remaining key bytes (less than 4 bytes)
 			c.xLen = WordSize - partiallyUsed
-			c.used -= uint64(c.xLen)
 			byteorder.BEPutUint32(c.x[:], keyStream[numWords-1])
 			copy(c.x[:], c.x[partiallyUsed:])
 		}
 	}
+}
+
+func (c *eea) XORKeyStreamAt(dst, src []byte, offset uint64) {
+	c.seek(offset)
 	c.XORKeyStream(dst, src)
 }
diff --git a/zuc/eea_test.go b/zuc/eea_test.go
index 97bbd25..a192831 100644
--- a/zuc/eea_test.go
+++ b/zuc/eea_test.go
@@ -105,7 +105,7 @@ func TestXORStreamAt(t *testing.T) {
 		}
 	})
 
-	t.Run("Jump and forward (incomplete word)", func(t *testing.T) {
+	t.Run("Jump and forward (incomplete word): gap > xLen", func(t *testing.T) {
 		for i := 0; i < 4; i++ {
 			c.XORKeyStreamAt(dst[i:16], src[i:16], uint64(i))
 			c.XORKeyStreamAt(dst[32:64], src[32:64], 32)
@@ -122,7 +122,7 @@ func TestXORStreamAt(t *testing.T) {
 		}
 	})
 
-	t.Run("Jump and forward (incomplete word): offsetDiff <= xLen", func(t *testing.T) {
+	t.Run("Jump and forward (incomplete word): gap <= xLen", func(t *testing.T) {
 		c.XORKeyStreamAt(dst[:1], src[:1], 0)
 		c.XORKeyStreamAt(dst[3:16], src[3:16], 3)
 		if !bytes.Equal(dst[3:16], expected[3:16]) {
@@ -227,3 +227,33 @@ func BenchmarkEncrypt1K(b *testing.B) {
 func BenchmarkEncrypt8K(b *testing.B) {
 	benchmarkStream(b, make([]byte, almost8K))
 }
+
+func benchmarkSeek(b *testing.B, offset uint64) {
+	var key [16]byte
+	var iv [16]byte
+
+	stream, _ := NewCipher(key[:], iv[:])
+
+	eea, ok := stream.(*eea)
+	if !ok {
+		b.Fatal("not an eea")
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		eea.reset()
+		eea.seek(offset)
+	}
+}
+
+func BenchmarkSeek1K(b *testing.B) {
+	benchmarkSeek(b, 1024)
+}
+
+func BenchmarkSeek8K(b *testing.B) {
+	benchmarkSeek(b, 8*1024)
+}
+
+func BenchmarkSeek1M(b *testing.B) {
+	benchmarkSeek(b, 1024*1024)
+}