Refactor AES encryption/decryption

Schezuk · Nov 3, 2015 · a46db06 · a46db06
1 parent 3f884ed
commit a46db06
Show file tree

Hide file tree

Showing 7 changed files with 90 additions and 109 deletions.
diff --git a/common/crypto/aes.go b/common/crypto/aes.go
@@ -0,0 +1,62 @@
+package crypto
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"io"
+)
+
+func NewAesDecryptionStream(key []byte, iv []byte) (cipher.Stream, error) {
+	aesBlock, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	return cipher.NewCFBDecrypter(aesBlock, iv), nil
+}
+
+func NewAesEncryptionStream(key []byte, iv []byte) (cipher.Stream, error) {
+	aesBlock, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	return cipher.NewCFBEncrypter(aesBlock, iv), nil
+}
+
+type cryptionReader struct {
+	stream cipher.Stream
+	reader io.Reader
+}
+
+func NewCryptionReader(stream cipher.Stream, reader io.Reader) io.Reader {
+	return &cryptionReader{
+		stream: stream,
+		reader: reader,
+	}
+}
+
+func (this *cryptionReader) Read(data []byte) (int, error) {
+	nBytes, err := this.reader.Read(data)
+	if nBytes > 0 {
+		this.stream.XORKeyStream(data[:nBytes], data[:nBytes])
+	}
+	return nBytes, err
+}
+
+type cryptionWriter struct {
+	stream cipher.Stream
+	writer io.Writer
+}
+
+func NewCryptionWriter(stream cipher.Stream, writer io.Writer) io.Writer {
+	return &cryptionWriter{
+		stream: stream,
+		writer: writer,
+	}
+}
+
+func (this *cryptionWriter) Write(data []byte) (int, error) {
+	this.stream.XORKeyStream(data, data)
+	return this.writer.Write(data)
+}
diff --git a/common/io/aes.go b/common/io/aes.go
diff --git a/common/io/encryption.go b/common/io/encryption.go
diff --git a/proxy/vmess/protocol/vmess.go b/proxy/vmess/protocol/vmess.go
@@ -2,15 +2,13 @@
 package protocol
 
 import (
-	"crypto/aes"
-	"crypto/cipher"
 	"encoding/binary"
 	"hash/fnv"
 	"io"
 	"time"
 
 	"github.com/v2ray/v2ray-core/common/alloc"
-	v2io "github.com/v2ray/v2ray-core/common/io"
+	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
 	"github.com/v2ray/v2ray-core/common/log"
 	v2net "github.com/v2ray/v2ray-core/common/net"
 	proxyerrors "github.com/v2ray/v2ray-core/proxy/common/errors"
@@ -80,16 +78,12 @@ func (r *VMessRequestReader) Read(reader io.Reader) (*VMessRequest, error) {
 		return nil, proxyerrors.InvalidAuthentication
 	}
 
-	aesCipher, err := aes.NewCipher(userObj.ID().CmdKey())
+	aesStream, err := v2crypto.NewAesDecryptionStream(userObj.ID().CmdKey(), user.Int64Hash(timeSec))
 	if err != nil {
 		return nil, err
 	}
-	aesStream := cipher.NewCFBDecrypter(aesCipher, user.Int64Hash(timeSec))
-	decryptor := v2io.NewCryptionReader(aesStream, reader)
 
-	if err != nil {
-		return nil, err
-	}
+	decryptor := v2crypto.NewCryptionReader(aesStream, reader)
 
 	nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[:41])
 	if err != nil {
@@ -201,11 +195,10 @@ func (request *VMessRequest) ToBytes(idHash user.CounterHash, randomRangeInt64 u
 	buffer.AppendBytes(byte(fnvHash>>24), byte(fnvHash>>16), byte(fnvHash>>8), byte(fnvHash))
 	encryptionEnd += 4
 
-	aesCipher, err := aes.NewCipher(request.User.ID().CmdKey())
+	aesStream, err := v2crypto.NewAesEncryptionStream(request.User.ID().CmdKey(), user.Int64Hash(counter))
 	if err != nil {
 		return nil, err
 	}
-	aesStream := cipher.NewCFBEncrypter(aesCipher, user.Int64Hash(counter))
 	aesStream.XORKeyStream(buffer.Value[encryptionBegin:encryptionEnd], buffer.Value[encryptionBegin:encryptionEnd])
 
 	return buffer, nil

diff --git a/proxy/vmess/vmessin.go b/proxy/vmess/vmessin.go
@@ -8,7 +8,7 @@ import (
 
 	"github.com/v2ray/v2ray-core/app"
 	"github.com/v2ray/v2ray-core/common/alloc"
-	v2io "github.com/v2ray/v2ray-core/common/io"
+	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
 	"github.com/v2ray/v2ray-core/common/log"
 	v2net "github.com/v2ray/v2ray-core/common/net"
 	"github.com/v2ray/v2ray-core/common/retry"
@@ -98,12 +98,14 @@ func (handler *VMessInboundHandler) HandleConnection(connection *net.TCPConn) er
 	responseKey := md5.Sum(request.RequestKey)
 	responseIV := md5.Sum(request.RequestIV)
 
-	responseWriter, err := v2io.NewAesEncryptWriter(responseKey[:], responseIV[:], connection)
+	aesStream, err := v2crypto.NewAesEncryptionStream(responseKey[:], responseIV[:])
 	if err != nil {
-		log.Error("VMessIn: Failed to create encrypt writer: %v", err)
+		log.Error("VMessIn: Failed to create AES decryption stream: %v", err)
 		return err
 	}
 
+	responseWriter := v2crypto.NewCryptionWriter(aesStream, connection)
+
 	// Optimize for small response packet
 	buffer := alloc.NewLargeBuffer().Clear()
 	buffer.Append(request.ResponseHeader)
@@ -127,12 +129,12 @@ func handleInput(request *protocol.VMessRequest, reader io.Reader, input chan<-
 	defer close(input)
 	defer finish.Unlock()
 
-	requestReader, err := v2io.NewAesDecryptReader(request.RequestKey, request.RequestIV, reader)
+	aesStream, err := v2crypto.NewAesDecryptionStream(request.RequestKey, request.RequestIV)
 	if err != nil {
-		log.Error("VMessIn: Failed to create decrypt reader: %v", err)
+		log.Error("VMessIn: Failed to create AES decryption stream: %v", err)
 		return
 	}
-
+	requestReader := v2crypto.NewCryptionReader(aesStream, reader)
 	v2net.ReaderToChan(input, requestReader)
 }
 

diff --git a/proxy/vmess/vmessin_udp.go b/proxy/vmess/vmessin_udp.go
@@ -6,7 +6,7 @@ import (
 	"net"
 
 	"github.com/v2ray/v2ray-core/common/alloc"
-	v2io "github.com/v2ray/v2ray-core/common/io"
+	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
 	"github.com/v2ray/v2ray-core/common/log"
 	v2net "github.com/v2ray/v2ray-core/common/net"
 	"github.com/v2ray/v2ray-core/proxy/vmess/protocol"
@@ -54,12 +54,13 @@ func (handler *VMessInboundHandler) AcceptPackets(conn *net.UDPConn) {
 		}
 		log.Access(addr.String(), request.Address.String(), log.AccessAccepted, "")
 
-		cryptReader, err := v2io.NewAesDecryptReader(request.RequestKey, request.RequestIV, reader)
+		aesStream, err := v2crypto.NewAesDecryptionStream(request.RequestKey, request.RequestIV)
 		if err != nil {
-			log.Error("VMessIn: Failed to create decrypt reader: %v", err)
+			log.Error("VMessIn: Failed to AES decryption stream: %v", err)
 			buffer.Release()
 			continue
 		}
+		cryptReader := v2crypto.NewCryptionReader(aesStream, reader)
 
 		data := alloc.NewBuffer()
 		nBytes, err = cryptReader.Read(data.Value)
@@ -86,11 +87,13 @@ func (handler *VMessInboundHandler) handlePacket(conn *net.UDPConn, request *pro
 	buffer := alloc.NewBuffer().Clear()
 	defer buffer.Release()
 
-	responseWriter, err := v2io.NewAesEncryptWriter(responseKey[:], responseIV[:], buffer)
+	aesStream, err := v2crypto.NewAesEncryptionStream(responseKey[:], responseIV[:])
 	if err != nil {
-		log.Error("VMessIn: Failed to create encrypt writer: %v", err)
+		log.Error("VMessIn: Failed to create AES encryption stream: %v", err)
 		return
 	}
+	responseWriter := v2crypto.NewCryptionWriter(aesStream, buffer)
+
 	responseWriter.Write(request.ResponseHeader)
 
 	hasData := false

diff --git a/proxy/vmess/vmessout.go b/proxy/vmess/vmessout.go
@@ -9,7 +9,7 @@ import (
 	"sync"
 
 	"github.com/v2ray/v2ray-core/common/alloc"
-	v2io "github.com/v2ray/v2ray-core/common/io"
+	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
 	"github.com/v2ray/v2ray-core/common/log"
 	v2net "github.com/v2ray/v2ray-core/common/net"
 	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
@@ -114,11 +114,12 @@ func startCommunicate(request *protocol.VMessRequest, dest v2net.Destination, ra
 
 func handleRequest(conn net.Conn, request *protocol.VMessRequest, firstPacket v2net.Packet, input <-chan *alloc.Buffer, finish *sync.Mutex) {
 	defer finish.Unlock()
-	encryptRequestWriter, err := v2io.NewAesEncryptWriter(request.RequestKey[:], request.RequestIV[:], conn)
+	aesStream, err := v2crypto.NewAesEncryptionStream(request.RequestKey[:], request.RequestIV[:])
 	if err != nil {
-		log.Error("VMessOut: Failed to create encrypt writer: %v", err)
+		log.Error("VMessOut: Failed to create AES encryption stream: %v", err)
 		return
 	}
+	encryptRequestWriter := v2crypto.NewCryptionWriter(aesStream, conn)
 
 	buffer := alloc.NewBuffer().Clear()
 	buffer, err = request.ToBytes(user.NewTimeHash(user.HMACHash{}), user.GenerateRandomInt64InRange, buffer)
@@ -136,7 +137,7 @@ func handleRequest(conn net.Conn, request *protocol.VMessRequest, firstPacket v2
 	}
 
 	if firstChunk != nil {
-		encryptRequestWriter.Crypt(firstChunk.Value)
+		aesStream.XORKeyStream(firstChunk.Value, firstChunk.Value)
 		buffer.Append(firstChunk.Value)
 		firstChunk.Release()
 
@@ -160,11 +161,12 @@ func handleResponse(conn net.Conn, request *protocol.VMessRequest, output chan<-
 	responseKey := md5.Sum(request.RequestKey[:])
 	responseIV := md5.Sum(request.RequestIV[:])
 
-	decryptResponseReader, err := v2io.NewAesDecryptReader(responseKey[:], responseIV[:], conn)
+	aesStream, err := v2crypto.NewAesDecryptionStream(responseKey[:], responseIV[:])
 	if err != nil {
-		log.Error("VMessOut: Failed to create decrypt reader: %v", err)
+		log.Error("VMessOut: Failed to create AES encryption stream: %v", err)
 		return
 	}
+	decryptResponseReader := v2crypto.NewCryptionReader(aesStream, conn)
 
 	buffer, err := v2net.ReadFrom(decryptResponseReader, nil)
 	if err != nil {