internal/zuc,zuc: eea seakable stream support zuc states cache per bucket #321

internal/zuc/eea.go

@@ -12,20 +12,21 @@ const (
 	RoundWords = 32
 	// number of bytes in a word
 	WordSize = 4
-	WordMask = WordSize - 1
 	// number of bytes in a round
 	RoundBytes = RoundWords * WordSize
 )
 
 type eea struct {
 	zucState32
-	x         [WordSize]byte // remaining bytes buffer
+	x          [RoundBytes]byte // remaining bytes buffer
-	xLen      int            // number of remaining bytes
+	xLen       int              // number of remaining bytes
-	initState zucState32     // initial state for reset
+	used       uint64           // number of key bytes processed, current offset
-	used      uint64         // number of key bytes processed, current offset
+	states     []*zucState32    // internal states for seek
+	stateIndex int              // current state index, for test usage
+	bucketSize int
 }
 
-// NewCipher create a stream cipher based on key and iv aguments.
+// NewCipher creates a stream cipher based on key and iv aguments.
 // The key must be 16 bytes long and iv must be 16 bytes long for zuc 128;
 // or the key must be 32 bytes long and iv must be 23 bytes long for zuc 256;
 // otherwise, an error will be returned.
@@ -36,22 +37,48 @@ func NewCipher(key, iv []byte) (*eea, error) {
 	}
 	c := new(eea)
 	c.zucState32 = *s
-	c.initState = *s
+	c.states = append(c.states, s)
 	c.used = 0
+	c.bucketSize = 0
+	c.stateIndex = 0
 	return c, nil
 }
 
