From ddb5b69b53fac374440bdb52d99f9ff98d391106 Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Wed, 4 Dec 2024 18:54:19 +0800
Subject: [PATCH] fix #284

---
 zuc/eea.go      |  9 ++++-----
 zuc/eea_test.go | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/zuc/eea.go b/zuc/eea.go
index 52fd4eb..e60d587 100644
--- a/zuc/eea.go
+++ b/zuc/eea.go
@@ -78,17 +78,16 @@ func (c *eea) XORKeyStream(dst, src []byte) {
 			return
 		}
 	}
-	words := (len(src) + WordSize - 1) / WordSize
-	rounds := words / RoundWords
 	var keyBytes [RoundBytes]byte
-	for i := 0; i < rounds; i++ {
+	for len(src) >= RoundBytes {
 		genKeyStreamRev32(keyBytes[:], &c.zucState32)
 		subtle.XORBytes(dst, src, keyBytes[:])
 		dst = dst[RoundBytes:]
 		src = src[RoundBytes:]
 	}
-	if processedWords := rounds * RoundWords; processedWords < words {
-		byteLen := WordSize * (words - processedWords)
+	if len(src) > 0 {
+		words := (len(src) + WordSize - 1) / WordSize
+		byteLen := WordSize * words
 		genKeyStreamRev32(keyBytes[:byteLen], &c.zucState32)
 		n := subtle.XORBytes(dst, src, keyBytes[:])
 		// save remaining key bytes
diff --git a/zuc/eea_test.go b/zuc/eea_test.go
index 3320a0b..64b0b55 100644
--- a/zuc/eea_test.go
+++ b/zuc/eea_test.go
@@ -160,6 +160,28 @@ func TestXORStreamAt(t *testing.T) {
 	})
 }
 
+func TestIssue284(t *testing.T) {
+	key, err := hex.DecodeString(zucEEATests[0].key)
+	if err != nil {
+		t.Error(err)
+	}
+	c, err := NewEEACipher(key, zucEEATests[0].count, zucEEATests[0].bearer, zucEEATests[0].direction)
+	if err != nil {
+		t.Error(err)
+	}
+	src := make([]byte, 200)
+	expected := make([]byte, 200)
+	dst := make([]byte, 200)
+	c.XORKeyStream(expected, src)
+
+	for i := 124; i <= 200; i++ {
+		c.XORKeyStreamAt(dst, src[:i], 0)
+		if !bytes.Equal(expected[:i], dst[:i]) {
+			t.Fatalf("failed for len %v", i)
+		}
+	}
+}
+
 func benchmarkStream(b *testing.B, buf []byte) {
 	b.SetBytes(int64(len(buf)))