simplify BytesGenerator

common/crypto/auth.go

@@ -9,44 +9,36 @@ import (
 	"v2ray.com/core/common/protocol"
 )
 
-type BytesGenerator interface {
-	Next() []byte
-}
-
-type NoOpBytesGenerator struct {
-	buffer [1]byte
-}
-
-func (v NoOpBytesGenerator) Next() []byte {
-	return v.buffer[:0]
-}
-
-type StaticBytesGenerator struct {
-	Content []byte
-}
+type BytesGenerator func() []byte
 
-func (v StaticBytesGenerator) Next() []byte {
-	return v.Content
-}
-
-type IncreasingAEADNonceGenerator struct {
-	nonce []byte
+func GenerateEmptyBytes() BytesGenerator {
+	var b [1]byte
+	return func() []byte {
+		return b[:0]
+	}
 }
 
-func NewIncreasingAEADNonceGenerator() *IncreasingAEADNonceGenerator {
-	return &IncreasingAEADNonceGenerator{
-		nonce: []byte{0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF},
+func GenerateStaticBytes(content []byte) BytesGenerator {
+	return func() []byte {
+		return content
 	}
 }
 
-func (g *IncreasingAEADNonceGenerator) Next() []byte {
-	for i := range g.nonce {
-		g.nonce[i]++
-		if g.nonce[i] != 0 {
-			break
+func GenerateIncreasingNonce(nonce []byte) BytesGenerator {
+	c := append([]byte(nil), nonce...)
+	return func() []byte {
+		for i := range c {
+			c[i]++
+			if c[i] != 0 {
+				break
+			}
 		}
+		return c
 	}
-	return g.nonce
+}
+
+func GenerateInitialAEADNonce() BytesGenerator {
+	return GenerateIncreasingNonce([]byte{0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF})
 }
 
 type Authenticator interface {
@@ -63,27 +55,27 @@ type AEADAuthenticator struct {
 }
 
 func (v *AEADAuthenticator) Open(dst, cipherText []byte) ([]byte, error) {
-	iv := v.NonceGenerator.Next()
+	iv := v.NonceGenerator()
 	if len(iv) != v.AEAD.NonceSize() {
 		return nil, newError("invalid AEAD nonce size: ", len(iv))
 	}
 
 	var additionalData []byte
 	if v.AdditionalDataGenerator != nil {
-		additionalData = v.AdditionalDataGenerator.Next()
+		additionalData = v.AdditionalDataGenerator()
 	}
 	return v.AEAD.Open(dst, iv, cipherText, additionalData)
 }
 
 func (v *AEADAuthenticator) Seal(dst, plainText []byte) ([]byte, error) {
-	iv := v.NonceGenerator.Next()
+	iv := v.NonceGenerator()
 	if len(iv) != v.AEAD.NonceSize() {
 		return nil, newError("invalid AEAD nonce size: ", len(iv))
 	}
 
 	var additionalData []byte
 	if v.AdditionalDataGenerator != nil {
-		additionalData = v.AdditionalDataGenerator.Next()
+		additionalData = v.AdditionalDataGenerator()
 	}
 	return v.AEAD.Seal(dst, iv, plainText, additionalData), nil
 }

common/crypto/auth_test.go

@@ -37,11 +37,9 @@ func TestAuthenticationReaderWriter(t *testing.T) {
 	rand.Read(iv)
 
 	writer := NewAuthenticationWriter(&AEADAuthenticator{
-		AEAD: aead,
-		NonceGenerator: &StaticBytesGenerator{
-			Content: iv,
-		},
-		AdditionalDataGenerator: &NoOpBytesGenerator{},
+		AEAD:                    aead,
+		NonceGenerator:          GenerateStaticBytes(iv),
+		AdditionalDataGenerator: GenerateEmptyBytes(),
 	}, PlainChunkSizeParser{}, cache, protocol.TransferTypeStream)
 
 	assert(writer.WriteMultiBuffer(buf.NewMultiBufferValue(payload)), IsNil)
@@ -49,11 +47,9 @@ func TestAuthenticationReaderWriter(t *testing.T) {
 	assert(writer.WriteMultiBuffer(buf.MultiBuffer{}), IsNil)
 
 	reader := NewAuthenticationReader(&AEADAuthenticator{
-		AEAD: aead,
-		NonceGenerator: &StaticBytesGenerator{
-			Content: iv,
-		},
-		AdditionalDataGenerator: &NoOpBytesGenerator{},
+		AEAD:                    aead,
+		NonceGenerator:          GenerateStaticBytes(iv),
+		AdditionalDataGenerator: GenerateEmptyBytes(),
 	}, PlainChunkSizeParser{}, cache, protocol.TransferTypeStream)
 
 	var mb buf.MultiBuffer
@@ -91,11 +87,9 @@ func TestAuthenticationReaderWriterPacket(t *testing.T) {
 	rand.Read(iv)
 
 	writer := NewAuthenticationWriter(&AEADAuthenticator{
-		AEAD: aead,
-		NonceGenerator: &StaticBytesGenerator{
-			Content: iv,
-		},
-		AdditionalDataGenerator: &NoOpBytesGenerator{},
+		AEAD:                    aead,
+		NonceGenerator:          GenerateStaticBytes(iv),
+		AdditionalDataGenerator: GenerateEmptyBytes(),
 	}, PlainChunkSizeParser{}, cache, protocol.TransferTypePacket)
 
 	var payload buf.MultiBuffer
@@ -113,11 +107,9 @@ func TestAuthenticationReaderWriterPacket(t *testing.T) {
 	assert(err, IsNil)
 
 	reader := NewAuthenticationReader(&AEADAuthenticator{
-		AEAD: aead,
-		NonceGenerator: &StaticBytesGenerator{
-			Content: iv,
-		},
-		AdditionalDataGenerator: &NoOpBytesGenerator{},
+		AEAD:                    aead,
+		NonceGenerator:          GenerateStaticBytes(iv),
+		AdditionalDataGenerator: GenerateEmptyBytes(),
 	}, PlainChunkSizeParser{}, cache, protocol.TransferTypePacket)
 
 	mb, err := reader.ReadMultiBuffer()

proxy/shadowsocks/config.go

@@ -169,7 +169,7 @@ func (c *AEADCipher) IVSize() int32 {
 }
 
 func (c *AEADCipher) createAuthenticator(key []byte, iv []byte) *crypto.AEADAuthenticator {
-	nonce := crypto.NewIncreasingAEADNonceGenerator()
+	nonce := crypto.GenerateInitialAEADNonce()
 	subkey := make([]byte, c.KeyBytes)
 	hkdfSHA1(key, iv, subkey)
 	return &crypto.AEADAuthenticator{

proxy/vmess/encoding/client.go

@@ -120,8 +120,8 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 			}
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(NoOpAuthenticator),
-				NonceGenerator:          crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 			return crypto.NewAuthenticationWriter(auth, sizeParser, writer, protocol.TransferTypePacket)
 		}
@@ -133,8 +133,8 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 		if request.Option.Has(protocol.RequestOptionChunkStream) {
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(FnvAuthenticator),
-				NonceGenerator:          crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 			return crypto.NewAuthenticationWriter(auth, sizeParser, cryptionWriter, request.Command.TransferType())
 		}
@@ -145,24 +145,18 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), c.requestBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(c.requestBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
 		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.requestBodyKey))
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), c.requestBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(c.requestBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())
 	default:
@@ -219,8 +213,8 @@ func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, read
 
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(NoOpAuthenticator),
-				NonceGenerator:          crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 
 			return crypto.NewAuthenticationReader(auth, sizeParser, reader, protocol.TransferTypePacket)
@@ -231,8 +225,8 @@ func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, read
 		if request.Option.Has(protocol.RequestOptionChunkStream) {
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(FnvAuthenticator),
-				NonceGenerator:          crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 			return crypto.NewAuthenticationReader(auth, sizeParser, c.responseReader, request.Command.TransferType())
 		}
@@ -243,24 +237,18 @@ func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, read
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), c.responseBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(c.responseBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
 		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.responseBodyKey))
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), c.responseBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(c.responseBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())
 	default:
@@ -268,14 +256,12 @@ func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, read
 	}
 }
 
-type ChunkNonceGenerator struct {
-	Nonce []byte
-	Size  int
-	count uint16
-}
-
-func (g *ChunkNonceGenerator) Next() []byte {
-	serial.Uint16ToBytes(g.count, g.Nonce[:0])
-	g.count++
-	return g.Nonce[:g.Size]
+func GenerateChunkNonce(nonce []byte, size uint32) crypto.BytesGenerator {
+	c := append([]byte(nil), nonce...)
+	count := uint16(0)
+	return func() []byte {
+		serial.Uint16ToBytes(count, c[:0])
+		count++
+		return c[:size]
+	}
 }

proxy/vmess/encoding/server.go

@@ -9,13 +9,13 @@ import (
 	"sync"
 	"time"
 
-	"v2ray.com/core/common/dice"
-
 	"golang.org/x/crypto/chacha20poly1305"
+
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/bitmask"
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/crypto"
+	"v2ray.com/core/common/dice"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/protocol"
 	"v2ray.com/core/common/serial"
@@ -238,8 +238,8 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(NoOpAuthenticator),
-				NonceGenerator:          crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 			return crypto.NewAuthenticationReader(auth, sizeParser, reader, protocol.TransferTypePacket)
 		}
@@ -251,8 +251,8 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 		if request.Option.Has(protocol.RequestOptionChunkStream) {
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(FnvAuthenticator),
-				NonceGenerator:          crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 			return crypto.NewAuthenticationReader(auth, sizeParser, cryptionReader, request.Command.TransferType())
 		}
@@ -263,24 +263,18 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), s.requestBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
 		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.requestBodyKey))
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), s.requestBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(s.requestBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationReader(auth, sizeParser, reader, request.Command.TransferType())
 	default:
@@ -319,8 +313,8 @@ func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writ
 
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(NoOpAuthenticator),
-				NonceGenerator:          &crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 			return crypto.NewAuthenticationWriter(auth, sizeParser, writer, protocol.TransferTypePacket)
 		}
