From 400147ed3c9e7182393bc7958b7926feb0d80b34 Mon Sep 17 00:00:00 2001
From: emmansun <emman.sun@foxmail.com>
Date: Fri, 29 Sep 2023 09:57:27 +0800
Subject: [PATCH] zuc: eea, fix can't src/dst can't use same buffer issue #169

---
 zuc/core.go        |  8 --------
 zuc/eea.go         | 26 +++++++++++++++-----------
 zuc/eea_asm.go     | 22 +++-------------------
 zuc/eea_generic.go |  4 ++--
 zuc/eea_test.go    |  5 +++--
 5 files changed, 23 insertions(+), 42 deletions(-)

diff --git a/zuc/core.go b/zuc/core.go
index ee89b7c..148e19b 100644
--- a/zuc/core.go
+++ b/zuc/core.go
@@ -237,11 +237,3 @@ func (s *zucState32) genKeywords(words []uint32) {
 	}
 	genKeyStream(words, s)
 }
-
-func genKeyStreamRev32(keyStream []byte, pState *zucState32) {
-	for len(keyStream) >= 4 {
-		z := genKeyword(pState)
-		binary.BigEndian.PutUint32(keyStream, z)
-		keyStream = keyStream[4:]
-	}
-}
diff --git a/zuc/eea.go b/zuc/eea.go
index 2b64df5..6442dd7 100644
--- a/zuc/eea.go
+++ b/zuc/eea.go
@@ -25,7 +25,21 @@ func NewEEACipher(key []byte, count, bearer, direction uint32) (cipher.Stream, e
 	return newZUCState(key, iv)
 }
 
-func xorKeyStreamGeneric(c *zucState32, dst, src []byte) {
+func genKeyStreamRev32Generic(keyStream []byte, pState *zucState32) {
+	for len(keyStream) >= 4 {
+		z := genKeyword(pState)
+		binary.BigEndian.PutUint32(keyStream, z)
+		keyStream = keyStream[4:]
+	}
+}
+
+func (c *zucState32) XORKeyStream(dst, src []byte) {
+	if len(dst) < len(src) {
+		panic("zuc: output smaller than input")
+	}
+	if alias.InexactOverlap(dst[:len(src)], src) {
+		panic("zuc: invalid buffer overlap")
+	}
 	words := (len(src) + 3) / 4
 	rounds := words / RoundWords
 	var keyBytes [RoundWords * 4]byte
@@ -40,13 +54,3 @@ func xorKeyStreamGeneric(c *zucState32, dst, src []byte) {
 		subtle.XORBytes(dst, src, keyBytes[:])
 	}
 }
-
-func (c *zucState32) XORKeyStream(dst, src []byte) {
-	if len(dst) < len(src) {
-		panic("zuc: output smaller than input")
-	}
-	if alias.InexactOverlap(dst[:len(src)], src) {
-		panic("zuc: invalid buffer overlap")
-	}
-	xorKeyStream(c, dst, src)
-}
diff --git a/zuc/eea_asm.go b/zuc/eea_asm.go
index 1f06425..72e45bb 100644
--- a/zuc/eea_asm.go
+++ b/zuc/eea_asm.go
@@ -3,29 +3,13 @@
 
 package zuc
 
-import (
-	"github.com/emmansun/gmsm/internal/subtle"
-)
-
 //go:noescape
 func genKeyStreamRev32Asm(keyStream []byte, pState *zucState32)
 
-func xorKeyStream(c *zucState32, dst, src []byte) {
+func genKeyStreamRev32(keyStream []byte, pState *zucState32) {
 	if supportsAES {
-		words := len(src) / 4
-		// handle complete words first
-		if words > 0 {
-			dstWords := dst[:words*4]
-			genKeyStreamRev32Asm(dstWords, c)
-			subtle.XORBytes(dst, src, dstWords)
-		}
-		// handle remain bytes
-		if words*4 < len(src) {
-			var singleWord [4]byte
-			genKeyStreamRev32Asm(singleWord[:], c)
-			subtle.XORBytes(dst[words*4:], src[words*4:], singleWord[:])
-		}
+		genKeyStreamRev32Asm(keyStream, pState)
 	} else {
-		xorKeyStreamGeneric(c, dst, src)
+		genKeyStreamRev32Generic(keyStream, pState)
 	}
 }
diff --git a/zuc/eea_generic.go b/zuc/eea_generic.go
index 641ee1b..2fb9c22 100644
--- a/zuc/eea_generic.go
+++ b/zuc/eea_generic.go
@@ -3,6 +3,6 @@
 
 package zuc
 
-func xorKeyStream(c *zucState32, dst, src []byte) {
-	xorKeyStreamGeneric(c, dst, src)
+func genKeyStreamRev32(keyStream []byte, pState *zucState32) {
+	genKeyStreamRev32Generic(keyStream, pState)
 }
diff --git a/zuc/eea_test.go b/zuc/eea_test.go
index a822c37..b99ede9 100644
--- a/zuc/eea_test.go
+++ b/zuc/eea_test.go
@@ -50,11 +50,12 @@ func Test_EEA(t *testing.T) {
 			t.Error(err)
 		}
 		in, err := hex.DecodeString(test.in)
-		out := make([]byte, len(in))
 		if err != nil {
 			t.Error(err)
 		}
-		c.XORKeyStream(out, in)
+		out := make([]byte, len(in))
+		copy(out, in)
+		c.XORKeyStream(out, out)
 		if hex.EncodeToString(out) != test.out {
 			t.Errorf("case %d, expected=%s, result=%s\n", i+1, test.out, hex.EncodeToString(out))
 		}