-// NewEEACipher create a stream cipher based on key, count, bearer and direction arguments according specification.
+// NewCipherWithBucketSize creates a new instance of the eea cipher with the specified
-// The key must be 16 bytes long and iv must be 16 bytes long, otherwise, an error will be returned.
+// bucket size. The bucket size is rounded up to the nearest multiple of RoundBytes.
-// The count is the 32-bit counter value, the bearer is the 5-bit bearer identity and the direction is the 1-bit
+func NewCipherWithBucketSize(key, iv []byte, bucketSize int) (*eea, error) {
-// transmission direction flag.
+	c, err := NewCipher(key, iv)
-func NewEEACipher(key []byte, count, bearer, direction uint32) (*eea, error) {
+	if err != nil {
+		return nil, err
+	}
+	if bucketSize > 0 {
+		c.bucketSize = ((bucketSize + RoundBytes - 1) / RoundBytes) * RoundBytes
+	}
+	return c, nil
+}
+
+func construcIV4EEA(count, bearer, direction uint32) []byte {
 	iv := make([]byte, 16)
 	byteorder.BEPutUint32(iv, count)
 	copy(iv[8:12], iv[:4])
 	iv[4] = byte(((bearer << 1) | (direction & 1)) << 2)
 	iv[12] = iv[4]
-	return NewCipher(key, iv)
+	return iv
+}
+
+// NewEEACipher creates a stream cipher based on key, count, bearer and direction arguments according specification.
+// The key must be 16 bytes long and iv must be 16 bytes long, otherwise, an error will be returned.
+// The count is the 32-bit counter value, the bearer is the 5-bit bearer identity and the direction is the 1-bit
+// transmission direction flag.
+func NewEEACipher(key []byte, count, bearer, direction uint32) (*eea, error) {
+	return NewCipher(key, construcIV4EEA(count, bearer, direction))
+}
+
+// NewEEACipherWithBucketSize creates a new instance of the EEA cipher with a specified bucket size.
+// It initializes the cipher using the provided key, count, bearer, and direction parameters,
+// and adjusts the bucket size to be a multiple of RoundBytes.
+func NewEEACipherWithBucketSize(key []byte, count, bearer, direction uint32, bucketSize int) (*eea, error) {
+	return NewCipherWithBucketSize(key, construcIV4EEA(count, bearer, direction), bucketSize)
 }
 
 func genKeyStreamRev32Generic(keyStream []byte, pState *zucState32) {
@@ -62,6 +89,11 @@ func genKeyStreamRev32Generic(keyStream []byte, pState *zucState32) {
 	}
 }
 
+func (c *eea) appendState() {
+	state := c.zucState32
+	c.states = append(c.states, &state)
+}
+
 func (c *eea) XORKeyStream(dst, src []byte) {
 	if len(dst) < len(src) {
 		panic("zuc: output smaller than input")
@@ -69,43 +101,55 @@ func (c *eea) XORKeyStream(dst, src []byte) {
 	if alias.InexactOverlap(dst[:len(src)], src) {
 		panic("zuc: invalid buffer overlap")
 	}
-	used := len(src)
 	if c.xLen > 0 {
 		// handle remaining key bytes
 		n := subtle.XORBytes(dst, src, c.x[:c.xLen])
 		c.xLen -= n
+		c.used += uint64(n)
 		dst = dst[n:]
 		src = src[n:]
 		if c.xLen > 0 {
 			copy(c.x[:], c.x[n:c.xLen+n])
-			c.used += uint64(used)
 			return
 		}
 	}
 	var keyBytes [RoundBytes]byte
+	stepLen := uint64(RoundBytes)
+	nextBucketOffset := c.bucketSize * len(c.states)
 	for len(src) >= RoundBytes {
 		genKeyStreamRev32(keyBytes[:], &c.zucState32)
 		subtle.XORBytes(dst, src, keyBytes[:])
 		dst = dst[RoundBytes:]
 		src = src[RoundBytes:]
-	}
+		c.used += stepLen
-	if len(src) > 0 {
+		if c.bucketSize > 0 && int(c.used) >= nextBucketOffset {
-		byteLen := (len(src) + WordMask) &^ WordMask
+			c.appendState()
-		genKeyStreamRev32(keyBytes[:byteLen], &c.zucState32)
+			nextBucketOffset += c.bucketSize
-		n := subtle.XORBytes(dst, src, keyBytes[:])
-		// save remaining key bytes
-		c.xLen = byteLen - n
-		if c.xLen > 0 {
-			copy(c.x[:], keyBytes[n:byteLen])
 		}
 	}
-	c.used += uint64(used)
+	remaining := len(src)
+	if remaining > 0 {
+		genKeyStreamRev32(keyBytes[:], &c.zucState32)
+		subtle.XORBytes(dst, src, keyBytes[:])
+		c.xLen = RoundBytes - remaining
+		copy(c.x[:], keyBytes[remaining:])
+		if c.bucketSize > 0 && int(c.used)+RoundBytes >= nextBucketOffset {
+			c.appendState()
+		}
+		c.used += uint64(remaining)
+	}
 }
 
-func (c *eea) reset() {
+func (c *eea) reset(offset uint64) {
-	c.zucState32 = c.initState
+	var n uint64
+	if c.bucketSize > 0 {
+		n = offset / uint64(c.bucketSize)
+	}
+	// due to offset < c.used, n must be less than len(c.states)
+	c.stateIndex = int(n)
+	c.zucState32 = *c.states[n]
 	c.xLen = 0
-	c.used = 0
+	c.used = n * uint64(c.bucketSize)
 }
 
 // seek sets the offset for the next XORKeyStream operation.
@@ -116,7 +160,7 @@ func (c *eea) reset() {
 // Note: This method is not thread-safe.
 func (c *eea) seek(offset uint64) {
 	if offset < c.used {
-		c.reset()
+		c.reset(offset)
 	}
 	if offset == c.used {
 		return
@@ -140,24 +184,28 @@ func (c *eea) seek(offset uint64) {
 	}
 
 	// forward the state to the offset
-	c.used += gap
+	nextBucketOffset := c.bucketSize * len(c.states)
 	stepLen := uint64(RoundBytes)
 	var keyStream [RoundWords]uint32
 	for gap >= stepLen {
 		genKeyStream(keyStream[:], &c.zucState32)
 		gap -= stepLen
+		c.used += stepLen
+		if c.bucketSize > 0 && int(c.used) >= nextBucketOffset {
+			c.appendState()
+			nextBucketOffset += c.bucketSize
+		}
 	}
 
 	if gap > 0 {
-		numWords := (gap + WordMask) / WordSize
+		var keyBytes [RoundBytes]byte
-		genKeyStream(keyStream[:numWords], &c.zucState32)
+		genKeyStreamRev32(keyBytes[:], &c.zucState32)
-		partiallyUsed := int(gap & WordMask)
+		c.xLen = RoundBytes - int(gap)
-		if partiallyUsed > 0 {
+		copy(c.x[:], keyBytes[gap:])
-			// save remaining key bytes (less than 4 bytes)
+		if c.bucketSize > 0 && int(c.used)+RoundBytes >= nextBucketOffset {
-			c.xLen = WordSize - partiallyUsed
+			c.appendState()
-			byteorder.BEPutUint32(c.x[:], keyStream[numWords-1])
-			copy(c.x[:], c.x[partiallyUsed:])
 		}
+		c.used += uint64(gap)
 	}
 }
 

internal/zuc/eea_test.go

@@ -203,6 +203,96 @@ func TestIssue284(t *testing.T) {
 	}
 }
 
+func TestEEAXORKeyStreamAtWithBucketSize(t *testing.T) {
+	key, err := hex.DecodeString(zucEEATests[0].key)
+	if err != nil {
+		t.Error(err)
+	}
+	noBucketCipher, err := NewEEACipher(key, zucEEATests[0].count, zucEEATests[0].bearer, zucEEATests[0].direction)
+	if err != nil {
+		t.Error(err)
+	}
+	src := make([]byte, 10000)
+	expected := make([]byte, 10000)
+	dst := make([]byte, 10000)
+	stateCount := 1 + (10000 + RoundBytes -1) / RoundBytes
+	noBucketCipher.XORKeyStream(expected, src)
+
+	t.Run("Make sure the cached states are used once backward", func(t *testing.T) {
+		bucketCipher, err := NewEEACipherWithBucketSize(key, zucEEATests[0].count, zucEEATests[0].bearer, zucEEATests[0].direction, 128)
+		if err != nil {
+			t.Error(err)
+		}
+		bucketCipher.XORKeyStream(dst, src)
+		if !bytes.Equal(expected, dst) {
+			t.Fatalf("expected=%x, result=%x\n", expected, dst)
+		}
+		clear(dst)
+		if len(bucketCipher.states) != stateCount {
+			t.Fatalf("expected=%d, result=%d\n", stateCount, len(bucketCipher.states))
+		}
+		// go backward to offset 128
+		bucketCipher.XORKeyStreamAt(dst[128:256], src[128:256], 128)
+		if bucketCipher.stateIndex != 1 {
+			t.Fatalf("expected=%d, result=%d\n", 1, bucketCipher.stateIndex)
+		}
+		if !bytes.Equal(expected[128:256], dst[128:256]) {
+			t.Fatalf("expected=%x, result=%x\n", expected, dst[128:256])
+		}
+		// go backward to offset 130
+		bucketCipher.XORKeyStreamAt(dst[130:258], src[130:258], 130)
+		if bucketCipher.stateIndex != 1 {
+			t.Fatalf("expected=%d, result=%d\n", 1, bucketCipher.stateIndex)
+		}
+		if !bytes.Equal(expected[130:258], dst[130:258]) {
+			t.Fatalf("expected=%x, result=%x\n", expected[130:258], dst[130:258])
+		}
+		if len(bucketCipher.states) != stateCount {
+			t.Fatalf("expected=%d, result=%d\n", stateCount, len(bucketCipher.states))
+		}
+	})
+
+	t.Run("Forward to offset", 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)
+		bucketCipher.XORKeyStreamAt(dst[256:512], src[256:512], 256)
+		if bucketCipher.stateIndex != 0 {
+			t.Fatalf("expected=%d, result=%d\n", 0, bucketCipher.stateIndex)
+		}
+		if len(bucketCipher.states) != 5 {
+			t.Fatalf("expected=%d, result=%d\n", 5, len(bucketCipher.states))
+		}
+		if !bytes.Equal(expected[256:512], dst[256:512]) {
+			t.Fatalf("expected=%x, result=%x\n", expected[256:512], dst[256:512])
+		}
+		clear(dst)
+		bucketCipher.XORKeyStreamAt(dst[513:768], src[513:768], 513)
+		if bucketCipher.stateIndex != 0 {
+			t.Fatalf("expected=%d, result=%d\n", 0, bucketCipher.stateIndex)
+		}
+		if len(bucketCipher.states) != 7 {
+			t.Fatalf("expected=%d, result=%d\n", 7, len(bucketCipher.states))
+		}
+		if !bytes.Equal(expected[513:768], dst[513:768]) {
+			t.Fatalf("expected=%x, result=%x\n", expected[513:768], dst[513:768])
+		}
+		clear(dst)
+		bucketCipher.XORKeyStreamAt(dst[512:768], src[512:768], 512)
+		if bucketCipher.stateIndex != 4 {
+			t.Fatalf("expected=%d, result=%d\n", 0, bucketCipher.stateIndex)
+		}
+		if len(bucketCipher.states) != 7 {
+			t.Fatalf("expected=%d, result=%d\n", 7, len(bucketCipher.states))
+		}
+		if !bytes.Equal(expected[512:768], dst[512:768]) {
+			t.Fatalf("expected=%x, result=%x\n", expected[512:768], dst[512:768])
+		}
+	})
+}
+
 func benchmarkStream(b *testing.B, buf []byte) {
 	b.SetBytes(int64(len(buf)))
 
@@ -236,7 +326,7 @@ func benchmarkSeek(b *testing.B, offset uint64) {
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		eea.reset()
+		eea.reset(0)
 		eea.seek(offset)
 	}
 }

zuc/eea.go

@@ -8,15 +8,15 @@ import (
 
 const (
 	// IV size in bytes for zuc 128
-	IVSize128 = 16
+	IVSize128 = zuc.IVSize128
 	// IV size in bytes for zuc 256
-	IVSize256 = 23
+	IVSize256 = zuc.IVSize256
 	// number of words in a round
-	RoundWords = 32
+	RoundWords = zuc.RoundWords
 	// number of bytes in a word
-	WordSize = 4
+	WordSize = zuc.WordSize
 	// number of bytes in a round
-	RoundBytes = RoundWords * WordSize
+	RoundBytes = zuc.RoundBytes
 )
 
 // NewCipher create a stream cipher based on key and iv aguments.
@@ -34,3 +34,22 @@ func NewCipher(key, iv []byte) (cipher.SeekableStream, error) {
 func NewEEACipher(key []byte, count, bearer, direction uint32) (cipher.SeekableStream, error) {
 	return zuc.NewEEACipher(key, count, bearer, direction)
 }
+
+// NewCipherWithBucketSize create a new instance of the eea cipher with the specified
+// bucket size. The bucket size is rounded up to the nearest multiple of RoundBytes.
+//
+// The implementation of this function is used for XORKeyStreamAt function optimization, which will keep states
+// for seekable stream cipher once the bucketSize is greater than 0.
+func NewCipherWithBucketSize(key, iv []byte, bucketSize int) (cipher.SeekableStream, error) {
+	return zuc.NewCipherWithBucketSize(key, iv, bucketSize)
+}
+
+// NewEEACipherWithBucketSize creates a new instance of a seekable stream cipher
+// for the EEA encryption algorithm with a specified bucket size. This function
+// is typically used in mobile communication systems for secure data encryption.
+//
+// The implementation of this function is used for XORKeyStreamAt function optimization, which will keep states
+// for seekable stream cipher once the bucketSize is greater than 0.
+func NewEEACipherWithBucketSize(key []byte, count, bearer, direction uint32, bucketSize int) (cipher.SeekableStream, error) {
+	return zuc.NewEEACipherWithBucketSize(key, count, bearer, direction, bucketSize)
+}