@@ -330,8 +324,8 @@ func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writ
 		if request.Option.Has(protocol.RequestOptionChunkStream) {
 			auth := &crypto.AEADAuthenticator{
 				AEAD:                    new(FnvAuthenticator),
-				NonceGenerator:          crypto.NoOpBytesGenerator{},
-				AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+				NonceGenerator:          crypto.GenerateEmptyBytes(),
+				AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 			}
 			return crypto.NewAuthenticationWriter(auth, sizeParser, s.responseWriter, request.Command.TransferType())
 		}
@@ -342,24 +336,18 @@ func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writ
 		aead, _ := cipher.NewGCM(block)
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), s.responseBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())
 	case protocol.SecurityType_CHACHA20_POLY1305:
 		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(s.responseBodyKey))
 
 		auth := &crypto.AEADAuthenticator{
-			AEAD: aead,
-			NonceGenerator: &ChunkNonceGenerator{
-				Nonce: append([]byte(nil), s.responseBodyIV...),
-				Size:  aead.NonceSize(),
-			},
-			AdditionalDataGenerator: crypto.NoOpBytesGenerator{},
+			AEAD:                    aead,
+			NonceGenerator:          GenerateChunkNonce(s.responseBodyIV, uint32(aead.NonceSize())),
+			AdditionalDataGenerator: crypto.GenerateEmptyBytes(),
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType())
 	default: