From 81c99681885603c9604574a4c9937fdbd92cd2bd Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Wed, 12 Sep 2018 15:43:36 +0200
Subject: [PATCH] simplify crypto related code

---
 common/crypto/aes.go           |  8 ++++++++
 proxy/vmess/encoding/client.go | 18 ++++++------------
 proxy/vmess/encoding/server.go | 14 ++++----------
 3 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/common/crypto/aes.go b/common/crypto/aes.go
index 5a002561..94a49735 100644
--- a/common/crypto/aes.go
+++ b/common/crypto/aes.go
@@ -28,3 +28,11 @@ func NewAesStreamMethod(key []byte, iv []byte, f func(cipher.Block, []byte) ciph
 func NewAesCTRStream(key []byte, iv []byte) cipher.Stream {
 	return NewAesStreamMethod(key, iv, cipher.NewCTR)
 }
+
+func NewAesGcm(key []byte) cipher.AEAD {
+	block, err := aes.NewCipher(key)
+	common.Must(err)
+	aead, err := cipher.NewGCM(block)
+	common.Must(err)
+	return aead
+}
diff --git a/proxy/vmess/encoding/client.go b/proxy/vmess/encoding/client.go
index df080038..30ea48e3 100644
--- a/proxy/vmess/encoding/client.go
+++ b/proxy/vmess/encoding/client.go
@@ -1,8 +1,6 @@
 package encoding
 
 import (
-	"crypto/aes"
-	"crypto/cipher"
 	"crypto/md5"
 	"crypto/rand"
 	"hash/fnv"
@@ -105,10 +103,8 @@ func (c *ClientSession) EncodeRequestHeader(header *protocol.RequestHeader, writ
 		common.Must(buffer.AppendSupplier(serial.WriteHash(fnv1a)))
 	}
 
-	timestampHash := md5.New()
-	common.Must2(timestampHash.Write(hashTimestamp(timestamp)))
-	iv := timestampHash.Sum(nil)
-	aesStream := crypto.NewAesEncryptionStream(account.ID.CmdKey(), iv)
+	iv := md5.Sum(hashTimestamp(timestamp))
+	aesStream := crypto.NewAesEncryptionStream(account.ID.CmdKey(), iv[:])
 	aesStream.XORKeyStream(buffer.Bytes(), buffer.Bytes())
 	common.Must2(writer.Write(buffer.Bytes()))
 	return nil
@@ -153,9 +149,7 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 
 		return &buf.SequentialWriter{Writer: cryptionWriter}
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(c.requestBodyKey[:])
-		aead, _ := cipher.NewGCM(block)
-
+		aead := crypto.NewAesGcm(c.requestBodyKey[:])
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
 			NonceGenerator:          GenerateChunkNonce(c.requestBodyIV[:], uint32(aead.NonceSize())),
@@ -163,7 +157,8 @@ func (c *ClientSession) EncodeRequestBody(request *protocol.RequestHeader, write
 		}
 		return crypto.NewAuthenticationWriter(auth, sizeParser, writer, request.Command.TransferType(), padding)
 	case protocol.SecurityType_CHACHA20_POLY1305:
-		aead, _ := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.requestBodyKey[:]))
+		aead, err := chacha20poly1305.New(GenerateChacha20Poly1305Key(c.requestBodyKey[:]))
+		common.Must(err)
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
@@ -250,8 +245,7 @@ func (c *ClientSession) DecodeResponseBody(request *protocol.RequestHeader, read
 
 		return buf.NewReader(c.responseReader)
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(c.responseBodyKey[:])
-		aead, _ := cipher.NewGCM(block)
+		aead := crypto.NewAesGcm(c.responseBodyKey[:])
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
diff --git a/proxy/vmess/encoding/server.go b/proxy/vmess/encoding/server.go
index e5736f5e..55442b7c 100644
--- a/proxy/vmess/encoding/server.go
+++ b/proxy/vmess/encoding/server.go
@@ -1,8 +1,6 @@
 package encoding
 
 import (
-	"crypto/aes"
-	"crypto/cipher"
 	"crypto/md5"
 	"hash/fnv"
 	"io"
@@ -147,12 +145,10 @@ func (s *ServerSession) DecodeRequestHeader(reader io.Reader) (*protocol.Request
 		return nil, newError("invalid user")
 	}
 
-	timestampHash := md5.New()
-	common.Must2(timestampHash.Write(hashTimestamp(timestamp)))
-	iv := timestampHash.Sum(nil)
+	iv := md5.Sum(hashTimestamp(timestamp))
 	vmessAccount := user.Account.(*vmess.InternalAccount)
 
-	aesStream := crypto.NewAesDecryptionStream(vmessAccount.ID.CmdKey(), iv)
+	aesStream := crypto.NewAesDecryptionStream(vmessAccount.ID.CmdKey(), iv[:])
 	decryptor := crypto.NewCryptionReader(aesStream, reader)
 
 	if err := buffer.Reset(buf.ReadFullFrom(decryptor, 38)); err != nil {
@@ -263,8 +259,7 @@ func (s *ServerSession) DecodeRequestBody(request *protocol.RequestHeader, reade
 
 		return buf.NewReader(cryptionReader)
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(s.requestBodyKey[:])
-		aead, _ := cipher.NewGCM(block)
+		aead := crypto.NewAesGcm(s.requestBodyKey[:])
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,
@@ -341,8 +336,7 @@ func (s *ServerSession) EncodeResponseBody(request *protocol.RequestHeader, writ
 
 		return &buf.SequentialWriter{Writer: s.responseWriter}
 	case protocol.SecurityType_AES128_GCM:
-		block, _ := aes.NewCipher(s.responseBodyKey[:])
-		aead, _ := cipher.NewGCM(block)
+		aead := crypto.NewAesGcm(s.responseBodyKey[:])
 
 		auth := &crypto.AEADAuthenticator{
 			AEAD:                    aead,